Views
No views yet
1import torch
2import transformers
3use_cuda = torch.cuda.is_available()
4device = torch.device("cuda" if use_cuda else "cpu")
5t5_tokenizer = transformers.GPT2Tokenizer.from_pretrained("SiberiaSoft/SiberianPersonaFred-2")
6t5_model = transformers.T5ForConditionalGeneration.from_pretrained("SiberiaSoft/SiberianPersonaFred-2")
7while True:
8 print('-'*80)
9 dialog = []
10 while True:
11 msg = input('H:> ').strip()
12 if len(msg) == 0:
13 break
14 msg = msg[0].upper() + msg[1:]
15 dialog.append('Ты: ' + msg)
16 # В начале ставится промпт персонажа.
17 prompt = '<SC6>Я парень, консультант по разным вопросам. Я очень умный. Я люблю помогать собеседнику. Недавно, у меня был следующий диалог:' + '\n'.join(dialog) + '\nЯ: <extra_id_0>'
18 input_ids = t5_tokenizer(prompt, return_tensors='pt').input_ids
19 out_ids = t5_model.generate(input_ids=input_ids.to(device), do_sample=True, temperature=0.9, max_new_tokens=512, top_p=0.85,
20 top_k=2, repetition_penalty=1.2)
21 t5_output = t5_tokenizer.decode(out_ids[0][1:])
22 if '</s>' in t5_output:
23 t5_output = t5_output[:t5_output.find('</s>')].strip()
24 t5_output = t5_output.replace('<extra_id_0>', '').strip()
25 t5_output = t5_output.split('Собеседник')[0].strip()
26 print('B:> {}'.format(t5_output))
27 dialog.append('Я: ' + t5_output)