Update app.py

app.py

@@ -1,74 +1,97 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import matplotlib.pyplot as plt
-import pdfplumber
+import fitz  # PyMuPDF
 import io
+import base64
 
-# Load MedAlpaca (slow tokenizer to avoid SentencePiece issue)
-model_name = "TheBloke/medalpaca-7B-GPTQ"
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
+# Load Mistral model (or any other open-access instruct model)
+model_id = "mistralai/Mistral-7B-Instruct-v0.2"
+
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
 
+# Function to extract text from PDF
 def extract_text_from_pdf(pdf_file):
+    text = ""
     try:
-        with pdfplumber.open(pdf_file) as pdf:
-            return "\n".join(page.extract_text() for page in pdf.pages if page.extract_text())
+        with fitz.open(stream=pdf_file.read(), filetype="pdf") as doc:
+            for page in doc:
+                text += page.get_text()
     except Exception as e:
-        return f"Error reading PDF: {e}"
+        text = f"Error reading PDF: {str(e)}"
+    return text
 
+# Function to generate answer using model
+def ask_model(input_text, context):
+    prompt = f"""[INST] Given the following context, answer the question:\n\nContext:\n{context}\n\nQuestion:\n{input_text} [/INST]"""
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.7)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# Create a flag for abnormal values in chart
 def generate_chart(text):
-    # Dummy abnormal value check
-    keywords = ['fever', 'pain', 'high', 'low']
-    is_abnormal = any(word in text.lower() for word in keywords)
-    
-    labels = ['Normal', 'Abnormal']
-    values = [1, 0]
-    colors = ['green', 'red'] if is_abnormal else ['green', 'gray']
+    # Dummy logic to extract and visualize values
+    lines = [line.strip() for line in text.split('\n') if ':' in line]
+    labels, values, flags = [], [], []
+
+    for line in lines:
+        try:
+            key, val = line.split(':')
+            num = float(val.strip().split()[0])
+            labels.append(key.strip())
+            values.append(num)
+            flags.append(num > 100 or num < 20)  # just dummy abnormal range
+        except:
+            continue
+
+    if not labels:
+        return "No numerical data found for plotting."
 
-    if is_abnormal:
-        values = [0.2, 0.8]
+    colors = ['red' if flag else 'green' for flag in flags]
 
     fig, ax = plt.subplots()
-    bars = ax.bar(labels, values, color=colors)
-    
-    for bar in bars:
-        height = bar.get_height()
-        ax.text(bar.get_x() + bar.get_width() / 2.0, height, f'{height:.2f}', ha='center', va='bottom')
+    ax.barh(labels, values, color=colors)
+    ax.set_xlabel('Values')
+    ax.set_title('📊 Test Results (Green = Normal, Red = Abnormal)')
+    plt.tight_layout()
 
-    ax.set_ylim([0, 1])
-    ax.set_title("Normal vs Abnormal Indicator")
     buf = io.BytesIO()
-    plt.savefig(buf, format="png")
+    plt.savefig(buf, format='png')
     buf.seek(0)
+    encoded = base64.b64encode(buf.read()).decode('utf-8')
     plt.close(fig)
-    return buf
 
-def medalpaca_response(input_text, pdf=None):
-    if pdf:
-        extracted_text = extract_text_from_pdf(pdf)
-        input_text += "\n\n" + extracted_text
+    return f"data:image/png;base64,{encoded}"
 
-    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
-    with torch.no_grad():
-        outputs = model.generate(**inputs, max_new_tokens=200)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    chart = generate_chart(input_text)
+# Main Gradio interface
+def process_input(user_question, pdf_file):
+    if pdf_file is None:
+        return "Please upload a medical report (PDF).", None
+
+    context = extract_text_from_pdf(pdf_file)
+    response = ask_model(user_question, context)
+    chart = generate_chart(context)
     return response, chart
 
-demo = gr.Interface(
-    fn=medalpaca_response,
-    inputs=[
-        gr.Textbox(label="Enter Symptoms or Question"),
-        gr.File(label="Upload Medical PDF (optional)", type="file")
-    ],
-    outputs=[
-        gr.Textbox(label="MedAlpaca Response"),
-        gr.Image(label="Normal vs Abnormal Chart")
-    ],
-    title="🧠 MedAlpaca Medical Assistant",
-    description="Ask medical questions or upload a PDF. Get AI-powered responses with visual abnormality flags."
-)
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("# 🩺 Medical Report Analyzer\nUpload a PDF report, ask questions, and see abnormalities visualized.")
+    
+    with gr.Row():
+        user_question = gr.Textbox(label="Ask a medical question", placeholder="e.g. What does the report say about cholesterol?")
+        pdf_input = gr.File(label="Upload medical report (PDF)", file_types=['.pdf'])
+    
+    submit_btn = gr.Button("Analyze")
+    output_text = gr.Textbox(label="Answer", lines=8)
+    output_img = gr.Image(label="📊 Chart")
+
+    submit_btn.click(fn=process_input, inputs=[user_question, pdf_input], outputs=[output_text, output_img])
 
+# Launch the app
 demo.launch()