| import json | |
| import bz2 | |
| import os | |
| import csv | |
| path = '/home/icml01/multi_rag/RAG/Search-in-the-Chain/data/enwiki-20171001-pages-meta-current-withlinks-abstracts' | |
| def load_data(path=None): | |
| assert path | |
| corpus = [] | |
| count = 0 | |
| with open( | |
| '/home/icml01/multi_rag/RAG/Search-in-the-Chain/data/enwiki-20171001-pages-meta-current-withlinks-abstracts.tsv', | |
| 'w', encoding='utf-8-sig') as w: | |
| writer = csv.writer(w, delimiter="\t") | |
| for root, dirs, files in os.walk(path): | |
| for file in files: | |
| file_path = root + '/' + file | |
| with bz2.open(file_path, 'r') as fin: | |
| for line in fin.readlines(): | |
| data = json.loads(line) | |
| row_text = '' | |
| for text in data['text']: | |
| row_text = row_text + ' ' + text | |
| writer.writerow([count, 'Title: '+data['title']+' Text: '+row_text]) #pid \t passage text | |
| count += 1 | |
| print(count) | |
| return corpus | |
| load_data(path) | |