Update app.py

app.py

@@ -1,15 +1,33 @@
 import gradio as gr
-from transformers import pipeline
-import torch
+from smolagents import CodeAgent, HfApiModel
 
-# Load the instruct version of the model
-physician = pipeline("text-generation", model="HuggingFaceTB/SmolLM2-1.7B-Instruct")
-patient = pipeline("text-generation", model="HuggingFaceTB/SmolLM2-1.7B-Instruct")
+# Define system prompts for the agents
+patient_system_prompt = """
+You are a patient describing your symptoms to a physician. You are here to talk about a health issue.
+Be concise and provide relevant information about your symptoms.
+"""
 
-# System prompts to define roles (not included in the input prompt)
-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."
+physician_system_prompt = """
+You are a physician responding to a patient's symptoms. 
+Ask relevant questions to understand the patient's condition and provide appropriate advice.
+"""
 
+# Load the models for the agents
+patient_model = HfApiModel(model_id="HuggingFaceTB/SmolLM2-1.7B-Instruct")
+physician_model = HfApiModel(model_id="HuggingFaceTB/SmolLM2-1.7B-Instruct")
+
+# Initialize the agents
+patient_agent = CodeAgent(
+    model=patient_model,
+    system_prompt=patient_system_prompt,
+    planning_interval=1  # Allow the agent to plan after each turn
+)
+
+physician_agent = CodeAgent(
+    model=physician_model,
+    system_prompt=physician_system_prompt,
+    planning_interval=1  # Allow the agent to plan after each turn
+)
 
 def generate_conversation(topic, turns):
     conversation = []
@@ -18,16 +36,9 @@ def generate_conversation(topic, turns):
     patient_tokens = 0
 
     # Initial prompt for the patient
-    patient_input = f"Patient: I'm here to talk about {topic}."
+    patient_input = f"I'm here to talk about {topic}."
     print(f"Patient Initial Input: {patient_input}")  # Debugging
-    patient_response = patient(
-        patient_input,
-        max_new_tokens=50,  # Allow the model to generate up to 50 new tokens
-        num_return_sequences=1,
-        truncation=True,  # Explicitly enable truncation
-        do_sample=True,  # Enable sampling
-        temperature=0.7  # Control randomness
-    )[0]['generated_text']
+    patient_response = patient_agent.run(patient_input)
     print(f"Patient Response: {patient_response}")  # Debugging
     patient_tokens += len(patient_response.split())
     conversation.append({"role": "patient", "message": patient_response, "tokens": len(patient_response.split())})
@@ -35,28 +46,14 @@ def generate_conversation(topic, turns):
     for turn in range(turns):
         # Physician's turn
         print(f"Physician Turn {turn} Prompt: {patient_response}")  # Debugging
-        physician_response = physician(
-            f"Physician: {patient_response}",
-            max_new_tokens=50,  # Allow the model to generate up to 50 new tokens
-            num_return_sequences=1,
-            truncation=True,  # Explicitly enable truncation
-            do_sample=True,  # Enable sampling
-            temperature=0.7  # Control randomness
-        )[0]['generated_text']
+        physician_response = physician_agent.run(patient_response)
         print(f"Physician Response: {physician_response}")  # Debugging
         physician_tokens += len(physician_response.split())
         conversation.append({"role": "physician", "message": physician_response, "tokens": len(physician_response.split())})
 
         # Patient's turn
         print(f"Patient Turn {turn} Prompt: {physician_response}")  # Debugging
-        patient_response = patient(
-            f"Patient: {physician_response}",
-            max_new_tokens=50,  # Allow the model to generate up to 50 new tokens
-            num_return_sequences=1,
-            truncation=True,  # Explicitly enable truncation
-            do_sample=True,  # Enable sampling
-            temperature=0.7  # Control randomness
-        )[0]['generated_text']
+        patient_response = patient_agent.run(physician_response)
         print(f"Patient Response: {patient_response}")  # Debugging
         patient_tokens += len(patient_response.split())
         conversation.append({"role": "patient", "message": patient_response, "tokens": len(patient_response.split())})
@@ -102,4 +99,4 @@ with gr.Blocks() as demo:
         outputs=output_json
     )
 
-demo.launch()
+demo.launch()