remove User: Assistant: from response

app.py

@@ -75,20 +75,23 @@ def generate_response(user_input, model_id):
     prompt = user_input
     inputs = tokenizer([prompt], return_tensors="pt")
 
-    generation_config = GenerationConfig(
+    generation_config = GenerationConfig(         
         penalty_alpha=0.6,
         do_sample=True,
-        top_k=5,
-        temperature=0.6,
+        top_p=0.2,
+        top_k=50,
+        temperature=0.3,
         repetition_penalty=1.2,
-        max_new_tokens=60,  # Adjust as necessary
+        max_new_tokens=60,
         pad_token_id=tokenizer.eos_token_id,
-        stop_sequences=["User:", "Assistant:"],
+        stop_sequences=["User:", "Assistant:", "\n"],
     )
 
     outputs = model.generate(**inputs, generation_config=generation_config)
     response = tokenizer.decode(outputs[:, inputs['input_ids'].shape[-1]:][0], skip_special_tokens=True)
-    return response.strip().split("Assistant:")[-1].strip()
+    cleaned_response = response.replace("User:", "").replace("Assistant:", "").strip() 
+    return cleaned_response.strip().split("\n")[0]  # Keep only the first line of response
+    #return response.strip().split("Assistant:")[-1].strip()
      
 def formatted_prompt(question) -> str:
     return f"<|startoftext|>User: {question}\nAssistant:"