app.py

import gradio as gr
from transformers import pipeline

MODEL_ID = "borisMI/ChGK_NER" 

pipe = pipeline("token-classification", model=MODEL_ID, aggregation_strategy="simple")


def ner_prediction(text):
    entities = pipe(text)
    if not entities:
        return {"text": text, "entities": []}    
    return {"text": text, "entities": [
        {"entity": e["entity_group"], "start": e["start"], "end": e["end"]} for e in entities
    ]}

title = "Демонстрация NER-модели для извлечения сущностей из вопросов 'Что? Где? Когда?'"
description = """
Этот проект демонстрирует кастомную NER-модель, обученную на базе ruBert-large. 
Модель извлекает 5 типов сущностей: AUTHOR, CHARACTER, LOCATION, PERSON, WORK_OF_ART.
Введите текст вопроса или любой другой текст на русском языке, чтобы увидеть модель в действии.
Модель лучше всего работает с текстами, похожими на те, на которых она обучалась.
"""

iface = gr.Interface(
    fn=ner_prediction,
    inputs=gr.Textbox(lines=5, label="Текст для анализа"),
    outputs=gr.HighlightedText(label="Извлеченные сущности"),
    title=title,
    description=description,
    examples=[
        ["В романе Льва Толстого 'Война и мир' одним из главных героев является Пьер Безухов, который много времени проводит в Москве."],
        ["Картина 'Утро в сосновом лесу' была написана Иваном Шишкиным в соавторстве с Константином Савицким."],
        ["Согласно одной из версий, персонаж Шерлок Холмс, созданный Артуром Конан Дойлом, проживал в Лондоне на Бейкер-стрит."]
    ]
)

iface.launch()