Add multiclass and shrink LLM model size

app.py

@@ -5,10 +5,10 @@ from transformers import T5Tokenizer, T5ForConditionalGeneration
 
 te_tokenizer = BartTokenizer.from_pretrained('facebook/bart-large-mnli')
 te_model = BartForSequenceClassification.from_pretrained('facebook/bart-large-mnli')
-qa_tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base")
-qa_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-base", device_map="auto")
+qa_tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small")
+qa_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-small", device_map="auto")
 
-def predict(context, intent):
+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
     opposite_output = qa_tokenizer.decode(qa_model.generate(input_ids, max_length=2)[0])
@@ -39,19 +39,19 @@ def predict(context, intent):
     pn_tensor[0] = pn_tensor[0] * outputs[2][1]
 
     pn_tensor = F.normalize(pn_tensor, p=1, dim=0)
-
-    pn_tensor = pn_tensor.softmax(dim=0)
+    if (multi_class):
+        pn_tensor = F.normalize(pn_tensor, p=1, dim=0)
+    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]}
 
 gradio_app = gr.Interface(
     predict,
-    inputs=["text", "text"],
+    inputs=[gr.Text("Sentence"), gr.Text("Class"), gr.Checkbox("Allow multiple true classes", default=True)],
     outputs=[gr.Label(num_top_classes=3)],
     title="Intent Analysis",
 )
 
-print(predict("The cat is short.", "long"))
-
 if __name__ == "__main__":
     gradio_app.launch()