Update text_gen.py

text_gen.py

@@ -1,9 +1,9 @@
 import transformers
 import string
 
-model_names = ['microsoft/GODEL-v1_1-base-seq2seq',
+model_names = ['microsoft/GODEL-v1_1-large-seq2seq
                'facebook/blenderbot-1B-distill',
-               'microsoft/DialoGPT-medium']
+               'satvikag/chatbot']
 
 tokenizers = [transformers.AutoTokenizer.from_pretrained(model_names[0]),
               transformers.BlenderbotTokenizer.from_pretrained(model_names[1]),
@@ -15,15 +15,16 @@ model = [transformers.AutoModelForSeq2SeqLM.from_pretrained(model_names[0]),
 
 
 def generate_text(text, context, model_name, model, tokenizer, minimum=15, maximum=300):
-    text = f'{context} {text}'
     if 'GODEL' in model_name:
-        text = 'Instruction: you need to response discreetly. [CONTEXT] ' + text
+        text = f'Instruction: you need to response discreetly. [CONTEXT] {context} {text}'
+        text.replace('\t', ' EOS ')
     else:
-        text = text.replace(' EOS ', tokenizer.eos_token) + tokenizer.eos_token
+        text = f'{context} {text}'
+        text = text.replace('\t', '\n')
     input_ids = tokenizer(text, return_tensors="pt").input_ids
     outputs = model.generate(input_ids, max_new_tokens=maximum, min_new_tokens=minimum, top_p=0.9, do_sample=True)
     output = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return capitalization(output)
+    return model_name + capitalization(output)
 
 
 def capitalization(line):