Update app.py

app.py

@@ -43,13 +43,17 @@ model, tokenizer = load_model()
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def get_prediction(prompt):
-    messages = [{"role": "user", "content": prompt}]
-    inputs = tokenizer.apply_chat_template(
-        messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
-    ).to(device)
+    st.write(f"Received prompt: {prompt}")  # Log the prompt received
+    # Tokenize the input prompt
+    inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
+    st.write(f"Tokenized input: {inputs}")  # Log the tokenized inputs
+    # Generate output from the model
     output = model.generate(inputs, max_new_tokens=150, temperature=0.7, top_p=0.95)
+    st.write(f"Output: {output}")  # Log the raw output from the model
+    # Decode the output to readable text
     decoded = tokenizer.decode(output[0], skip_special_tokens=True)
-    return decoded.split("###")[-1].strip()
+    st.write(f"Decoded output: {decoded}")  # Log the decoded output
+    return decoded.strip()
 
 
 # UI Header
@@ -108,3 +112,5 @@ with tab2:
 
                 csv_output = df.to_csv(index=False).encode("utf-8")
                 st.download_button("📤 Download Predictions", data=csv_output, file_name="predictions.csv")
+
+