Update app.py

app.py

@@ -11,52 +11,205 @@ import io
 from transformers.tokenization_utils_base import BatchEncoding
 
 # Login using Hugging Face token
-login(token=os.getenv("HUGGINGFACEHUB_TOKEN"))
+try:
+    login(token=os.getenv("HUGGINGFACEHUB_TOKEN"))
+except Exception as e:
+    st.error(f"Error logging in to Hugging Face: {str(e)}")
+    st.stop()
 
 st.set_page_config(page_title="AnthroBot", page_icon="🤖", layout="centered")
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load model and tokenizer
+# Load model & tokenizer
 @st.cache_resource
 def load_model():
-    model = AutoModelForCausalLM.from_pretrained("SallySims/AnthroBot_Model_Lora").to(device)
-    tokenizer = AutoTokenizer.from_pretrained("SallySims/AnthroBot_Model_Lora")
-    return model, tokenizer
+    try:
+        peft_config = PeftConfig.from_pretrained("SallySims/AnthroBot_Model_Lora")
+        base_model = AutoModelForCausalLM.from_pretrained(
+            peft_config.base_model_name_or_path,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            trust_remote_code=True,
+            token=True
+        )
+        model = PeftModel.from_pretrained(base_model, "SallySims/AnthroBot_Model_Lora")
+        model.eval()
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            peft_config.base_model_name_or_path,
+            trust_remote_code=True,
+            token=True
+        )
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id  # Set pad_token_id to eos_token_id (128001)
+
+        st.write("✅ Model and tokenizer loaded successfully.")
+        return model, tokenizer
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        raise e
 
 model, tokenizer = load_model()
 
-st.title("🧠 Health Metric Estimator")
-st.markdown("Enter your details below to get an AI-generated estimation.")
-
-# Input fields
-age = st.number_input("Age", min_value=1, max_value=120, value=30)
-sex = st.selectbox("Sex", options=["male", "female"])
-height = st.number_input("Height (cm)", min_value=50.0, max_value=250.0, value=170.0)
-weight = st.number_input("Weight (kg)", min_value=10.0, max_value=300.0, value=70.0)
-wc = st.number_input("Waist Circumference (cm)", min_value=20.0, max_value=200.0, value=80.0)
-
-if st.button("Estimate Metrics"):
-    prompt = f"Age: {age}, Sex: {sex}, Height: {height} cm, Weight: {weight} kg, WC: {wc} cm"
-    st.write(f"📝 Prompt Sent to Model: `{prompt}`")
-
-    messages = [{"role": "user", "content": prompt}]
-
-    # Tokenize
-    inputs = tokenizer.apply_chat_template(
-        messages,
-        tokenize=True,
-        add_generation_prompt=True,
-        return_tensors="pt"
-    ).to("cuda")
-
-    # Generate response with streaming
-    st.write("🤖 Model response:")
-    with st.empty():
-        text_streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-        _ = model.generate(
-            inputs,
-            streamer=text_streamer,
-            max_new_tokens=250,
-            use_cache=True
-        )
+# Initialize session state for prediction history
+if 'history' not in st.session_state:
+    st.session_state.history = []
+
+# Prediction function
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def generate_response(age, sex, height_cm, weight_kg, wc_cm):
+    try:
+        # Create prompt
+        prompt = f"Age: {age}, Sex: {sex}, Height: {height_cm} cm, Weight: {weight_kg} kg, WC: {wc_cm} cm"
+        st.write(f"📝 Prompt Sent to Model: `{prompt}`")
+
+        # Create message structure
+        messages = [{"role": "user", "content": prompt}]
+
+        # Tokenize the input
+        try:
+            inputs = tokenizer.apply_chat_template(
+                messages,
+                tokenize=True,
+                add_generation_prompt=True,
+                return_tensors="pt",
+                max_length=512,
+                truncation=True,
+                return_dict=True
+            )
+        except Exception as e:
+            st.warning(f"apply_chat_template failed: {str(e)}. Falling back to manual tokenization.")
+            inputs = tokenizer(
+                prompt,
+                return_tensors="pt",
+                max_length=512,
+                truncation=True,
+                padding=False,
+                return_attention_mask=True
+            )
+
+        # Debug: Log inputs structure
+        st.write(f"Inputs type: {type(inputs)}")
+        st.write(f"Inputs keys: {list(inputs.keys()) if isinstance(inputs, (dict, BatchEncoding)) else 'N/A'}")
+
+        # Handle inputs
+        if isinstance(inputs, (dict, BatchEncoding)):
+            input_ids = inputs['input_ids']
+            attention_mask = inputs.get('attention_mask', torch.ones_like(input_ids))
+        elif isinstance(inputs, torch.Tensor):
+            input_ids = inputs
+            attention_mask = torch.ones_like(input_ids)
+        else:
+            st.error(f"Unexpected inputs format: {type(inputs)}")
+            return None
+
+        # Ensure 2D tensors
+        if len(input_ids.shape) == 1:
+            input_ids = input_ids.unsqueeze(0)
+            attention_mask = attention_mask.unsqueeze(0)
+        elif len(input_ids.shape) > 2:
+            input_ids = input_ids.squeeze()
+            attention_mask = attention_mask.squeeze()
+            if len(input_ids.shape) == 1:
+                input_ids = input_ids.unsqueeze(0)
+                attention_mask = attention_mask.unsqueeze(0)
+
+        st.write(f"Input IDs shape: {input_ids.shape}")
+        st.write(f"Attention mask shape: {attention_mask.shape}")
+
+        # Move to device
+        input_ids = input_ids.to(device)
+        attention_mask = attention_mask.to(device)
+
+        # 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
+            )
+
+        # Decode the output
+        decoded = tokenizer.decode(output[0], skip_special_tokens=False)
+        st.write(f"Decoded output: {decoded}")
+
+        # Update history
+        st.session_state.history.append((prompt, decoded))
+        return decoded
+
+    except Exception as e:
+        st.error(f"Error during generation: {str(e)}")
+        return None
+
+# UI Header
+st.title("🧠 AnthroBot")
+st.markdown("Enter your anthropometric details to receive an AI-generated summary of health metrics.")
+
+# Tabs for input method
+tab1, tab2 = st.tabs(["🧍 Manual Input", "📄 CSV Upload"])
+
+with tab1:
+    st.subheader("Manual Entry")
+    age = st.number_input("Age", min_value=1, max_value=120, value=30)
+    sex = st.selectbox("Sex", options=["male", "female"])
+    height = st.number_input("Height (cm)", min_value=50.0, max_value=250.0, value=170.0)
+    weight = st.number_input("Weight (kg)", min_value=10.0, max_value=300.0, value=70.0)
+    wc = st.number_input("Waist Circumference (cm)", min_value=20.0, max_value=200.0, value=80.0)
+
+    if st.button("Estimate Metrics"):
+        prediction = generate_response(age, sex, height, weight, wc)
+        if prediction:
+            st.success("Prediction:")
+            st.write(prediction)
+
+    # Display history
+    st.subheader("Prediction History")
+    for prompt, response in st.session_state.history:
+        st.markdown(f"**Input**: {prompt}")
+        st.markdown(f"**Output**: {response}")
+
+with tab2:
+    st.subheader("Batch Upload via CSV")
+    sample_csv = pd.DataFrame({
+        "Age": [30],
+        "Sex": ["male"],
+        "Height": [170.0],
+        "Weight": [70.0],
+        "WC": [80.0]
+    })
+
+    st.download_button("📥 Download Sample CSV", sample_csv.to_csv(index=False), file_name="sample_input.csv")
+
+    uploaded_file = st.file_uploader("Upload your CSV file", type=["csv"])
+
+    if uploaded_file:
+        df = pd.read_csv(uploaded_file)
+        if not all(col in df.columns for col in ["Age", "Sex", "Height", "Weight", "WC"]):
+            st.error("CSV must contain columns: Age, Sex, Height, Weight, WC")
+        else:
+            outputs = []
+            with st.spinner("Generating predictions..."):
+                for _, row in df.iterrows():
+                    prediction = generate_response(row['Age'], row['Sex'], row['Height'], row['Weight'], row['WC'])
+                    outputs.append(prediction if prediction else "Error")
+
+                df["Prediction"] = outputs
+                st.success("Here are your predictions:")
+                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")
+
+# Clear history button
+if st.button("Clear History"):
+    st.session_state.history = []
+    st.rerun()
+