Update app.py

app.py

@@ -4,6 +4,7 @@ import streamlit as st
 import pandas as pd
 import torch
 import os
+import re
 from huggingface_hub import login
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 from peft import PeftModel, PeftConfig
@@ -54,6 +55,18 @@ model, tokenizer = load_model()
 if 'history' not in st.session_state:
     st.session_state.history = []
 
+# Placeholder MET prediction function (replace with actual model)
+def predict_met(age, sex, weight_kg, bfp):
+    try:
+        # Simulated MET prediction (replace with your model)
+        # Example: Linear regression based on AntDatpromt.csv features
+        lbm = weight_kg * (1 - (bfp / 100))
+        met = 3.5 + 0.1 * lbm - 0.05 * age + (0.3 if sex == "male" else 0.0)  # Placeholder formula
+        return round(met, 2)
+    except Exception as e:
+        st.warning(f"MET prediction error: {str(e)}")
+        return None
+
 # Prediction function
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -123,34 +136,51 @@ def generate_response(age, sex, height_cm, weight_kg, wc_cm):
 
         # Generate output
         st.write("🤖 Model response:")
-        with st.empty():
-            text_streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)
-            output = model.generate(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                max_new_tokens=250,
-                temperature=0.7,
-                top_p=0.95,
-                do_sample=True,
-                pad_token_id=tokenizer.eos_token_id,
-                use_cache=True,
-                streamer=text_streamer
-            )
+        output_container = st.empty()
+        text_streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        output = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=250,
+            temperature=0.7,
+            top_p=0.95,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+            use_cache=True,
+            streamer=text_streamer
+        )
 
         # Decode the output
         decoded = tokenizer.decode(output[0], skip_special_tokens=False)
         st.write(f"Decoded output: {decoded}")
 
+        # Extract assistant response (remove system header and tokens)
+        assistant_response = re.search(r"<|start_header_id|>assistant<|end_header_id>\n(.*)<|eot_id>", decoded, re.DOTALL)
+        clean_response = assistant_response.group(1).strip() if assistant_response else decoded
+
+        # MET prediction
+        try:
+            bfp_match = re.search(r"BFP:\s*(\d+\.\d+)%", clean_response)
+            bfp = float(bfp_match.group(1)) if bfp_match else 18.36  # Fallback
+            met_pred = predict_met(age, sex, weight_kg, bfp)
+            if met_pred is not None:
+                clean_response += f", Predicted MET: {met_pred}"
+        except Exception as e:
+            st.warning(f"MET prediction skipped: {str(e)}")
+
         # Update history
-        st.session_state.history.append((prompt, decoded))
-        return decoded
+        st.session_state.history.append((prompt, clean_response))
+
+        # Display clean response
+        output_container.write(clean_response)
+        return clean_response
 
     except Exception as e:
         st.error(f"Error during generation: {str(e)}")
         return None
 
 # UI Header
-st.title("🧠 AnthroBot")
+st.title("🧠 Health Metric Estimator")
 st.markdown("Enter your anthropometric details to receive an AI-generated summary of health metrics.")
 
 # Tabs for input method
@@ -211,5 +241,4 @@ with tab2:
 # Clear history button
 if st.button("Clear History"):
     st.session_state.history = []
-    st.rerun()
-
+    st.rerun()