Update app.py

app.py

@@ -1,47 +1,349 @@
 import time
+import re
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification
 
-TASK = "token-classification"
-MODEL_NAME = "tanaos/tanaos-NER-v1"
-pipe = pipeline(TASK, model=MODEL_NAME)
+# ================== КОНФИГУРАЦИЯ ==================
+# Можно легко добавить новые модели
+MODELS = {
+    "Davlan/xlm-roberta-large-ner-hrl": "Многоязычная (50+ языков, 7 типов сущностей)",
+    "Babelscape/wikineural-multilingual-ner": "Многоязычная (википедия, 4 типа)",
+    "ai-forever/ruBert-base-ner": "Русская (4 типа сущностей)",
+}
 
-MAX_CHARS = 2000
+# Выбранная модель по умолчанию
+DEFAULT_MODEL = "Davlan/xlm-roberta-large-ner-hrl"
 
-def run(text: str):
-    if text is None or not text.strip():
-        return "Ошибка: пустой ввод.", None, None
-    t0 = time.time()
+# Цвета для разных типов сущностей (для красивого отображения)
+ENTITY_COLORS = {
+    "PER": "#FF6B6B",  # персона
+    "ORG": "#4ECDC4",  # организация
+    "LOC": "#FFD166",  # локация
+    "MISC": "#06D6A0", # прочее
+    "PERSON": "#FF6B6B",
+    "ORGANIZATION": "#4ECDC4",
+    "LOCATION": "#FFD166",
+    "DATE": "#118AB2",
+    "TIME": "#073B4C",
+}
+
+MAX_CHARS = 2000  # ограничение длины текста
+
+# ================== ИНИЦИАЛИЗАЦИЯ ==================
+def load_model(model_name):
+    """Загрузка модели и токенизатора"""
     try:
-        result = pipe(text)
-        latency = round((time.time() - t0) * 1000, 1)
-        return "OK", result, f"{latency} ms"
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForTokenClassification.from_pretrained(model_name)
+        nlp_pipeline = pipeline(
+            "ner",
+            model=model,
+            tokenizer=tokenizer,
+            aggregation_strategy="simple",  # группируем токены
+            device=-1  # CPU (для Space)
+        )
+        return nlp_pipeline
     except Exception as e:
-        return f"Ошибка: {type(e).__name__}: {e}", None, 
-None
-
-with gr.Blocks() as demo:
-    gr.Markdown(f"""
-    # NER-приложение (Hugging Face Spaces + Gradio)
-    **Задача:** {TASK}  
-    **Модель:** {MODEL_NAME}
-    """)
+        raise Exception(f"Ошибка загрузки модели: {str(e)}")
 
-    inp = gr.Textbox(label="Введите текст", lines=6, placeholder="Скопируйте сюда текст...")
-    btn = gr.Button("Обработать")
-    status = gr.Textbox(label="Статус")
-    out = gr.JSON(label="Результат модели")
-    latency = gr.Textbox(label="Время ответа")
+# Загружаем модель по умолчанию
+try:
+    pipe = load_model(DEFAULT_MODEL)
+    current_model_name = DEFAULT_MODEL
+except Exception as e:
+    print(f"Warning: {e}")
+    pipe = None
+    current_model_name = None
 
-    btn.click(run, inputs=inp, outputs=[status, out, latency])
+# ================== ОСНОВНАЯ ФУНКЦИЯ ==================
+def extract_entities(text, model_choice):
+    global pipe, current_model_name
+    
+    # Проверка ввода
+    if not text or not text.strip():
+        return "⚠️ Введите текст для анализа", None, None, None, None
+    
+    text = text.strip()
+    
+    # Ограничение длины
+    if len(text) > MAX_CHARS:
+        text = text[:MAX_CHARS]
+    
+    # Загрузка новой модели при необходимости
+    if model_choice != current_model_name:
+        try:
+            pipe = load_model(model_choice)
+            current_model_name = model_choice
+        except Exception as e:
+            return f"❌ Ошибка загрузки модели: {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:
+            formatted_result = "Сущности не обнаружены"
+            html_output = "<p>Сущности не обнаружены</p>"
+        else:
+            # 1. Структурированный список
+            formatted_result = []
+            for entity in entities:
+                entity_info = {
+                    "Текст": entity['word'],
+                    "Тип": entity['entity_group'],
+                    "Уверенность": round(entity['score'], 3),
+                    "Позиция": f"{entity['start']}-{entity['end']}"
+                }
+                formatted_result.append(entity_info)
+            
+            # 2. Подсветка в тексте (HTML)
+            html_parts = []
+            last_end = 0
+            
+            for entity in sorted(entities, key=lambda x: x['start']):
+                # Текст до сущности
+                if entity['start'] > last_end:
+                    html_parts.append(text[last_end:entity['start']])
+                
+                # Сущность с подсветкой
+                color = ENTITY_COLORS.get(entity['entity_group'], "#CCCCCC")
+                html_parts.append(
+                    f'<span style="background-color: {color}; padding: 2px 4px; '
+                    f'border-radius: 3px; margin: 2px;" title="{entity["entity_group"]} '
+                    f'(уверенность: {entity["score"]:.2f})">{text[entity["start"]:entity["end"]]}</span>'
+                )
+                
+                last_end = entity['end']
+            
+            # Остаток текста
+            if last_end < len(text):
+                html_parts.append(text[last_end:])
+            
+            html_output = '<div style="line-height: 1.8; font-size: 16px;">' + ''.join(html_parts) + '</div>'
+            
+            # 3. Статистика
+            stats = {}
+            for entity in entities:
+                etype = entity['entity_group']
+                stats[etype] = stats.get(etype, 0) + 1
+            
+            stats_text = " | ".join([f"{k}: {v}" for k, v in stats.items()])
+        
+        return (
+            "✅ Анализ завершен",
+            formatted_result,
+            html_output,
+            stats_text,
+            f"{latency} мс"
+        )
+        
+    except Exception as e:
+        return f"❌ Ошибка: {str(e)}", None, None, None, None
 
-    gr.Examples(
-        examples=[
-            ["I love this product! It works great."],
-            ["This is the worst experience ever."],
-            ["It's okay, nothing special."]
-        ],
-        inputs=inp
-    )
+# ================== ДОПОЛНИТЕЛЬНЫЕ ФУНКЦИИ ==================
+def anonymize_text(text, entities):
+    """Базовая анонимизация текста (для демонстрации)"""
+    if not entities:
+        return text
+    
+    result = text
+    # Заменяем в обратном порядке, чтобы не сбивались индексы
+    for entity in sorted(entities, key=lambda x: x['start'], reverse=True):
+        if entity['entity_group'] in ['PER', 'PERSON']:
+            replacement = '[ЛИЦО]'
+        elif entity['entity_group'] in ['ORG', 'ORGANIZATION']:
+            replacement = '[ОРГАНИЗАЦИЯ]'
+        elif entity['entity_group'] in ['LOC', 'LOCATION']:
+            replacement = '[МЕСТО]'
+        else:
+            replacement = f'[{entity["entity_group"]}]'
+        
+        result = result[:entity['start']] + replacement + result[entity['end']:]
+    
+    return result
+
+def batch_process(files):
+    """Обработка нескольких файлов"""
+    if not files:
+        return "⚠️ Загрузите файлы", []
+    
+    results = []
+    for file_info in files:
+        try:
+            with open(file_info.name, 'r', encoding='utf-8') as f:
+                text = f.read(MAX_CHARS)  # ограничиваем длину
+                _, entities, _, stats, _ = extract_entities(text, DEFAULT_MODEL)
+                results.append({
+                    "Файл": file_info.name,
+                    "Сущности": len(entities) if entities else 0,
+                    "Статистика": stats
+                })
+        except Exception as e:
+            results.append({
+                "Файл": file_info.name,
+                "Ошибка": str(e)
+            })
+    
+    return "✅ Обработка завершена", results
+
+# ================== ИНТЕРФЕЙС GRADIO ==================
+with gr.Blocks(title="NER — Извлечение сущностей", theme=gr.themes.Soft()) as demo:
+    
+    # Заголовок
+    gr.Markdown("""
+    # 🔍 Извлечение именованных сущностей (NER)
+    **Распознавание имен, организаций, локаций и других сущностей в тексте**
+    
+    *Используются модели из Hugging Face Hub*
+    """)
+    
+    # Основные вкладки
+    with gr.Tabs():
+        # Вкладка 1: Основной анализ
+        with gr.TabItem("📝 Анализ текста"):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    # Выбор модели
+                    model_dropdown = gr.Dropdown(
+                        choices=list(MODELS.keys()),
+                        value=DEFAULT_MODEL,
+                        label="Выберите модель",
+                        info="Разные модели поддерж��вают разные языки и типы сущностей"
+                    )
+                    
+                    # Поле ввода
+                    text_input = gr.Textbox(
+                        label="Введите текст для анализа",
+                        placeholder="Пример: Компания Microsoft, основанная Биллом Гейтсом, находится в Редмонде, штат Вашингтон.",
+                        lines=8,
+                        max_length=MAX_CHARS
+                    )
+                    
+                    # Кнопки
+                    with gr.Row():
+                        analyze_btn = gr.Button("🔎 Анализировать", variant="primary")
+                        clear_btn = gr.Button("🗑️ Очистить")
+                    
+                    # Примеры
+                    gr.Examples(
+                        examples=[
+                            ["Apple решила открыть новый офис в Париже, где Тим Кук встретится с президентом Франции."],
+                            ["Вчера в Москве прошла встреча представителей Google и Яндекса, обсудили ИИ с Илоном Маском."],
+                            ["Президент США Джо Байден выступил в Белом доме на встрече с генеральным директором Tesla."],
+                            ["The Eiffel Tower in Paris is visited by millions of tourists every year from all over the world."]
+                        ],
+                        inputs=text_input,
+                        label="Примеры текстов"
+                    )
+                
+                with gr.Column(scale=3):
+                    # Статус
+                    status = gr.Textbox(label="Статус")
+                    
+                    # Результаты в разных форматах
+                    with gr.Tab("📊 Структурированный"):
+                        result_json = gr.JSON(label="Найденные сущности")
+                    
+                    with gr.Tab("🎨 Визуализация"):
+                        result_html = gr.HTML(label="Текст с подсветкой сущностей")
+                    
+                    # Статистика
+                    stats_output = gr.Textbox(label="Статистика")
+                    
+                    # Производительность
+                    with gr.Row():
+                        latency_output = gr.Textbox(label="Время обработки")
+                    
+                    # Анонимизация
+                    with gr.Accordion("🛡️ Анонимизация текста", open=False):
+                        anonymized_text = 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, result_json, result_html, stats_output, latency_output, anonymized_text]
+            )
+            
+            anonymize_btn.click(
+                fn=lambda text, entities: anonymize_text(text, entities) if entities else "Сначала выполните анализ",
+                inputs=[text_input, result_json],
+                outputs=anonymized_text
+            )
+        
+        # Вкладка 2: Пакетная обработка
+        with gr.TabItem("📁 Пакетная обработка"):
+            gr.Markdown("### Загрузите текстовые файлы (.txt)")
+            
+            file_input = gr.File(
+                label="Выберите файлы",
+                file_count="multiple",
+                file_types=[".txt"]
+            )
+            
+            batch_btn = gr.Button("🚀 Обработать файлы", variant="primary")
+            
+            batch_status = gr.Textbox(label="Статус обработки")
+            batch_results = gr.JSON(label="Результаты")
+            
+            batch_btn.click(
+                fn=batch_process,
+                inputs=file_input,
+                outputs=[batch_status, batch_results]
+            )
+        
+        # Вкладка 3: О приложении
+        with gr.TabItem("ℹ️ О приложении"):
+            gr.Markdown(f"""
+            ## Информация о приложении
+            
+            **Поддерживаемые модели:**
+            {chr(10).join([f"- **{k}**: {v}" for k, v in MODELS.items()])}
+            
+            **Типы сущностей:**
+            - **PER/PERSON**: Персона, человек
+            - **ORG/ORGANIZATION**: Организация, компания
+            - **LOC/LOCATION**: Локация, место
+            - **MISC**: Прочие сущности
+            - **DATE/TIME**: Даты и время
+            
+            **Ограничения:**
+            - Максимальная длина текста: {MAX_CHARS} символов
+            - Работает на CPU (бесплатно, но медленнее)
+            - Модели могут ошибаться, особенно на сложных текстах
+            
+            **Рекомендации по использованию:**
+            1. Для русских текстов используйте `ai-forever/ruBert-base-ner`
+            2. Для многоязычных текстов используйте `Davlan/xlm-roberta-large-ner-hrl`
+            3. Текст должен быть грамотно написан для лучшего распознавания
+            """)
+    
+    # Футер
+    gr.Markdown("---")
+    gr.Markdown("""
+    <div style="text-align: center; color: #666;">
+    <small>NLP приложение для извлечения сущностей | Hugging Face Spaces + Gradio</small>
+    </div>
+    """)
 
-demo.launch()
+# ================== ЗАПУСК ==================
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        share=False
+    )