Update text_gen.py

text_gen.py

@@ -3,15 +3,15 @@ import string
 
 model_names = ['microsoft/GODEL-v1_1-large-seq2seq',
                'facebook/blenderbot-1B-distill',
-               'satvikag/chatbot']
+               'facebook/blenderbot_small-90M']
 
 tokenizers = [transformers.AutoTokenizer.from_pretrained(model_names[0]),
               transformers.BlenderbotTokenizer.from_pretrained(model_names[1]),
-              transformers.GPT2Tokenizer.from_pretrained(model_names[2])]
+              transformers.BlenderbotSmallTokenizer.from_pretrained(model_names[2])]
 
 model = [transformers.AutoModelForSeq2SeqLM.from_pretrained(model_names[0]),
          transformers.BlenderbotForConditionalGeneration.from_pretrained(model_names[1]),
-         transformers.GPT2LMHeadModel.from_pretrained(model_names[2])]
+         transformers.BlenderbotSmallForConditionalGeneration.from_pretrained(model_names[2])]
 
 
 def generate_text(text, context, model_name, model, tokenizer, minimum=15, maximum=300):
@@ -24,7 +24,7 @@ def generate_text(text, context, model_name, model, tokenizer, minimum=15, maxim
     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 model_name + capitalization(output)
+    return capitalization(output)
 
 
 def capitalization(line):