Update text_gen.py

text_gen.py

@@ -2,16 +2,16 @@ import transformers
 import string
 
 model_names = ['microsoft/GODEL-v1_1-base-seq2seq',
-               'facebook/blenderbot-400M-distill',
-               'facebook/blenderbot_small-90M']
+               'facebook/blenderbot-3B',
+               'PygmalionAI/pygmalion-6b']
 
 tokenizers = [transformers.AutoTokenizer.from_pretrained(model_names[0]),
               transformers.BlenderbotTokenizer.from_pretrained(model_names[1]),
-              transformers.BlenderbotSmallTokenizer.from_pretrained(model_names[2])]
+              transformers.GPTNeoXTokenizerFast.from_pretrained(model_names[2])]
 
 model = [transformers.AutoModelForSeq2SeqLM.from_pretrained(model_names[0]),
          transformers.BlenderbotForConditionalGeneration.from_pretrained(model_names[1]),
-         transformers.BlenderbotSmallForConditionalGeneration.from_pretrained(model_names[2])]
+         transformers.GPTNeoXForCausalLM.from_pretrained(model_names[2])]
 
 
 def generate_text(text, context, model_name, model, tokenizer, minimum=15, maximum=300):