Update app.py

app.py

@@ -4,44 +4,44 @@ import tokenizers
 import gradio as gr
 import re
 
-from PIL import Image
-
 
 def get_model_gpt(model_name,tokenizer_name):
     tokenizer = transformers.GPT2Tokenizer.from_pretrained(tokenizer_name)
     model = transformers.GPT2LMHeadModel.from_pretrained(model_name)
-    model.eval()
     return model, tokenizer
 
 def get_model_t5(model_name,tokenizer_name):
     tokenizer = transformers.T5Tokenizer.from_pretrained(tokenizer_name)
     model = transformers.T5ForConditionalGeneration.from_pretrained(model_name)
-    model.eval()
     return model, tokenizer
 
 
 def predict_gpt(text, model, tokenizer, temperature=1.0):
     input_ids = tokenizer.encode(text+" \n Описание:", return_tensors="pt")
+
+    model.eval()
     with torch.no_grad():
         out = model.generate(input_ids, 
                             do_sample=True,
-                            num_beams=4,
-                            temperature= temperature,
-                            top_p=0.65,
-                            max_length=512,
-                            length_penalty = 2.5,
+                            num_beams=3,
+                            temperature=temperature,
+                            top_p=0.75,
+                            max_length=1024,
                             eos_token_id = tokenizer.eos_token_id,
                             pad_token_id = tokenizer.pad_token_id,
+                            repetition_penalty = 2.5,
                             num_return_sequences = 1,
-                            output_attentions = True,
-                            return_dict_in_generate=True,
+                             output_attentions = True,
+                             return_dict_in_generate=True,
                             )
     decode = lambda x : tokenizer.decode(x, skip_special_tokens=True)
-    generated_text = list(map(decode, out['sequences']))[0].replace(text,'')
-    return generated_text
+    generated_text = list(map(decode, out['sequences']))
+    return generated_text[0].split('Описание :')[1]
 
 def predict_t5(text, model, tokenizer, temperature=1.2):
     input_ids = tokenizer.encode(text, return_tensors="pt")
+    
+    model.eval()
     with torch.no_grad():
         out = model.generate(input_ids, 
                         do_sample=True,