zs

app.py

@@ -1,19 +1,24 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import BartForSequenceClassification, BartTokenizer
 
-# def greet(name):
-#     return "Hello " + name + "!!"
 
-# iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-# iface.launch()
+# model = pipeline("text-generation")
 
-model = pipeline("text-generation")
+# following https://joeddav.github.io/blog/2020/05/29/ZSL.html
+tokenizer_bart = BartTokenizer.from_pretrained('facebook/bart-large-mnli')
+model_bart_sq = BartForSequenceClassification.from_pretrained('facebook/bart-large-mnli')
 
 
-def predict(prompt):
-    completion = model(prompt)[0]["generated_text"]
-    return completion
+# def predict(prompt):
+#     completion = model(prompt)[0]["generated_text"]
+#     return completion
 
-predict("My favorite programming language is")
+def zs(premise,hypothesis):
+    input_ids = tokenizer_bart.encode(premise, hypothesis, return_tensors='pt')
+    logits = model_bart_sq(input_ids)[0]
+    entail_contradiction_logits = logits[:,[0,2]]
+    probs = entail_contradiction_logits.softmax(dim=1)
+    true_prob = probs[:,1].item() * 100
+    return true_prob
 
-gr.Interface(fn=predict, inputs="text", outputs="text").launch()
+gr.Interface(fn=zs, inputs="text", outputs="text").launch()