Update app.py

app.py

@@ -6,6 +6,10 @@ import torch
 physician = pipeline("text-generation", model="HuggingFaceTB/SmolLM2-1.7B-Instruct")
 patient = pipeline("text-generation", model="HuggingFaceTB/SmolLM2-1.7B-Instruct")
 
+# System prompts to define roles
+patient_system_prompt = "You are a patient describing your symptoms to a physician."
+physician_system_prompt = "You are a physician responding to a patient's symptoms."
+
 def generate_conversation(topic, turns):
     conversation = []
     total_tokens = 0
@@ -13,10 +17,10 @@ def generate_conversation(topic, turns):
     patient_tokens = 0
 
     # Initial prompt for the patient
-    patient_prompt = f"Imagine you are a patient. You are here to talk about {topic}."
+    patient_prompt = f"I'm here to talk about {topic}."
     print(f"Patient Initial Prompt: {patient_prompt}")  # Debugging
     patient_response = patient(
-        patient_prompt,
+        f"{patient_system_prompt} {patient_prompt}",
         max_length=50,  # Reduce max_length for faster responses
         num_return_sequences=1,
         truncation=True,  # Explicitly enable truncation
@@ -29,10 +33,9 @@ def generate_conversation(topic, turns):
 
     for turn in range(turns):
         # Physician's turn
-        physician_prompt = f"Imagine you are a physician. How would you respond to: {patient_response}"
-        print(f"Physician Turn {turn} Prompt: {physician_prompt}")  # Debugging
+        print(f"Physician Turn {turn} Prompt: {patient_response}")  # Debugging
         physician_response = physician(
-            physician_prompt,
+            f"{physician_system_prompt} Patient says: {patient_response}",
             max_length=50,  # Reduce max_length for faster responses
             num_return_sequences=1,
             truncation=True,  # Explicitly enable truncation
@@ -44,10 +47,9 @@ def generate_conversation(topic, turns):
         conversation.append({"role": "physician", "message": physician_response, "tokens": len(physician_response.split())})
 
         # Patient's turn
-        patient_prompt = f"Imagine you are a patient. How would you respond to: {physician_response}"
-        print(f"Patient Turn {turn} Prompt: {patient_prompt}")  # Debugging
+        print(f"Patient Turn {turn} Prompt: {physician_response}")  # Debugging
         patient_response = patient(
-            patient_prompt,
+            f"{patient_system_prompt} Physician says: {physician_response}",
             max_length=50,  # Reduce max_length for faster responses
             num_return_sequences=1,
             truncation=True,  # Explicitly enable truncation