app.py

import gradio as gr
from keybert import KeyBERT
from sentence_transformers import SentenceTransformer
from langdetect import detect
import re

# Инициализация мультиязычной модели
# (Можно заменить на другую модель, если нужна ещё более быстрая или иная точность)
model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
kw_model = KeyBERT(model)

# Стоп-слова для разных языков
STOP_WORDS = {
    "en": ["the", "and", "for", "of", "a", "in", "on", "at", "to", "is"],
    "es": ["en", "de", "y", "el", "la", "un", "una", "con", "es", "por"],
    "fr": ["le", "la", "un", "une", "et", "de", "en", "avec", "est"],
    "pt": ["em", "de", "e", "o", "a", "com", "é", "por"],
    "de": ["der", "die", "das", "und", "ein", "eine", "mit", "ist", "zu"],
    "ar": ["في", "من", "و", "على", "مع", "إلى", "هو", "عن"],
    # Добавьте другие языки, если необходимо
}

def detect_language(text: str) -> str:
    """
    Определение языка текста с помощью langdetect.
    Если язык определить не удалось, возвращаем 'en' (английский) по умолчанию.
    """
    try:
        return detect(text)
    except:
        return "en"

def clean_text(text: str, lang: str) -> str:
    """
    Удаление из текста стоп-слов на основе определённого языка.
    """
    stop_words = STOP_WORDS.get(lang, [])
    words = text.split()
    cleaned_text = " ".join([word for word in words if word.lower() not in stop_words])
    return cleaned_text

def segment_text(text: str, max_segment_length: int = 2000) -> list:
    """
    Разделение текста на сегменты по предложениям, чтобы KeyBERT
    мог обрабатывать длинные тексты по частям.
    Увеличили max_segment_length до 2000 (ранее было 1000),
    чтобы уменьшить общее число сегментов и улучшить скорость.
    """
    sentences = re.split(r'(?<=[.!?]) +', text)
    segments = []
    current_segment = ""

    for sentence in sentences:
        if len(current_segment) + len(sentence) <= max_segment_length:
            current_segment += " " + sentence
        else:
            segments.append(current_segment.strip())
            current_segment = sentence

    if current_segment:
        segments.append(current_segment.strip())

    return segments

def postprocess_keywords(keywords: list, min_score: float = 0.5) -> list:
    """
    Постобработка ключевых слов:
      - убираем дубли и слишком низкий score,
      - НЕ исключаем однословные фразы (ранее исключали),
      - при желании можно смягчить логику удаления "вложенных" фраз.
    """
    unique_keywords = []
    seen = set()

    # Сортируем по длине фразы (от самых длинных к коротким),
    # чтобы при проверке "вложенности" не выкидывать большие фразы первыми.
    for kw in sorted(keywords, key=lambda x: len(x[0]), reverse=True):
        phrase, score = kw[0], kw[1]
        if phrase not in seen and score >= min_score:
            # Если хотим менее агрессивно исключать частичные совпадения,
            # можно убрать проверку "phrase in u_kw".
            if not any(phrase in u_kw["keyword"] for u_kw in unique_keywords):
                unique_keywords.append({"keyword": phrase, "score": score})
                seen.add(phrase)

    # Раньше здесь отсеивали короткие фразы:
    # filtered_keywords = [kw for kw in unique_keywords if len(kw["keyword"].split()) > 1]
    # Теперь разрешаем и однословные:
    filtered_keywords = unique_keywords

    return filtered_keywords

def evaluate_keywords(keywords: list, text_length: int) -> dict:
    """
    Оценка качества ключевых слов:
      - average_score: средний рейтинг ключевых слов
      - coverage: кол-во ключевых слов на 100 символов
    """
    if not keywords:
        return {"average_score": 0, "coverage": 0}
    avg_score = sum([kw["score"] for kw in keywords]) / len(keywords)
    coverage = len(keywords) / (text_length / 100.0)
    return {"average_score": avg_score, "coverage": coverage}

def calculate_keyword_density(keywords, text):
    """
    Рассчитать плотность ключевых слов (тошноту).
    :param keywords: Список ключевых слов [{"keyword": ..., "score": ...}, ...]
    :param text: Исходный (оригинальный) текст.
    :return: Список, где к каждому ключу добавляются поля count и density.
    """
    # Общее число слов (разделяем по пробелу)
    text_length = len(text.split())
    keyword_density = []

    for kw in keywords:
        keyword = kw["keyword"]
        count = text.lower().count(keyword.lower())
        # Плотность (в процентах): (количество вхождений / общее число слов) * 100
        density = 0.0
        if text_length > 0:
            density = (count / text_length) * 100

        keyword_density.append({
            "keyword": keyword,
            "count": count,
            "density": round(density, 2),  # округляем до 2 знаков
            "score": kw["score"]          # оставляем оригинальный score
        })

    return keyword_density

def extract_keywords_interface(text: str,
                              top_n: int = 20,
                              diversity: float = 0.3,
                              max_ngram: int = 3) -> dict:
    """
    Извлечение ключевых слов через Gradio интерфейс.
    По умолчанию:
      - top_n=20 (было 15), чтобы извлекать больше ключей
      - diversity=0.3 (было 0.4)
      - max_segment_length=2000
      - min_score=0.5 для постфильтра
    """
    if not text or not text.strip():
        return {
            "keywords": [],
            "metrics": {
                "average_score": 0,
                "coverage": 0
            }
        }

    # 1. Определяем язык
    lang = detect_language(text)

    # 2. Очищаем текст от базовых стоп-слов
    cleaned_text = clean_text(text, lang)

    # 3. Сегментируем текст на большие блоки
    segments = segment_text(cleaned_text, max_segment_length=2000)

    # 4. Извлекаем ключи из каждого сегмента
    all_keywords = []
    for segment in segments:
        keywords = kw_model.extract_keywords(
            segment,
            keyphrase_ngram_range=(1, max_ngram),
            top_n=top_n,
            diversity=diversity
        )
        all_keywords.extend(keywords)

    # 5. Постобработка ключевых слов (снижение min_score и прочее)
    processed_keywords = postprocess_keywords(all_keywords, min_score=0.5)

    # 6. Метрики: средний score + coverage
    metrics = evaluate_keywords(processed_keywords, len(text))

    # 7. Рассчитываем «тошноту» (Keyword Density) по оригинальному тексту
    keyword_density = calculate_keyword_density(processed_keywords, text)

    # Формируем итоговый ответ
    return {
        "keywords": [
            {
                "keyword": kw["keyword"],
                "score": kw["score"],
                "count": kw["count"],
                "density": kw["density"]
            }
            for kw in keyword_density
        ],
        "metrics": metrics
    }

# Gradio интерфейс
iface = gr.Interface(
    fn=extract_keywords_interface,
    inputs=[
        gr.Textbox(label="Введите текст", lines=10, placeholder="Введите текст для извлечения ключевых слов"),
        gr.Slider(5, 50, value=20, step=1, label="Количество ключевых слов (top_n)"),
        gr.Slider(0.0, 1.0, value=0.3, step=0.1, label="Разнообразие ключевых слов (diversity)"),
        gr.Slider(1, 3, value=3, step=1, label="Максимальная длина фраз (n-gram)")
    ],
    outputs="json",
    title="KeyBERT Extractor",
    description="Извлечение ключевых слов мирового уровня с поддержкой мультиязычности и расчётом «тошноты»"
)

if __name__ == "__main__":
    iface.launch()