ui.py

"""
Gradio интерфейс для NER-приложения.
"""

import gradio as gr

from config import MAX_CHARS, MAX_BATCH_ROWS, HISTORY_SIZE, MODELS, COLOR_MAP
from processing import (
    process_single_text,
    compare_models,
    process_batch,
    get_history_df,
    clear_history
)


def create_interface():
    """Создание Gradio интерфейса."""

    with gr.Blocks(
        theme=gr.themes.Soft(),
        title="Russian NER — Извлечение сущностей",
        css="""
        .entity-legend {
            display: flex;
            gap: 20px;
            margin: 10px 0;
            flex-wrap: wrap;
        }
        .entity-legend-item {
            display: flex;
            align-items: center;
            gap: 5px;
        }
        .entity-color {
            width: 16px;
            height: 16px;
            border-radius: 3px;
        }
        """
    ) as demo:

        # Заголовок
        gr.Markdown("""
        # Russian NER — Извлечение именованных сущностей

        Приложение для автоматического извлечения именованных сущностей из русского текста:
        **персоны (ФИО)**, **организации**, **локации (города, страны)** и **прочее**.

        ---
        """)

        # Легенда цветов
        gr.HTML("""
        <div class="entity-legend">
            <div class="entity-legend-item">
                <div class="entity-color" style="background-color: #3b82f6;"></div>
                <span><b>PER</b> — Персоны</span>
            </div>
            <div class="entity-legend-item">
                <div class="entity-color" style="background-color: #22c55e;"></div>
                <span><b>ORG</b> — Организации</span>
            </div>
            <div class="entity-legend-item">
                <div class="entity-color" style="background-color: #f97316;"></div>
                <span><b>LOC</b> — Локации</span>
            </div>
            <div class="entity-legend-item">
                <div class="entity-color" style="background-color: #a855f7;"></div>
                <span><b>MISC</b> — Прочее</span>
            </div>
        </div>
        """)

        with gr.Tabs():
            # ==================== ВКЛАДКА 1: АНАЛИЗ ТЕКСТА ====================
            with gr.Tab("Анализ текста"):
                gr.Markdown("### Введите текст для извлечения сущностей")

                with gr.Row():
                    with gr.Column(scale=2):
                        model_dropdown = gr.Dropdown(
                            choices=list(MODELS.keys()),
                            value=list(MODELS.keys())[0],
                            label="Выберите модель NER",
                            interactive=True
                        )
                    with gr.Column(scale=1):
                        latency_box = gr.Textbox(
                            label="Время обработки",
                            value="—",
                            interactive=False
                        )

                text_input = gr.Textbox(
                    label=f"Текст для анализа (максимум {MAX_CHARS} символов)",
                    placeholder="Введите или вставьте текст на русском языке...",
                    lines=5
                )

                process_btn = gr.Button("Обработать", variant="primary", size="lg")

                status_box = gr.Textbox(label="Статус", interactive=False)

                gr.Markdown("### Результат с подсветкой сущностей")
                highlighted_text = gr.HighlightedText(
                    label="Найденные сущности в тексте",
                    combine_adjacent=True,
                    color_map=COLOR_MAP
                )

                gr.Markdown("### Таблица извлечённых сущностей")
                entities_table = gr.Dataframe(
                    headers=["Текст", "Тип", "Описание", "Уверенность"],
                    label="Извлечённые сущности",
                    wrap=True
                )

                # Обработчик
                process_btn.click(
                    fn=process_single_text,
                    inputs=[text_input, model_dropdown],
                    outputs=[status_box, highlighted_text, entities_table, latency_box]
                )

                # Примеры
                gr.Markdown("### Примеры текстов")
                gr.Examples(
                    examples=[
                        ["Владимир Путин встретился с президентом Франции Эммануэлем Макроном в Москве для обсуждения вопросов безопасности."],
                        ["Компания Яндекс открыла новый офис в Санкт-Петербурге рядом с Невским проспектом."],
                        ["Сбербанк и ВТБ объявили о запуске совместного проекта в Казани при поддержке Министерства финансов."],
                        ["Иван Петров работает программистом в компании Mail.ru Group в Москве с 2020 года."],
                        ["Александр Сергеевич Пушкин родился в Москве в 1799 году и стал величайшим русским поэтом."]
                    ],
                    inputs=text_input,
                    label="Нажмите на пример для автозаполнения"
                )

            # ==================== ВКЛАДКА 2: СРАВНЕНИЕ МОДЕЛЕЙ ====================
            with gr.Tab("Сравнение моделей"):
                gr.Markdown("""
                ### Сравнение результатов двух моделей
                Введите текст, чтобы увидеть, как разные модели распознают сущности.
                """)

                compare_input = gr.Textbox(
                    label=f"Текст для сравнения (максимум {MAX_CHARS} символов)",
                    placeholder="Введите текст для сравнения моделей...",
                    lines=4
                )

                compare_btn = gr.Button("Сравнить модели", variant="primary", size="lg")
                compare_status = gr.Textbox(label="Статус сравнения", interactive=False)

                model_keys = list(MODELS.keys())

                with gr.Row():
                    with gr.Column():
                        gr.Markdown(f"#### {model_keys[0]}")
                        highlight_1 = gr.HighlightedText(
                            label="Результат модели 1",
                            color_map=COLOR_MAP
                        )
                        table_1 = gr.Dataframe(label="Сущности (модель 1)")
                        latency_1 = gr.Textbox(label="Время", interactive=False)

                    with gr.Column():
                        gr.Markdown(f"#### {model_keys[1]}")
                        highlight_2 = gr.HighlightedText(
                            label="Результат модели 2",
                            color_map=COLOR_MAP
                        )
                        table_2 = gr.Dataframe(label="Сущности (модель 2)")
                        latency_2 = gr.Textbox(label="Время", interactive=False)

                compare_btn.click(
                    fn=compare_models,
                    inputs=[compare_input],
                    outputs=[compare_status, highlight_1, table_1, latency_1, highlight_2, table_2, latency_2]
                )

            # ==================== ВКЛАДКА 3: ПАКЕТНАЯ ОБРАБОТКА ====================
            with gr.Tab("Пакетная обработка"):
                gr.Markdown(f"""
                ### Массовая обработка текстов из файла

                Загрузите файл **CSV** (с колонкой `text`) или **TXT** (каждая строка — отдельный текст).

                **Ограничения:** максимум {MAX_BATCH_ROWS} строк, {MAX_CHARS} символов на текст.
                """)

                with gr.Row():
                    batch_model = gr.Dropdown(
                        choices=list(MODELS.keys()),
                        value=list(MODELS.keys())[0],
                        label="Модель для обработки"
                    )
                    batch_file = gr.File(
                        label="Загрузите CSV или TXT файл",
                        file_types=[".csv", ".txt"]
                    )

                batch_btn = gr.Button("Обработать файл", variant="primary", size="lg")
                batch_status = gr.Textbox(label="Статус обработки", interactive=False)

                batch_results = gr.Dataframe(
                    label="Результаты обработки",
                    wrap=True
                )

                batch_download = gr.Textbox(
                    label="CSV для скачивания (скопируйте содержимое)",
                    lines=5,
                    visible=True
                )

                batch_btn.click(
                    fn=process_batch,
                    inputs=[batch_file, batch_model],
                    outputs=[batch_status, batch_results, batch_download]
                )

            # ==================== ВКЛАДКА 4: ИСТОРИЯ ====================
            with gr.Tab("История запросов"):
                gr.Markdown(f"""
                ### История последних {HISTORY_SIZE} запросов

                Здесь отображаются ваши недавние запросы с результатами.
                """)

                refresh_btn = gr.Button("Обновить историю", variant="secondary")
                clear_btn = gr.Button("Очистить историю", variant="stop")

                history_status = gr.Textbox(label="Статус", interactive=False, visible=False)
                history_table = gr.Dataframe(
                    label="История запросов",
                    headers=["Время", "Модель", "Текст", "Найдено", "Типы", "Latency"],
                    wrap=True
                )

                refresh_btn.click(
                    fn=get_history_df,
                    outputs=[history_table]
                )

                clear_btn.click(
                    fn=clear_history,
                    outputs=[history_table, history_status]
                )

        # Футер
        gr.Markdown("""
        ---

        **Модели:**
        - [Babelscape/wikineural-multilingual-ner](https://huggingface.co/Babelscape/wikineural-multilingual-ner) — мультиязычная модель NER
        - [Gherman/bert-base-NER-Russian](https://huggingface.co/Gherman/bert-base-NER-Russian) — BERT для русского NER

        **Ограничения:** CPU-режим, максимум 2000 символов на текст.

        **Внимание:** Не вводите реальные персональные данные в демонстрационных целях.
        """)

    return demo