app.py

"""Gradio app that showcases Scandinavian zero-shot text classification models."""

from typing import Dict, Tuple
import gradio as gr
from transformers import pipeline
from luga import language as detect_language
import re


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"])}


def main():

    # Load the zero-shot classification pipeline
    global classifier
    classifier = pipeline(
        "zero-shot-classification", model="alexandrainst/scandi-nli-large"
    )

    # 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
        gr.Markdown("# Scandinavian Zero-shot Text Classification")
        gr.Markdown("""
            Classify text in Danish, Swedish or Norwegian into categories, without
            finetuning on any training data!

            Note that the models will most likely not work as well as a finetuned model
            on your specific data, but they can be used as a starting point for your
            own classification task ✨

            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 gr.Row():

            # Input column
            with gr.Column():

                # Create a dropdown menu for the task
                dropdown = gr.inputs.Dropdown(
                    label="Task",
                    choices=[
                        "Sentiment classification",
                        "News topic classification",
                        "Spam detection",
                        "Product feedback detection",
                        "Define your own task!",
                    ],
                    default="Sentiment classification",
                )

                with gr.Row(variant="compact"):
                    da_hypothesis_template = gr.inputs.Textbox(
                        label="Danish hypothesis template",
                        default="Dette eksempel er {}.",
                    )
                    da_candidate_labels = gr.inputs.Textbox(
                        label="Danish candidate labels (comma separated)",
                        default="positivt, negativt, neutralt",
                    )

                with gr.Row(variant="compact"):
                    sv_hypothesis_template = gr.inputs.Textbox(
                        label="Swedish hypothesis template",
                        default="Detta exempel är {}.",
                    )
                    sv_candidate_labels = gr.inputs.Textbox(
                        label="Swedish candidate labels (comma separated)",
                        default="positivt, negativt, neutralt",
                    )

                with gr.Row(variant="compact"):
                    no_hypothesis_template = gr.inputs.Textbox(
                        label="Norwegian hypothesis template",
                        default="Dette eksemplet er {}.",
                    )
                    no_candidate_labels = gr.inputs.Textbox(
                        label="Norwegian candidate labels (comma separated)",
                        default="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 gr.Column():

                # Create a text box for the input text
                input_textbox = gr.inputs.Textbox(
                    label="Input text", default="Jeg er helt vild med fodbolden 😊"
                )

                with gr.Row():
                    clear_btn = gr.Button(value="Clear", width=0.5)
                    submit_btn = gr.Button(value="Submit", width=0.5, 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 gr.Column():

                # Create output text box
                output_textbox = gr.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()


if __name__ == "__main__":
    main()