Update app.py

app.py

@@ -1,5 +1,5 @@
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoConfig
-
+import re
 import gradio as gr
 from torch.nn import functional as F
 import seaborn
@@ -60,7 +60,9 @@ def generate(model_name, text):
     input_ids = tokenizer.encode("AFA:{}".format(text), return_tensors="pt")
     outputs = model.generate(input_ids, max_length=200, num_beams=2, repetition_penalty=2.5, top_k=50, top_p=0.98, length_penalty=1.0, early_stopping=True)
     output = tokenizer.decode(outputs[0])
-    return ".".join(output.split(".")[:-1]) + "."
+    #return ".".join(output.split(".")[:-1]) + "."
+    sent = ".".join(output.split(".")[:-1]) + "."
+    return re.match(r'<pad> ([^<>]*)', sent).group(1)
 
 
 output_text = gr.outputs.Textbox()