Update app.py

app.py

@@ -1,64 +1,69 @@
 # %%writefile deployment_files/app.py
 import streamlit as st
-# REMOVED: import os (we still need it for safety later, but the primary ENV setting goes to Docker)
 import time
+# REMOVED: import os (Not needed as it's in Dockerfile)
 from transformers import T5Tokenizer, TFT5ForConditionalGeneration 
 
-# --- REMOVED: os.environ["HF_HOME"] = "/tmp/hf_cache" ---
-
-# --- Configuration ---
+# --- Configuration (Unchanged) ---
 MODEL_NAME = "google/flan-t5-small" 
 
 # -------------------- Model Logic --------------------
 
-sys_prompt = """
-... (rest of sys_prompt content unchanged) ...
-"""
+# CRITICAL FIX: Simplified and highly directive prompt for the smallest model
+sys_prompt = "Classify the sentiment of the following customer review as either 'positive', 'negative', or 'neutral'. Respond with only one word." 
 
 @st.cache_resource
 def load_llm():
-    """Loads the model and tokenizer without PyTorch."""
-    # Since we removed torch, device is always 'cpu' conceptually
+    # ... (load_llm function remains identical) ...
     device = "CPU (TensorFlow)"
-    
     try:
         with st.spinner(f"Loading tokenizer and model ({MODEL_NAME}) on {device}..."):
             st.info(f"Using device: **{device}**. Starting model download...")
-            
             start_time = time.time()
             tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME)
-            
-            # CRUCIAL CHANGE: Load the TensorFlow version of the model
             model = TFT5ForConditionalGeneration.from_pretrained(MODEL_NAME)
-            
             end_time = time.time()
             st.success(f"Model {MODEL_NAME} loaded successfully in {end_time - start_time:.2f} seconds!")
-            
             return tokenizer, model, device
-            
     except Exception as e:
         st.error(f"FATAL ERROR LOADING MODEL: {e}")
-        st.info("Model load failed. Check the logs for missing dependencies.")
+        st.info("Model load failed.")
         return None, None, None
-# ... (rest of app.py functions and UI logic are identical) ...
 
 def llm_response(tokenizer, model, device, prompt):
     if tokenizer is None or model is None:
         return "Model not initialized due to previous error."
     
-    # Process input (TensorFlow models automatically handle device placement on CPU)
-    # TF models use the 'tf' argument for return tensors
     input_ids = tokenizer(prompt, return_tensors="tf").input_ids
+    # Set max_length=1 to force a single token output if possible, but 2 is safer for labels
+    outputs = model.generate(input_ids, max_length=3, do_sample=False) 
     
-    # Generate response
-    outputs = model.generate(input_ids, max_length=150, do_sample=False)
-    
-    # Decode and clean the output
     return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
 
 def predict_review_sentiment(tokenizer, model, device, review):
-    full_prompt = f"{sys_prompt}\nReview text: '{review}'"
-    return llm_response(tokenizer, model, device, full_prompt)
+    """
+    CLEANED PROMPT FORMATTING.
+    The final prompt sent to the model is simple: 
+    "Classify the sentiment... Respond with only one word. Review: {review text}"
+    """
+    # FIX: Combine the strict system prompt and the review text clearly
+    full_prompt = f"{sys_prompt} Review: '{review}'" 
+    
+    # Run the prediction and convert the output to standard casing
+    response = llm_response(tokenizer, model, device, full_prompt)
+    
+    # Attempt to normalize the model's output to the three categories
+    normalized_response = response.lower().strip()
+    
+    if "positive" in normalized_response:
+        return "Positive"
+    elif "negative" in normalized_response:
+        return "Negative"
+    elif "neutral" in normalized_response:
+        return "Neutral"
+    else:
+        # For non-classification outputs like 'hi', return the raw response
+        return response
 
 
 # -------------------- Streamlit UI --------------------