"""Gradio app that showcases Scandinavian zero-shot text classification models.""" from typing import Dict, Tuple import gradio as gr from gradio.components import Dropdown, Textbox, Row, Column, Button, Label, Markdown from types import MethodType from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer from luga import language as detect_language import torch import re import os import torch._dynamo def main(): # Disable tokenizers parallelism os.environ["TOKENIZERS_PARALLELISM"] = "false" # Load the zero-shot classification pipeline global classifier, model, tokenizer model_id = "alexandrainst/scandi-nli-large" model = AutoModelForSequenceClassification.from_pretrained(model_id) tokenizer = AutoTokenizer.from_pretrained(model_id) model = torch.compile(model=model, backend="aot_eager") model.eval() classifier = pipeline("zero-shot-classification", model=model, tokenizer=tokenizer) classifier.get_inference_context = MethodType( lambda self: torch.no_grad, classifier ) # Create dictionary of descriptions for each task, containing the hypothesis template # and candidate labels task_configs: Dict[str, Tuple[str, str, str, str, str, str]] = { "Sentiment classification": ( "Dette eksempel er {}.", "positivt, negativt, neutralt", "Detta exempel är {}.", "positivt, negativt, neutralt", "Dette eksemplet er {}.", "positivt, negativt, nøytralt", ), "News topic classification": ( "Denne nyhedsartikel handler primært om {}.", "krig, politik, uddannelse, sundhed, økonomi, mode, sport", "Den här nyhetsartikeln handlar främst om {}.", "krig, politik, utbildning, hälsa, ekonomi, mode, sport", "Denne nyhetsartikkelen handler først og fremst om {}.", "krig, politikk, utdanning, helse, økonomi, mote, sport", ), "Spam detection": ( "Denne e-mail ligner {}.", "en spam e-mail, ikke en spam e-mail", "Det här e-postmeddelandet ser {}.", "ut som ett skräppostmeddelande, inte ut som ett skräppostmeddelande", "Denne e-posten ser {}.", "ut som en spam-e-post, ikke ut som en spam-e-post", ), "Product feedback detection": ( "Denne kommentar er {}.", "en anmeldelse af et produkt, ikke en anmeldelse af et produkt", "Den här kommentaren är {}.", "en recension av en produkt, inte en recension av en produkt", "Denne kommentaren er {}.", "en anmeldelse av et produkt, ikke en anmeldelse av et produkt", ), "Define your own task!": ( "Dette eksempel er {}.", "", "Detta exempel är {}.", "", "Dette eksemplet er {}.", "", ), } def set_task_setup(task: str) -> Tuple[str, str, str, str, str, str]: return task_configs[task] with gr.Blocks() as demo: # Create title and description Markdown("# Scandinavian Zero-shot Text Classification") Markdown(""" Classify text in Danish, Swedish or Norwegian into categories, without finetuning on any training data! Select one of the tasks from the dropdown menu on the left, and try entering some input text (in Danish, Swedish or Norwegian) in the input text box and press submit, to see the model in action! The labels are generated by putting in each candidate label into the hypothesis template, and then running the classifier on each label separately. Feel free to change the "hypothesis template" and "candidate labels" on the left as you please as well, and try to come up with your own tasks too 😊 _Also, be patient, as this demo is running on a CPU!_ """) with Row(): # Input column with Column(): # Create a dropdown menu for the task dropdown = Dropdown( label="Task", choices=[ "Sentiment classification", "News topic classification", "Spam detection", "Product feedback detection", "Define your own task!", ], value="Sentiment classification", ) with Row(variant="compact"): da_hypothesis_template = Textbox( label="Danish hypothesis template", value="Dette eksempel er {}.", ) da_candidate_labels = Textbox( label="Danish candidate labels (comma separated)", value="positivt, negativt, neutralt", ) with Row(variant="compact"): sv_hypothesis_template = Textbox( label="Swedish hypothesis template", value="Detta exempel är {}.", ) sv_candidate_labels = Textbox( label="Swedish candidate labels (comma separated)", value="positivt, negativt, neutralt", ) with Row(variant="compact"): no_hypothesis_template = Textbox( label="Norwegian hypothesis template", value="Dette eksemplet er {}.", ) no_candidate_labels = Textbox( label="Norwegian candidate labels (comma separated)", value="positivt, negativt, nøytralt", ) # When a new task is chosen, update the description dropdown.change( fn=set_task_setup, inputs=dropdown, outputs=[ da_hypothesis_template, da_candidate_labels, sv_hypothesis_template, sv_candidate_labels, no_hypothesis_template, no_candidate_labels, ], ) # Output column with Column(): # Create a text box for the input text input_textbox = Textbox( label="Input text", value="Jeg er helt vild med fodbolden 😊" ) with Row(): clear_btn = Button(value="Clear") submit_btn = Button(value="Submit", variant="primary") # When the clear button is clicked, clear the input text box clear_btn.click( fn=lambda _: "", inputs=input_textbox, outputs=input_textbox ) with Column(): # Create output text box output_textbox = Label(label="Result") # When the submit button is clicked, run the classifier on the input text # and display the result in the output text box submit_btn.click( fn=classification, inputs=[ input_textbox, da_hypothesis_template, da_candidate_labels, sv_hypothesis_template, sv_candidate_labels, no_hypothesis_template, no_candidate_labels, ], outputs=output_textbox, ) # Run the app demo.launch(width=.5) @torch.compile() def classification( doc: str, da_hypothesis_template: str, da_candidate_labels: str, sv_hypothesis_template: str, sv_candidate_labels: str, no_hypothesis_template: str, no_candidate_labels: str, ) -> Dict[str, float]: """Classify text into categories. Args: doc (str): Text to classify. da_hypothesis_template (str): Template for the hypothesis to be used for Danish classification. da_candidate_labels (str): Comma-separated list of candidate labels for Danish classification. sv_hypothesis_template (str): Template for the hypothesis to be used for Swedish classification. sv_candidate_labels (str): Comma-separated list of candidate labels for Swedish classification. no_hypothesis_template (str): Template for the hypothesis to be used for Norwegian classification. no_candidate_labels (str): Comma-separated list of candidate labels for Norwegian classification. Returns: dict of str to float: The predicted label and the confidence score. """ # Detect the language of the text language = detect_language(doc.replace('\n', ' ')).name # Set the hypothesis template and candidate labels based on the detected language if language == "sv": hypothesis_template = sv_hypothesis_template candidate_labels = re.split(r', *', sv_candidate_labels) elif language == "no": hypothesis_template = no_hypothesis_template candidate_labels = re.split(r', *', no_candidate_labels) else: hypothesis_template = da_hypothesis_template candidate_labels = re.split(r', *', da_candidate_labels) # Run the classifier on the text result = classifier( doc, candidate_labels=candidate_labels, hypothesis_template=hypothesis_template, ) print(result) # Return the predicted label return {lbl: score for lbl, score in zip(result["labels"], result["scores"])} if __name__ == "__main__": main()