Update text_gen.py

text_gen.py

@@ -3,24 +3,21 @@ import string
 
 model_names = ['microsoft/GODEL-v1_1-base-seq2seq',
                'facebook/blenderbot-1B-distill',
-               'PygmalionAI/pygmalion-1.3b']
+               'microsoft/DialoGPT-medium']
 
 tokenizers = [transformers.AutoTokenizer.from_pretrained(model_names[0]),
               transformers.BlenderbotTokenizer.from_pretrained(model_names[1]),
-              transformers.GPTNeoXTokenizerFast.from_pretrained(model_names[2])]
+              transformers.GPT2Tokenizer.from_pretrained(model_names[2])]
 
 model = [transformers.AutoModelForSeq2SeqLM.from_pretrained(model_names[0]),
          transformers.BlenderbotForConditionalGeneration.from_pretrained(model_names[1]),
-         transformers.GPTNeoXForCausalLM.from_pretrained(model_names[2])]
+         transformers.GPT2Model.from_pretrained(model_names[2])]
 
 
 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
-    elif 'pygmalion' in model_name:
-        text = 'ROBOTS\'s Persona: Discreet and polite humanoid robot. ' + text
-        text = text.replace(' EOS ', '\n')
     else:
         text = text.replace(' EOS ', '\n')
     input_ids = tokenizer(text, return_tensors="pt").input_ids