Update app.py

app.py

@@ -1,24 +1,21 @@
 ## Deploying on HuggingFace
 
-
 import streamlit as st
 import pandas as pd
 import torch
 import os
 from huggingface_hub import login
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel, PeftConfig
 import io
 
 # Login using Hugging Face token stored in Space secrets
 login(token=os.getenv("HUGGINGFACEHUB_TOKEN"))
 
-
 st.set_page_config(page_title="AnthroBot", page_icon="🤖", layout="centered")
 
 # Load model & tokenizer
 @st.cache_resource
-
 def load_model():
     try:
         peft_config = PeftConfig.from_pretrained("SallySims/AnthroBot_Model_Lora")
@@ -35,14 +32,13 @@ def load_model():
 
         st.write("✅ Model and tokenizer loaded successfully.")
         return model, tokenizer
-
     except Exception as e:
         st.error(f"Error loading model: {str(e)}")
-        raise e 
+        raise e
+
 model, tokenizer = load_model()
 
 # Prediction function
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def get_prediction(prompt):
@@ -52,44 +48,70 @@ def get_prediction(prompt):
     messages = [{"role": "user", "content": prompt}]
     
     # Tokenize the input
-    inputs = tokenizer.apply_chat_template(
-        messages,
-        tokenize=True,
-        add_generation_prompt=True,  # This is needed for generation
-        return_tensors="pt",
-    ).to(device)
-
-    # Log the tokenized input
-    st.write(f"Tokenized input: {inputs}")
-
-    # Verify the shape of the tokenized input
-    st.write(f"Shape of tokenized input: {inputs['input_ids'].shape}")
-
-    # Ensure that input_ids has the correct shape
-    input_ids = inputs["input_ids"].squeeze(0)  # Remove the batch dimension if it's there
-    st.write(f"Corrected tokenized input shape: {input_ids.shape}")
-
-    # Generate output using the model
-    output = model.generate(
-        input_ids,  # Use the tokenized input
-        max_new_tokens=150,  # Limit the number of tokens
-        temperature=0.7,  # Control randomness of output
-        top_p=0.95,  # Sampling parameter
-        do_sample=True,  # Ensure sampling for diverse output
-    )
+    try:
+        inputs = tokenizer.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_tensors="pt",
+            max_length=512,
+            truncation=True
+        ).to(device)
+    except Exception as e:
+        st.error(f"Error during tokenization: {str(e)}")
+        return None
+
+    # Debug: Log inputs structure
+    st.write(f"Inputs type: {type(inputs)}")
+    if isinstance(inputs, dict):
+        st.write(f"Inputs keys: {list(inputs.keys())}")
+        if 'input_ids' in inputs:
+            st.write(f"Input IDs shape: {inputs['input_ids'].shape}")
+        else:
+            st.error("No 'input_ids' in tokenized inputs")
+            return None
+    else:
+        st.error(f"Unexpected inputs format: {type(inputs)}")
+        return None
+
+    # Extract input_ids safely
+    input_ids = inputs['input_ids']
+    if len(input_ids.shape) == 3 and input_ids.shape[0] == 1:
+        input_ids = input_ids.squeeze(0)  # Remove batch dimension if 3D
+    elif len(input_ids.shape) == 2:
+        pass  # Already 2D, no squeeze needed
+    else:
+        st.error(f"Invalid input_ids shape: {input_ids.shape}")
+        return None
+
+    st.write(f"Final input_ids shape: {input_ids.shape}")
+
+    # Generate output
+    try:
+        output = model.generate(
+            input_ids=input_ids,
+            max_new_tokens=150,
+            temperature=0.7,
+            top_p=0.95,
+            do_sample=True,
+            pad_token_id=tokenizer.pad_token_id
+        )
+    except Exception as e:
+        st.error(f"Error during generation: {str(e)}")
+        return None
 
     # Decode the output
-    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
-
-    # Log the decoded output
-    st.write(f"Decoded output: {decoded}")
-
-    return decoded
-
+    try:
+        decoded = tokenizer.decode(output[0], skip_special_tokens=True)
+        st.write(f"Decoded output: {decoded}")
+        return decoded
+    except Exception as e:
+        st.error(f"Error decoding output: {str(e)}")
+        return None
 
 # UI Header
 st.title("🧠 AnthroBot")
-st.write("Enter your anthropometric estimates to receive an interpreted summary inputs — manually or via CSV upload.")
+st.write("Enter your anthropometric estimates to receive an interpreted summary — manually or via CSV upload.")
 
 # Tabs for input method
 tab1, tab2 = st.tabs(["🧍 Manual Input", "📄 CSV Upload"])
@@ -105,8 +127,9 @@ with tab1:
     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)
-        st.success("Prediction:")
-        st.write(prediction)
+        if prediction:
+            st.success("Prediction:")
+            st.write(prediction)
 
 with tab2:
     st.subheader("Batch Upload via CSV")
@@ -135,13 +158,11 @@ with tab2:
                         f"Weight: {row['Weight']} kg, WC: {row['WC']} cm\n\n###"
                     )
                     prediction = get_prediction(prompt)
-                    outputs.append(prediction)
+                    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")
-
-
+                st.download_button("📤 Download Predictions", data=csv_output, file_name="predictions.csv")