import json

from tqdm import tqdm

### read: cast 22 eval json
eval_json_path="/root/Corpus/CAsT22_eval_queries/cqr_inferred_results.json"

# read data
with open(eval_json_path, 'r') as fr:
    data = json.load(fr)

# write: queries.tsv (format: {qid}\t{query})
eval_queries_path = "/root/Corpus/CAsT22_eval_queries/queries_cqr_result.tsv"
with open(eval_queries_path, 'w') as fw:
    for sample in tqdm(data):
        conv_id = sample['number']
        for turn in sample['turn']:
            turn_id = turn['number']
            automatic_rewritten_utterance = turn['automatic_rewritten_utterance']

            q_id = f"{conv_id}_{turn_id}"
            
            fw.write(f"{q_id}\t{automatic_rewritten_utterance}\n")