Update text_gen.py

text_gen.py

@@ -17,7 +17,10 @@ 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 = '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