Fix multi-class

app.py

@@ -1,7 +1,9 @@
 import gradio as gr
 from transformers import BartForSequenceClassification, BartTokenizer
-import torch.nn.functional as F
 from transformers import T5Tokenizer, T5ForConditionalGeneration
+import torch
+
+device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
 
 te_tokenizer = BartTokenizer.from_pretrained('facebook/bart-large-mnli')
 te_model = BartForSequenceClassification.from_pretrained('facebook/bart-large-mnli')
@@ -10,15 +12,15 @@ qa_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-small", de
 
 def predict(context, intent, multi_class):
     input_text = "In one word, what is the opposite of: " + intent + "?"
-    input_ids = qa_tokenizer(input_text, return_tensors="pt").input_ids
+    input_ids = qa_tokenizer(input_text, return_tensors="pt").input_ids.to(device)
     opposite_output = qa_tokenizer.decode(qa_model.generate(input_ids, max_length=2)[0])
     input_text = "In one word, what is the following describing: " + context
-    input_ids = qa_tokenizer(input_text, return_tensors="pt").input_ids
+    input_ids = qa_tokenizer(input_text, return_tensors="pt").input_ids.to(device)
     object_output = qa_tokenizer.decode(qa_model.generate(input_ids, max_length=2)[0])
     batch = ['I think the ' + object_output + ' are long.', 'I think the ' + object_output + ' are ' + opposite_output, 'I think the ' + object_output + ' are the perfect']
     outputs = []
     for i, hypothesis in enumerate(batch):
-        input_ids = te_tokenizer.encode(context, hypothesis, return_tensors='pt')
+        input_ids = te_tokenizer.encode(context, hypothesis, return_tensors='pt').to(device)
         # -> [contradiction, neutral, entailment]
         logits = te_model(input_ids)[0][0]
 
@@ -38,20 +40,19 @@ def predict(context, intent, multi_class):
     pn_tensor[2] = pn_tensor[2] * outputs[2][1]
     pn_tensor[0] = pn_tensor[0] * outputs[2][1]
 
-    pn_tensor = F.normalize(pn_tensor, p=1, dim=0)
     if (multi_class):
-        pn_tensor = F.normalize(pn_tensor, p=1, dim=0)
+        pn_tensor = torch.sigmoid(pn_tensor)
     else:
         pn_tensor = pn_tensor.softmax(dim=0)
     pn_tensor = pn_tensor.tolist()
-    return {"entailment": pn_tensor[0], "neutral": pn_tensor[1], "contradiction": pn_tensor[2]}
+    return {"agree": pn_tensor[0], "neutral": pn_tensor[1], "disagree": pn_tensor[2]}
 
 gradio_app = gr.Interface(
     predict,
-    inputs=[gr.Text("Sentence"), gr.Text("Class"), gr.Checkbox("Allow multiple true classes")],
+    inputs=[gr.Text(label="Sentence"), gr.Text(label="Class"), gr.Checkbox(label="Allow multiple true classes")],
     outputs=[gr.Label(num_top_classes=3)],
     title="Intent Analysis",
+    description="This model predicts whether or not the **class** describes the **object described in the sentence.**"
 )
 
-if __name__ == "__main__":
-    gradio_app.launch()
+gradio_app.launch()