Update app.py

app.py

@@ -4,7 +4,7 @@ import pandas as pd
 import torch
 import os
 from huggingface_hub import login
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 from peft import PeftModel, PeftConfig
 import io
 
@@ -28,6 +28,7 @@ def load_model():
 
         tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path)
         tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id  # Explicitly set pad_token_id
 
         st.write("✅ Model and tokenizer loaded successfully.")
         return model, tokenizer
@@ -37,105 +38,17 @@ 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 fat"
-        elif bfp <= 20:
-            bfp_category = "Athlete"
-        elif bfp <= 24:
-            bfp_category = "Fit"
-        elif bfp <= 31:
-            bfp_category = "Average"
-        else:
-            bfp_category = "Obese"
-    else:
-        if bfp <= 5:
-            bfp_category = "Essential fat"
-        elif bfp <= 13:
-            bfp_category = "Athlete"
-        elif bfp <= 17:
-            bfp_category = "Fit"
-        elif bfp <= 24:
-            bfp_category = "Average"
-        else:
-            bfp_category = "Obese"
-
-    # Interpretation
-    if bfp_category=="Essential fat":
-        interpretation ='Minimum fat required for basic physiological functions (e.g., hormone production, insulation). Females require higher essential fat due to reproductive functions.'
-    elif bfp_category=='Athletes':
-        interpretation='Typical for competitive athletes with high muscle mass and low fat (e.g., runners, bodybuilders).'
-    elif bfp_category=='Fit':
-        interpretation='Healthy range for active individuals who exercise regularly but aren’t competitive athletes.'
-    elif bfp_category=='Average':
-        interpretation='Common for the general population, still within healthy limits'
-    elif bfp_category=='Obese':
-        interpretation='Associated with increased health risks (e.g., diabetes, heart disease and other CVDs)'
-
-
-DataAll
-    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,
-        "Interpretation": interpretation
-    }
-
 # Prediction function
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 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"Interpretation: {metrics['Interpretation']}.\n###"
-    )
+    # Create prompt matching test code
+    prompt = f"Age: {age}, Sex: {sex}, Height: {height_cm} cm, Weight: {weight_kg} kg, WC: {wc_cm} cm"
     st.write(f"Received prompt: {prompt}")
 
     # Create message structure
     messages = [{"role": "user", "content": prompt}]
-    
+
     # Tokenize the input
     try:
         inputs = tokenizer.apply_chat_template(
@@ -144,7 +57,8 @@ def get_prediction(age, sex, height_cm, weight_kg, wc_cm):
             add_generation_prompt=True,
             return_tensors="pt",
             max_length=512,
-            truncation=True
+            truncation=True,
+            return_dict=True  # Ensure dictionary output
         )
     except Exception as e:
         st.warning(f"apply_chat_template failed: {str(e)}. Falling back to manual tokenization.")
@@ -153,45 +67,53 @@ def get_prediction(age, sex, height_cm, weight_kg, wc_cm):
             return_tensors="pt",
             max_length=512,
             truncation=True,
-            padding=False
+            padding=False,
+            return_attention_mask=True
         )
-    
+
     # Debug: Log inputs structure
     st.write(f"Inputs type: {type(inputs)}")
-    
+
     # Handle inputs (tensor or dict)
     if isinstance(inputs, torch.Tensor):
+        # Assume tensor is input_ids, create attention_mask
         input_ids = inputs
         if len(input_ids.shape) == 1:
-            input_ids = input_ids.unsqueeze(0)  # [sequence_length] -> [1, sequence_length]
-        elif len(input_ids.shape) > 2:
-            input_ids = input_ids.squeeze()
-            if len(input_ids.shape) == 1:
-                input_ids = input_ids.unsqueeze(0)
+            input_ids = input_ids.unsqueeze(0)
+        attention_mask = torch.ones_like(input_ids)  # Create attention_mask
     elif isinstance(inputs, dict) and 'input_ids' in inputs:
         input_ids = inputs['input_ids']
+        attention_mask = inputs.get('attention_mask', torch.ones_like(input_ids))
         if len(input_ids.shape) == 3 and input_ids.shape[0] == 1:
             input_ids = input_ids.squeeze(0)
+            attention_mask = attention_mask.squeeze(0)
         elif len(input_ids.shape) == 1:
             input_ids = input_ids.unsqueeze(0)
+            attention_mask = attention_mask.unsqueeze(0)
     else:
         st.error(f"Unexpected inputs format: {type(inputs)}")
         return None
 
     st.write(f"Input IDs shape: {input_ids.shape}")
+    st.write(f"Attention mask shape: {attention_mask.shape}")
 
-    # Ensure input_ids is on the correct device
+    # Move to device
     input_ids = input_ids.to(device)
+    attention_mask = attention_mask.to(device)
 
     # Generate output
     try:
+        text_streamer = TextStreamer(tokenizer)
         output = model.generate(
             input_ids=input_ids,
-            max_new_tokens=150,
+            attention_mask=attention_mask,
+            max_new_tokens=250,
             temperature=0.7,
             top_p=0.95,
             do_sample=True,
-            pad_token_id=tokenizer.pad_token_id
+            pad_token_id=tokenizer.eos_token_id,
+            use_cache=True,
+            streamer=text_streamer
         )
     except Exception as e:
         st.error(f"Error during generation: {str(e)}")
@@ -199,7 +121,7 @@ def get_prediction(age, sex, height_cm, weight_kg, wc_cm):
 
     # Decode the output
     try:
-        decoded = tokenizer.decode(output[0], skip_special_tokens=False)  # Preserve special tokens
+        decoded = tokenizer.decode(output[0], skip_special_tokens=False)
         st.write(f"Decoded output: {decoded}")
         return decoded
     except Exception as e:
@@ -257,4 +179,5 @@ with tab2:
                 st.dataframe(df)
 
                 csv_output = df.to_csv(index=False).encode("utf-8")
-                st.download_button("📤 Download Predictions", data=csv_output, file_name="predictions.csv")
+                st.download_button("📤 Download Predictions", data=csv_output, file_name="predictions.csv")
+