Add chat model

app.py

@@ -1,13 +1,14 @@
 import gradio as gr
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification  #for classifer model
+
 import torch
+from transformers import T5Tokenizer, T5ForConditionalGeneration  #for chat model
 
-# Load model + extractor
+# ------------------ Load classifier model + extractor ------------------
 model_name = "Aalaa/Fine_tuned_Vit_trash_classification"
 
 feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
 model = AutoModelForImageClassification.from_pretrained(model_name)
-
 # Label mapping
 id2label = model.config.id2label
 
@@ -29,12 +30,41 @@ def predict(image):
 
     return result
 
+#------------------ Load flan-t5 chat model------------------
+tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base")
+chat_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-base")
+
+def explain_recycling(class_label):
+    prompt = f"""
+        Waste category: {class_label}
+
+        Explain the followings:
+        1. What recycling type this is (e.g., plastic, organic, glass, paper, metal).
+        2. How to properly dispose of or recycle it.
+
+        Use this format:
+
+        Recycling type: <type>
+        Recommendation: <instruction>
+        """
+
+    inputs = tokenizer(prompt, return_tensors="pt").input_ids
+
+    outputs = chat_model.generate(
+        inputs,
+        max_length=180,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9
+    )
+
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
     outputs=gr.Label(num_top_classes=3),
-    title="AI Waste Classifier (ViT)"
+    title="AI Waste Classifier"
 )
 
 demo.launch()