main.py

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

import fasttext
from fasttext.FastText import load_model, _FastText
import gradio as gr
from gradio import inputs, outputs
from langid.langid import LanguageIdentifier, model

from project_settings import project_path, temp_directory


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--language_identification_md_file",
        default=(project_path / "language_identification.md").as_posix(),
        type=str
    )
    parser.add_argument(
        "--lang_id_examples_file",
        default=(project_path / "lang_id_examples.json").as_posix(),
        type=str
    )
    parser.add_argument(
        "--fasttext_model",
        default=(project_path / "pretrained_models/lid.176.bin").as_posix(),
        type=str
    )
    args = parser.parse_args()
    return args


lang_id_identifier: LanguageIdentifier = None
fasttext_model: _FastText = None


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

    text = str(text).strip()

    if model_name == "langid":
        label, prob = lang_id_identifier.classify(text)
    elif model_name == "fasttext":
        label, prob = fasttext_model.predict(text, k=1)
        label = label[0][9:]
        prob = prob[0]
    else:
        label = "model_name not available."
        prob = -1
    return label, str(round(prob, 4))


def main():
    args = get_args()

    brief_description = """
    Language Identification
    """

    # description
    with open(args.language_identification_md_file, "r", encoding="utf-8") as f:
        description = f.read()

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

    global lang_id_identifier
    global fasttext_model
    lang_id_identifier = LanguageIdentifier.from_modelstring(model, norm_probs=True)
    fasttext_model = fasttext.load_model(args.fasttext_model)

    blocks = gr.Interface(
        click_lang_id_button,
        inputs=[
            inputs.Textbox(lines=3, label="text"),
            inputs.Textbox(label="ground_true"),
            inputs.Dropdown(choices=["langid", "fasttext", "qgyd_lang_id_1"], default="langid", label="model_name"),
        ],
        outputs=[
            outputs.Textbox(label="label"),
            outputs.Textbox(label="prob"),
        ],
        examples=lang_id_examples,
        description=brief_description,
        title="Language Identification",
        server_name="127.0.0.1" if platform.system() == "Windows" else "0.0.0.0",
        server_port=7860
    )

    blocks.launch(
        share=False if platform.system() == "Windows" else False,
    )
    return


if __name__ == "__main__":
    main()