Views
No views yet
1from transformers import PreTrainedTokenizerFast, GPT2LMHeadModel
2import torch
3
4
5device = 'cuda' if torch.cuda.is_available() else 'cpu'
6
7tokenizer = PreTrainedTokenizerFast.from_pretrained('byeongal/Ko-DialoGPT')
8model = GPT2LMHeadModel.from_pretrained('byeongal/Ko-DialoGPT').to(device)
9
10past_user_inputs = []
11generated_responses = []
12
13while True:
14 user_input = input(">> User:")
15 if user_input == 'bye':
16 break
17 text_idx = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
18 for i in range(len(generated_responses)-1, len(generated_responses)-3, -1):
19 if i < 0:
20 break
21 encoded_vector = tokenizer.encode(generated_responses[i] + tokenizer.eos_token, return_tensors='pt')
22 if text_idx.shape[-1] + encoded_vector.shape[-1] < 1000:
23 text_idx = torch.cat([encoded_vector, text_idx], dim=-1)
24 else:
25 break
26 encoded_vector = tokenizer.encode(past_user_inputs[i] + tokenizer.eos_token, return_tensors='pt')
27 if text_idx.shape[-1] + encoded_vector.shape[-1] < 1000:
28 text_idx = torch.cat([encoded_vector, text_idx], dim=-1)
29 else:
30 break
31 text_idx = text_idx.to(device)
32 inference_output = model.generate(
33 text_idx,
34 max_length=1000,
35 num_beams=5,
36 top_k=20,
37 no_repeat_ngram_size=4,
38 length_penalty=0.65,
39 repetition_penalty=2.0,
40 )
41 inference_output = inference_output.tolist()
42 bot_response = tokenizer.decode(inference_output[0][text_idx.shape[-1]:], skip_special_tokens=True)
43 print(f"Bot: {bot_response}")
44 past_user_inputs.append(user_input)
45 generated_responses.append(bot_response)