app.py

from typing import Dict, Union
from gliner import GLiNER
import gradio as gr

# Available models
AVAILABLE_MODELS = {
    "nvidia/gliner-PII": "nvidia/gliner-PII",
    "urchade/gliner_multi_pii-v1": "urchade/gliner_multi_pii-v1",
    # "BaoLocTown/gliner-vn-demo": "BaoLocTown/gliner-vn-demo",
    # "BaoLocTown/gliner-bi-large-v1.0-tuned-v3-50k": "BaoLocTown/gliner-bi-large-v1.0-tuned-v3-50k",
    "GreenNode/gliner-pii-base-v1.0-tuned": "GreenNode/gliner-pii-base-v1.0-tuned",
    # "GreenNode/gliner_multi_pii-v1-tuned": "GreenNode/gliner_multi_pii-v1-tuned",
    # "GreenNode/gliner-bi-large-v1.0-tuned": "GreenNode/gliner-bi-large-v1.0-tuned",
    "GreenNode/gliner_multi_pii-v1-tuned-2048": "GreenNode/gliner_multi_pii-v1-tuned-2048",
    "GreenNode/gliner-bi-large-v1.0-tuned-2048": "GreenNode/gliner-bi-large-v1.0-tuned-2048",
    
}

# Cache for loaded models
model_cache = {}

def get_model(model_name: str) -> GLiNER:
    """Get model from cache or load it if not cached."""
    if model_name not in model_cache:
        print(f"Loading model {model_name}")
        try:
            model_cache[model_name] = GLiNER.from_pretrained(model_name)
        except Exception as e:
            print(f"Error loading model {model_name}: {e}")
            raise e
    return model_cache[model_name]

examples = [
    [
        "Công ty TNHH XYZ tuyển dụng vị trí **Nhân viên Kinh doanh** với yêu cầu tốt nghiệp đại học, kỹ năng giao tiếp tốt và đam mê xây dựng mối quan hệ khách hàng. Quyền lợi hấp dẫn, lương cạnh tranh, thưởng theo doanh số và cơ hội thăng tiến. Liên hệ: **0903 123 456** hoặc **[tuyendung@xyz.com](mailto:tuyendung@xyz.com)**.",
        "person, company, phone, job title, email",
        0.3,
        False,
    ],
    [
        "Số hộ chiếu của ông Phạm Hùng Cường là A1234567 và ông đang cư trú tại địa chỉ 123 Đường Láng, Hà Nội. Ông cũng sở hữu số tài khoản ngân hàng 9876543210 tại Ngân hàng VietinBank.",
        "person, company, phone, job title, location, email, credit card number, date of birth, passport number, bank number",
        0.3,
        False
    ],
    [
        "Mã số bảo hiểm y tế của ông Trần Quang Minh là <b>987654321</b>, và số điện thoại của ông là 0987-654-321. Ông đang điều trị bệnh tại bệnh viện <i>Chợ Rẫy</i>, TP.HCM.",
        "person, company, phone, job title, location, email, credit card number, date of birth, passport number, bank number, student id number, organization, social security number, health insurance id number",
        0.3,
        False
    ],
    [
        "Chị **Hoàng Thị Lan**, mã số sinh viên **20237865**, đã đăng ký tham gia khóa học “Data Science” tại **Đại học FPT**, và đã cung cấp số tài khoản ngân hàng ***1234567890123456*** tại **Vietcombank**. Cô cũng đã thanh toán qua số thẻ tín dụng ***4321-8765-1234-5678*** và cung cấp số hộ chiếu <b>C1234567</b>. Chị Hoàng Thị Lan sinh ngày <i>15/08/1998</i> và có email là ***hoangthilan123@fpt.edu.vn***. Cô cũng đã cung cấp số điện thoại ***0909123456*** để nhận thông báo.",
        "person, company, phone, job title, location, email, credit card number, date of birth, passport number, bank number, student id number, organization, social security number, health insurance id number",
        0.3,
        False
    ]
]


def ner(
    text, labels: str, threshold: float, nested_ner: bool, model_name: str = "BaoLocTown/gliner-vn-demo"
) -> Dict[str, Union[str, int, float]]:
    labels = labels.split(",")
    model = get_model(model_name)
    return {
        "text": text,
        "entities": [
            {
                "entity": entity["label"],
                "word": entity["text"],
                "start": entity["start"],
                "end": entity["end"],
                "score": 0,
            }
            for entity in model.predict_entities(
                text, labels, flat_ner=not nested_ner, threshold=threshold
            )
        ],
    }


with gr.Blocks(title="VN-GLiNER") as demo:
    gr.Markdown(
        """
        # VN-GLiNER (Vietnamese Named Entity Recognition)

        GLiNER is a Named Entity Recognition (NER) model capable of identifying any entity type using a bidirectional transformer encoder (BERT-like). It provides a practical alternative to traditional NER models, which are limited to predefined entities, and Large Language Models (LLMs) that, despite their flexibility, are costly and large for resource-constrained scenarios.

        The model has been trained by fine-tuning urchade/gliner_multi-v2.1 on the urchade/synthetic-pii-ner-mistral-v1 dataset.
        
        ## Links

        * All VN-GLiNER models: https://huggingface.co/collections/BaoLocTown/vn-gliner
        * Paper: https://arxiv.org/abs/2311.08526
        * Repository: https://github.com/urchade/GLiNER
        """
    )
    with gr.Accordion("How to run this model locally", open=False):
        gr.Markdown(
            """
            ## Installation
            To use this model, you must install the GLiNER Python library:
            ```
            !pip install gliner
            ```
         
            ## Usage
            Once you've downloaded the GLiNER library, you can import the GLiNER class. You can then load this model using `GLiNER.from_pretrained` and predict entities with `predict_entities`.
            """
        )
        gr.Code(
            '''
from gliner import GLiNER

model = GLiNER.from_pretrained("urchade/gliner_multi_pii-v1")

text = """
Harilala Rasoanaivo, un homme d'affaires local d'Antananarivo, a enregistré une nouvelle société nommée "Rasoanaivo Enterprises" au Lot II M 92 Antohomadinika. Son numéro est le +261 32 22 345 67, et son adresse électronique est harilala.rasoanaivo@telma.mg. Il a fourni son numéro de sécu 501-02-1234 pour l'enregistrement.
"""

labels = ["work", "booking number", "personally identifiable information", "driver licence", "person", "book", "full address", "company", "actor", "character", "email", "passport number", "Social Security Number", "phone number"]
entities = model.predict_entities(text, labels)

for entity in entities:
    print(entity["text"], "=>", entity["label"])
            ''',
            language="python",
        )
        gr.Code(
            """
Harilala Rasoanaivo => person
Rasoanaivo Enterprises => company
Lot II M 92 Antohomadinika => full address
+261 32 22 345 67 => phone number
harilala.rasoanaivo@telma.mg => email
501-02-1234 => Social Security Number
            """
        )

    model_dropdown = gr.Dropdown(
        choices=list(AVAILABLE_MODELS.keys()),
        value="BaoLocTown/gliner-vn-demo",
        label="Model",
        info="Select the GLiNER model to use",
    )
    input_text = gr.Textbox(
        value=examples[0][0], label="Text input", placeholder="Enter your text here"
    )
    with gr.Row() as row:
        labels = gr.Textbox(
            value=examples[0][1],
            label="Labels",
            placeholder="Enter your labels here (comma separated)",
            scale=2,
        )
        threshold = gr.Slider(
            0,
            1,
            value=0.3,
            step=0.01,
            label="Threshold",
            info="Lower the threshold to increase how many entities get predicted.",
            scale=1,
        )
        nested_ner = gr.Checkbox(
            value=examples[0][2],
            label="Nested NER",
            info="Allow for nested NER?",
            scale=0,
        )
    output = gr.HighlightedText(label="Predicted Entities")
    submit_btn = gr.Button("Submit")
    examples = gr.Examples(
        examples,
        fn=ner,
        inputs=[input_text, labels, threshold, nested_ner, model_dropdown],
        outputs=output,
        cache_examples=True,
    )

    # Submitting
    input_text.submit(
        fn=ner, inputs=[input_text, labels, threshold, nested_ner, model_dropdown], outputs=output
    )
    labels.submit(
        fn=ner, inputs=[input_text, labels, threshold, nested_ner, model_dropdown], outputs=output
    )
    threshold.release(
        fn=ner, inputs=[input_text, labels, threshold, nested_ner, model_dropdown], outputs=output
    )
    submit_btn.click(
        fn=ner, inputs=[input_text, labels, threshold, nested_ner, model_dropdown], outputs=output
    )
    nested_ner.change(
        fn=ner, inputs=[input_text, labels, threshold, nested_ner, model_dropdown], outputs=output
    )
    model_dropdown.change(
        fn=ner, inputs=[input_text, labels, threshold, nested_ner, model_dropdown], outputs=output
    )

demo.queue()
demo.launch(debug=True)