prompt model change

app.py

@@ -9,9 +9,8 @@ import torch
 from transformers import (
     AutoFeatureExtractor,
     AutoModelForImageClassification,
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    pipeline     # <-- added
+    T5Tokenizer,
+    T5ForConditionalGeneration
 )
 
 # ------------------ LOAD CLASSIFIER ------------------
@@ -35,71 +34,24 @@ def classify_image(image):
     return {id2label[i]: float(probs[i]) for i in top3}
 
 
-# ------------------ LOAD TinyLlama ------------------
-tiny_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+# ------------------ LOAD FLAN-T5 ------------------
+tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
+chat_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large")
 
-tokenizer = AutoTokenizer.from_pretrained(tiny_model)
-chat_model = AutoModelForCausalLM.from_pretrained(
-    tiny_model,
-    torch_dtype=torch.float32)
+def explain_recycling(class_label):
+    prompt = f"""
+    Give recycling instructions for item {class_label}.
+    """
 
-# ✅ Create pipeline (this was missing)
-pipe = pipeline(
-    "text-generation",
-    model=chat_model,
-    tokenizer=tokenizer,
-    device_map="auto",
-    max_new_tokens=200
-)
-
-def clean_chat_output(full_text):
-    # Split at the last assistant tag (TinyLlama uses <|assistant|>)
-    if "<|assistant|>" in full_text:
-        full_text = full_text.split("<|assistant|>")[-1]
-
-    # Remove any leftover "Item:" echoes
-    lines = full_text.strip().split("\n")
-    if lines[0].lower().startswith("item:"):
-        lines = lines[1:]
-
-    return "\n".join(lines).strip()
-
-
-def explain_recycling(label):
-    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: <one clear correct sentence>\n"
-            "No extra text, no introductions, no explanations."
-        )
-    }
-
-    user_msg = {
-        "role": "user",
-        "content": f"Item: {label}\nReturn the two bullet points now."
-    }
-
-    messages = [system_msg, user_msg]
-
-    # Chat template
-    prompt = tokenizer.apply_chat_template(
-        messages, 
-        tokenize=False, 
-        add_generation_prompt=True
-    )
+    inputs = tokenizer(prompt, return_tensors="pt").input_ids
 
-    # Use the pipeline (fixed)
-    outputs = pipe(
-        prompt,
-        do_sample=False   # deterministic to avoid repeating instructions
+    outputs = chat_model.generate(
+        inputs,
+        max_length=180,
+        do_sample=False
     )
 
-    raw = outputs[0]["generated_text"]
-    return clean_chat_output(raw)
-
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # ------------------ PIPELINE ------------------
 def full_pipeline(image):
@@ -121,7 +73,7 @@ with gr.Blocks() as demo:
     explain_output = gr.Textbox(
         label="Detailed Recycling & Disposal Advice",
         elem_id="explainbox",
-        lines=4
+        lines=18
     )
 
     analyze_btn = gr.Button("Analyze", variant="primary")