Fix: CPU-safe model for HF Space

app.py

@@ -1,12 +1,22 @@
 import streamlit as st
-import requests
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import re
 from pydub import AudioSegment
 import speech_recognition as sr
 import io
 
-# Backend API endpoint
-API_URL = "http://10.24.6.170:8000/predict"  # Change to your server IP
+# Load model and tokenizer from local fine-tuned directory
+MODEL_PATH = "Tufan1/BioMedLM-Cardio-Fold4-CPU"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+#model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map=None, low_cpu_mem_usage=True, torch_dtype=torch.float32)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_PATH,
+    device_map="auto" if torch.cuda.is_available() else None,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+).to(device)
+#model.to(device)
 
 # Dictionaries to decode user inputs
 gender_map = {1: "Female", 2: "Male"}
@@ -14,16 +24,31 @@ cholesterol_map = {1: "Normal", 2: "Elevated", 3: "Peak"}
 glucose_map = {1: "Normal", 2: "High", 3: "Extreme"}
 binary_map = {0: "No", 1: "Yes"}
 
-# Function to call the backend API
-def get_prediction_from_api(data):
-    try:
-        response = requests.post(API_URL, json=data)
-        if response.status_code == 200:
-            return response.json().get("diagnosis", "No diagnosis returned.")
-        else:
-            return f"Error {response.status_code}: {response.text}"
-    except Exception as e:
-        return f"API request failed: {e}"
+# Function to predict diagnosis using the LLM
+def get_prediction(age, gender, height, weight, ap_hi, ap_lo,
+                   cholesterol, glucose, smoke, alco, active):
+    input_text = f"""Patient Record:
+- Age: {age} years
+- Gender: {gender_map[gender]}
+- Height: {height} cm
+- Weight: {weight} kg
+- Systolic BP: {ap_hi} mmHg
+- Diastolic BP: {ap_lo} mmHg
+- Cholesterol Level: {cholesterol_map[cholesterol]}
+- Glucose Level: {glucose_map[glucose]}
+- Smokes: {binary_map[smoke]}
+- Alcohol Intake: {binary_map[alco]}
+- Physically Active: {binary_map[active]}
+
+Diagnosis:"""
+
+    inputs = tokenizer(input_text, return_tensors="pt").to(device)
+    model.eval()
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_new_tokens=4)
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    diagnosis = decoded.split("Diagnosis:")[-1].strip()
+    return diagnosis
 
 # Function to extract patient features from a phrase or transcribed audio
 def extract_details_from_text(text):
@@ -61,20 +86,8 @@ if input_mode == "Manual Input":
     active = st.radio("Physically Active?", [("No", 0), ("Yes", 1)], format_func=lambda x: x[0])[1]
 
     if st.button("Predict Diagnosis"):
-        payload = {
-            "age": age,
-            "gender": gender,
-            "height": height,
-            "weight": weight,
-            "ap_hi": ap_hi,
-            "ap_lo": ap_lo,
-            "cholesterol": cholesterol,
-            "glucose": glucose,
-            "smoke": smoke,
-            "alco": alco,
-            "active": active
-        }
-        diagnosis = get_prediction_from_api(payload)
+        diagnosis = get_prediction(age, gender, height, weight, ap_hi, ap_lo,
+                                   cholesterol, glucose, smoke, alco, active)
         st.success(f"🩺 **Predicted Diagnosis:** {diagnosis}")
 
 elif input_mode == "Text Phrase":
@@ -83,9 +96,7 @@ elif input_mode == "Text Phrase":
         try:
             values = extract_details_from_text(phrase)
             if all(v is not None for v in values):
-                keys = ["age", "gender", "height", "weight", "ap_hi", "ap_lo", "cholesterol", "glucose", "smoke", "alco", "active"]
-                payload = dict(zip(keys, values))
-                diagnosis = get_prediction_from_api(payload)
+                diagnosis = get_prediction(*values)
                 st.success(f"🩺 **Predicted Diagnosis:** {diagnosis}")
             else:
                 st.warning("Couldn't extract all fields from the text. Please revise.")
@@ -93,7 +104,7 @@ elif input_mode == "Text Phrase":
             st.error(f"Error: {e}")
 
 elif input_mode == "Audio Upload":
-    uploaded_file = st.file_uploader("Upload audio file (WAV, MP3, M4A, MPEG)", type=["wav", "mp3", "m4a", "mpeg"])
+    uploaded_file = st.file_uploader("Upload audio file (WAV, MP3, M4A, MPEG)", type=["wav", "mp3", "m4a","mpeg"])
     if uploaded_file:
         st.audio(uploaded_file, format='audio/wav')
         audio = AudioSegment.from_file(uploaded_file)
@@ -110,9 +121,7 @@ elif input_mode == "Audio Upload":
             st.markdown(f"**Transcribed Text:** _{text}_")
             values = extract_details_from_text(text)
             if all(v is not None for v in values):
-                keys = ["age", "gender", "height", "weight", "ap_hi", "ap_lo", "cholesterol", "glucose", "smoke", "alco", "active"]
-                payload = dict(zip(keys, values))
-                diagnosis = get_prediction_from_api(payload)
+                diagnosis = get_prediction(*values)
                 st.success(f"🩺 **Predicted Diagnosis:** {diagnosis}")
             else:
                 st.warning("Could not extract complete information from audio.")