Create app.py

app.py

@@ -0,0 +1,130 @@
+import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import os
+
+# Get the Hugging Face token from the environment variables
+hf_token = os.environ.get("HF_TOKEN")
+
+# Initialize the tokenizer and model
+# We are now using MedGemma, a 4 billion parameter instruction-tuned model
+# specialized for the medical domain.
+model_id = "google/medgemma-4b-it"
+
+# Check for GPU availability and set the data type accordingly
+if torch.cuda.is_available():
+    dtype = torch.bfloat16
+else:
+    dtype = torch.float32
+
+# Load the tokenizer and model from Hugging Face
+try:
+    tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        token=hf_token,
+        torch_dtype=dtype,
+        device_map="auto",
+    )
+    model_loaded = True
+except Exception as e:
+    print(f"Error loading model: {e}")
+    model_loaded = False
+
+# This is the core function that will take the clinical text and return a code
+def get_clinical_code(clinical_text):
+    """
+    Generates a clinical code from unstructured clinical text using the MedGemma model.
+    """
+    if not model_loaded:
+        return "Error: The model could not be loaded. Please check the logs."
+        
+    if not clinical_text:
+        return "Please enter some clinical text."
+
+    # This is our prompt template. It's designed to guide the model
+    # to perform the specific task of clinical coding.
+    # We are asking for an ICD-10 code, which is a common standard.
+    prompt = f"""
+    <start_of_turn>user
+    You are an expert medical coder. Your task is to analyze the following clinical text and determine the most appropriate ICD-10 code. Provide only the ICD-10 code and a brief description.
+
+    Clinical Text: "{clinical_text}"
+
+    Provide the ICD-10 code and a brief description.
+    <end_of_turn>
+    <start_of_turn>model
+    """
+
+    # Prepare the input for the model
+    input_ids = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    # Generate the output from the model
+    # We are using a max length of 256 tokens which should be sufficient
+    # for a code and a short description.
+    outputs = model.generate(
+        **input_ids,
+        max_new_tokens=256,
+        do_sample=True,
+        temperature=0.7, # A lower temperature makes the output more deterministic
+    )
+    
+    # Decode the output and clean it up
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Extract the relevant part of the response
+    # The model will output the prompt as well, so we need to remove it.
+    model_response_start = response.find("<start_of_turn>model") + len("<start_of_turn>model")
+    clean_response = response[model_response_start:].strip()
+    
+    return clean_response
+
+# Create the Gradio Interface
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # Clinical Code Generator with Google MedGemma
+        Enter a piece of unstructured clinical text below, and the app will suggest an ICD-10 clinical code.
+        *Disclaimer: This is a demonstration and not for professional medical use.*
+        """
+    )
+    
+    with gr.Row():
+        # Input Textbox
+        input_text = gr.Textbox(
+            label="Unstructured Clinical Text",
+            placeholder="e.g., Patient presents with a severe headache and photophobia...",
+            lines=10
+        )
+        
+        # Output Textbox
+        output_text = gr.Textbox(
+            label="Suggested Clinical Code (ICD-10)",
+            interactive=False,
+            lines=5
+        )
+
+    # Submit Button
+    submit_button = gr.Button("Get Clinical Code", variant="primary")
+
+    # Connect the button to the function
+    submit_button.click(
+        fn=get_clinical_code,
+        inputs=input_text,
+        outputs=output_text
+    )
+    
+    gr.Examples(
+        examples=[
+            ["The patient complains of a persistent cough and fever for the past three days. Chest X-ray shows signs of pneumonia."],
+            ["45-year-old male with a history of hypertension presents with chest pain radiating to the left arm."],
+            ["The patient has a history of type 2 diabetes and is here for a routine check-up. Blood sugar levels are elevated."]
+        ],
+        inputs=input_text,
+        outputs=output_text,
+        fn=get_clinical_code
+    )
+
+# Launch the Gradio app
+if __name__ == "__main__":
+    demo.launch()