import sys sys.path.extend(['.','..']) import os import re import torch import pandas as pd import numpy as np import ujson from rich import progress import pyarrow.parquet as pq from model.infer import ChatBot from logger import Logger from config import PROJECT_ROOT, InferConfig from utils.raw_data_process import delete_file log = Logger('data_process', save2file=True, file_name=PROJECT_ROOT + '/logs/raw_data_process.log') def process_alpaca_gpt4_data(max_len: int=512) -> None: '''' 处理RM高质量回答部分 数据集：https://huggingface.co/datasets/c-s-ale/alpaca-gpt4-data-zh ''' read_file = PROJECT_ROOT + '/data/raw_data/alpaca_gpt4_data_zh.json' save_file = PROJECT_ROOT + '/data/alpaca_gpt4_data_zh.json' max_len += 8 my_data = [] with open(read_file, 'r', encoding='utf-8') as f: data = ujson.load(f) print('length of {} is {}'.format(read_file, len(data))) for item in progress.track(data): prompt = item['instruction'] inputs = item['input'] response = item['output'] if len(response) > max_len: continue # 超长的不要 if len(inputs.strip()) > 0: prompt = f"{prompt}，{inputs}" if len(prompt) > max_len: continue if len(prompt) == 0 or len(response) == 0: continue my_data.append( { 'prompt': prompt, 'chosen': response } ) print('length of {} is {}'.format(save_file, len(my_data))) with open(save_file, 'w', encoding='utf-8') as f: ujson.dump(my_data, f, indent=4, ensure_ascii=False) def generate_alpaca_gpt4_reject_response(groups_cnt: int=50000, max_len: int=320, batch_size: int=32) -> None: '''生成不是很满意的回答回答 ''' print('load model...') # load config infer_config = InferConfig() chatbot = ChatBot(infer_config) model = chatbot.model tokenizer = chatbot.tokenizer device = 'cuda' if torch.cuda.is_available() else 'cpu' finetune_file = PROJECT_ROOT + '/data/alpaca_gpt4_data_zh.json' save_rw_json_file = PROJECT_ROOT + '/data/my_dpo_alpaca_gpt4_data_zh.json' # save_rw_parquet_file = PROJECT_ROOT + '/data/my_rlhf_dataset.parquet' data = [] with open(finetune_file, 'r', encoding='utf-8') as f: data = ujson.load(f) log.info('length of {} is {}'.format(save_rw_json_file, len(data)), save_to_file=True) model_outs = [] batch_prompt = [] process_item = [] for i, item in progress.track(enumerate(data), total=len(data)): # 模型生成的答案为拒绝答案 batch_prompt.append(f"{item['prompt']}[EOS]") process_item.append(item) if i % 500 == 0: print('process {} items.'.format(i)) if len(batch_prompt) >= batch_size or i == len(data) - 1: encoded = tokenizer.batch_encode_plus(batch_prompt, truncation=False, padding=True) with torch.no_grad(): input_ids = torch.LongTensor(encoded.input_ids).to(device) attention_mask = torch.LongTensor(encoded.attention_mask).to(device) outputs = model.my_generate( input_ids=input_ids, attention_mask=attention_mask, max_seq_len=infer_config.max_seq_len, search_type='greedy', ) outputs = tokenizer.batch_decode(outputs.cpu().numpy(), clean_up_tokenization_spaces=True, skip_special_tokens=True) model_outs.extend(outputs) batch_prompt = [] if len(model_outs) % 2000 == 0: for i in range(len(model_outs)): process_item[i]['reject'] = model_outs[i] try: with open(PROJECT_ROOT + '/data/outs.ckp.json', 'w', encoding='utf-8') as f: ujson.dump(process_item, f, indent=4, ensure_ascii=False) except Exception as e: print(e) for i in range(len(model_outs)): process_item[i]['reject'] = model_outs[i] with open(save_rw_json_file, 'w', encoding='utf-8') as f: ujson.dump(process_item, f, indent=4, ensure_ascii=False) # df = pd.DataFrame(data) # write_single_parquet_file(save_rw_parquet_file, df) def replace_line(s: str) -> str: '''将双斜杠替换为单斜杠，既是 \\n 替换为 \n ''' return re.sub('\\\\n', '\n', s) def merge_rlhf_data(max_len: int=512) -> None: '''' 处理RM高质量回答部分 数据集：https://huggingface.co/datasets/Skepsun/huozi_rlhf_data_json https://huggingface.co/datasets/beyond/rlhf-reward-single-round-trans_chinese ''' my_data = [] read_files = [ PROJECT_ROOT + '/data/raw_data/huozi_rlhf_data.json', PROJECT_ROOT + '/data/my_dpo_alpaca_gpt4_data_zh.json', ] save_file = PROJECT_ROOT + '/data/my_dpo_data.json' if os.path.exists(save_file): assert delete_file(save_file) max_len += 8 # for eos token for read_file in read_files: items = [] with open(read_file, 'r', encoding='utf-8') as f: items = ujson.load(f) for item in progress.track(items): prompt, chosen, reject = item['prompt'], item['chosen'], item['reject'] if len(prompt) > max_len or len(chosen) > max_len or len(reject) > max_len: continue # reject.strip() == chosen.strip()，这两个相同的也不要 if len(prompt) == 0 or len(chosen) == 0 or len(reject) == 0 or reject.strip() == chosen.strip(): continue my_data.append({ 'prompt': replace_line(prompt), 'chosen': replace_line(chosen), 'rejected': replace_line(reject), }) read_files = [ PROJECT_ROOT + '/data/raw_data/train-00000-of-00001-789dc5dece0f1fc1.parquet', PROJECT_ROOT + '/data/raw_data/test-00000-of-00001-8ecd46436fadcf7f.parquet', ] for read_file in read_files: pf = pq.read_table(read_file) for prompt, chosen, rejected in progress.track(zip(pf['prompt'], pf['chosen'], pf['rejected']), total=pf.num_rows): prompt, chosen, rejected = prompt.as_py(), chosen.as_py(), rejected.as_py() if len(prompt) > max_len or len(chosen) > max_len or len(rejected) > max_len: continue if len(prompt) == 0 or len(chosen) == 0 or len(rejected) == 0 or rejected.strip() == chosen.strip(): continue my_data.append({ 'prompt': replace_line(prompt), 'chosen': replace_line(chosen), 'rejected': replace_line(rejected), }) print('length of {} is {}'.format(save_file, len(my_data))) with open(save_file, 'w', encoding='utf-8') as f: ujson.dump(my_data, f, indent=4, ensure_ascii=False) def split_train_eval_dataset() -> None: '''划分数据集 ''' rw_json_file = PROJECT_ROOT + '/data/my_dpo_data.json' train_file = PROJECT_ROOT + '/data/my_dpo_train.json' eval_file = PROJECT_ROOT + '/data/my_dpo_eval.json' data = [] with open(rw_json_file, 'r', encoding='utf-8') as f: data = ujson.load(f) np.random.shuffle(data) split_idx = int(len(data) * 0.99) train_data = data[0: split_idx] eval_data = data[split_idx: ] log.info('train size: {}, eval size:{}'.format(len(train_data), len(eval_data)), save_to_file=True) with open(train_file, 'w', encoding='utf-8') as f: ujson.dump(train_data, f, indent=4, ensure_ascii=False) with open(eval_file, 'w', encoding='utf-8') as f: ujson.dump(eval_data, f, indent=4, ensure_ascii=False) if __name__ == '__main__': # 1. 处理chosen文本 # process_alpaca_gpt4_data() # 2. 生成rejected文本 # generate_alpaca_gpt4_reject_response() # 合并数据集 merge_rlhf_data() # 3. split train and eval dataset # split_train_eval_dataset()