UI change

app.py

@@ -6,6 +6,7 @@ except RuntimeError:
 
 import gradio as gr
 from transformers import (
+    AutoTokenizer, AutoModelForCausalLM,
     T5Tokenizer,
     T5ForConditionalGeneration,
     pipeline
@@ -26,26 +27,43 @@ def classify_image(image):
     return results
 
 
-# ------------------ LOAD FLAN-T5 ------------------
-tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small")
-chat_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-small")
+# ------------------ LOAD CHAT MODEL ------------------
+model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+chat_model = AutoModelForCausalLM.from_pretrained(model_name)
 
 def explain_recycling(class_label):
-    prompt = f"""
-    Provide recycling Instructions for {class_label}:
-    Provide waste reduction tips for {class_label}:
-    """
-
-    inputs = tokenizer(prompt, return_tensors="pt").input_ids
+    system_msg = {
+        "role": "system",
+        "content": (
+            "You are an expert in waste sorting. "
+            "You ALWAYS answer using exactly two bullet points:\n"
+            "• Recycling type: <one short category>\n"
+            "• Disposal: <clear correct sentence>\n"
+            "No extra text, no introductions, no explanations."
+        )
+    }
+    user_msg = {
+        "role": "user",
+        "content": f"Item: {class_label}\nReturn the two bullet points now."
+    }
+    messages = [system_msg, user_msg]
+    
+    prompt = tokenizer.apply_chat_template(
+        messages, 
+        tokenize=False, 
+        add_generation_prompt=True
+    )
 
     outputs = chat_model.generate(
-        inputs,
-        max_length=250,
+        prompt,
         do_sample=True,
-        top_p=0.9,
+        temperature=0.3,
+        top_p=0.9
     )
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    return tokenizer.decode(outputs[0])
+    return outputs[0]["generated_text"]
 
 
 # ------------------ PIPELINE ------------------