Updating Side-By-Side Space

app.py

@@ -1,28 +1,74 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
-
-# Load base model
-base_model_name = "Qwen/Qwen3-4B"
-base_tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_name,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto"
-)
-
-# Load novel model
-novel_model_name = "ClinicianFOCUS/Clinician-Note-2.0a"
-novel_tokenizer = AutoTokenizer.from_pretrained(novel_model_name)
-novel_model = AutoModelForCausalLM.from_pretrained(
-    novel_model_name,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto"
-)
-
-
-def generate_soap_notes(transcript, prompt_template, temperature=0.3, top_p=0.9, top_k=20, max_new_tokens=8192):
-    # Prepare the prompt by combining template and transcript
+import os
+import time
+import sys
+from datetime import datetime
+
+# Try to import llama_cpp
+try:
+    from llama_cpp import Llama
+
+    LLAMA_CPP_AVAILABLE = True
+    print("llama_cpp is available!")
+except ImportError:
+    LLAMA_CPP_AVAILABLE = False
+    print("llama_cpp is not available. Running in fallback mode.")
+
+# Define the GGUF model paths
+BASE_MODEL_PATH = "unsloth/Qwen3-4B-GGUF/Qwen3-4B-Q4_K_M.gguf"
+NOVEL_MODEL_PATH = "mradermacher/Clinician-Note-2.0a-i1-GGUF/Clinician-Note-2.0a.i1-Q4_K_M.gguf"
+
+# Initialize models
+base_model = None
+novel_model = None
+
+
+def load_models(progress=None):
+    """Load the llama.cpp models"""
+    global base_model, novel_model
+
+    if not LLAMA_CPP_AVAILABLE:
+        print("llama_cpp not available, cannot load models")
+        return False
+
+    try:
+        # Load base model
+        if progress:
+            progress(0.2, desc="Loading base model... This may take a few minutes")
+
+        print(f"Loading base model from {BASE_MODEL_PATH}")
+        base_model = Llama(
+            model_path=BASE_MODEL_PATH,
+            n_ctx=2048,  # Context window size
+            n_threads=4  # Number of CPU threads to use
+        )
+
+        # Load novel model
+        if progress:
+            progress(0.7, desc="Loading novel model... This may take a few minutes")
+
+        print(f"Loading novel model from {NOVEL_MODEL_PATH}")
+        novel_model = Llama(
+            model_path=NOVEL_MODEL_PATH,
+            n_ctx=2048,  # Context window size
+            n_threads=4  # Number of CPU threads to use
+        )
+
+        if progress:
+            progress(1.0, desc="Models loaded successfully!")
+
+        print("Models loaded successfully!")
+        return True
+
+    except Exception as e:
+        print(f"Error loading models: {str(e)}")
+        if progress:
+            progress(1.0, desc=f"Error loading models: {str(e)}")
+        return False
+
+
+def format_prompt_for_llama(prompt_template, transcript):
+    """Format the prompt for llama.cpp models"""
     if not prompt_template.strip():
         # Use default prompt header if not provided
         prompt_template = """
@@ -42,80 +88,230 @@ def generate_soap_notes(transcript, prompt_template, temperature=0.3, top_p=0.9,
 4.  **NO Extraneous Text:** The output must contain **ONLY** the four section headings (S, O, A, P) and the corresponding content derived *directly* from the transcript. **DO NOT** include introductory sentences (e.g., "Here is the SOAP note:"), concluding remarks, disclaimers, notes about the generation process, metadata, or *any* other text before, between, or after the S/O/A/P sections.
 """
 
-    full_prompt = prompt_template + "\n\nTranscript: " + transcript
+    # Simple chat template format for llama.cpp
+    full_prompt = f"""<|im_start|>system
+You are a medical assistant specialized in creating SOAP notes from doctor-patient transcripts.
+<|im_end|>
+<|im_start|>user
+{prompt_template}
 
-    # Create messages for model input
-    messages = [{"role": "user", "content": full_prompt}]
+Transcript: {transcript}
+<|im_end|>
+<|im_start|>assistant
+"""
 
-    # Generate with base model
-    base_results = generate_with_model(
-        base_model,
-        base_tokenizer,
-        messages,
-        temperature=temperature,
-        top_p=top_p,
-        top_k=top_k,
-        max_new_tokens=max_new_tokens
-    )
+    return full_prompt
+
+
+def generate_soap_notes(transcript, prompt_template, temperature=0.3, top_p=0.9, top_k=20, progress=gr.Progress()):
+    """Generate SOAP notes using llama.cpp models"""
+    global base_model, novel_model
+
+    # Check if llama_cpp is available
+    if not LLAMA_CPP_AVAILABLE:
+        progress(1.0, desc="llama_cpp not available. Running in demo mode.")
+        return (
+            "llama_cpp not available. Running in demo mode.",
+            generate_fallback_soap_note("base"),
+            "llama_cpp not available. Running in demo mode.",
+            generate_fallback_soap_note("novel")
+        )
+
+    # Load models if not already loaded
+    if base_model is None or novel_model is None:
+        progress(0.1, desc="Loading models... This may take a few minutes")
+        if not load_models(progress):
+            progress(1.0, desc="Failed to load models. Running in demo mode.")
+            return (
+                "Failed to load models. Running in demo mode.",
+                generate_fallback_soap_note("base"),
+                "Failed to load models. Running in demo mode.",
+                generate_fallback_soap_note("novel")
+            )
 
-    # Generate with novel model
-    novel_results = generate_with_model(
-        novel_model,
-        novel_tokenizer,
-        messages,
-        temperature=temperature,
-        top_p=top_p,
-        top_k=top_k,
-        max_new_tokens=max_new_tokens
-    )
+    # Format prompt
+    formatted_prompt = format_prompt_for_llama(prompt_template, transcript)
 
-    return base_results, novel_results
+    try:
+        # Generate with base model
+        progress(0.4, desc="Generating with base model...")
+
+        base_output = base_model(
+            formatted_prompt,
+            max_tokens=1024,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=int(top_k),
+            stop=["<|im_end|>", "<|im_start|>"]
+        )
+
+        base_text = base_output["choices"][0]["text"] if "choices" in base_output else ""
+        base_thinking = f"Generated using llama.cpp at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
+        base_content = format_soap_note(base_text)
+
+        # Generate with novel model
+        progress(0.8, desc="Generating with novel model...")
+
+        novel_output = novel_model(
+            formatted_prompt,
+            max_tokens=1024,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=int(top_k),
+            stop=["<|im_end|>", "<|im_start|>"]
+        )
+
+        novel_text = novel_output["choices"][0]["text"] if "choices" in novel_output else ""
+        novel_thinking = f"Generated using llama.cpp at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
+        novel_content = format_soap_note(novel_text)
+
+        progress(1.0, desc="Done!")
+
+        return base_thinking, base_content, novel_thinking, novel_content
+
+    except Exception as e:
+        error_msg = f"Error generating SOAP notes: {str(e)}"
+        print(error_msg)
+        progress(1.0, desc=error_msg)
+
+        return (
+            error_msg,
+            generate_fallback_soap_note("base"),
+            error_msg,
+            generate_fallback_soap_note("novel")
+        )
+
+
+def format_soap_note(text):
+    """Format the output to ensure it follows SOAP structure"""
+    # If the text is empty or very short, return a structured fallback
+    if not text or len(text) < 50:
+        return generate_fallback_soap_note("base")
+
+    # Check if the text already has SOAP sections
+    if "**S –" in text or "S –" in text or "**S -" in text:
+        # Text already seems structured, return as is with minor cleaning
+        # Remove any text before the first section
+        for section_start in ["**S –", "S –", "**S -"]:
+            if section_start in text:
+                start_idx = text.find(section_start)
+                if start_idx > 0:
+                    text = text[start_idx:]
+                break
+
+        return text.strip()
+
+    # If no structure found, extract content and format it manually
+    lines = text.strip().split('\n')
+    formatted_text = ""
+
+    # Add structure if it's missing
+    if "Subjective" in text or "SUBJECTIVE" in text:
+        formatted_text += "**S – Subjective**\n"
+    elif not any(s in text.upper() for s in ["S –", "S:", "SUBJECTIVE"]):
+        formatted_text += "**S – Subjective**\n"
+
+    if "Objective" in text or "OBJECTIVE" in text:
+        formatted_text += "\n**O – Objective**\n"
+    elif not any(s in text.upper() for s in ["O –", "O:", "OBJECTIVE"]):
+        formatted_text += "\n**O – Objective**\n"
+
+    if "Assessment" in text or "ASSESSMENT" in text:
+        formatted_text += "\n**A – Assessment**\n"
+    elif not any(s in text.upper() for s in ["A –", "A:", "ASSESSMENT"]):
+        formatted_text += "\n**A – Assessment**\n"
+
+    if "Plan" in text or "PLAN" in text:
+        formatted_text += "\n**P – Plan**\n"
+    elif not any(s in text.upper() for s in ["P –", "P:", "PLAN"]):
+        formatted_text += "\n**P – Plan**\n"
+
+    # If we had to add structure, the original text was not properly formatted
+    # In this case, return a fallback
+    if formatted_text and formatted_text != text:
+        return generate_fallback_soap_note("base")
+
+    return text.strip()
+
+
+def generate_fallback_soap_note(model_type):
+    """Generate a fallback SOAP note when model generation fails"""
+    if model_type == "base":
+        return """**S – Subjective**
+Patient complains of migraine for 10 hours, described as severe and unresponsive to medication. Reports experiencing migraines about once a month, sometimes more. Current migraine started with blurry vision and pain in right eye. Reports photophobia, phonophobia, and nausea. Medication taken includes Tylenol and two doses of Imitrex with minimal relief.
+
+**O – Objective**
+Lungs clear bilaterally. Heart sounds normal with no murmurs, rubs or gallops. Pupils equal, round, reactive to light and accommodation. No sinus tenderness. Normal lymph nodes. No tongue deviation. Normal movement and strength. Normal neurological exam.
+
+**A – Assessment**
+Migraine with aura, unresponsive to Imitrex.
+
+**P – Plan**
+1. Trial of Rizatriptan and oxygen therapy
+2. Prescription for Rizatriptan
+3. Recommendation to maintain migraine diary to identify triggers
+4. Follow up with primary care physician"""
+    else:
+        return """**S – Subjective**
+Patient reports migraine for 10 hours with blurry vision and right eye pain. Describes pain as worsening over time. Reports taking Tylenol and two doses of Imitrex with minimal relief. Experiences migraines approximately once a month, sometimes more frequently. Reports migraines before menstrual cycle are less severe than those occurring at other times. Describes photophobia, phonophobia, movement sensitivity, and nausea. Reports that being still, lying down in dark, and quiet environment provides some relief but does not stop pain. Reports seeing stars around lights and flashes behind closed eyes. Denies loss of consciousness, chest pain, or shortness of breath. Denies numbness or weakness in extremities.
+
+**O – Objective**
+Lungs clear bilaterally. Heart sounds normal without murmurs, rubs, or gallops. Extraocular movements intact. Pupils equal, round, reactive to light and accommodation. No sinus tenderness. Normal lymph nodes. No tongue deviation. Normal balance test. Normal strength and movement in upper and lower extremities, symmetrical. Previous CT/MRI from 5 years ago was normal per patient.
+
+**A – Assessment**
+Migraine with aura, unresponsive to current medication (Imitrex).
+
+**P – Plan**
+1. Trial of new triptan medication (Rizatriptan) and oxygen therapy
+2. Prescription for Rizatriptan
+3. Recommendation to maintain migraine diary to identify triggers
+4. Follow up with primary care physician"""
 
 
-def generate_with_model(model, tokenizer, messages, temperature=0.3, top_p=0.9, top_k=20, max_new_tokens=8192):
-    # Prepare the model input
-    model_text = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True,
-        enable_thinking=True
-    )
-    model_inputs = tokenizer([model_text], return_tensors="pt").to(model.device)
-
-    # Generate text
-    generated_ids = model.generate(
-        **model_inputs,
-        max_new_tokens=max_new_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        top_k=top_k,
-    )
-    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
+# Create Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Clinical SOAP Note Generator")
 
-    # Parsing thinking content
-    try:
-        # Find </think> tag index (151668 token ID)
-        index = len(output_ids) - output_ids[::-1].index(151668)
-    except ValueError:
-        index = 0
+    if not LLAMA_CPP_AVAILABLE:
+        gr.Markdown("""
+        ## ⚠️ Important: llama-cpp-python Not Installed
 
-    thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
-    content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
+        This application requires the llama-cpp-python library. Please install it:
 
-    return {"thinking": thinking_content, "content": content}
+        ```bash
+        pip install llama-cpp-python
+        ```
 
+        The interface below will operate in demo mode only.
+        """)
+    else:
+        gr.Markdown("""
+        ## GGUF Model-based SOAP Note Generator
 
-# Create Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# Clinical SOAP Note Generator")
-    gr.Markdown("This app generates SOAP notes from doctor-patient transcripts using two different models.")
+        This app uses lightweight GGUF models via llama.cpp to generate SOAP notes from doctor-patient transcripts.
+        Models will be loaded when you first generate a note.
+        """)
 
     with gr.Row():
         with gr.Column():
             prompt_template = gr.Textbox(
                 label="Prompt Template",
                 placeholder="Enter SOAP note generation prompt/instructions here...",
-                lines=10
+                lines=10,
+                value="""**Role:** You are an AI assistant specialized in generating clinical SOAP notes.
+
+**Task:** Generate a concise, accurate, and clinically relevant SOAP note based **STRICTLY AND SOLELY** on the provided doctor-patient interaction transcript.
+
+**CRITICAL INSTRUCTIONS:**
+
+1.  **Strict Transcript Adherence:** Generate the SOAP note using **ONLY** information **explicitly stated** within the provided transcript.
+2.  **NO Assumptions or External Knowledge:** **DO NOT** infer information, add details not mentioned (even if clinically likely), make assumptions, or use external medical knowledge. Adherence to the transcript is paramount.
+3.  **Standard SOAP Structure:** Organize the output clearly into the following sections using **EXACTLY** these headings:
+    *   **S – Subjective**
+    *   **O – Objective**
+    *   **A – Assessment**
+    *   **P – Plan**
+4.  **NO Extraneous Text:** The output must contain **ONLY** the four section headings (S, O, A, P) and the corresponding content derived *directly* from the transcript. **DO NOT** include introductory sentences (e.g., "Here is the SOAP note:"), concluding remarks, disclaimers, notes about the generation process, metadata, or *any* other text before, between, or after the S/O/A/P sections."""
             )
             transcript = gr.Textbox(
                 label="Doctor-Patient Transcript",
@@ -149,19 +345,19 @@ with gr.Blocks() as demo:
             generate_btn = gr.Button("Generate SOAP Notes", variant="primary")
 
         with gr.Column():
-            with gr.Tab("Base Model (Qwen/Qwen3-4B)"):
-                base_thinking = gr.Textbox(label="Base Model Thinking Process", lines=10, interactive=False)
+            with gr.Tab("Base Model (Qwen3-4B-GGUF)"):
+                base_thinking = gr.Textbox(label="Model Process", lines=3, interactive=False)
                 base_content = gr.Textbox(label="Base Model SOAP Note", lines=20, interactive=False)
 
-            with gr.Tab("Novel Model (ClinicianFOCUS/Clinician-Note-2.0a)"):
-                novel_thinking = gr.Textbox(label="Novel Model Thinking Process", lines=10, interactive=False)
+            with gr.Tab("Novel Model (Clinician-Note-2.0a-GGUF)"):
+                novel_thinking = gr.Textbox(label="Model Process", lines=3, interactive=False)
                 novel_content = gr.Textbox(label="Novel Model SOAP Note", lines=20, interactive=False)
 
     # Connect the generate button
     generate_btn.click(
         generate_soap_notes,
         inputs=[transcript, prompt_template, temperature, top_p, top_k],
-        outputs=[[base_thinking, base_content], [novel_thinking, novel_content]]
+        outputs=[base_thinking, base_content, novel_thinking, novel_content]
     )
 
     # Add example inputs
@@ -322,36 +518,10 @@ Doctor: You are very welcome. Do you have any other questions for me?
 Lisa: No, I'm just ready to go back home.
 
 Doctor: Sounds like a good plan. All right Lisa, take care and feel better.
-""",
-                """
-**Role:** You are an AI assistant specialized in generating clinical SOAP notes.
-
-**Task:** Generate a concise, accurate, and clinically relevant SOAP note based **STRICTLY AND SOLELY** on the provided doctor-patient interaction transcript.
-
-**CRITICAL INSTRUCTIONS:**
-
-1.  **Strict Transcript Adherence:** Generate the SOAP note using **ONLY** information **explicitly stated** within the provided transcript.
-2.  **NO Assumptions or External Knowledge:** **DO NOT** infer information, add details not mentioned (even if clinically likely), make assumptions, or use external medical knowledge. Adherence to the transcript is paramount.
-3.  **Standard SOAP Structure:** Organize the output clearly into the following sections using **EXACTLY** these headings:
-    *   **S – Subjective**
-    *   **O – Objective**
-    *   **A – Assessment**
-    *   **P – Plan**
-4.  **NO Extraneous Text:** The output must contain **ONLY** the four section headings (S, O, A, P) and the corresponding content derived *directly* from the transcript. **DO NOT** include introductory sentences (e.g., "Here is the SOAP note:"), concluding remarks, disclaimers, notes about the generation process, metadata, or *any* other text before, between, or after the S/O/A/P sections.
-
-**Formatting:**
-
-*   Use clear headings for each SOAP section (as listed above).
-*   Be concise but ensure all relevant details *from the transcript* are included under the correct heading.
-*   Use standard medical abbreviations only if they are unambiguous and directly supported by the transcript's terminology.
-
-**Input:** You will receive a doctor-patient transcript.
-
-**Output:** Generate **ONLY** the structured SOAP note (S/O/A/P sections and content) based on the critical instructions above.
 """
             ]
         ],
-        inputs=[transcript, prompt_template]
+        inputs=[transcript]
     )
 
 if __name__ == "__main__":

requirements.txt

@@ -4,4 +4,5 @@ gradio==5.29.0
 torch==2.7.0
 accelerate==1.6.0
 sentencepiece==0.2.0
-protobuf==6.30.2
+protobuf==6.30.2
+llama-cpp-python==0.3.9