main.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
https://huggingface.co/spaces/sayakpaul/demo-docker-gradio
"""
import argparse
import json
import platform
from typing import Tuple

import gradio as gr
import langid
from langid.langid import LanguageIdentifier, model
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image

from project_settings import project_path, temp_directory


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--lang_id_examples_file",
        default=(project_path / "lang_id_examples.json").as_posix(),
        type=str
    )
    args = parser.parse_args()
    return args


lang_id_identifier = LanguageIdentifier.from_modelstring(model, norm_probs=True)


def click_lang_id_button(text: str, ground_true: str, model_name: str):
    global lang_id_identifier

    if model_name == "langid":
        label, prob = lang_id_identifier.classify(text)
    else:
        label = "model_name not available."
        prob = 0.0
    return label, round(prob, 4)


def main():
    args = get_args()

    brief_description = """
    ## Language Identification

    langid 识别 97 种语言。
    https://github.com/saffsd/langid.py
    """

    # examples
    with open(args.lang_id_examples_file, "r", encoding="utf-8") as f:
        lang_id_examples = json.load(f)

    # ui
    with gr.Blocks() as blocks:
        gr.Markdown(value=brief_description)

        with gr.Row():
            with gr.Column(scale=5):
                with gr.Tabs():
                    with gr.TabItem("lang_id"):
                        gr.Markdown(value="")

                        with gr.Row():
                            with gr.Column(scale=1):
                                lang_id_text = gr.Textbox(lines=2, max_lines=50, label="text")
                                lang_id_ground_true = gr.Textbox(label="ground_true")

                                lang_id_model_name = gr.Dropdown(choices=["langid"], value="langid", label="model_name")
                                lang_id_button = gr.Button("run", variant="primary")

                            with gr.Column(scale=1):
                                lang_id_label = gr.Textbox(label="label")
                                lang_id_prob = gr.Number(label="prob")

                        gr.Examples(
                            examples=lang_id_examples,
                            inputs=[
                                lang_id_text,
                                lang_id_ground_true,
                                lang_id_model_name,
                            ],
                            outputs=[lang_id_label, lang_id_prob],
                            fn=click_lang_id_button
                        )

                        # click event
                        lang_id_button.click(
                            click_lang_id_button,
                            inputs=[
                                lang_id_text,
                                lang_id_ground_true,
                                lang_id_model_name,
                            ],
                            outputs=[lang_id_label, lang_id_prob],
                        )

    blocks.queue().launch(
        share=False if platform.system() == "Windows" else False,
        server_name="127.0.0.1" if platform.system() == "Windows" else "0.0.0.0",
        server_port=7860
    )
    return


if __name__ == "__main__":
    main()