DJCheng/embedding_data

1import os2import pandas as pd3import json4from tqdm import tqdm5from glob import glob6from collections import defaultdict7 8# root_path = '/etc/ssd1/jiangzhongtao/baai_embedding_tune/data/all_collect/mteb-classification'9# fnames = os.listdir(root_path)10# dataset2size = defaultdict(int)11# for fname in fnames:12#     _, name, lang = fname.split('_')13#     fname = os.path.join(root_path, fname)14#     with open(fname, 'r') as f:15#         for line in f:16#             dataset2size[name] += 117 18src = 'mteb-Clustering'19# postfix = '_hn'20 21with open(f'/etc/ssd1/jiangzhongtao/baai_embedding_tune/data/all_collect/instructions/{src}_instructions.json') as f:22    instructions = json.load(f)23 24# src += postfix25root_path = f'/etc/ssd1/jiangzhongtao/baai_embedding_tune/data/all_collect/{src}'26fnames = os.listdir(root_path)27pbar = tqdm()28for fname in tqdm(fnames):29    _, name, lang = fname[:-6].split('_')30    instruction = instructions[name]31    f_in = open(os.path.join(root_path, fname), 'r')32    f_out = open(os.path.join(root_path, fname) + '.inst', 'w')33    for line in f_in:34        line = json.loads(line)35        line['instruction'] = instruction36        line = json.dumps(line, ensure_ascii=False)37        f_out.write(line + '\n')38        pbar.update(1)39    f_in.close()40    f_out.close()41 42for fname in tqdm(glob(f'/etc/ssd1/jiangzhongtao/baai_embedding_tune/data/all_collect/{src}/*.jsonl')):43    os.remove(fname)44 45for fname in tqdm(glob(f'/etc/ssd1/jiangzhongtao/baai_embedding_tune/data/all_collect/{src}/*.jsonl.inst')):46    os.rename(fname, fname.replace('.inst', ''))47 48