app.py

import time
import gradio as gr
from transformers import pipeline

# ================== КОНФИГУРАЦИЯ ==================
# ГАРАНТИРОВАННО РАБОЧИЕ МОДЕЛИ НА CPU
MODELS = {
    "dslim/bert-base-NER": "Английская NER (4 типа сущностей)",
    "Davlan/xlm-roberta-large-ner-hrl": "Многоязычная (крупная, медленная)",
    "Babelscape/wikineural-multilingual-ner": "Многоязычная (9 языков)",
}

# Основная модель - легкая и гарантированно рабочая
DEFAULT_MODEL = "dslim/bert-base-NER"

# Цвета для типов сущностей
ENTITY_COLORS = {
    "PER": "#FF6B6B",      # персона
    "ORG": "#4ECDC4",      # организация
    "LOC": "#FFD166",      # локация
    "MISC": "#06D6A0",     # прочее
}

MAX_CHARS = 1500  # уменьшил для стабильности

# ================== ИНИЦИАЛИЗАЦИЯ ==================
# Инициализируем сразу, чтобы проверить
try:
    print(f"Загрузка модели: {DEFAULT_MODEL}")
    pipe = pipeline(
        "ner",
        model=DEFAULT_MODEL,
        aggregation_strategy="simple",  # группируем сущности
        device=-1  # CPU
    )
    current_model_name = DEFAULT_MODEL
    print("Модель успешно загружена")
except Exception as e:
    print(f"Ошибка загрузки основной модели: {e}")
    print("Попытка загрузки резервной модели...")
    try:
        pipe = pipeline(
            "ner",
            model="Babelscape/wikineural-multilingual-ner",
            aggregation_strategy="simple",
            device=-1
        )
        current_model_name = "Babelscape/wikineural-multilingual-ner"
        print("Резервная модель загружена")
    except Exception as e2:
        print(f"Все модели не загрузились: {e2}")
        pipe = None
        current_model_name = None

# ================== ОСНОВНАЯ ФУНКЦИЯ ==================
def extract_entities(text, model_choice=None):
    global pipe, current_model_name
    
    # Если модель не загрузилась при запуске
    if pipe is None:
        return "❌ Модель не загрузилась. Проверьте логи Space.", None, None, None, None
    
    # Проверка ввода
    if not text or not text.strip():
        return "⚠️ Введите текст для анализа", None, None, None, None
    
    text = text.strip()
    
    # Ограничение длины
    if len(text) > MAX_CHARS:
        text = text[:MAX_CHARS]
        length_warning = f" (обрезано до {MAX_CHARS} символов)"
    else:
        length_warning = ""
    
    # Если выбрана новая модель - перезагружаем
    if model_choice and model_choice != current_model_name:
        try:
            new_pipe = pipeline(
                "ner",
                model=model_choice,
                aggregation_strategy="simple",
                device=-1
            )
            pipe = new_pipe
            current_model_name = model_choice
        except Exception as e:
            return f"❌ Не удалось загрузить модель {model_choice}: {str(e)}", None, None, None, None
    
    # Измерение времени
    start_time = time.time()
    
    try:
        # Выполнение NER
        entities = pipe(text)
        latency = round((time.time() - start_time) * 1000, 1)
        
        if not entities:
            result_text = "Сущности не обнаружены"
            html_output = f"<p>Сущности не обнаружены{length_warning}</p>"
            stats = "Нет сущностей"
        else:
            # 1. Структурированный результат
            formatted_result = []
            for entity in entities:
                formatted_result.append({
                    "Текст": entity['word'],
                    "Тип": entity['entity_group'],
                    "Уверенность": round(entity['score'], 3)
                })
            
            # 2. Подсветка в тексте
            html_parts = []
            last_end = 0
            
            # Сортируем по началу
            sorted_entities = sorted(entities, key=lambda x: x['start'])
            
            for entity in sorted_entities:
                # Текст до сущности
                if entity['start'] > last_end:
                    html_parts.append(text[last_end:entity['start']])
                
                # Сущность с подсветкой
                color = ENTITY_COLORS.get(entity['entity_group'], "#CCCCCC")
                entity_text = text[entity['start']:entity['end']]
                html_parts.append(
                    f'<mark style="background-color: {color}; padding: 2px 4px; '
                    f'border-radius: 3px; margin: 1px; border: 1px solid {color}80;" '
                    f'title="{entity["entity_group"]} (уверенность: {entity["score"]:.2f})">'
                    f'{entity_text}<sup>{entity["entity_group"]}</sup></mark>'
                )
                
                last_end = entity['end']
            
            # Остаток текста
            if last_end < len(text):
                html_parts.append(text[last_end:])
            
            html_output = '<div style="line-height: 1.6; padding: 10px; background: #f5f5f5; border-radius: 5px;">' + \
                          ''.join(html_parts) + f'<br><small>{length_warning}</small></div>'
            
            # 3. Статистика
            stats_dict = {}
            for entity in entities:
                etype = entity['entity_group']
                stats_dict[etype] = stats_dict.get(etype, 0) + 1
            
            stats = " | ".join([f"<b>{k}</b>: {v}" for k, v in stats_dict.items()])
        
        return (
            f"✅ Анализ завершен ({len(entities) if entities else 0} сущностей)",
            formatted_result,
            html_output,
            stats,
            f"{latency} мс"
        )
        
    except Exception as e:
        return f"❌ Ошибка при обработке: {str(e)}", None, None, None, None

# ================== ДОПОЛНИТЕЛЬНЫЕ ФУНКЦИИ ==================
def anonymize_text(text, entities_json):
    """Простая анонимизация текста"""
    if not text or not entities_json:
        return "Сначала выполните анализ текста"
    
    result = text
    # Для безопасности заменяем с конца
    if isinstance(entities_json, list):
        for entity in sorted(entities_json, key=lambda x: x.get('Позиция', ''), reverse=True):
            etype = entity.get('Тип', '')
            if etype in ['PER', 'PERSON']:
                replacement = '[ЛИЦО]'
            elif etype in ['ORG', 'ORGANIZATION']:
                replacement = '[ОРГАНИЗАЦИЯ]'
            elif etype in ['LOC', 'LOCATION']:
                replacement = '[МЕСТО]'
            else:
                replacement = f'[{etype}]'
            
            # Находим позицию в тексте
            entity_text = entity.get('Текст', '')
            if entity_text in result:
                result = result.replace(entity_text, replacement, 1)
    
    return result

def create_example(text):
    """Создание примера работы"""
    return extract_entities(text, DEFAULT_MODEL)

# ================== ИНТЕРФЕЙС GRADIO ==================
with gr.Blocks(title="NER App", theme=gr.themes.Soft()) as demo:
    
    # Заголовок
    gr.Markdown("""
    # 🔍 Извлечение именованных сущностей (NER)
    **Распознавание персон, организаций, локаций и других сущностей в тексте**
    """)
    
    # Предупреждение если модель не загрузилась
    if pipe is None:
        gr.Error("⚠️ Модель не загрузилась. Проверьте файл requirements.txt и логи Space.")
    
    with gr.Row():
        with gr.Column(scale=1):
            # Выбор модели
            model_dropdown = gr.Dropdown(
                choices=list(MODELS.keys()),
                value=DEFAULT_MODEL,
                label="Модель для анализа",
                info="Для английского текста используйте bert-base-NER"
            )
            
            # Поле ввода
            text_input = gr.Textbox(
                label="Введите текст",
                placeholder="Пример: Apple announced new iPhone at Steve Jobs Theater in Cupertino.",
                lines=6,
                max_length=MAX_CHARS
            )
            
            with gr.Row():
                analyze_btn = gr.Button("Анализировать", variant="primary")
                clear_btn = gr.Button("Очистить")
            
            # Примеры
            gr.Examples(
                examples=[
                    ["Apple CEO Tim Cook announced new products in California."],
                    ["Microsoft was founded by Bill Gates and Paul Allen in Albuquerque."],
                    ["Elon Musk is the CEO of Tesla and SpaceX, both based in the United States."],
                    ["The president of France will visit Berlin next month for EU summit."]
                ],
                inputs=text_input,
                label="Примеры текстов"
            )
            
            gr.Markdown(f"""
            **Ограничения:**
            - Макс. длина: {MAX_CHARS} символов
            - Только текст (без файлов)
            - CPU-режим (может быть медленно)
            """)
        
        with gr.Column(scale=2):
            # Статус
            status = gr.Textbox(label="Статус", interactive=False)
            
            # Результаты
            with gr.Tab("📊 Результаты"):
                result_json = gr.JSON(label="Найденные сущности")
            
            with gr.Tab("🎨 Визуализация"):
                result_html = gr.HTML(label="Текст с выделением сущностей")
            
            with gr.Tab("📈 Статистика"):
                stats_output = gr.HTML(label="Статистика по типам")
                latency_output = gr.Textbox(label="Время обработки")
            
            with gr.Accordion("🛡️ Анонимизация", open=False):
                anonymized_output = gr.Textbox(
                    label="Анонимизированный текст",
                    lines=4,
                    interactive=False
                )
                anonymize_btn = gr.Button("Анонимизировать текст")
    
    # Обработчики событий
    analyze_btn.click(
        fn=extract_entities,
        inputs=[text_input, model_dropdown],
        outputs=[status, result_json, result_html, stats_output, latency_output]
    )
    
    clear_btn.click(
        fn=lambda: ["", None, None, None, None, ""],
        outputs=[text_input, status, result_json, result_html, stats_output, anonymized_output]
    )
    
    anonymize_btn.click(
        fn=anonymize_text,
        inputs=[text_input, result_json],
        outputs=anonymized_output
    )
    
    # Информация
    gr.Markdown("""
    ---
    ### 🎯 Как пользоваться:
    1. Введите текст в поле слева
    2. Нажмите **"Анализировать"**
    3. Смотрите результаты во вкладках
    
    ### 🏷️ Обозначения цветов:
    - <mark style="background-color: #FF6B6B; padding: 2px 6px; border-radius: 3px;">PER</mark> — Персона
    - <mark style="background-color: #4ECDC4; padding: 2px 6px; border-radius: 3px;">ORG</mark> — Организация
    - <mark style="background-color: #FFD166; padding: 2px 6px; border-radius: 3px;">LOC</mark> — Локация
    - <mark style="background-color: #06D6A0; padding: 2px 6px; border-radius: 3px;">MISC</mark> — Прочее
    
    ### ⚠️ Важно:
    - Модель `dslim/bert-base-NER` работает только с **английским** текстом
    - Для русского текста выберите другую модель в выпадающем списке
    - Большие тексты обрабатываются дольше
    """)

# Запуск
if __name__ == "__main__":
    demo.launch(debug=False)