Update app.py

app.py

@@ -63,8 +63,6 @@ def generate_medical_response(message, history):
     
     prompt = f"<|system|>\nYou are a highly accurate and helpful medical assistant.</s>\n<|user|>\n{message}</s>\n<|assistant|>\n"
     inputs = tokenizer(prompt, return_tensors="pt")
-
-    input_length = inputs.input_ids.shape[1]
     
     with torch.no_grad():
         outputs = model.generate(
@@ -78,13 +76,9 @@ def generate_medical_response(message, history):
             pad_token_id=tokenizer.eos_token_id
         )
     
-    generated_tokens = outputs[0][input_length:]
-
-    final_answer = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    if "user:" in final_answer.lower():
-        final_answer = final_answer.split("user:")[0].strip()
-        final_answer = final_answer.split("User:")[0].strip()
+    final_answer = generated_text.split("<|assistant|>\n")[-1].strip()
 
     if not final_answer:
         final_answer = "I apologize, but I am unable to generate a confident medical response to that exact phrasing. Could you please rephrase your question?"