import os
import json

dir = "/home/tiger/nlpcc/clean"

for split in ['train', 'dev', 'test']:
    documents = []
    summaries = []
    with open(os.path.join(dir, "{}.simple.label.jsonl".format(split)), 'r') as fin:
        for line in fin:
            info = json.loads(line.strip())
            documents.append(info['article'])
            summaries.append(info['summary'])
    with open(os.path.join(dir, "{}.doc".format(split)), 'w') as fout:
        for doc in documents:
            fout.write(' <q> '.join(doc) + '\n')

    with open(os.path.join(dir, "{}.sum".format(split)), 'w') as fout:
        for sum in summaries:
            fout.write(sum + '\n')
    