switching to an image classifier

app.py

@@ -2,17 +2,19 @@
 # gradio demo                                                                                                     
 #                                                                                                                 
 import gradio as gr
-from transformers import pipeline
+from transformers import ViTFeatureExtractor, ViTModel
 
-model = pipeline("text-generation", model="openai-community/gpt2")
+feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
+model = ViTModel.from_pretrained('akahana/vit-base-cats-vs-dogs')
 
-title = "openai-community/gpt2 prompt completion"
-description = "Basic demo from the NLP course"
+title = "Sandbox"
+description = "Place to try various models"
 
-def predict(prompt):
-    completion = model(prompt)[0]["generated_text"]
-    return completion
-
-demo = gr.Interface(fn=predict, inputs="text", outputs="text")
+def classify(image):
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+    return "dog"
+    
+demo = gr.Interface(fn=classify, inputs="image", outputs="text")
 
 demo.launch()