Update README.md

README.md

@@ -18,41 +18,35 @@ use_cuda = torch.cuda.is_available()
 device = torch.device("cuda" if use_cuda else "cpu")
 
 
-t5_tokenizer = transformers.GPT2Tokenizer.from_pretrained("Den4ikAI/FRED-T5-XL-chitchat")
-t5_model = transformers.T5ForConditionalGeneration.from_pretrained("Den4ikAI/FRED-T5-XL-chitchat")
-t5_model.to(device)
-t5_model.eval()
+t5_tokenizer = transformers.GPT2Tokenizer.from_pretrained(model_dir)
+t5_model = transformers.T5ForConditionalGeneration.from_pretrained(model_dir)
 
 while True:
     print('-'*80)
     dialog = []
     while True:
-        msg = input('Вопрос :> ').strip()
+        msg = input('H:> ').strip()
         if len(msg) == 0:
             break
 
-        msg = msg[0].upper() + msg[1:]
-        dialog.append('человек: ' + msg)
-        prompt = '<SC1>' + '\n'.join(dialog) + '\nчатбот: <extra_id_0>'
-
-        input_ids = t5_tokenizer(prompt, return_tensors='pt').input_ids
-        out_ids = t5_model.generate(input_ids=input_ids.to(device),
+        dialog.append('- ' + msg)
+        dialog.append('- <extra_id_0>')
+        input_ids = t5_tokenizer('<SC1>'+'\n'.join(dialog), return_tensors='pt').input_ids
+        out_ids = t5_model.generate(input_ids=input_ids,
                                     max_length=200,
                                     eos_token_id=t5_tokenizer.eos_token_id,
                                     early_stopping=True,
                                     do_sample=True,
                                     temperature=1.0,
                                     top_k=0,
-                                    top_p=0.95)
-
-        t5_output = t5_tokenizer.decode(out_ids[0][1:])
+                                    top_p=0.85)
+        dialog.pop(-1)
+        t5_output = t5_tokenizer.decode(out_ids[0][1:]).replace('<extra_id_0>','')
         if '</s>' in t5_output:
             t5_output = t5_output[:t5_output.find('</s>')].strip()
 
-        t5_output = t5_output.replace('<extra_id_0>', '').strip()
-
-        print('Ответ :> {}'.format(t5_output))
-        dialog.append('чатбот: ' + t5_output)
+        print('B:> {}'.format(t5_output))
+        dialog.append('- '+t5_output)
 ```
 # Citation
 ```