Update app.py

app.py

@@ -1,15 +1,29 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
 
-pipeline = pipeline(task="image-classification", model="akansha2k2/Burger_sandwich_pizza")
+tokenizer = AutoTokenizer.from_pretrained("akhooli/mistral-7B-llm")
+model = AutoModelForSequenceClassification.from_pretrained("akhooli/mistral-7B-llm")
 
-def predict(image):
-    predictions = pipeline(image)
-    return {p["label"]: p["score"] for p in predictions}
+def predict(image_path):
+    # Load and preprocess the image
+    with open(image_path, "rb") as f:
+        image_bytes = f.read()
+
+    # Tokenize and predict
+    inputs = tokenizer(image_bytes, return_tensors="pt", padding=True, truncation=True)
+    outputs = model(**inputs)
+    predicted_class_idx = torch.argmax(outputs.logits)
+    
+    # In this example, we are assuming the labels are ['pizza', 'burger', 'sandwich']
+    labels = ['pizza', 'burger', 'sandwich']
+    predicted_label = labels[predicted_class_idx]
+
+    return predicted_label
 
 gr.Interface(
     predict,
-    inputs=gr.Image(label="Upload junk food (sandwich,pizza,burgur) candidate", type="filepath"),
-    outputs=gr.Label(num_top_classes=2),
-    title="pizza or burgur or sandwich?",
-).launch()
+    inputs=gr.Image(label="Upload junk food (sandwich, pizza, burger) candidate", type="file"),
+    outputs=gr.Label(num_top_classes=3),
+    title="Pizza, Burger, or Sandwich?",
+).launch()