chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,213 @@
|
||||
"""
|
||||
python split_data_by_length.py \
|
||||
--input_path train_data \
|
||||
--output_dir train_data_split \
|
||||
--cache_dir .cache \
|
||||
--log_name .split_log \
|
||||
--length_list 0 500 1000 2000 3000 4000 5000 6000 7000 \
|
||||
--model_name_or_path BAAI/bge-m3 \
|
||||
--num_proc 16 \
|
||||
--overwrite False
|
||||
"""
|
||||
import os
|
||||
import json
|
||||
import math
|
||||
import time
|
||||
import argparse
|
||||
import datasets
|
||||
from tqdm import tqdm
|
||||
from pprint import pprint
|
||||
from transformers import AutoTokenizer
|
||||
from datasets import load_dataset, Features, Value, Sequence
|
||||
|
||||
|
||||
def get_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--input_path', type=str, required=True, help='the path of input datas')
|
||||
parser.add_argument('--output_dir', type=str, required=True, help='the dir of output datas')
|
||||
parser.add_argument('--cache_dir', type=str, default=None, help='the cache dir')
|
||||
parser.add_argument('--log_name', type=str, default='.split_log', help='the name of log file, default: `.split_log`, which will be saved to `output_dir`')
|
||||
parser.add_argument('--length_list', type=int, default=[0, 500, 1000, 2000, 3000, 4000, 5000, 6000, 7000], nargs='+', help='the length list to split')
|
||||
parser.add_argument('--model_name_or_path', type=str, default='BAAI/bge-m3', help='the model name or path of the tokenizer')
|
||||
parser.add_argument('--num_proc', type=int, default=16, help='the number of process, default: 16')
|
||||
parser.add_argument('--overwrite', action='store_true', default=False, help='whether to overwrite the output file, default: False')
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
class SplitByLengthHandler:
|
||||
def __init__(self,
|
||||
model_name_or_path: str,
|
||||
cache_dir: str=None,
|
||||
num_proc: int=16,
|
||||
length_list: list=[0, 500, 1000, 2000, 3000, 4000, 5000, 6000, 7000],
|
||||
overwrite: bool=False):
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
|
||||
self.cache_dir = cache_dir
|
||||
self.num_proc = num_proc
|
||||
self.length_ranges_list = self._get_length_ranges_list(length_list)
|
||||
self.overwrite = overwrite
|
||||
|
||||
pprint(self.length_ranges_list)
|
||||
|
||||
def _map_func(examples):
|
||||
results = {}
|
||||
results['idx'] = []
|
||||
results['max_length'] = []
|
||||
for i in range(len(examples['query'])):
|
||||
idx = examples['idx'][i]
|
||||
query = examples['query'][i]
|
||||
pos, neg = examples['pos'][i], examples['neg'][i]
|
||||
all_texts = [query] + pos + neg
|
||||
|
||||
max_len = 0
|
||||
for x in all_texts:
|
||||
tokenized_x = self.tokenizer(x)['input_ids']
|
||||
if len(tokenized_x) > max_len:
|
||||
max_len = len(tokenized_x)
|
||||
|
||||
results['idx'].append(idx)
|
||||
results['max_length'].append(max_len)
|
||||
return results
|
||||
|
||||
self._map_func = _map_func
|
||||
|
||||
@staticmethod
|
||||
def _get_length_ranges_list(length_list: list):
|
||||
length_ranges_list = []
|
||||
length_list = sorted(length_list)
|
||||
for i in range(len(length_list)):
|
||||
length_l = length_list[i]
|
||||
if i == len(length_list) - 1:
|
||||
length_r = math.inf
|
||||
else:
|
||||
length_r = length_list[i + 1]
|
||||
assert 0 <= length_l < length_r
|
||||
length_ranges_list.append((length_l, length_r))
|
||||
|
||||
return length_ranges_list
|
||||
|
||||
def _process_dir(self, dir_path: str, output_dir: str):
|
||||
assert os.path.isdir(dir_path)
|
||||
log_info_list = []
|
||||
for file in tqdm(os.listdir(dir_path), desc=f'processing {dir_path}'):
|
||||
file_path = os.path.join(dir_path, file)
|
||||
if not file_path.endswith('.jsonl'):
|
||||
print(f"skip {file_path} ...")
|
||||
continue
|
||||
|
||||
output_path = os.path.join(output_dir, '.'.join(file.split('.')[:-1]))
|
||||
log_info = self._process_file(file_path, output_path)
|
||||
log_info_list.append(log_info)
|
||||
return log_info_list
|
||||
|
||||
def _process_file(self, file_path: str, output_path: str):
|
||||
assert not os.path.isdir(file_path)
|
||||
|
||||
start_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
|
||||
|
||||
features = Features({
|
||||
'query': Value('string'),
|
||||
'pos': Sequence(Value('string')),
|
||||
'neg': Sequence(Value('string'))
|
||||
})
|
||||
kd_features = Features({
|
||||
'query': Value('string'),
|
||||
'pos': Sequence(Value('string')),
|
||||
'neg': Sequence(Value('string')),
|
||||
'pos_scores': Sequence(Value('float')),
|
||||
'neg_scores': Sequence(Value('float'))
|
||||
})
|
||||
try:
|
||||
dataset = load_dataset('json', data_files=file_path, cache_dir=self.cache_dir, features=features)['train']
|
||||
except:
|
||||
dataset = load_dataset('json', data_files=file_path, cache_dir=self.cache_dir, features=kd_features)['train']
|
||||
|
||||
dataset_with_idx_list = []
|
||||
for i, data in enumerate(dataset):
|
||||
data['idx'] = i
|
||||
dataset_with_idx_list.append(data)
|
||||
dataset_with_idx = datasets.Dataset.from_list(dataset_with_idx_list)
|
||||
|
||||
mapped_dataset = dataset_with_idx.map(self._map_func, batched=True, num_proc=self.num_proc)
|
||||
|
||||
split_info_dict = {}
|
||||
for length_l, length_r in self.length_ranges_list:
|
||||
save_path = output_path + f'_len-{length_l}-{length_r}.jsonl'
|
||||
if os.path.exists(save_path) and not self.overwrite:
|
||||
print(f'{save_path} exists, skip')
|
||||
continue
|
||||
|
||||
idxs = mapped_dataset.filter(lambda x: length_l <= x['max_length'] < length_r, num_proc=self.num_proc)
|
||||
split_dataset = dataset_with_idx.select(idxs['idx'])
|
||||
split_dataset = split_dataset.remove_columns('idx')
|
||||
|
||||
split_info_dict[f'len-{length_l}-{length_r}'] = len(split_dataset)
|
||||
|
||||
if len(split_dataset) > 0:
|
||||
split_dataset.to_json(save_path, force_ascii=False)
|
||||
|
||||
end_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
|
||||
|
||||
size = len(dataset)
|
||||
avg_length = sum(mapped_dataset['max_length']) / size
|
||||
log_info = {
|
||||
'file_name': os.path.basename(file_path),
|
||||
'size': size,
|
||||
'avg_length': avg_length,
|
||||
'file_path': file_path,
|
||||
'start_time': start_time,
|
||||
'end_time': end_time,
|
||||
'split_info': split_info_dict
|
||||
}
|
||||
return log_info
|
||||
|
||||
def run(self, input_path: str, output_dir: str, log_name: str=None):
|
||||
if not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir)
|
||||
|
||||
if log_name is None:
|
||||
log_path = os.path.join(output_dir, '.split_log')
|
||||
else:
|
||||
log_path = os.path.join(output_dir, log_name)
|
||||
|
||||
log_info_list = []
|
||||
|
||||
if os.path.isdir(input_path):
|
||||
log_info_list = self._process_dir(input_path, output_dir)
|
||||
else:
|
||||
file_name = os.path.basename(input_path)
|
||||
output_path = os.path.join(output_dir, '.'.join(file_name.split('.')[:-1]))
|
||||
log_info = self._process_file(input_path, output_path)
|
||||
log_info_list.append(log_info)
|
||||
|
||||
with open(log_path, 'a', encoding='utf-8') as f:
|
||||
for log_info in log_info_list:
|
||||
json.dump(log_info, f, ensure_ascii=False)
|
||||
f.write('\n')
|
||||
|
||||
|
||||
def main(args):
|
||||
input_path = args.input_path
|
||||
output_dir = args.output_dir
|
||||
log_name = args.log_name
|
||||
|
||||
handler = SplitByLengthHandler(
|
||||
model_name_or_path=args.model_name_or_path,
|
||||
cache_dir=args.cache_dir,
|
||||
num_proc=args.num_proc,
|
||||
length_list=args.length_list if isinstance(args.length_list, list) else [args.length_list],
|
||||
overwrite=args.overwrite
|
||||
)
|
||||
|
||||
handler.run(
|
||||
input_path=input_path,
|
||||
output_dir=output_dir,
|
||||
log_name=log_name
|
||||
)
|
||||
print('\nDONE!')
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = get_args()
|
||||
main(args)
|
||||
Reference in New Issue
Block a user