Update app.py

app.py

@@ -6,6 +6,7 @@ import time
 
 zero_shot = pipeline("zero-shot-classification")
 distilbert = ktrain.load_predictor("models/distilbert-base-uncased-finetuned-internet-provider")
+distilbert_v2 = ktrain.load_predictor("models/distilbert-base-uncased-finetuned-internet-provider")
 
 def zero_shot_predict(text):
     labels = ["Slow Connection", "Billing", "Setup", "No Connectivity"]
@@ -17,17 +18,24 @@ def distilbert_predict(text):
     preds = distilbert.predict_proba(text)
     return {label: float(pred) for label, pred in zip(labels, preds)}
 
+def distilbert_v2_predict(text):
+    labels = distilbert_v2.get_classes()
+    preds = distilbert_v2.predict_proba(text)
+    return {label: float(pred) for label, pred in zip(labels, preds)}
+
 def predict(text):
-    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
         zero_shot_future = executor.submit(zero_shot_predict, text)
         distilbert_future = executor.submit(distilbert_predict, text)
-    concurrent.futures.wait([zero_shot_future, distilbert_future])
+        distilbert_v2_future = executor.submit(distilbert_v2_predict, text)
+    concurrent.futures.wait([zero_shot_future, distilbert_future, distilbert_v2_future])
     zero_shot_preds = zero_shot_future.result()
     distilbert_preds = distilbert_future.result()
-    return zero_shot_preds, distilbert_preds
+    distilbert_v2_preds = distilbert_v2_future.result()
+    return zero_shot_preds, distilbert_preds, distilbert_v2_preds
 
 input = gr.inputs.Textbox(label="Customer Sentence")
-outputs = [gr.outputs.Label(num_top_classes=4, label="Zero-Shot-Classification"), gr.outputs.Label(num_top_classes=4, label="DistilBERT")]
+outputs = [gr.outputs.Label(num_top_classes=4, label="Zero-Shot-Classification"), gr.outputs.Label(num_top_classes=4, label="DistilBERT"), gr.outputs.Label(num_top_classes=4, label="DistilBERT v2")]
 title = "Case Classification"
 description = "Comparison of Zero-Shot-Classification and a fine-tuned DistilBERT."
 gr.Interface(predict, input, outputs, live=False, live_update=False, title=title, analytics_enabled=False,