nothing

app.py

@@ -7,7 +7,7 @@ from pydantic import BaseModel
 
 import random
 import torch
-from transformers import (
+from transformers import ( T5Tokenizer, T5ForConditionalGeneration,
     AutoTokenizer, AutoModelForCausalLM,
     pipeline
 )
@@ -40,72 +40,43 @@ def classify_image(image):
     return results
 
 
-# =========================
-# LOAD CHAT MODEL
-# =========================
-tiny_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+# ------------------ LOAD CHAT MODEL
+tiny_model = "google/flan-t5-small"
 
-tokenizer = AutoTokenizer.from_pretrained(tiny_model)
-chat_model = AutoModelForCausalLM.from_pretrained(
+tokenizer = T5Tokenizer.from_pretrained(tiny_model)
+chat_model = T5ForConditionalGeneration.from_pretrained(
     tiny_model,
-    dtype=torch.bfloat16,
-    device_map="auto",
-    low_cpu_mem_usage=True
+    device_map="cpu"
 )
 
 pipe = pipeline(
-    "text-generation",
+    "text2text-generation",
     model=chat_model,
     tokenizer=tokenizer,
-    device_map="auto",
     max_new_tokens=80
 )
 
-
-def clean_chat_output(full_text):
-    if "<|assistant|>" in full_text:
-        full_text = full_text.split("<|assistant|>")[-1]
-
-    lines = full_text.strip().split("\n")
-    if lines and lines[0].lower().startswith("item:"):
-        lines = lines[1:]
-
-    return "\n".join(lines).strip()
-
-
 def explain_recycling(class_label):
-    system_msg = {
-        "role": "system",
-        "content": (
-            "You are an expert in waste sorting. "
-            "You ALWAYS answer using exactly two bullet points:\n"
-            "• Recycling type: <Item category>\n"
-            "• Disposal: <clear, detailed 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
+    prompt = (
+        "You are an expert in waste sorting. "
+        "You ALWAYS answer using exactly two bullet points:\n"
+        "• Recycling type: <Item category>\n"
+        "• Disposal: <clear, detailed correct sentence>\n"
+        f"Item: {class_label}\n"
+        "Return the two bullet points now."
     )
 
-    outputs = pipe(
-        prompt,
-        max_new_tokens=80,
-        do_sample=True,
-        top_p=0.9,
-        temperature=0.3
-    )
+    output = pipe(prompt)[0]["generated_text"]
+    return output.strip()
+
+pipe = pipeline(
+    "text-generation",
+    model=chat_model,
+    tokenizer=tokenizer,
+    device_map="auto",
+    max_new_tokens=80
+)
 
-    raw = outputs[0]["generated_text"]
-    return clean_chat_output(raw)
 
 
 # =========================