src/summarization_manager.py - master - IvanMysin/Topics

src/summarization_manager.py
333 строки15 KB
ivan
AI summarization is added
07 окт 2025, 15:47
07 окт 2025, 15:477baa21b
import sqlite3
import pandas as pd
import requests
import time
import logging
from typing import List, Dict, Optional, Tuple
from tqdm.auto import tqdm
import os
import re

from config import DB_CONFIG, SUMMARIZATION_CONFIG

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class SummarizationManager:
    """Менеджер для суммаризации научных статей с использованием Ollama"""

    def __init__(self):
        self.db_path = DB_CONFIG["db_path"]
        self.table_name = DB_CONFIG["table_name"]
        self.text_column = DB_CONFIG["text_column"]
        self.summary_column = SUMMARIZATION_CONFIG["database"]["summary_column"]
        self.ollama_config = SUMMARIZATION_CONFIG["ollama"]
        self.summarization_config = SUMMARIZATION_CONFIG["summarization"]

        # Загружаем шаблоны промптов
        self.prompt_templates = self._load_prompt_templates()

    def _load_prompt_templates(self) -> Dict[str, str]:
        """Загружает шаблоны промптов из файлов"""
        templates_dir = os.path.join(os.path.dirname(__file__), "prompts_templates")
        templates = {}

        try:
            # Основной шаблон суммаризации
            with open(os.path.join(templates_dir, "summarization_prompt.txt"),
                      "r", encoding="utf-8") as f:
                templates["summarization"] = f.read().strip()

            # Детализированный шаблон
            detailed_path = os.path.join(templates_dir, "detailed_summarization_prompt.txt")
            if os.path.exists(detailed_path):
                with open(detailed_path, "r", encoding="utf-8") as f:
                    templates["detailed"] = f.read().strip()
            else:
                templates["detailed"] = templates["summarization"]

        except Exception as e:
            logger.error(f"Ошибка загрузки шаблонов промптов: {e}")
            # Резервные шаблоны
            templates["summarization"] = """
Создай краткое резюме следующей научной статьи на русском языке:

{text}

Резюме должно содержать основную идею, методы, результаты и выводы. Объем: 300-500 слов.

Резюме:
"""
            templates["detailed"] = templates["summarization"]

        return templates

    def check_summary_column(self) -> bool:
        """Проверяет существование колонки для суммаризации, создает если нет"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            # Получаем информацию о колонках таблицы
            cursor.execute(f"PRAGMA table_info({self.table_name})")
            columns = [column[1] for column in cursor.fetchall()]

            if self.summary_column not in columns:
                logger.info(f"Создание колонки {self.summary_column} в таблице {self.table_name}")
                cursor.execute(f"ALTER TABLE {self.table_name} ADD COLUMN {self.summary_column} TEXT")
                conn.commit()
                logger.info("✅ Колонка успешно создана")
            else:
                logger.info(f"✅ Колонка {self.summary_column} уже существует")

            conn.close()
            return True

        except Exception as e:
            logger.error(f"❌ Ошибка при работе с базой данных: {e}")
            return False

    def get_articles_for_summarization(self) -> pd.DataFrame:
        """Возвращает статьи, требующие суммаризации"""
        try:
            conn = sqlite3.connect(self.db_path)

            if self.summarization_config["update_existing"]:
                # Берем все статьи с полным текстом
                query = f"""
                SELECT {DB_CONFIG['id_column']}, {DB_CONFIG['title_column']}, 
                       {self.text_column}, {self.summary_column}
                FROM {self.table_name}
                WHERE {self.text_column} IS NOT NULL 
                AND TRIM({self.text_column}) != ''
                """
            else:
                # Берем только статьи без существующей суммаризации
                query = f"""
                SELECT {DB_CONFIG['id_column']}, {DB_CONFIG['title_column']}, 
                       {self.text_column}, {self.summary_column}
                FROM {self.table_name}
                WHERE {self.text_column} IS NOT NULL 
                AND TRIM({self.text_column}) != ''
                AND ({self.summary_column} IS NULL OR TRIM({self.summary_column}) = '')
                """

            df = pd.read_sql_query(query, conn)
            conn.close()

            logger.info(f"📋 Найдено статей для суммаризации: {len(df)}")
            return df

        except Exception as e:
            logger.error(f"❌ Ошибка при загрузке статей: {e}")
            return pd.DataFrame()

    def _prepare_text_for_summarization(self, text: str) -> str:
        """Подготавливает текст для суммаризации"""
        if pd.isna(text) or text == "":
            return ""

        # Очищаем текст
        text_clean = str(text).strip()

        # Обрезаем если слишком длинный
        max_length = self.ollama_config["max_input_length"]
        if len(text_clean) > max_length:
            logger.warning(f"Текст обрезан с {len(text_clean)} до {max_length} символов")
            text_clean = text_clean[:max_length]

        return text_clean

    def _create_summarization_prompt(self, text: str, template_type: str = "summarization") -> str:
        """Создает промпт для суммаризации на основе шаблона"""
        template = self.prompt_templates.get(template_type, self.prompt_templates["summarization"])

        # Заменяем плейсхолдеры
        prompt = template.replace("{text}", text)

        return prompt

    def summarize_text(self, text: str, title: str = "") -> Tuple[str, bool]:
        """Суммаризирует один текст с помощью Ollama"""
        if not text or text.strip() == "":
            return "Текст для суммаризации отсутствует", False

        prepared_text = 'title: ' + title + '\n\n' + self._prepare_text_for_summarization(text)
        if not prepared_text:
            return "Не удалось подготовить текст для суммаризации", False

        # Выбираем шаблон в зависимости от сложности текста
        template_type = "detailed" if len(prepared_text) > 2000 else "summarization"
        prompt = self._create_summarization_prompt(prepared_text, template_type)

        url = f"{self.ollama_config['base_url']}/api/generate"
        payload = {
            "model": self.ollama_config["model"],
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": 0.3,  # Низкая температура для более детерминированных результатов
                "top_p": 0.9,
                "num_predict": self.summarization_config["max_summary_length"]
            }
        }

        try:
            response = requests.post(url, json=payload, timeout=self.ollama_config["timeout"])

            if response.status_code == 200:
                result = response.json()
                summary = result.get("response", "").strip()

                # Очищаем результат от возможных артефактов
                summary = self._clean_summary(summary)

                return summary, True
            else:
                error_msg = f"Ошибка API Ollama: {response.status_code}"
                logger.error(error_msg)
                return error_msg, False

        except requests.exceptions.Timeout:
            error_msg = "Таймаут при обращении к Ollama"
            logger.error(error_msg)
            return error_msg, False
        except Exception as e:
            error_msg = f"Ошибка при суммаризации: {e}"
            logger.error(error_msg)
            return error_msg, False

    def _clean_summary(self, summary: str) -> str:
        """Очищает сгенерированную суммаризацию от артефактов"""
        if not summary:
            return ""

        # Удаляем повторяющиеся фразы промпта
        cleaned = summary.strip()

        # Удаляем возможные префиксы типа "Резюме:" если они есть в начале
        prefixes = ["Резюме:", "РЕЗЮМЕ:", "Summary:", "SUMMARY:"]
        for prefix in prefixes:
            if cleaned.startswith(prefix):
                cleaned = cleaned[len(prefix):].strip()

        # Удаляем лишние переносы строк
        cleaned = re.sub(r'\n\s*\n', '\n\n', cleaned)

        return cleaned

    def update_article_summary(self, article_id: int, summary: str) -> bool:
        """Обновляет суммаризацию статьи в базе данных"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            query = f"""
            UPDATE {self.table_name} 
            SET {self.summary_column} = ? 
            WHERE {DB_CONFIG['id_column']} = ?
            """

            cursor.execute(query, (summary, article_id))
            conn.commit()
            conn.close()

            return True

        except Exception as e:
            logger.error(f"❌ Ошибка при обновлении статьи {article_id}: {e}")
            return False

    def process_batch_summarization(self, batch_size: int = None) -> Dict:
        """Обрабатывает суммаризацию батчами"""
        if batch_size is None:
            batch_size = self.summarization_config["batch_size"]

        # Проверяем и создаем колонку если нужно
        if not self.check_summary_column():
            return {"success": False, "error": "Не удалось создать колонку для суммаризации"}

        # Загружаем статьи для обработки
        articles_df = self.get_articles_for_summarization()

        if len(articles_df) == 0:
            return {"success": True, "message": "Нет статей для суммаризации", "processed": 0}

        logger.info(f"🚀 Начало суммаризации {len(articles_df)} статей")

        # Статистика
        stats = {
            "total": len(articles_df),
            "processed": 0,
            "successful": 0,
            "failed": 0,
            "errors": []
        }

        # Обрабатываем батчами
        for i in tqdm(range(0, len(articles_df), batch_size), desc="Суммаризация статей"):
            batch = articles_df.iloc[i:i + batch_size]

            for _, article in batch.iterrows():
                article_id = article[DB_CONFIG['id_column']]
                title = article[DB_CONFIG['title_column']]
                text = article[self.text_column]

                logger.info(f"📝 Суммаризация статьи {article_id}: {title[:50]}...")

                # Суммаризируем текст
                summary, success = self.summarize_text(text, title)

                if success:
                    # Сохраняем в базу
                    if self.update_article_summary(article_id, summary):
                        stats["successful"] += 1
                        logger.info(f"✅ Статья {article_id} успешно суммаризирована")
                    else:
                        stats["failed"] += 1
                        stats["errors"].append(f"Не удалось сохранить суммаризацию для статьи {article_id}")
                        logger.error(f"❌ Не удалось сохранить суммаризацию для статьи {article_id}")
                else:
                    stats["failed"] += 1
                    stats["errors"].append(f"Ошибка суммаризации статьи {article_id}: {summary}")
                    logger.error(f"❌ Ошибка суммаризации статьи {article_id}: {summary}")

                stats["processed"] += 1

                # Небольшая пауза между запросами чтобы не перегружать Ollama
                # time.sleep(1)

        logger.info(f"✅ Суммаризация завершена. Успешно: {stats['successful']}, Ошибок: {stats['failed']}")

        return {
            "success": True,
            "stats": stats,
            "message": f"Обработано {stats['processed']} статей, успешно: {stats['successful']}"
        }


class SummarizationValidator:
    """Валидатор для проверки качества суммаризации"""

    @staticmethod
    def validate_summary_length(summary: str, min_length: int = 50, max_length: int = 800) -> bool:
        """Проверяет длину суммаризации"""
        if not summary:
            return False
        return min_length <= len(summary) <= max_length

    @staticmethod
    def validate_summary_content(summary: str) -> bool:
        """Проверяет содержание суммаризации на наличие ключевых элементов"""
        if not summary:
            return False

        # Проверяем наличие ключевых слов, указывающих на содержательность
        key_phrases =  ["research", "method", "result", "conclusion", "analysis", "study", "исследование", "метод", "результат", "вывод", "анализ", "изучение"]
        summary_lower = summary.lower()

        # Считаем сколько ключевых фраз присутствует
        matches = sum(1 for phrase in key_phrases if phrase in summary_lower)

        return matches >= 2  # Минимум 2 ключевых элемента