Spaces:
Paused
Paused
File size: 4,809 Bytes
6b29344 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 |
import os
import torch
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr
# ——— CONFIG ———
REPO_ID = "CodCodingCode/llama-3.1-8b-clinical"
SUBFOLDER = "checkpoint-45000"
HF_TOKEN = os.environ["HUGGINGFACE_HUB_TOKEN"] # set in Settings→Secrets
# ——— SNAPSHOT & LOAD ———
# This will grab all .json and .safetensors under checkpoint-45000:
local_dir = snapshot_download(
repo_id=REPO_ID,
subfolder=SUBFOLDER,
token=HF_TOKEN,
allow_patterns=["*.json", "*.safetensors"],
)
# Now point at that folder:
MODEL_DIR = local_dir # e.g. ~/.cache/huggingface/…/checkpoint-45000
# Load tokenizer & model from the real files you just pulled:
tokenizer = AutoTokenizer.from_pretrained(
MODEL_DIR,
use_fast=False,
trust_remote_code=True,
)
model = AutoModelForCausalLM.from_pretrained(
MODEL_DIR,
device_map="auto",
torch_dtype=torch.float16,
trust_remote_code=True,
)
model.eval()
# === Role Agent with instruction/input/output format ===
class RoleAgent:
def __init__(self, role_instruction):
self.role_instruction = role_instruction
def act(self, input_text):
prompt = (
f"Instruction: {self.role_instruction}\n"
f"Input: {input_text}\n"
f"Output:"
)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(
**inputs,
max_new_tokens=256,
do_sample=True,
temperature=0.7,
pad_token_id=tokenizer.eos_token_id,
)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
# extract THINKING / ANSWER if present
thinking, answer = "", response
if "THINKING:" in response and "ANSWER:" in response and "END" in response:
block = response.split("THINKING:")[1].split("END")[0]
thinking = block.split("ANSWER:")[0].strip()
answer = block.split("ANSWER:")[1].strip()
return {
"instruction": f"You are {self.role_instruction}.",
"input": input_text,
"thinking": thinking,
"output": answer,
}
# === Agents ===
summarizer = RoleAgent(
"You are a clinical summarizer trained to extract structured vignettes from doctor–patient dialogues."
)
diagnoser = RoleAgent(
"You are a board-certified diagnostician that diagnoses patients."
)
questioner = RoleAgent("You are a physician asking questions to diagnose a patient.")
treatment_agent = RoleAgent(
"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."
)
# === Inference State ===
conversation_history = []
summary = ""
diagnosis = ""
# === Gradio Inference ===
def simulate_interaction(user_input, iterations=1):
history = [f"Doctor: What brings you in today?", f"Patient: {user_input}"]
summary, diagnosis = "", ""
for i in range(iterations):
# Summarize
sum_in = "\n".join(history) + f"\nPrevious Vignette: {summary}"
sum_out = summarizer.act(sum_in)
summary = sum_out["output"]
# Diagnose
diag_out = diagnoser.act(summary)
diagnosis = diag_out["output"]
# Question
q_in = f"Vignette: {summary}\nCurrent Estimated Diagnosis: {diag_out['thinking']} {diagnosis}"
q_out = questioner.act(q_in)
history.append(f"Doctor: {q_out['output']}")
history.append("Patient: (awaiting response)")
# Treatment
treatment_out = treatment_agent.act(
f"Diagnosis: {diagnosis}\nVignette: {summary}"
)
return {
"summary": sum_out,
"diagnosis": diag_out,
"question": q_out,
"treatment": treatment_out,
"conversation": "\n".join(history),
}
# === Gradio UI ===
def ui_fn(user_input):
res = simulate_interaction(user_input)
return f"""📋 Vignette Summary:
💭 THINKING: {res['summary']['thinking']}
ANSWER: {res['summary']['output']}
🩺 Diagnosis:
💭 THINKING: {res['diagnosis']['thinking']}
ANSWER: {res['diagnosis']['output']}
T
❓ Follow-up Question:
💭 THINKING: {res['question']['thinking']}
ANSWER: {res['question']['output']}
💊 Treatment Plan:
{res['treatment']['output']}
💬 Conversation:
{res['conversation']}
"""
demo = gr.Interface(
fn=ui_fn,
inputs=gr.Textbox(label="Patient Response"),
outputs=gr.Textbox(label="Doctor Simulation Output"),
title="🧠 AI Doctor Multi-Agent Reasoning",
)
if __name__ == "__main__":
demo.launch(share=True)
|