Update text_gen.py

text_gen.py

@@ -15,11 +15,12 @@ model = [transformers.AutoModelForSeq2SeqLM.from_pretrained(model_names[0]),
 
 
 def generate_text(text, context, model_name, model, tokenizer, minimum=15, maximum=300):
+    if context:
+        text = f'{context} {text}'
     if 'GODEL' in model_name:
-        text = f'Instruction: you need to response discreetly. [CONTEXT] {context} {text}'
+        text = f'Instruction: you need to response discreetly. [CONTEXT] {text}'
         text.replace('\t', ' EOS ')
     else:
-        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)