Update app.py

app.py

@@ -15,172 +15,46 @@ login(token=os.getenv("HUGGINGFACEHUB_TOKEN"))
 
 st.set_page_config(page_title="AnthroBot", page_icon="🤖", layout="centered")
 
-# Load model & tokenizer
+# Load model and tokenizer
 @st.cache_resource
 def load_model():
-    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"
-        )
-        model = PeftModel.from_pretrained(base_model, "SallySims/AnthroBot_Model_Lora")
-        model.eval()
-
-        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
-    except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
-        raise e
+    model = AutoModelForCausalLM.from_pretrained("SallySims/AnthroBot_Model_Lora").to("cuda")
+    tokenizer = AutoTokenizer.from_pretrained("SallySims/AnthroBot_Model_Lora")
+    return model, tokenizer
 
 model, tokenizer = load_model()
 
-# Prediction function
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-def get_prediction(age, sex, height_cm, weight_kg, wc_cm):
-    # 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(
-            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 content: {inputs}")
+st.title("🧠 Health Metric Estimator")
+st.markdown("Enter your details below to get an AI-generated estimation.")
 
-    # Handle inputs (tensor, dict, or BatchEncoding)
-    if isinstance(inputs, torch.Tensor):
-        input_ids = inputs
-        attention_mask = torch.ones_like(input_ids)
-        if len(input_ids.shape) == 1:
-            input_ids = input_ids.unsqueeze(0)
-            attention_mask = attention_mask.unsqueeze(0)
-    elif isinstance(inputs, (dict, BatchEncoding)):
-        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
+# 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)
 
-    st.write(f"Input IDs shape: {input_ids.shape}")
-    st.write(f"Attention mask shape: {attention_mask.shape}")
+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}`")
 
-    # Move to device
-    input_ids = input_ids.to(device)
-    attention_mask = attention_mask.to(device)
+    messages = [{"role": "user", "content": prompt}]
 
-    # Generate output
-    try:
-        text_streamer = TextStreamer(tokenizer)
-        output = model.generate(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
+    # 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,
-            temperature=0.7,
-            top_p=0.95,
-            do_sample=True,
-            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)}")
-        return None
-
-    # Decode the output
-    try:
-        decoded = tokenizer.decode(output[0], skip_special_tokens=False)
-        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 — manually or via CSV upload.")
-
-# Tabs for input method
-tab1, tab2 = st.tabs(["🧍 Manual Input", "📄 CSV Upload"])
-
-with tab1:
-    st.subheader("Manual Entry")
-    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"):
-        prediction = get_prediction(age, sex, height, weight, wc)
-        if prediction:
-            st.success("Prediction:")
-            st.write(prediction)
-
-with tab2:
-    st.subheader("Batch Upload via CSV")
-    sample_csv = pd.DataFrame({
-        "Age": [16],
-        "Sex": ["female"],
-        "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")
-
-    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 = get_prediction(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")
-
+            use_cache=True
+        )