Update app.py

app.py

@@ -8,7 +8,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel, PeftConfig
 import io
 
-# Login using Hugging Face token stored in Space secrets
+# Login using Hugging Face token
 login(token=os.getenv("HUGGINGFACEHUB_TOKEN"))
 
 st.set_page_config(page_title="AnthroBot", page_icon="🤖", layout="centered")
@@ -37,13 +37,88 @@ def load_model():
 
 model, tokenizer = load_model()
 
+# Calculate anthropometric metrics
+def calculate_metrics(age, sex, height_cm, weight_kg, wc_cm):
+    # Convert height and WC to meters
+    height_m = height_cm / 100
+    wc_m = wc_cm / 100
+    
+    # BMI
+    bmi = weight_kg / (height_m ** 2)
+    
+    # WHTR
+    whtr = wc_m / height_m
+    
+    # BFP (Boer's formula approximation)
+    if sex.lower() == "male":
+        bfp = (1.20 * bmi) + (0.23 * age) - 16.2
+    else:
+        bfp = (1.20 * bmi) + (0.23 * age) - 5.4
+    
+    # LBM
+    lbm = weight_kg * (1 - (bfp / 100))
+    
+    # BMI Category (WHO)
+    if bmi < 18.5:
+        bmi_category = "Underweight"
+    elif bmi <= 24.9:
+        bmi_category = "Normal"
+    elif bmi <= 29.9:
+        bmi_category = "Overweight"
+    else:
+        bmi_category = "Obese"
+    
+    # BFP Category (ACE)
+    if sex.lower() == "female":
+        if bfp <= 13:
+            bfp_category = "Essential"
+        elif bfp <= 20:
+            bfp_category = "Athlete"
+        elif bfp <= 24:
+            bfp_category = "Fitness"
+        elif bfp <= 31:
+            bfp_category = "Average"
+        else:
+            bfp_category = "Obese"
+    else:
+        if bfp <= 5:
+            bfp_category = "Essential"
+        elif bfp <= 13:
+            bfp_category = "Athlete"
+        elif bfp <= 17:
+            bfp_category = "Fitness"
+        elif bfp <= 24:
+            bfp_category = "Average"
+        else:
+            bfp_category = "Obese"
+    
+    return {
+        "BMI": round(bmi, 2),
+        "WHTR": round(whtr, 2),
+        "BFP": round(bfp, 2),
+        "LBM": round(lbm, 2),
+        "BMI_Category": bmi_category,
+        "BFP_Category": bfp_category
+    }
+
 # Prediction function
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-def get_prediction(prompt):
+def get_prediction(age, sex, height_cm, weight_kg, wc_cm):
+    # Calculate metrics
+    metrics = calculate_metrics(age, sex, height_cm, weight_kg, wc_cm)
+    
+    # Create prompt with metrics
+    prompt = (
+        f"Age: {age}, Sex: {sex}, Height: {height_cm} cm, Weight: {weight_kg} kg, WC: {wc_cm} cm\n"
+        f"BMI: {metrics['BMI']} kg/m2, WHTR: {metrics['WHTR']} m, BFP: {metrics['BFP']}%, "
+        f"LBM: {metrics['LBM']} kg, BMI Category: {metrics['BMI_Category']}, "
+        f"BFP Category: {metrics['BFP_Category']}\n"
+        f"Provide an interpretation of these anthropometric metrics.\n###"
+    )
     st.write(f"Received prompt: {prompt}")
 
-    # Create a message structure
+    # Create message structure
     messages = [{"role": "user", "content": prompt}]
     
     # Tokenize the input
@@ -58,7 +133,6 @@ def get_prediction(prompt):
         )
     except Exception as e:
         st.warning(f"apply_chat_template failed: {str(e)}. Falling back to manual tokenization.")
-        # Fallback: Manual tokenization
         inputs = tokenizer(
             prompt,
             return_tensors="pt",
@@ -72,12 +146,11 @@ def get_prediction(prompt):
     
     # Handle inputs (tensor or dict)
     if isinstance(inputs, torch.Tensor):
-        # Direct tensor (likely input_ids)
         input_ids = inputs
         if len(input_ids.shape) == 1:
-            input_ids = input_ids.unsqueeze(0)  # Add batch dimension: [sequence_length] -> [1, sequence_length]
+            input_ids = input_ids.unsqueeze(0)  # [sequence_length] -> [1, sequence_length]
         elif len(input_ids.shape) > 2:
-            input_ids = input_ids.squeeze()  # Remove extra dimensions if any
+            input_ids = input_ids.squeeze()
             if len(input_ids.shape) == 1:
                 input_ids = input_ids.unsqueeze(0)
     elif isinstance(inputs, dict) and 'input_ids' in inputs:
@@ -111,7 +184,7 @@ def get_prediction(prompt):
 
     # Decode the output
     try:
-        decoded = tokenizer.decode(output[0], skip_special_tokens=True)
+        decoded = tokenizer.decode(output[0], skip_special_tokens=False)  # Preserve special tokens
         st.write(f"Decoded output: {decoded}")
         return decoded
     except Exception as e:
@@ -127,15 +200,14 @@ tab1, tab2 = st.tabs(["🧍 Manual Input", "📄 CSV Upload"])
 
 with tab1:
     st.subheader("Manual Entry")
-    age = st.number_input("Age", 0, 100, 30)
-    sex = st.selectbox("Sex", ["male", "female"])
-    height = st.number_input("Height (cm)", 100.0, 250.0, 150.5)
-    weight = st.number_input("Weight (kg)", 30.0, 200.0, 75.3)
-    wc = st.number_input("Waist Circumference (cm)", 30.0, 150.0, 68.0)
+    age = st.number_input("Age", 0, 100, 16)
+    sex = st.selectbox("Sex", ["male", "female"], index=1)
+    height = st.number_input("Height (cm)", 100.0, 250.0, 153.0)
+    weight = st.number_input("Weight (kg)", 30.0, 200.0, 51.1)
+    wc = st.number_input("Waist Circumference (cm)", 30.0, 150.0, 64.0)
 
     if st.button("Get Prediction"):
-        prompt = f"Age: {age}, Sex: {sex}, Height: {height} cm, Weight: {weight} kg, WC: {wc} cm\n\n###"
-        prediction = get_prediction(prompt)
+        prediction = get_prediction(age, sex, height, weight, wc)
         if prediction:
             st.success("Prediction:")
             st.write(prediction)
@@ -143,11 +215,11 @@ with tab1:
 with tab2:
     st.subheader("Batch Upload via CSV")
     sample_csv = pd.DataFrame({
-        "Age": [30],
+        "Age": [16],
         "Sex": ["female"],
-        "Height": [150.5],
-        "Weight": [75.3],
-        "WC": [68.0]
+        "Height": [153.0],
+        "Weight": [51.1],
+        "WC": [64.0]
     })
 
     st.download_button("📥 Download Sample CSV", sample_csv.to_csv(index=False), file_name="sample_input.csv")
@@ -162,11 +234,7 @@ with tab2:
             outputs = []
             with st.spinner("Generating predictions..."):
                 for _, row in df.iterrows():
-                    prompt = (
-                        f"Age: {row['Age']}, Sex: {row['Sex']}, Height: {row['Height']} cm, "
-                        f"Weight: {row['Weight']} kg, WC: {row['WC']} cm\n\n###"
-                    )
-                    prediction = get_prediction(prompt)
+                    prediction = get_prediction(row['Age'], row['Sex'], row['Height'], row['Weight'], row['WC'])
                     outputs.append(prediction if prediction else "Error")
 
                 df["Prediction"] = outputs