Update app.py

app.py

@@ -5,6 +5,8 @@ from transformers import pipeline
 from joblib import load
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch.nn.functional as F
+from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
+import matplotlib.pyplot as plt
 
 # global variables to load models
 lr_model = load("lr_model.joblib")
@@ -35,6 +37,14 @@ def predict_sentiment(text, model):
     elif model == "custom BERT":
         pred = F.softmax(bert_model(**bert_tokenizer(text, return_tensors="pt")).logits[0], dim=0).tolist()
         return {"neg": pred[0], "pos": pred[1]}
+
+def plot():
+    actual = ["pos", "pos", "neg", "neg", "pos"]
+    pred = ["pos", "neg", "pos", "neg", "pos"]
+    cm = confusion_matrix(y_test, predictions)
+    disp = ConfusionMatrixDisplay(confusion_matrix=cm)
+    disp.plot()
+    return plt.gca()
         
 
 demo = gr.Blocks()
@@ -72,5 +82,12 @@ with demo:
                     ["Sad frown", "custom BERT"],
                 ]
             )
+        with gr.TabItem("Multiple Inputs"):
+            gr.Markdown("A more complex demo showing a plot and two outputs")
+            interface = gr.Interface(
+                plot,
+                [],
+                "image"
+            )
 
 demo.launch()