DJCheng/embedding_data

1import json2import os3 4lines = []5with open('/etc/ssd1/jiangzhongtao/baai_embedding_tune/data/all_collect/p+a_hn+aug_en_task_files.txt', 'r') as f:6    for line in f:7        lines.append(line.strip().split('\t'))8 9with open('/etc/ssd1/jiangzhongtao/baai_embedding_tune/data/task_meta/casual/all.json', 'r') as f:10    meta_info = json.load(f)11 12with open('/etc/ssd1/jiangzhongtao/baai_embedding_tune/data/all_collect/p+a_hn+aug_en_retrieval_task_files.txt', 'w') as f:13    for file, train_size in lines:14        _, task_name, _ = os.path.split(file)[-1].split('_')15        if meta_info[task_name]['task_type'] == 'retrieval':16            f.write(f'{file}\t{train_size}\n')