Update app.py

app.py

@@ -1,12 +1,22 @@
-from transformers import pipeline
+# from transformers import pipeline
+from transformers import BartForSequenceClassification, BartTokenizer
 import gradio as grad
 
-zero_shot_classifier = pipeline("zero-shot-classification")
+# zero_shot_classifier = pipeline("zero-shot-classification")
+bart_tkn = BartTokenizer.from_pretrained('facebook/bart-large-mnli')
+mdl = BartForSequenceClassification.from_pretrained('facebook/bart-large-mnli')
 
-def classify(text, labels):
-    classifier_labels = labels.split(",")
-    #["software", "politics", "love", "movies", "emergency", "advertisment", "sports"]
-    response = zero_shot_classifier(text, classifier_labels)
+# def classify(text, labels):
+def classify(text, label):
+    # classifier_labels = labels.split(",")
+    # #["software", "politics", "love", "movies", "emergency", "advertisment", "sports"]
+    # response = zero_shot_classifier(text, classifier_labels)
+
+    tkn_ids = bart_tkn.encode(text, label, return_tensors = "pt")
+    tkn_lgts = mdl(tkn_ids)[0]
+    entail_contra_tkn_lgts = tkn_lgts[:, [0, 2]]
+    probab = entail_contra_tkn_lgts.softmax(dim = 1)
+    response = probab[:, 1].item() * 100
 
     return response