prompt model change

app.py

@@ -1,23 +1,16 @@
-import asyncio
-try:
-    asyncio.get_event_loop()
-except RuntimeError:
-    asyncio.set_event_loop(asyncio.new_event_loop())
-
 import gradio as gr
 import torch
 from transformers import (
     AutoFeatureExtractor,
     AutoModelForImageClassification,
-    T5Tokenizer,
-    T5ForConditionalGeneration
+    AutoTokenizer,
+    AutoModelForCausalLM
 )
 
 # ------------------ LOAD CLASSIFIER ------------------
 cls_model_name = "Aalaa/Fine_tuned_Vit_trash_classification"
 feature_extractor = AutoFeatureExtractor.from_pretrained(cls_model_name)
 cls_model = AutoModelForImageClassification.from_pretrained(cls_model_name)
-
 id2label = cls_model.config.id2label
 
 
@@ -27,63 +20,69 @@ def classify_image(image):
     with torch.no_grad():
         outputs = cls_model(**inputs)
 
-    logits = outputs.logits
-    probs = torch.nn.functional.softmax(logits, dim=-1)[0]
-
+    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
     top3 = probs.topk(3).indices.tolist()
+
     return {id2label[i]: float(probs[i]) for i in top3}
 
 
-# ------------------ LOAD FLAN-T5 ------------------
-tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base")
-chat_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-base")
+# ------------------ LOAD TinyLlama ------------------
+tiny_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
-def explain_recycling(class_label):
+tokenizer = AutoTokenizer.from_pretrained(tiny_model)
+chat_model = AutoModelForCausalLM.from_pretrained(
+    tiny_model,
+    torch_dtype=torch.float32
+)
+
+
+def explain_recycling(label):
     prompt = f"""
-    Give short recycling instructions for the item {class_label}."""
-    inputs = tokenizer(prompt, return_tensors="pt").input_ids
+You are an expert in waste sorting.
+Return ONLY the following 2 bullet points:
+
+• Recycling type: (one short category)
+• Disposal: (1 clear and correct sentence)
 
+Example:
+Item: Glass bottle
+• Recycling type: Glass recycling
+• Disposal: Rinse and put in the glass-recycling bin.
+
+Now answer for:
+Item: {label}
+"""
+
+    inputs = tokenizer(prompt, return_tensors="pt")
     outputs = chat_model.generate(
-        inputs,
-        max_length=80,
-        do_sample=False,
-        repetition_penalty=2.0,  
+        **inputs,
+        max_new_tokens=80,
+        temperature=0.3,
+        do_sample=True,
+        top_p=0.9
     )
 
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+
 # ------------------ PIPELINE ------------------
 def full_pipeline(image):
-    predictions = classify_image(image)
-    top_label = max(predictions, key=predictions.get)
-    explanation = explain_recycling(top_label)
-    return predictions, explanation
+    preds = classify_image(image)
+    label = max(preds, key=preds.get)
+    explanation = explain_recycling(label)
+    return preds, explanation
 
 
 # ------------------ GRADIO UI ------------------
 with gr.Blocks() as demo:
-    gr.Markdown("<h1 style='text-align:center;'>♻️ AI Waste Classifier + Eco Advisor</h1>")
-
-    with gr.Row():
-        img_input = gr.Image(type="pil", label="Upload waste image")
+    gr.Markdown("## ♻️ Waste Classifier + TinyLlama Advisor")
 
-    cls_output = gr.Label(num_top_classes=3, label="Classifier Prediction")
+    img = gr.Image(type="pil")
+    preds = gr.Label(num_top_classes=3, label="Classifier")
+    advice = gr.Textbox(lines=6, label="Recycling Advice")
 
-    explain_output = gr.Textbox(
-        label="Detailed Recycling & Disposal Advice",
-        elem_id="explainbox",
-        lines=4
-    )
-
-    analyze_btn = gr.Button("Analyze", variant="primary")
+    btn = gr.Button("Analyze")
 
-    analyze_btn.click(
-        full_pipeline,
-        inputs=img_input,
-        outputs=[cls_output, explain_output]
-    )
+    btn.click(full_pipeline, inputs=img, outputs=[preds, advice])
 
-demo.launch(
-    theme=gr.themes.Soft(primary_hue="green"),
-    css="#explainbox {height: 330px; font-size: 15px;}"
-)
+demo.launch()