python json文件数据恢复_python处理json文件(Yelp数据集)

python脚本处理yelp数据集

import sys

import json

import re

import os

import time

if __name__ == '__main__':

dataset_sizes = {'sample': (40, 10), 'small': (1280, 320), 'medium': (32000, 8000), 'large': (2000000, 400000)}

dataset_name = sys.argv[1]

num_train, num_test = dataset_sizes[dataset_name]

os.chdir('yelp_dataset')

os.system('head -n ' + str(num_train + num_test) + ' review.json > temp.json')

os.system('mv temp.json ..')

os.chdir('..')

with open('temp.json', 'r') as f:

dataset = f.read().split('\n')

dataset.remove('')

print("Done loading in dataset")

try:

os.mkdir('datasets')

except:

pass

os.chdir('datasets')

try:

os.mkdir(dataset_name)

except:

pass

os.chdir(dataset_name)

train = open('train.txt', 'w')

test = open('test.txt', 'w')

reviews = [json.loads(review) for review in dataset]

count = 1

total = len(reviews)

for review in reviews:

review_id = review['review_id']

num_stars = review['stars']

if num_stars > 4.0:

num_stars = 5

elif num_stars > 2.0:

num_stars = 3

else:

num_stars = 1

review_text = review['text']

review_text = re.sub(r'[^\w\s]','', review_text).replace("\n", "").replace("\r", "").lower()

if num_train > 0:

train.write(str(review_id) + ' ' + str(num_stars) + ' ' + str(review_text))

train.write('\n')

num_train -= 1

elif num_test > 0:

test.write(str(review_id) + ' ' + str(num_stars) + ' ' + str(review_text))

test.write('\n')

num_test -= 1

bar_len = 30

filled_len = int(round(bar_len * count / float(total)))

percents = round(100.0 * count / float(total), 1)

bar = '=' * filled_len + '-' * (bar_len - filled_len)

sys.stdout.write("\r[%s] %s%s ... %d/%d reviews created" % (bar, percents, '%', count, total))

sys.stdout.flush()

count += 1

sys.stdout.write("\n")

train.close()

test.close()

os.chdir('../../')

os.system('rm -f temp.json')