Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
-import time
 from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
+import time
 
 model_name = "ethanrom/a2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -11,39 +11,26 @@ pretrained_tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name)
 pretrained_model = pipeline("zero-shot-classification", model=pretrained_model_name, tokenizer=pretrained_tokenizer)
 candidate_labels = ["negative", "positive", "no impact", "mixed"]
 
-accuracy_scores = {"Fine-tuned": 0.0, "Pretrained": 0.0}
-model_sizes = {"Fine-tuned": 0, "Pretrained": 0}
-inference_times = {"Fine-tuned": 0.0, "Pretrained": 0.0}
 
 def predict_sentiment(text_input, model_selection):
-    global accuracy_scores, model_sizes, inference_times
-    
-    start_time = time.time()
     if model_selection == "Fine-tuned":
         inputs = tokenizer.encode_plus(text_input, return_tensors='pt')
+        start_time = time.time()
         outputs = model(**inputs)
+        end_time = time.time()
         logits = outputs.logits.detach().cpu().numpy()[0]
         predicted_class = int(logits.argmax())
-        accuracy_scores[model_selection] += 1 if candidate_labels[predicted_class] == "positive" else 0
-        model_sizes[model_selection] = model.num_parameters()
+        inference_time = end_time - start_time
+        model_size = model.num_parameters()
+        return candidate_labels[predicted_class], inference_time, model_size
     else:
+        start_time = time.time()
         result = pretrained_model(text_input, candidate_labels)
+        end_time = time.time()
         predicted_class = result["labels"][0]
-        accuracy_scores[model_selection] += 1 if predicted_class == 1 else 0
-        model_sizes[model_selection] = pretrained_model.model.num_parameters()
-    end_time = time.time()
-    inference_times[model_selection] = end_time - start_time
-
-    return candidate_labels[predicted_class]
-
-def accuracy(model_selection):
-    return accuracy_scores[model_selection]/10
-
-def model_size(model_selection):
-    return str(model_sizes[model_selection]//(1024*1024)) + " MB"
-
-def inference_time(model_selection):
-    return str(inference_times[model_selection]*1000) + " ms"
+        inference_time = end_time - start_time
+        model_size = pretrained_tokenizer.model_max_length + pretrained_model.model.num_parameters()
+        return candidate_labels[predicted_class], inference_time, model_size
 
 inputs = [
     gr.inputs.Textbox("Enter text"),
@@ -52,13 +39,8 @@ inputs = [
 
 outputs = [
     gr.outputs.Textbox(label="Predicted Sentiment"),
-    gr.outputs.Label(label="Accuracy:"),
-    gr.outputs.Label(label="Model Size:"),
-    gr.outputs.Label(label="Inference Time:")
+    gr.outputs.Textbox(label="Inference Time (s)"),
+    gr.outputs.Textbox(label="Model Size (params)"),
 ]
 
-gr.Interface(fn=predict_sentiment, inputs=inputs, outputs=outputs, 
-             title="Sentiment Analysis", description="Compare the output of two models",
-             live=True,
-             examples=[["on us lift up the light", "Fine-tuned"], ["max laid his hand upon the old man's arm", "Pretrained"]]
-            ).launch();
+gr.Interface(fn=predict_sentiment, inputs=inputs, outputs=outputs, title="Sentiment Analysis", description="Compare the output, inference time, and model size of two models").launch();