Update model references to Qwen2.5-1.5B-Instruct; adjust download and preload functions accordingly

download_models.py

@@ -1,11 +1,13 @@
 from huggingface_hub import snapshot_download
 import os
 
+# Keep the download set aligned with the active inference pipeline.
+# The app now uses the lighter Qwen2.5-1.5B-Instruct path for reasoning/report synthesis
+# to minimize latency on small GPU instances.
 MODELS = [
     "yasserrmd/GemmaECG-Vision",
     "google/tabfm-1.0.0-pytorch",
-    "Qwen/Qwen3-0.6B",
-    "BioMistral/BioMistral-7B"
+    "Qwen/Qwen2.5-1.5B-Instruct",
 ]
 
 def download():

main.py

@@ -2,7 +2,7 @@ import torch
 from fastapi import FastAPI, UploadFile, File, Form
 from pydantic import BaseModel
 from typing import Optional
-from transformers import AutoProcessor, AutoModelForMultimodalLM, AutoTokenizer, AutoModelForCausalLM, AutoModel, BitsAndBytesConfig
+from transformers import AutoProcessor, AutoModelForMultimodalLM, AutoTokenizer, AutoModel, BitsAndBytesConfig, pipeline
 from PIL import Image
 import io
 import uvicorn
@@ -80,6 +80,21 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 models = {}
 
+def extract_generated_text(response):
+    """Normalize chat-pipeline output to a plain string."""
+    if not response:
+        return ""
+
+    generated = response[0].get("generated_text", "")
+    if isinstance(generated, list):
+        if generated and isinstance(generated[-1], dict):
+            return generated[-1].get("content", "")
+        return " ".join(part.get("content", "") for part in generated if isinstance(part, dict))
+    if isinstance(generated, str):
+        return generated
+    return str(generated)
+
+
 def load_models():
     if models:
         return
@@ -123,23 +138,14 @@ def load_models():
         print(f"Warning: Could not load TabFM model: {e}. Using mock for pipeline demonstration.")
         models['tabfm_clf'] = None
 
-    print("Loading Qwen3-0.6B...")
-    models['qwen_tokenizer'] = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
-    models['qwen_model'] = AutoModelForCausalLM.from_pretrained(
-        "Qwen/Qwen3-0.6B",
+    print("Loading Qwen2.5-1.5B-Instruct via pipeline...")
+    models['qwen_pipeline'] = pipeline(
+        "text-generation",
+        model="Qwen/Qwen2.5-1.5B-Instruct",
+        torch_dtype=DTYPE,
         device_map="auto",
-        quantization_config=quant_config,
-        low_cpu_mem_usage=True
     )
 
-    print("Loading BioMistral-7B...")
-    models['mistral_tokenizer'] = AutoTokenizer.from_pretrained("BioMistral/BioMistral-7B")
-    models['mistral_model'] = AutoModelForCausalLM.from_pretrained(
-        "BioMistral/BioMistral-7B",
-        device_map="auto",
-        quantization_config=quant_config,
-        low_cpu_mem_usage=True
-    )
     print("All models loaded successfully.")
 
 class PatientMetrics(BaseModel):
@@ -215,8 +221,8 @@ async def predict(
         tabfm_risk_probability = 68.2
         tabfm_classification_label = "Positive for Heart Disease"
 
-    # 3. Reasoning Layer via Qwen3-0.6B
-    qwen_system_prompt = "You are a clinical reasoning engine. Analyze the provided patient metrics and ECG findings to provide a detailed, step-by-step diagnostic reasoning chain. Think through the correlations between demographics, vitals, and electrical cardiac activity."
+    # 3. Reasoning Layer via Qwen2.5-1.5B-Instruct
+    qwen_system_prompt = "You are a clinical reasoning engine powered by Qwen2.5-1.5B-Instruct. Analyze the provided patient metrics and ECG findings to provide a detailed, step-by-step diagnostic reasoning chain. Think through the correlations between demographics, vitals, and electrical cardiac activity."
     qwen_user_content = f"""
     PATIENT DATA:
     - Age: {patient_metrics.age}, Sex: {patient_metrics.sex}, Ethnicity: {patient_metrics.ethnicity}
@@ -234,35 +240,26 @@ async def predict(
     """
 
     qwen_messages = [{"role": "user", "content": f"{qwen_system_prompt}\n\n{qwen_user_content}"}]
-    qwen_text = models['qwen_tokenizer'].apply_chat_template(
+    qwen_outputs = models['qwen_pipeline'](
         qwen_messages,
-        tokenize=False,
-        add_generation_prompt=True,
-        enable_thinking=True
-    )
-    qwen_inputs = models['qwen_tokenizer'](qwen_text, return_tensors="pt").to(models['qwen_model'].device)
-    qwen_outputs = models['qwen_model'].generate(
-        **qwen_inputs,
-        max_new_tokens=1024,
-        do_sample=True,
-        temperature=0.6,
-        top_p=0.95,
-        top_k=20,
-        min_p=0,
-        repetition_penalty=1.2
+        max_new_tokens=300,
+        do_sample=False,
+        temperature=0.2,
+        top_p=0.9,
+        repetition_penalty=1.1,
     )
-    qwen_reasoning = models['qwen_tokenizer'].decode(qwen_outputs[0][qwen_inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+    qwen_reasoning = extract_generated_text(qwen_outputs)
 
     bmi_bp_interaction = 1.24 # Derived output
     health_risk_score = 0.78   # Derived output
 
-    system_prompt = """You are BioMistral-7B, a specialized clinical reasoning and medical decision-support model. You are part of an end-to-end multimodal diagnostic pipeline that integrates three upstream specialized models:
+    system_prompt = """You are a specialized clinical decision-support model powered by Qwen2.5-1.5B-Instruct. You are part of an end-to-end multimodal diagnostic pipeline that integrates three upstream specialized models:
 
 1. Google TabFM (google/tabfm-1.0.0-pytorch): Processes tabular clinical, demographic, metabolic, and engineered metrics.
 2. GemmaECG-Vision (yasserrmd/GemmaECG-Vision via image-text-to-text pipeline): Processes 12-lead ECG image waveforms.
-3. Qwen3-0.6B (Qwen/Qwen3-0.6B): A reasoning engine that generates a diagnostic chain-of-thought.
+3. Qwen2.5-1.5B-Instruct: A reasoning engine that generates a diagnostic chain-of-thought.
 
-Your job is to ingest all patient metrics along with the outputs of GemmaECG-Vision, Google TabFM, and the reasoning generated by Qwen3-0.6B to generate a unified, explainable medical evaluation report.
+Your job is to ingest all patient metrics along with the outputs of GemmaECG-Vision, Google TabFM, and the reasoning generated by Qwen2.5-1.5B-Instruct to generate a unified, explainable medical evaluation report.
 
 OPERATIONAL MANDATES:
 1. Complete Parameter Coverage: You MUST account for every single parameter in the input data block (Demographics, Vitals, Labs, Cardiac Features, Lifestyle, and Engineered Metrics). Do not omit any parameter.
@@ -316,31 +313,24 @@ OUTPUT HEADINGS:
 - TabFM Risk Probability: {tabfm_risk_probability}%
 - TabFM Classification Label: {tabfm_classification_label}
 
-### 4. UPSTREAM MODEL 3 OUTPUT: Qwen3-0.6B (Qwen/Qwen3-0.6B)
+### 4. UPSTREAM MODEL 3 OUTPUT: Qwen2.5-1.5B-Instruct
 - Clinical Reasoning Chain: {qwen_reasoning}
 """
 
-    messages = [
-        {"role": "user", "content": f"{system_prompt}\n\n{user_content}"},
+    final_messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_content},
     ]
 
-    inputs = models['mistral_tokenizer'].apply_chat_template(
-        messages,
-        add_generation_prompt=True,
-        tokenize=True,
-        return_dict=True,
-        return_tensors="pt",
-    ).to(models['mistral_model'].device)
-
-    outputs = models['mistral_model'].generate(
-        **inputs,
-        max_new_tokens=1024,
-        do_sample=True,
-        temperature=0.7,
+    final_output = models['qwen_pipeline'](
+        final_messages,
+        max_new_tokens=350,
+        do_sample=False,
+        temperature=0.2,
         top_p=0.9,
-        repetition_penalty=1.2
+        repetition_penalty=1.1,
     )
-    final_report = models['mistral_tokenizer'].decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+    final_report = extract_generated_text(final_output)
 
     return {
         "clinical_report": final_report,
@@ -359,7 +349,7 @@ OUTPUT HEADINGS:
 async def analyze_ecg(file: UploadFile = File(...)):
     """
     Analyzes a 12-lead ECG image and generates a detailed clinical report using
-    GemmaECG-Vision for visual analysis and BioMistral-7B for clinical synthesis.
+    GemmaECG-Vision for visual analysis and Qwen2.5-1.5B-Instruct for clinical synthesis.
     """
     image_data = await file.read()
     image = Image.open(io.BytesIO(image_data)).convert("RGB")
@@ -386,31 +376,22 @@ async def analyze_ecg(file: UploadFile = File(...)):
     ecg_outputs = models['ecg_model'].generate(**ecg_inputs, max_new_tokens=100)
     raw_findings = models['ecg_processor'].tokenizer.decode(ecg_outputs[0][ecg_inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
 
-    # 2. Reasoning Layer via Qwen3-0.6B
-    qwen_system_prompt = "You are a clinical reasoning engine. Analyze the raw AI-extracted ECG findings and provide a detailed, step-by-step clinical interpretation."
+    # 2. Reasoning Layer via Qwen2.5-1.5B-Instruct
+    qwen_system_prompt = "You are a clinical reasoning engine powered by Qwen2.5-1.5B-Instruct. Analyze the raw AI-extracted ECG findings and provide a detailed, step-by-step clinical interpretation."
     qwen_user_content = f"RAW AI ECG FINDINGS: {raw_findings}"
 
     qwen_messages = [{"role": "user", "content": f"{qwen_system_prompt}\n\n{qwen_user_content}"}]
-    qwen_text = models['qwen_tokenizer'].apply_chat_template(
+    qwen_outputs = models['qwen_pipeline'](
         qwen_messages,
-        tokenize=False,
-        add_generation_prompt=True,
-        enable_thinking=True
-    )
-    qwen_inputs = models['qwen_tokenizer'](qwen_text, return_tensors="pt").to(models['qwen_model'].device)
-    qwen_outputs = models['qwen_model'].generate(
-        **qwen_inputs,
-        max_new_tokens=1024,
-        do_sample=True,
-        temperature=0.6,
-        top_p=0.95,
-        top_k=20,
-        min_p=0,
-        repetition_penalty=1.2
+        max_new_tokens=200,
+        do_sample=False,
+        temperature=0.2,
+        top_p=0.9,
+        repetition_penalty=1.1,
     )
-    qwen_reasoning = models['qwen_tokenizer'].decode(qwen_outputs[0][qwen_inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+    qwen_reasoning = extract_generated_text(qwen_outputs)
 
-    # 3. Clinical Synthesis via BioMistral-7B
+    # 3. Clinical Synthesis via Qwen2.5-1.5B-Instruct
     ecg_report_system_prompt = """You are a Board-Certified Cardiologist. Your task is to take raw visual findings from an AI-based ECG analysis tool and a reasoning chain from a clinical reasoning engine, and synthesize them into a professional, structured clinical ECG report.
 
     Your report MUST include:
@@ -428,27 +409,20 @@ async def analyze_ecg(file: UploadFile = File(...)):
     {qwen_reasoning}
     """
 
-    messages = [
-        {"role": "user", "content": f"{ecg_report_system_prompt}\n\n{user_content}"},
+    final_messages = [
+        {"role": "system", "content": ecg_report_system_prompt},
+        {"role": "user", "content": user_content},
     ]
 
-    inputs = models['mistral_tokenizer'].apply_chat_template(
-        messages,
-        add_generation_prompt=True,
-        tokenize=True,
-        return_dict=True,
-        return_tensors="pt",
-    ).to(models['mistral_model'].device)
-
-    outputs = models['mistral_model'].generate(
-        **inputs,
-        max_new_tokens=1024,
-        do_sample=True,
-        temperature=0.7,
+    final_output = models['qwen_pipeline'](
+        final_messages,
+        max_new_tokens=250,
+        do_sample=False,
+        temperature=0.2,
         top_p=0.9,
-        repetition_penalty=1.2
+        repetition_penalty=1.1,
     )
-    final_report = models['mistral_tokenizer'].decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+    final_report = extract_generated_text(final_output)
 
     return {
         "ecg_report": final_report,

preload_models.py

@@ -1,8 +1,9 @@
 import torch
-from transformers import AutoProcessor, AutoModelForMultimodalLM, AutoTokenizer, AutoModelForCausalLM, AutoModel, BitsAndBytesConfig
+from transformers import AutoProcessor, AutoModelForMultimodalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
 
 DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
 
+
 def preload():
     print("Preloading models to disk...")
 
@@ -17,9 +18,9 @@ def preload():
     AutoProcessor.from_pretrained("yasserrmd/GemmaECG-Vision")
     AutoModelForMultimodalLM.from_pretrained(
         "yasserrmd/GemmaECG-Vision",
-        device_map="cpu", # Use CPU for build time to avoid needing GPU in Docker build
+        device_map="cpu",  # Use CPU for build time to avoid needing GPU in Docker build
         quantization_config=quant_config,
-        low_cpu_mem_usage=True
+        low_cpu_mem_usage=True,
     )
 
     print("Preloading Google TabFM...")
@@ -29,24 +30,17 @@ def preload():
     except Exception as e:
         print(f"Warning: Could not preload TabFM model: {e}")
 
-    print("Preloading Qwen3-0.6B...")
-    AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
-    AutoModelForCausalLM.from_pretrained(
-        "Qwen/Qwen3-0.6B",
-        device_map="cpu", # Use CPU for build time to avoid needing GPU in Docker build
-        quantization_config=quant_config,
-        low_cpu_mem_usage=True
-    )
-
-    print("Preloading BioMistral-7B...")
-    AutoTokenizer.from_pretrained("BioMistral/BioMistral-7B")
-    AutoModelForCausalLM.from_pretrained(
-        "BioMistral/BioMistral-7B",
+    print("Preloading Qwen2.5-1.5B-Instruct...")
+    AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct")
+    pipeline(
+        "text-generation",
+        model="Qwen/Qwen2.5-1.5B-Instruct",
+        torch_dtype=DTYPE,
         device_map="cpu",
-        quantization_config=quant_config,
-        low_cpu_mem_usage=True
     )
+
     print("All models preloaded successfully.")
 
+
 if __name__ == "__main__":
     preload()