app.py

import gradio as gr
import os
import time
import sys
from datetime import datetime

# Import model configuration
try:
    from model_paths import (
        BASE_MODEL_REPO,
        NOVEL_MODEL_REPO,
        BASE_MODEL_FILENAME,
        NOVEL_MODEL_FILENAME,
        MODELS_DIR
    )
except ImportError:
    # Default values if import fails
    BASE_MODEL_REPO = "unsloth/Qwen3-4B-GGUF"
    NOVEL_MODEL_REPO = "mradermacher/Clinician-Note-2.0a-i1-GGUF"
    BASE_MODEL_FILENAME = "Qwen3-4B-Q4_K_M.gguf"
    NOVEL_MODEL_FILENAME = "Clinician-Note-2.0a.i1-Q4_K_M.gguf"
    MODELS_DIR = "models"

    # Create models directory
    os.makedirs(MODELS_DIR, exist_ok=True)

# 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.")

# Set up initial model paths
BASE_MODEL_PATH = os.path.join(MODELS_DIR, BASE_MODEL_FILENAME)
NOVEL_MODEL_PATH = os.path.join(MODELS_DIR, NOVEL_MODEL_FILENAME)

# Initialize models
base_model = None
novel_model = None


def download_model(model_repo_id, model_filename, progress=None, progress_start=0.0, progress_end=0.1):
    """Download a model from Hugging Face Hub if it doesn't exist locally"""
    try:
        from huggingface_hub import hf_hub_download
        import os

        # Create models directory if it doesn't exist
        os.makedirs("models", exist_ok=True)

        # Define the local path for the model
        local_path = os.path.join("models", model_filename)

        # Check if model already exists locally
        if os.path.exists(local_path):
            print(f"Model {model_filename} already exists at {local_path}")
            return local_path

        # If progress is provided, update it
        if progress is not None:
            progress(progress_start, desc=f"Downloading {model_filename}... This may take a while")

        print(f"Downloading {model_filename} from Hugging Face Hub...")

        # Download the model from Hugging Face
        downloaded_path = hf_hub_download(
            repo_id=model_repo_id,
            filename=model_filename,
            local_dir="models",
            local_dir_use_symlinks=False
        )

        # If progress is provided, update it
        if progress is not None:
            progress(progress_end, desc=f"Downloaded {model_filename}")

        print(f"Model downloaded to {downloaded_path}")
        return downloaded_path

    except Exception as e:
        print(f"Error downloading model {model_filename}: {str(e)}")
        if progress is not None:
            progress(progress_end, desc=f"Error downloading model: {str(e)}")
        return None


def load_models(progress=None):
    """Load the llama.cpp models"""
    global base_model, novel_model, BASE_MODEL_PATH, NOVEL_MODEL_PATH

    if not LLAMA_CPP_AVAILABLE:
        print("llama_cpp not available, cannot load models")
        return False

    try:
        # Download base model if needed
        base_model_repo = "unsloth/Qwen3-4B-GGUF"
        base_model_filename = "Qwen3-4B-Q4_K_M.gguf"

        if progress is not None:
            progress(0.1, desc="Checking for base model...")

        base_model_path = download_model(
            base_model_repo,
            base_model_filename,
            progress,
            0.1, 0.3
        )

        if not base_model_path:
            raise Exception(f"Failed to download base model {base_model_filename}")

        BASE_MODEL_PATH = base_model_path

        # Download novel model if needed
        novel_model_repo = "mradermacher/Clinician-Note-2.0a-i1-GGUF"
        novel_model_filename = "Clinician-Note-2.0a.i1-Q4_K_M.gguf"

        if progress is not None:
            progress(0.4, desc="Checking for novel model...")

        novel_model_path = download_model(
            novel_model_repo,
            novel_model_filename,
            progress,
            0.4, 0.6
        )

        if not novel_model_path:
            raise Exception(f"Failed to download novel model {novel_model_filename}")

        NOVEL_MODEL_PATH = novel_model_path

        # Load base model
        if progress is not None:
            progress(0.7, 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=8092,  # Context window size
            n_threads=4,  # Number of CPU threads to use
            temperature=0.3,  # Temperature for sampling
            top_p=0.9,  # Top-p sampling
            top_k=20,  # Top-k sampling
        )

        # Load novel model
        if progress is not None:
            progress(0.9, 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=8092,  # Context window size
            n_threads=4,  # Number of CPU threads to use
            temperature=0.3,  # Temperature for sampling
            top_p=0.9,  # Top-p sampling
            top_k=20,  # Top-k sampling
        )

        if progress is not None:
            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 is not None:
            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 = """
**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.
"""

    # 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}

Transcript: {transcript}
<|im_end|>
<|im_start|>assistant
"""

    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")
            )

    # Format prompt
    formatted_prompt = format_prompt_for_llama(prompt_template, transcript)

    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"""


# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Clinical SOAP Note Generator")

    if not LLAMA_CPP_AVAILABLE:
        gr.Markdown("""
        ## ⚠️ Important: llama-cpp-python Not Installed

        This application requires the llama-cpp-python library. Please install it:

        ```bash
        pip install llama-cpp-python
        ```

        The interface below will operate in demo mode only.
        """)
    else:
        gr.Markdown("""
        ## GGUF Model-based SOAP Note Generator

        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,
                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",
                placeholder="Enter the doctor-patient transcript here...",
                lines=15
            )

            with gr.Row():
                temperature = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.3,
                    step=0.1,
                    label="Temperature"
                )
                top_p = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.9,
                    step=0.05,
                    label="Top-p"
                )
                top_k = gr.Slider(
                    minimum=1,
                    maximum=50,
                    value=20,
                    step=1,
                    label="Top-k"
                )

            generate_btn = gr.Button("Generate SOAP Notes", variant="primary")

        with gr.Column():
            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 (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]
    )

    # Add example inputs
    gr.Examples(
        [
            [
                """
Doctor: Hi Lisa, I'm Dr. Petros. How can I help you today?

Lisa: Hi doctor, I've had this migraine for hours and it's really frustrating me. I've taken Tylenol and my migraine meds but they aren't helping.

Doctor: How long have you been hurting?

Lisa: It's been about 10 hours now but it feels like forever. These usually only last a few hours. It started with my normal blurry vision and pain in my right eye. The pain just keeps getting worse.

Doctor: I'm sorry to hear that. I'm going to ask you some questions to get some more information. I'll try to keep this brief since I know you don't feel well. I see in your chart that you have had migraines previously. Does this occur often?

Lisa: About once a month, sometimes more.

Doctor: Okay, and you have a prescription that you take when the headache gets bad?

Lisa: Yes, except this one isn't getting any better with my medicine.

Doctor: Well that's no good. Do you know the name of the medication you're prescribed?

Lisa: Yes, it's called Imitrex.

Doctor: Do you use the pill or the injection?

Lisa: Pills. I have taken two doses so far but the medicine has barely helped.

Doctor: Do you have any triggers? Anything that you know brings on your migraines?

Lisa: Not really, except I usually get one just before my period, but those aren't as bad as the others. The ones I get at the other times, I don't know what causes them. They're usually worse.

Doctor: Alright, are you experiencing any other symptoms related to your migraine right now?

Lisa: Light and sound make it worse, and any movement makes it worse too. I felt a little nauseous but I haven't thrown up any.

Doctor: Let's dim the lights in here a little to try to make you more comfortable. Are all of those symptoms normal for your migraine as well?

Lisa: Thanks, that does help a little. And yeah, all of that is typical. Sometimes I do throw up when I have a migraine, but not always.

Doctor: Okay, is there anything else I should know? Has anything made it better at all?

Lisa: Being still and lying down in the dark has helped some. Quiet helps, but nothing makes it stop. The pain hasn't gone away, just gotten a little better, then worse again.

Doctor: Have you noticed any vision changes or any numbness or weakness in your arms or legs?

Lisa: My vision is a little blurry, but that's normal during these episodes. I usually see stars around lights and flashes behind my eyes when I close them. I haven't noticed any weakness though.

Doctor: It sounds like you typically have migraines with aura. That's the technical term for what you describe as stars and flashes. Have you had any loss of consciousness, chest pain, or shortness of breath?

Lisa: No, none of that.

Doctor: Okay, so other than the extended length of time, this is a typical migraine for you? No new symptoms?

Lisa: Yeah, pretty much.

Doctor: Have you ever had a migraine that did not respond to the Imitrex before?

Lisa: Yeah, this happens every so often. It's really annoying.

Doctor: Do you have another treatment that usually works for you?

Lisa: I don't know. I think it can be different every time.

Doctor: Okay, that's fine. We will figure out a solution. Do you mind if I do a quick physical exam?

Lisa: No, go ahead.

Doctor: Thank you. I'm going to listen to your heart and lungs. Just breathe normally. Lungs are clear bilaterally. Heart sounds are normal with no murmurs, rubs or gallops. Now can you look at me and follow my finger please with just your eyes?

Lisa: [complies]

Doctor: Okay, great. Now I'm sorry to do this, but I need to briefly check your eyes with the light. It's going to be a little unpleasant.

Lisa: It's okay, I understand.

Doctor: I am looking for your pupils' response to light. I just want to be sure I don't see any signs pointing to anything other than a migraine here, like a stroke or aneurysm. Good, pupils are equal, round, and react normally to light and accommodation. Now I'm going to press on your sinuses. You should feel a little pressure, but tell me if you feel any pain.

Lisa: Okay. No pain.

Doctor: Good. Let me check your neck for a moment. Okay, your lymph nodes are normal as well. Now open your mouth and stick your tongue out.

Lisa: [complies]

Doctor: Good, no deviation there. Now, have you ever had a CT or MRI since you started having these headaches?

Lisa: Yes, I had one about five years ago when they started happening more regularly. It was normal.

Doctor: Okay, that's good. Can you stand up for me? I will stand behind you in case you lose your balance. Close your eyes and hold your arms out to your sides. I'm just checking your inner ear and proprioception.

Lisa: [complies]

Doctor: Good, I don't see any problems. You can sit back down. Now I want you to put both your hands against mine and push them down.

Lisa: [complies]

Doctor: Good. Now do the same with your legs, try to push my hands up.

Lisa: [complies]

Doctor: Good, thank you. Your movement and strength are normal and symmetrical, so that's a good sign. You have a normal neuro exam. Alright Lisa, I don't see anything too concerning. It seems you have a migraine that is simply unresponsive to your normal medication. Let's see if we can help you with this headache. I'm going to put in some orders and I'll be back in a few minutes to talk with you again.

Lisa: Okay, sure. Thank you.

[Later]

Doctor: All right Lisa, I have a sample here of a newer type of triptan medication. You can chew it or let it dissolve in your mouth so it acts a little more quickly. I'd like to try it and see if it is any more effective than the Imitrex. I'm also going to put some low flow oxygen on you. There's evidence that some other types of headaches can mimic migraine symptoms and some of them respond to oxygen therapy.

Lisa: I'm okay with that. Anything that makes this better.

Doctor: Okay, I will turn the lights off and send the nurse in with the medication and oxygen. I'd like you to lie back and rest until we see if it helps. Do you have any questions for now?

Lisa: No, I don't. Thank you.

[Later]

Doctor: Hi Lisa, are you feeling any better?

Lisa: Actually, I am. I started to feel the pain fade about 15 minutes ago. It's down to a dull ache now.

Doctor: That's great. How's your vision now?

Lisa: Almost normal. Nausea completely gone.

Doctor: Wonderful!

Lisa: You have no idea.

Doctor: I'm glad you're feeling better. I will write you a prescription for the newer medication. It's called Rizatriptan. Your most up-to-date pharmacy is in our system, correct?

Lisa: Yes, it is.

Doctor: Good. So now that you can think a little better, let's talk about what might be causing these.

Lisa: Okay.

Doctor: Have you ever done a migraine diary?

Lisa: No, not really.

Doctor: I think it would be a really good idea. You can keep a calendar just for this on your phone or something like that. Whenever you have a migraine, write down everything you can remember about that day - what you ate or drank, where you were, how you slept, that kind of thing. Hopefully you will start to see a pattern in what sets them off. It might be a food allergy, stress, or lack of sleep. Many women have hormonal migraines. I know you said you have them monthly, but that the ones unrelated to your menstrual cycle are different?

Lisa: Yes, they are.

Doctor: So maybe concentrate on those. Migraines are much easier to prevent than to treat. If you can discover your triggers, there are some medications you could try as a preventative measure, but I hate to go there if you can figure out what causes them and just avoid those triggers. Some of the preventative medications have significant side effects. At your age, I would rather avoid them if possible.

Lisa: I totally agree. I will start keeping track of the headaches and try to find out what is causing them.

Doctor: Great. I suggest you follow up on that with your primary care physician. If you don't have one, I will be happy to help you, but your best option is for consistent follow-up with your own doctor. Does this all sound okay to you?

Lisa: Yes, that works for me. Thank you so much.

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.
"""
            ]
        ],
        inputs=[transcript]
    )

if __name__ == "__main__":
    demo.launch()