Update app.py

app.py

@@ -1,4 +1,5 @@
 ## Deploying on HuggingFace
+
 import streamlit as st
 import pandas as pd
 import torch
@@ -7,6 +8,7 @@ from huggingface_hub import login
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 from peft import PeftModel, PeftConfig
 import io
+from transformers.tokenization_utils_base import BatchEncoding
 
 # Login using Hugging Face token
 login(token=os.getenv("HUGGINGFACEHUB_TOKEN"))
@@ -58,7 +60,7 @@ def get_prediction(age, sex, height_cm, weight_kg, wc_cm):
             return_tensors="pt",
             max_length=512,
             truncation=True,
-            return_dict=True  # Ensure dictionary output
+            return_dict=True
         )
     except Exception as e:
         st.warning(f"apply_chat_template failed: {str(e)}. Falling back to manual tokenization.")
@@ -73,15 +75,16 @@ def get_prediction(age, sex, height_cm, weight_kg, wc_cm):
 
     # Debug: Log inputs structure
     st.write(f"Inputs type: {type(inputs)}")
+    st.write(f"Inputs content: {inputs}")
 
-    # Handle inputs (tensor or dict)
+    # Handle inputs (tensor, dict, or BatchEncoding)
     if isinstance(inputs, torch.Tensor):
-        # Assume tensor is input_ids, create attention_mask
         input_ids = inputs
+        attention_mask = torch.ones_like(input_ids)
         if len(input_ids.shape) == 1:
             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:
+            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: