utils/dpo_data_process.py

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()