switched to new model name

app.py

@@ -5,8 +5,8 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 
 # ——— CONFIG ———
-REPO_ID = "CodCodingCode/llama-3.1-8b-clinical"
-SUBFOLDER = "checkpoint-45000"
+REPO_ID = "CodCodingCode/llama-3.1-8b-clinical-v1.1"
+SUBFOLDER = "checkpoint-2250"
 HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
 if not HF_TOKEN:
     raise RuntimeError("Missing HUGGINGFACE_HUB_TOKEN in env")
@@ -179,23 +179,80 @@ class RoleAgent:
 
 
 # === Agents ===
+# ——— Instantiate RoleAgents for each of your eight roles ———
+
 summarizer = RoleAgent(
-    role_instruction="You are a clinical summarizer trained to extract structured vignettes from doctor–patient dialogues.",
+    role_instruction=(
+        "“You are a clinical summarizer. Given a transcript of a doctor–patient dialogue, "
+        "extract a structured clinical vignette summarizing the key symptoms, relevant history, "
+        "and any diagnostic clues.”"
+    ),
     tokenizer=tokenizer,
     model=model,
 )
-diagnoser = RoleAgent(
-    role_instruction="You are a board-certified diagnostician that diagnoses patients.",
+
+treatment_agent = RoleAgent(
+    role_instruction=(
+        "You are a board-certified clinician. Based on the provided diagnosis and patient vignette, "
+        "propose a realistic, evidence-based treatment plan suitable for initiation by a primary care "
+        "physician or psychiatrist."
+    ),
     tokenizer=tokenizer,
     model=model,
 )
-questioner = RoleAgent(
-    role_instruction="You are a physician asking questions to diagnose a patient.",
+
+diagnoser_early = RoleAgent(
+    role_instruction=(
+        "You are a diagnostic reasoning model (Early Stage). Based on the patient vignette and "
+        "early-stage observations, generate a list of plausible diagnoses with reasoning. Focus on "
+        "broad differentials, considering common and uncommon conditions."
+    ),
     tokenizer=tokenizer,
     model=model,
 )
-treatment_agent = RoleAgent(
-    role_instruction="You are a board-certified clinician. Based on the diagnosis and patient vignette provided below, suggest a concise treatment plan that could realistically be initiated by a primary care physician or psychiatrist.",
+
+diagnoser_middle = RoleAgent(
+    role_instruction=(
+        "You are a diagnostic reasoning model (Middle Stage). Given the current vignette, prior dialogue, "
+        "and diagnostic hypothesis, refine the list of possible diagnoses with concise justifications for each. "
+        "Aim to reduce diagnostic uncertainty."
+    ),
+    tokenizer=tokenizer,
+    model=model,
+)
+
+diagnoser_late = RoleAgent(
+    role_instruction=(
+        "You are a diagnostic reasoning model (Late Stage). Based on the final patient vignette summary and full conversation, "
+        "provide the most likely diagnosis with structured reasoning. Confirm diagnostic certainty and include END if no more questioning is necessary."
+    ),
+    tokenizer=tokenizer,
+    model=model,
+)
+
+questioner_early = RoleAgent(
+    role_instruction=(
+        "You are a questioning agent (Early Stage). Your task is to propose highly relevant early-stage questions "
+        "that can open the differential diagnosis widely. Use epidemiology, demographics, and vague presenting symptoms as guides."
+    ),
+    tokenizer=tokenizer,
+    model=model,
+)
+
+questioner_middle = RoleAgent(
+    role_instruction=(
+        "You are a questioning agent (Middle Stage). Using the current diagnosis, past questions, and patient vignette, "
+        "generate a specific question to refine the current differential diagnosis. Return your reasoning and next question."
+    ),
+    tokenizer=tokenizer,
+    model=model,
+)
+
+questioner_late = RoleAgent(
+    role_instruction=(
+        "You are a questioning agent (Late Stage). Based on narrowed differentials and previous dialogue, "
+        "generate a focused question that would help confirm or eliminate the final 1-2 suspected diagnoses."
+    ),
     tokenizer=tokenizer,
     model=model,
 )
@@ -235,12 +292,12 @@ def simulate_interaction(user_input, conversation_history=None):
     summary = sum_out["output"]
 
     # Diagnose based on summary
-    diag_out = diagnoser.act(summary)
+    diag_out = diagnoser_middle.act(summary)
     diagnosis = diag_out["output"]
 
     # Generate next question based on current understanding
     q_in = f"Vignette: {summary}\nCurrent Estimated Diagnosis: {diagnosis}"
-    q_out = questioner.act(q_in)
+    q_out = questioner_middle.act(q_in)
 
     # Add doctor's response to history
     history.append(f"Doctor: {q_out['output']}")