Update app.py

app.py

@@ -1,5 +1,6 @@
 ## Deploying on HuggingFace
 
+
 import streamlit as st
 import pandas as pd
 import torch
@@ -57,28 +58,32 @@ def get_prediction(prompt):
         add_generation_prompt=True,  # This is needed for generation
         return_tensors="pt",
     ).to(device)
-    
+
     # Log the tokenized input
     st.write(f"Tokenized input: {inputs}")
 
-    # Initialize TextStreamer for real-time streaming
-    text_streamer = TextStreamer(tokenizer)
+    # Verify the shape of the tokenized input
+    st.write(f"Shape of tokenized input: {inputs['input_ids'].shape}")
 
-    # Generate output using the model with streaming
+    # 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(
-        inputs["input_ids"],  # Use the tokenized input
-        max_new_tokens=250,  # Limit the number of tokens
+        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
-        streamer=text_streamer,  # Use the TextStreamer for output
     )
 
-    st.write(f"Output: {output}")  # Log the raw output from the model
     # Decode the output
     decoded = tokenizer.decode(output[0], skip_special_tokens=True)
-    # Log decoded output
+
+    # Log the decoded output
     st.write(f"Decoded output: {decoded}")
+
     return decoded