Update app.py

app.py

@@ -1,168 +1,249 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+import re
 
-# Load your model from Hugging Face Hub
+# Load model
 MODEL_ID = "Muhammadidrees/my-gpt-oss"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto")
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID, 
+    device_map="auto",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+)
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+    model.config.pad_token_id = tokenizer.eos_token_id
+
+
+def clean_output(text):
+    """Remove reasoning artifacts and extract only the actual report"""
+    
+    # Remove common reasoning patterns
+    patterns_to_remove = [
+        r"Let's produce.*?(?=\*\*|$)",
+        r"We need to.*?(?=\*\*|$)",
+        r"We must.*?(?=\*\*|$)",
+        r"assistantfinal\*\*",
+        r"Note that.*?(?=\*\*|$)",
+        r"Use concise statements.*?(?=\*\*|$)",
+        r"Provide bullet points.*?(?=\*\*|$)",
+        r"✅ Medical Insights(?!\*\*)",
+    ]
+    
+    for pattern in patterns_to_remove:
+        text = re.sub(pattern, "", text, flags=re.DOTALL | re.IGNORECASE)
+    
+    # Find where actual report starts
+    start_markers = [
+        "**1. Executive Summary**",
+        "**Executive Summary**",
+        "1. Executive Summary",
+        "## 1. Executive Summary"
+    ]
+    
+    for marker in start_markers:
+        if marker in text:
+            idx = text.find(marker)
+            text = text[idx:]
+            break
+    
+    # Clean up extra whitespace and duplicate asterisks
+    text = re.sub(r'\*{3,}', '**', text)
+    text = re.sub(r'\n{3,}', '\n\n', text)
+    text = text.strip()
+    
+    return text
+
 
-# Function to build structured input and query the LLM
 def analyze(
     albumin, creatinine, glucose, crp, mcv, rdw, alp,
     wbc, lymph, age, gender, height, weight
 ):
     # Calculate BMI
     try:
-        height_m = height / 100  # cm → m
+        height_m = height / 100
         bmi = round(weight / (height_m ** 2), 2)
-    except Exception:
+    except:
         bmi = "N/A"
 
-    # Improved system prompt with clearer instructions
-    system_prompt = """You are a professional AI Medical Assistant analyzing patient biomarkers.
-
-CRITICAL: Generate a COMPLETE report following this EXACT structure. Do not stop mid-sentence.
-
-=== REQUIRED OUTPUT FORMAT ===
+    # Determine status for each biomarker
+    def get_status(value, normal_range):
+        low, high = normal_range
+        if value < low:
+            return "Low"
+        elif value > high:
+            return "High"
+        return "Normal"
 
-1. Executive Summary
-   - Top Priority Issues: [List 2-3 main concerns]
-   - Key Strengths: [List 2-3 positive findings]
+    albumin_status = get_status(albumin, (3.5, 5.5))
+    creatinine_status = get_status(creatinine, (0.7, 1.3))
+    glucose_status = get_status(glucose, (70, 100))
+    crp_status = get_status(crp, (0, 3))
+    mcv_status = get_status(mcv, (80, 100))
+    rdw_status = get_status(rdw, (11.5, 14.5))
+    alp_status = get_status(alp, (44, 147))
+    wbc_status = get_status(wbc, (4, 11))
+    lymph_status = get_status(lymph, (20, 40))
 
-2. System-Specific Analysis
-   - Blood Health (MCV, RDW, Lymphocytes, WBC)
-   - Protein & Liver Health (Albumin, ALP)
-   - Kidney Health (Creatinine)
-   - Metabolic Health (Glucose, CRP)
+    # Ultra-direct prompt with explicit instruction to skip reasoning
+    prompt = f"""You are a medical AI. Generate ONLY the final report. Do NOT include any reasoning, planning, or meta-commentary.
 
-3. Personalized Action Plan
-   - Medical: [Recommended tests/consultations]
-   - Nutrition: [Dietary recommendations and supplements]
-   - Lifestyle: [Exercise, hydration, sleep guidance]
-   - Testing: [Follow-up labs needed]
+OUTPUT ONLY THIS STRUCTURE:
 
-4. Interaction Alerts
-   [Explain how biomarkers interact and influence each other]
+**1. Executive Summary**
+- **Top Priority Issues**: List 2-3 main concerns based on abnormal values
+- **Key Strengths**: List 2-3 positive findings
 
-6. Tabular Mapping
-   | Biomarker | Value | Status | AI Insight | Client Message |
-   |-----------|-------|--------|------------|----------------|
-   [Complete table for all biomarkers]
+**2. System-Specific Analysis**
+- **Blood Health** (MCV {mcv} fL, RDW {rdw}%, Lymph {lymph}%, WBC {wbc} K/uL): Brief clinical interpretation
+- **Protein & Liver Health** (Albumin {albumin} g/dL, ALP {alp} U/L): Brief clinical interpretation
+- **Kidney Health** (Creatinine {creatinine} mg/dL): Brief clinical interpretation
+- **Metabolic Health** (Glucose {glucose} mg/dL, CRP {crp} mg/L): Brief clinical interpretation
 
-7. Enhanced AI Insights
-   - Subclinical Nutrient Analysis (Iron, B12, Folate status)
-   - ALP Interpretation (bone vs liver origin)
-   - Immune System Assessment (WBC & lymphocyte trends)
-   - Long-term Health Considerations
+**3. Personalized Action Plan**
+- **Medical**: Specific tests or consultations needed
+- **Nutrition**: Dietary recommendations and supplements
+- **Lifestyle**: Exercise, hydration, sleep recommendations
+- **Testing**: Follow-up labs (ferritin, Vitamin D, GGT, etc.)
 
-=== END FORMAT ===
+**4. Interaction Alerts**
+Explain how abnormal biomarkers may interact or compound risks.
 
-Now analyze the following patient data and provide a COMPLETE report:"""
+**5. Tabular Mapping**
+| Biomarker | Value | Status | Clinical Insight | Recommendation |
+|-----------|-------|--------|------------------|----------------|
+| Albumin | {albumin} g/dL | {albumin_status} | [insight] | [action] |
+| Creatinine | {creatinine} mg/dL | {creatinine_status} | [insight] | [action] |
+| Glucose | {glucose} mg/dL | {glucose_status} | [insight] | [action] |
+| CRP | {crp} mg/L | {crp_status} | [insight] | [action] |
+| MCV | {mcv} fL | {mcv_status} | [insight] | [action] |
+| RDW | {rdw}% | {rdw_status} | [insight] | [action] |
+| ALP | {alp} U/L | {alp_status} | [insight] | [action] |
+| WBC | {wbc} K/uL | {wbc_status} | [insight] | [action] |
+| Lymphocytes | {lymph}% | {lymph_status} | [insight] | [action] |
 
-    # Construct patient profile
-    patient_input = f"""
-Patient Profile:
-- Age: {age} years
-- Gender: {gender}
-- Height: {height} cm
-- Weight: {weight} kg
-- BMI: {bmi}
+**6. Enhanced AI Insights & Longitudinal Risk**
+- **Subclinical Nutrients**: Predicted Iron, B12, Folate, Copper status
+- **ALP Source**: Likely bone vs liver origin
+- **Immune Trends**: WBC and lymphocyte assessment
+- **Long-term Risk**: Predictive health considerations
 
-Laboratory Values:
-- Albumin: {albumin} g/dL
-- Creatinine: {creatinine} mg/dL
-- Glucose: {glucose} mg/dL
-- C-Reactive Protein (CRP): {crp} mg/L
-- Mean Cell Volume (MCV): {mcv} fL
-- Red Cell Distribution Width (RDW): {rdw} %
-- Alkaline Phosphatase (ALP): {alp} U/L
-- White Blood Cell Count (WBC): {wbc} K/uL
-- Lymphocyte Percentage: {lymph} %
+Patient: {gender}, Age {age}, BMI {bmi}
 
-Generate complete analysis now:"""
+START REPORT NOW (no reasoning or commentary):
 
-    prompt = system_prompt + "\n" + patient_input
+**1. Executive Summary**"""
 
     try:
-        # Increased max_new_tokens significantly for complete output
-        result = pipe(
-            prompt,
-            max_new_tokens=2500,  # INCREASED from 1000
-            do_sample=True,
-            temperature=0.7,  # INCREASED from 0.3 for better generation
-            top_p=0.92,
-            top_k=50,
-            repetition_penalty=1.1,  # Prevent repetition
-            return_full_text=False,
-            pad_token_id=tokenizer.eos_token_id
-        )
-
-        output_text = result[0]["generated_text"].strip()
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=2800,
+                min_new_tokens=1000,
+                temperature=0.75,
+                top_p=0.92,
+                top_k=40,
+                repetition_penalty=1.2,
+                do_sample=True,
+                early_stopping=False,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                no_repeat_ngram_size=4
+            )
+        
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract the actual report part
+        if "START REPORT NOW" in generated_text:
+            output_text = generated_text.split("START REPORT NOW")[-1]
+        elif "**1. Executive Summary**" in generated_text:
+            parts = generated_text.split("**1. Executive Summary**")
+            # Take the LAST occurrence (the actual output, not from prompt)
+            if len(parts) > 1:
+                output_text = "**1. Executive Summary**" + parts[-1]
+            else:
+                output_text = generated_text
+        else:
+            output_text = generated_text
+        
+        # Clean up the output
+        output_text = clean_output(output_text)
         
-        # Clean up output - remove any prompt leakage
-        if "Executive Summary" in output_text:
-            idx = output_text.find("Executive Summary")
-            output_text = output_text[idx:]
-        elif "1. Executive Summary" in output_text:
-            idx = output_text.find("1. Executive Summary")
-            output_text = output_text[idx:]
+        # Validation
+        required_sections = [
+            "Executive Summary",
+            "System-Specific Analysis", 
+            "Personalized Action Plan",
+            "Interaction Alerts",
+            "Tabular Mapping",
+            "Enhanced AI Insights"
+        ]
         
-        # If output seems incomplete, add a note
-        if len(output_text) < 500:
-            output_text += "\n\n⚠️ Note: Output may be incomplete. Consider re-running the analysis."
+        missing_sections = [s for s in required_sections if s not in output_text]
+        
+        if len(output_text) < 800 or len(missing_sections) >= 3:
+            warning = "\n\n⚠️ **Model Performance Issue Detected**\n\n"
+            warning += f"Generated text length: {len(output_text)} characters\n"
+            warning += f"Missing sections: {', '.join(missing_sections) if missing_sections else 'None'}\n\n"
+            warning += "**This model may not be suitable for this task. Consider:**\n"
+            warning += "1. Using Llama-3-8B-Instruct or Mistral-7B-Instruct\n"
+            warning += "2. Fine-tuning on medical report generation\n"
+            warning += "3. Using API-based models (GPT-4, Claude, etc.)\n"
+            output_text += warning
         
         return output_text
-    
+        
     except Exception as e:
-        return f"Error during analysis: {str(e)}\n\nPlease check your input values and try again."
+        return f"❌ **Error**: {str(e)}\n\nPlease verify all inputs and model availability."
 
 
-# Build Gradio UI with improved layout
+# Gradio Interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    
     gr.Markdown("""
-    # 🧪 Wellness Insights AI
-    ### Enter Patient Profile & Lab Values for Comprehensive Analysis
+    # 🏥 AI Medical Biomarker Analysis
+    ### Comprehensive wellness insights from lab values
     """)
-
+    
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=1):
             gr.Markdown("### 👤 Demographics")
-            age = gr.Number(label="Age (years)", value=30)
-            gender = gr.Dropdown(choices=["Male", "Female"], label="Gender", value="Male")
+            age = gr.Number(label="Age", value=35)
+            gender = gr.Dropdown(["Male", "Female"], label="Gender", value="Male")
             height = gr.Number(label="Height (cm)", value=175)
-            weight = gr.Number(label="Weight (kg)", value=70)
-        
-        with gr.Column():
-            gr.Markdown("### 🩸 Blood Health Markers")
-            wbc = gr.Number(label="White Blood Cell Count (K/uL)", value=7.0)
-            lymph = gr.Number(label="Lymphocyte Percentage (%)", value=30)
-            mcv = gr.Number(label="Mean Cell Volume (fL)", value=90)
-            rdw = gr.Number(label="Red Cell Distribution Width (%)", value=13)
-
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("### 🫀 Metabolic Markers")
-            glucose = gr.Number(label="Glucose (mg/dL)", value=95)
-            crp = gr.Number(label="C-Reactive Protein (mg/L)", value=1.5)
+            weight = gr.Number(label="Weight (kg)", value=75)
+            
+            gr.Markdown("### 🩸 Blood Panel")
+            wbc = gr.Number(label="WBC (K/uL)", value=7.0)
+            lymph = gr.Number(label="Lymphocytes (%)", value=30)
+            mcv = gr.Number(label="MCV (fL)", value=90)
+            rdw = gr.Number(label="RDW (%)", value=13)
         
-        with gr.Column():
-            gr.Markdown("### 🧬 Organ Function Markers")
+        with gr.Column(scale=1):
+            gr.Markdown("### 🧬 Chemistry Panel")
             albumin = gr.Number(label="Albumin (g/dL)", value=4.2)
             creatinine = gr.Number(label="Creatinine (mg/dL)", value=1.0)
-            alp = gr.Number(label="Alkaline Phosphatase (U/L)", value=70)
-
-    analyze_btn = gr.Button("🔎 Generate Comprehensive Analysis", variant="primary", size="lg")
+            glucose = gr.Number(label="Glucose (mg/dL)", value=95)
+            crp = gr.Number(label="CRP (mg/L)", value=1.5)
+            alp = gr.Number(label="ALP (U/L)", value=75)
+            
+            analyze_btn = gr.Button("🔬 Generate Report", variant="primary", size="lg")
     
-    gr.Markdown("### 📋 Analysis Report")
+    gr.Markdown("### 📊 Analysis Output")
     output = gr.Textbox(
-        label="AI-Generated Lab Report",
-        lines=25,
+        label="Medical Report",
+        lines=30,
         max_lines=50,
-        show_copy_button=True
+        show_copy_button=True,
+        placeholder="Results will appear here..."
     )
-
+    
     analyze_btn.click(
         fn=analyze,
         inputs=[albumin, creatinine, glucose, crp, mcv, rdw, alp,
@@ -170,10 +251,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         outputs=output
     )
     
-    gr.Markdown("""
-    ---
-    **Note:** This tool provides educational insights based on biomarker analysis. 
-    Always consult healthcare professionals for medical advice.
-    """)
+    gr.Markdown("*Disclaimer: For educational purposes only. Consult healthcare professionals.*")
 
-demo.launch(share=False)
+demo.launch()