Updated app with remote API inference and multi-input support

app.py

@@ -1,17 +1,12 @@
 import streamlit as st
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import requests
 import re
 from pydub import AudioSegment
 import speech_recognition as sr
 import io
 
-# Load model and tokenizer from local fine-tuned directory
-MODEL_PATH = "Tufan1/BioMedLM-Cardio-Fold1-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.to(device)
+# Backend API endpoint
+API_URL = "http://10.24.6.170:8000/predict"  # Change to your server IP
 
 # Dictionaries to decode user inputs
 gender_map = {1: "Female", 2: "Male"}
@@ -19,31 +14,16 @@ cholesterol_map = {1: "Normal", 2: "Elevated", 3: "Peak"}
 glucose_map = {1: "Normal", 2: "High", 3: "Extreme"}
 binary_map = {0: "No", 1: "Yes"}
 
-# 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 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 extract patient features from a phrase or transcribed audio
 def extract_details_from_text(text):
@@ -74,15 +54,27 @@ if input_mode == "Manual Input":
     weight = st.number_input("Weight (kg)", min_value=10, max_value=200)
     ap_hi = st.number_input("Systolic BP", min_value=80, max_value=250)
     ap_lo = st.number_input("Diastolic BP", min_value=40, max_value=150)
-    cholesterol = st.selectbox("Cholesterol", [("Normal", 1), ("High", 2), ("Extreme", 3)], format_func=lambda x: x[0])[1]
+    cholesterol = st.selectbox("Cholesterol", [("Normal", 1), ("Elevated", 2), ("Peak", 3)], format_func=lambda x: x[0])[1]
     glucose = st.selectbox("Glucose", [("Normal", 1), ("High", 2), ("Extreme", 3)], format_func=lambda x: x[0])[1]
     smoke = st.radio("Smoker?", [("No", 0), ("Yes", 1)], format_func=lambda x: x[0])[1]
     alco = st.radio("Alcohol Intake?", [("No", 0), ("Yes", 1)], format_func=lambda x: x[0])[1]
     active = st.radio("Physically Active?", [("No", 0), ("Yes", 1)], format_func=lambda x: x[0])[1]
 
     if st.button("Predict Diagnosis"):
-        diagnosis = get_prediction(age, gender, height, weight, ap_hi, ap_lo,
-                                   cholesterol, glucose, smoke, alco, active)
+        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)
         st.success(f"🩺 **Predicted Diagnosis:** {diagnosis}")
 
 elif input_mode == "Text Phrase":
@@ -91,7 +83,9 @@ elif input_mode == "Text Phrase":
         try:
             values = extract_details_from_text(phrase)
             if all(v is not None for v in values):
-                diagnosis = get_prediction(*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)
                 st.success(f"🩺 **Predicted Diagnosis:** {diagnosis}")
             else:
                 st.warning("Couldn't extract all fields from the text. Please revise.")
@@ -99,7 +93,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)", type=["wav", "mp3", "m4a"])
+    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)
@@ -116,7 +110,9 @@ 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):
-                diagnosis = get_prediction(*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)
                 st.success(f"🩺 **Predicted Diagnosis:** {diagnosis}")
             else:
                 st.warning("Could not extract complete information from audio.")

main.py

@@ -0,0 +1,24 @@
+# main.py
+from fastapi import FastAPI, Request
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+app = FastAPI()
+
+MODEL_PATH = "Tufan1/BioMedLM-Cardio-Fold1-CPU"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map="auto", torch_dtype=torch.float16)
+
+class PatientData(BaseModel):
+    input_text: str
+
+@app.post("/predict")
+def predict(data: PatientData):
+    inputs = tokenizer(data.input_text, return_tensors="pt").to("cuda")
+    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": diagnosis}