GPT-2-chinese-dialog

用法：

# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained("MelonFIsh/GPT-2-chinese-dialog")
model = AutoModelForCausalLM.from_pretrained("MelonFIsh/GPT-2-chinese-dialog")
model_dev = torch.device('cuda:0')

# 设置模型特殊符号
bos, eos, speaker1, speaker2, pad = tokenizer.convert_tokens_to_ids(["<|endoftext|>", "<|endoftext|>", "<|speaker1|>", "<|speaker2|>", '<|pad|>'])

# 定义相关函数
def find_last_index_from_list(ls, num):
    for i in range(len(ls)-1, 0, -1):
        if ls[i] == num:
            return i
    return -1

def convert_dialog_to_ids(bos,eos,speaker1,speaker2,pad, this_dialo, max_length, add_eos=True):
    input_ids = [bos]
    token_type_ids = [bos]
    lm_labels_first_pad_id = -100
    for i,l in enumerate(this_dialo):
        if i%2 == 0:
            speaker = speaker1
        else:
            speaker = speaker2

        token = tokenizer.encode(l, add_special_tokens=False)
        input_ids += [speaker]+token
        token_type_ids += [speaker]*(len(token)+1)

    #检查最后一个speaker后，下一个speaker应该是什么
    next_speaker = speaker1
    if speaker == speaker1: next_speaker = speaker2
        
    if add_eos:
        input_ids.append(eos)
        token_type_ids.append(speaker)
    else:
        input_ids.append(next_speaker)
        token_type_ids.append(next_speaker)
    if max_length!=0 and len(input_ids) > max_length:
        if add_eos:
            input_ids = input_ids[0:max_length-1]+[eos]
            token_type_ids = token_type_ids[0:max_length-1]+[speaker]
        else:
            input_ids = input_ids[0:max_length-1]+[next_speaker]
            token_type_ids = token_type_ids[0:max_length-1]+[next_speaker]
    
    if max_length!=0:
        attention_mask = [1]*len(input_ids) + [0]*(max_length-len(input_ids))
        input_ids += [pad]*(max_length-len(input_ids))
        token_type_ids += [pad]*(max_length-len(token_type_ids))
    else:
        attention_mask = [1]*len(input_ids)

    last_speak_token_index = find_last_index_from_list(input_ids, speaker)
    lm_labels = [lm_labels_first_pad_id]*(last_speak_token_index+1) + input_ids[last_speak_token_index+1:]
    if max_length!=0:
        lm_labels += [0]*(max_length-len(lm_labels))
    
    return input_ids, attention_mask, token_type_ids, lm_labels


## 多轮对话测试
max_steps = 10

model.to(model_dev)
history_list = []
while True:
    user_input = input('User >>> ')
    if user_input == '':
        break
    
    if user_input != 'skip':
        history_list.append(user_input)
    
    if len(history_list)>=max_steps:
        history_list = history_list[len(history_list)-max_steps:]
    
    input_ids, attention_mask, token_type_ids, lm_labels = convert_dialog_to_ids(bos,eos,speaker1,speaker2,pad, history_list, 0, add_eos=False)
    input_ids = torch.tensor(input_ids, dtype=torch.long, device=model_dev).unsqueeze(0)
    token_type_ids = torch.tensor(token_type_ids, dtype=torch.long, device=model_dev).unsqueeze(0)
    attention_mask = torch.tensor(attention_mask, dtype=torch.long, device=model_dev).unsqueeze(0)
    results = model.generate(input_ids,
        do_sample=True,
        attention_mask=attention_mask,
        token_type_ids = token_type_ids,
        top_k=20,
        temperature=0.6,
        max_length=200,
        top_p=0.45,
        length_penalty=-0.3,
        pad_token_id=pad
    )
    d = tokenizer.batch_decode(results, skip_special_tokens=True)[0]
    #print(d)
    #d = d.replace(''.join(history_list), '')
    d = d[len(''.join(history_list)):]
    history_list.append(d)
    if '�' in d:
        history_list = []
    print('Robot >>> ', d)