anisaallahdadi/ZeroShotTextClassification

1 2from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline3model = AutoModelForSequenceClassification.from_pretrained('facebook/bart-large-mnli')4tokenizer = AutoTokenizer.from_pretrained('facebook/bart-large-mnli')5 6classifier = pipeline("zero-shot-classification",7                      model=model, tokenizer=tokenizer)8 9import gradio as gr10 11def classify(input_query, input_classes, input_multi_class):12    13    input_candidate_classes = input_classes14 15    res = classifier(input_query, input_candidate_classes, multi_class=input_multi_class)16    res_dict = {res.get('labels')[i]: res.get('scores')[i]17                for i in range(len(res.get('labels')))}18 19    return res_dict20 21demo = gr.Interface(22  fn=classify,23  inputs=[gr.Text(label='Search Query'), gr.Text(label='Candidate Classes'), gr.Checkbox(label='Multi_class')],24  outputs=gr.Label(label='Prediction:')25)26 27demo.launch()