Update app.py

app.py

@@ -1,5 +1,4 @@
 ## Deploying on HuggingFace
-
 import streamlit as st
 import pandas as pd
 import torch
@@ -56,35 +55,45 @@ def get_prediction(prompt):
             return_tensors="pt",
             max_length=512,
             truncation=True
-        ).to(device)
+        )
     except Exception as e:
-        st.error(f"Error during tokenization: {str(e)}")
-        return None
-
+        st.warning(f"apply_chat_template failed: {str(e)}. Falling back to manual tokenization.")
+        # Fallback: Manual tokenization
+        inputs = tokenizer(
+            prompt,
+            return_tensors="pt",
+            max_length=512,
+            truncation=True,
+            padding=False
+        )
+    
     # 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
+    
+    # Handle inputs (tensor or dict)
+    if isinstance(inputs, torch.Tensor):
+        # Direct tensor (likely input_ids)
+        input_ids = inputs
+        if len(input_ids.shape) == 1:
+            input_ids = input_ids.unsqueeze(0)  # Add batch dimension: [sequence_length] -> [1, sequence_length]
+        elif len(input_ids.shape) > 2:
+            input_ids = input_ids.squeeze()  # Remove extra dimensions if any
+            if len(input_ids.shape) == 1:
+                input_ids = input_ids.unsqueeze(0)
+    elif isinstance(inputs, dict) and 'input_ids' in inputs:
+        input_ids = inputs['input_ids']
+        if len(input_ids.shape) == 3 and input_ids.shape[0] == 1:
+            input_ids = input_ids.squeeze(0)
+        elif len(input_ids.shape) == 1:
+            input_ids = input_ids.unsqueeze(0)
     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"Input IDs shape: {input_ids.shape}")
 
-    st.write(f"Final input_ids shape: {input_ids.shape}")
+    # Ensure input_ids is on the correct device
+    input_ids = input_ids.to(device)
 
     # Generate output
     try: