prompt model change

app.py

@@ -5,33 +5,25 @@ except RuntimeError:
     asyncio.set_event_loop(asyncio.new_event_loop())
 
 import gradio as gr
-import torch
 from transformers import (
-    AutoFeatureExtractor,
-    AutoModelForImageClassification,
     T5Tokenizer,
-    T5ForConditionalGeneration
+    T5ForConditionalGeneration,
+    pipeline
 )
 
-# ------------------ 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
-
+# ------------------ LOAD NEW CLASSIFIER ------------------
+cls_model_name = "yangy50/garbage-classification"
+classifier = pipeline("image-classification", model=cls_model_name)
 
 def classify_image(image):
-    inputs = feature_extractor(images=image, return_tensors="pt")
-
-    with torch.no_grad():
-        outputs = cls_model(**inputs)
-
-    logits = outputs.logits
-    probs = torch.nn.functional.softmax(logits, dim=-1)[0]
+    preds = classifier(image)
 
-    top3 = probs.topk(3).indices.tolist()
-    return {id2label[i]: float(probs[i]) for i in top3}
+    results = {}
+    for item in preds[:3]:
+        label = item["label"]
+        score = float(item["score"])
+        results[label] = score
+    return results
 
 
 # ------------------ LOAD FLAN-T5 ------------------
@@ -42,12 +34,11 @@ def explain_recycling(class_label):
     prompt = f"""
     You are a waste management expert.
 
-    The waste item is classified as: **{class_label}**
+    The waste item is classified as: **{class_label}**.
     Provide a 2-paragraph explanation including:
 
-    1. **How to dispose of {class_label} correctly**
-    2. **Optional tips for reducing waste or reusing {class_label}**
-
+    1. How to dispose of {class_label} correctly.
+    2. Optional tips for reducing waste or reusing {class_label}.
     """
 
     inputs = tokenizer(prompt, return_tensors="pt").input_ids
@@ -56,14 +47,16 @@ def explain_recycling(class_label):
         inputs,
         max_length=250,
         do_sample=True,
-        top_p = 0.9,)
+        top_p=0.9,
+    )
 
     return tokenizer.decode(outputs[0])
 
+
 # ------------------ PIPELINE ------------------
 def full_pipeline(image):
     predictions = classify_image(image)
-    top_label = max(predictions, key=predictions.get)
+    top_label = max(predictions, key=predictions.get)  # top-1
     explanation = explain_recycling(top_label)
     return predictions, explanation