src/thesis_extractor.py - master - IvanMysin/Topics

src/thesis_extractor.py
416 строк17 KB
ivan
Work on autoreview by theses
08 янв 2026, 20:58
08 янв 2026, 20:583ca3f82
import sqlite3
from langchain_ollama import OllamaLLM
from typing import List, Dict, Any
import logging
import time
from dataclasses import dataclass
from pprint import pprint
# Настройка логирования
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


THESES_EXCTRACTION = './prompts_templates/theses_exctraction.txt'
THESES_FORMATTING = './prompts_templates/theses_formatting.txt'

@dataclass
class Thesis:
    article_id: int
    text: str

import re
from typing import List, Optional

def parse_model_theses(text: str) -> List[str]:
    """
    Извлекает тезисы из вывода языковой модели.
    Учитывает возможные отклонения от ожидаемого формата.

    Args:
        text (str): Текст вывода модели

    Returns:
        List[str]: Список извлеченных тезисов
    """

    # Очистка текста от лишних пробелов
    text = text.strip()

    # Паттерны для поиска тезисов в разных возможных форматах
    patterns = [
        # Формат: ### Тезис [N]:
        r'###\s*Thesis\s*\[\d+\]:?(.*?)(?=###\s*Thesis\s*\[\d+\]:|### Thesis|\Z)',
        # Формат: Thesis [N]:
        r'Thesis\s*\[\d+\]:?(.*?)(?=Thesis\s*\[\d+\]:|Thesis|\Z)',
        # Формат: Thesis N:
        r'Thesis\s*\d+:?(.*?)(?=Thesis\s*\d+:|Thesis|\Z)',
        # Формат с маркерами * Wording:
        r'\*?\s*Wording:\s*(.*?)(?=\*?\s*Context|\*?\s*Rationale|\*?\s*Wording|###|Thesis|\Z)',
        # Просто нумерованный список
        r'\d+\.\s*(.*?)(?=\d+\.\s*|\n\n|\Z)',
        # Формат с дефисами
        r'-\s*(.*?)(?=\n-|\n\d+\.|\n###|\nThesis|\Z)'
    ]

    theses = []

    for pattern in patterns:
        if theses:  # Если уже нашли тезисы, выходим
            break

        matches = re.findall(pattern, text, re.DOTALL | re.IGNORECASE)
        if matches:
            for match in matches:
                thesis_text = match.strip()
                if thesis_text and _is_valid_thesis(thesis_text):
                    theses.append(thesis_text)

    # Если не нашли структурированных тезисов, попробуем извлечь по абзацам
    if not theses:
        paragraphs = [p.strip() for p in text.split('\n\n') if p.strip()]
        for paragraph in paragraphs:
            if (_is_valid_thesis(paragraph) and
                    len(paragraph) > 50 and  # Минимальная длина
                    not _is_instruction(paragraph)):
                theses.append(paragraph)

    return theses

def _is_valid_thesis(text: str) -> bool:
    """
    Проверяет, является ли текст валидным тезисом.
    """
    # Исключаем слишком короткие тексты
    if len(text) < 30:
        return False

    # Исключаем тексты, которые похожи на инструкции или мета-комментарии
    exclusion_patterns = [
        r'^(criteria|instructions|structure|format|task|role)',
        r' document.*(called|titled)',
        r' is the original one',
        r'output',
        r' languages.*model',
        r' is artificial.*Intelligence',
        r' prompt',
        r' header:',
        r'###',
        r'\*+\s*'
    ]

    for pattern in exclusion_patterns:
        if re.search(pattern, text, re.IGNORECASE):
            return False

    return True

def _is_instruction(text: str) -> bool:
    """
    Проверяет, является ли текст инструкцией, а не тезисом.
    """
    instruction_indicators = [
        'follow', 'instruction', 'criterion', 'format', 'structure',
        'required', 'requirement', 'make sure', 'please',
        'use', 'stick to', 'observe'
    ]

    text_lower = text.lower()
    return any(indicator in text_lower for indicator in instruction_indicators)

# Дополнительная функция для более детального разбора структурированных тезисов
def parse_structured_theses(text: str) -> List[dict]:
    """
    Парсит структурированные тезисы с выделением формулировки, контекста и обоснования.
    """
    structured_theses = []

    # Паттерн для структурированного тезиса
    thesis_pattern = r'###\s*Thesis\s*\[(\d+)\]:?(.*?)(?=###\s*Thesis\s*\[\d+\]:|### Thesis|\Z)'

    matches = re.findall(thesis_pattern, text, re.DOTALL | re.IGNORECASE)

    for number, content in matches:
        thesis_data = {
            'number': int(number),
            'formulation': '',
            'context': '',
            'justification': ''
        }

        # Извлекаем формулировку
        formulation_match = re.search(r'\*?\s*Wording:\s*(.*?)(?=\*?\s*Context|\*?\s*Rationale|\Z)', content, re.DOTALL | re.IGNORECASE)
        if formulation_match:
            thesis_data['formulation'] = formulation_match.group(1).strip()

        # Извлекаем контекст
        context_match = re.search(r'\*?\s*Context[^:]*:\s*(.*?)(?=\*?\s*Rationale|\*?\s*significance|\Z)', content, re.DOTALL | re.IGNORECASE)
        if context_match:
            thesis_data['context'] = context_match.group(1).strip()

        # Извлекаем обоснование
        justification_match = re.search(r'\*?\s*Rationale[^:]*:\s*(.*?)(?=\*?\s*Wording|\*?\s*Context|\Z)', content, re.DOTALL | re.IGNORECASE)
        if justification_match:
            thesis_data['justification'] = justification_match.group(1).strip()

        structured_theses.append(thesis_data['formulation'] + '\n' + thesis_data['context'] + '\n' + thesis_data['justification'])

    return structured_theses


class ThesisExtractor:
    def __init__(self, db_path: str, model: OllamaLLM ):
        self.db_path = db_path
        self.model = model
        self.setup_database()

    def setup_database(self):
        """Создает таблицу для тезисов, если она не существует"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS Thesis (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    article_id INTEGER NOT NULL,
                    text TEXT NOT NULL,
                    FOREIGN KEY (article_id) REFERENCES Articles (id),
                    UNIQUE(article_id, text)
                )
            ''')
            conn.commit()

    def get_articles(self) -> List[Dict[str, Any]]:
        """Получает все статьи из базы данных"""
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()
            cursor.execute('''
                SELECT id, doi, title, abstract, full_text, date 
                FROM Articles 
                WHERE full_text IS NOT NULL AND full_text != ''
            ''')
            return [dict(row) for row in cursor.fetchall()]

    def extract_theses_from_article(self, title: str, full_text: str) -> List[str]:
        """Извлекает выводы из статьи с помощью LLM"""

        # Подготовка промпта
        with open(THESES_EXCTRACTION, 'r') as file:
            extraction_prompt_template = file.read()

        with open(THESES_FORMATTING, 'r') as file:
            formatting_prompt_template = file.read()


        extraction_prompt = extraction_prompt_template.format(title=title, full_text=full_text)

        # try:
        if True:
            response = self.model.invoke(extraction_prompt)

            pprint(response)
            # Обрабатываем вызов функции
            theses = self.parse_theses_from_text(response)


            if not theses:

                formatting_prompt = formatting_prompt_template.format(extracted_theses=response)
                for attempt in range(5):

                    response = self.model.invoke(formatting_prompt)
                    theses = self.parse_theses_from_text(response)


                    if theses:
                        break
                    else:
                        print(f'Неудачная попытка извлечения №{attempt}, повторяем...')


            # Очищаем и фильтруем тезисы
            cleaned_theses = self.clean_theses(theses)
            logger.info(f"Извлечено {len(cleaned_theses)} выводов")

            return cleaned_theses

        # except Exception as e:
        #     logger.error(f"Ошибка при извлечении выводов: {e}")
        #     return []

    def parse_theses_from_text(self, text: str) -> List[str]:
        """
        Парсит тезисы из текстового ответа модели, ориентируясь на заданный markdown-формат:

        ### Title of thesis 1
        Context and explanation of thesis 1

        ### Title of thesis 2
        Context and explanation of thesis 2
        """
        theses = []

        # Разделяем текст по заголовкам ###, соответствующим формату
        # Используем регулярное выражение для поиска всех блоков ### <заголовок> + текст до следующего заголовка
        pattern = r'###\s*(.+?)(?=\n###|\Z)'
        matches = re.findall(pattern, text.strip(), re.DOTALL)

        for match in matches:
            # Каждый match содержит заголовок и тело (контекст)
            block = match.strip()
            if not block:
                continue

            # Разделяем первую строку (заголовок) и остальной текст (объяснение)
            lines = block.split('\n', 1)  # Делим только по первому переносу
            title = lines[0].strip()
            context = lines[1].strip() if len(lines) > 1 else ""

            # Формируем итоговый тезис: заголовок + контекст
            full_thesis = f"{title}\n{context}".strip()

            if len(full_thesis) > 50:  # Минимальная длина для валидного тезиса
                theses.append(full_thesis)


        return theses

    def clean_theses(self, theses: List[str]) -> List[str]:
        """Очищает и фильтрует список выводов"""
        cleaned = []
        for thesis in theses:
            if not thesis:
                continue

            # Очистка текста
            thesis = thesis.strip()
            thesis = thesis.replace('"', '').replace("'", "")

            # Удаляем слишком короткие или неинформативные тезисы
            if len(thesis) > 40:
                # and
                # not thesis.lower().startswith('conclusion') and
                # not thesis.lower().startswith('theses') and
                # not thesis.lower().startswith('thesis')):
                cleaned.append(thesis)

        return list(set(cleaned))  # Удаляем дубликаты

    def save_theses(self, article_id: int, theses: List[str]):
        """Сохраняет тезисы в базу данных"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            for thesis in theses:
                try:
                    cursor.execute(
                        'INSERT OR IGNORE INTO Thesis (article_id, text) VALUES (?, ?)',
                        (article_id, thesis)
                    )
                except sqlite3.Error as e:
                    logger.error(f"Ошибка при сохранении тезиса: {e}")
            conn.commit()

    def process_all_articles(self, batch_delay: int = 2):
        """Обрабатывает все статьи в базе данных"""
        articles = self.get_articles()
        logger.info(f"Найдено {len(articles)} статей для обработки")

        for i, article in enumerate(articles, 1):
            logger.info(f"Обработка статьи {i}/{len(articles)}: {article['title'][:50]}...")

            # Проверяем, не обрабатывали ли мы уже эту статью
            if self.is_article_processed(article['id']):
                logger.info(f"Статья {article['id']} уже обработана, пропускаем")
                continue

            # try:
            if True:
                theses = self.extract_theses_from_article(
                    article['title'],
                    article['full_text']
                )

                if theses:
                    self.save_theses(article['id'], theses)
                    logger.info(f"Сохранено {len(theses)} выводов для статьи {article['id']}")
                else:
                    logger.warning(f"Не удалось извлечь выводы для статьи {article['id']}")

                # Задержка между обработкой статей
                if i < len(articles):
                    time.sleep(batch_delay)

            # except Exception as e:
            #     logger.error(f"Ошибка при обработке статьи {article['id']}: {e}")
            #     continue

    def is_article_processed(self, article_id: int) -> bool:
        """Проверяет, была ли статья уже обработана"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute(
                'SELECT COUNT(*) FROM Thesis WHERE article_id = ?',
                (article_id,)
            )
            count = cursor.fetchone()[0]
            return count > 0

    def get_statistics(self) -> Dict[str, Any]:
        """Возвращает статистику по обработке"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()

            cursor.execute('SELECT COUNT(*) FROM Articles')
            total_articles = cursor.fetchone()[0]

            cursor.execute('SELECT COUNT(DISTINCT article_id) FROM Thesis')
            processed_articles = cursor.fetchone()[0]

            cursor.execute('SELECT COUNT(*) FROM Thesis')
            total_theses = cursor.fetchone()[0]

            cursor.execute('''
                SELECT COUNT(*) 
                FROM Thesis t 
                GROUP BY article_id 
                HAVING COUNT(*) >= 3
            ''')
            articles_with_multiple_theses = len(cursor.fetchall())

            return {
                'total_articles': total_articles,
                'processed_articles': processed_articles,
                'total_theses': total_theses,
                'articles_with_multiple_theses': articles_with_multiple_theses
            }

def main():
    # Конфигурация
    DB_PATH = "../data/interneurons.db"  # Укажите путь к вашей базе данных
    MODEL_NAME = 'mistral-nemo'  #"gpt-oss:20b" #"llama3.2"  #   Или другая модель, доступная в Ollama


    MODEL = OllamaLLM(model=MODEL_NAME, temperature=0.7, num_predict=4056)

    # Инициализация экстрактора
    extractor = ThesisExtractor(DB_PATH, MODEL)

    # Получение статистики до обработки
    stats_before = extractor.get_statistics()
    print("Статистика до обработки:")
    print(f"Всего статей: {stats_before['total_articles']}")
    print(f"Обработанных статей: {stats_before['processed_articles']}")
    print(f"Всего тезисов: {stats_before['total_theses']}")

    # Запуск обработки
    print("\nЗапуск обработки статей...")
    extractor.process_all_articles(batch_delay=0)  # 3 секунды между статьями

    # Статистика после обработки
    stats_after = extractor.get_statistics()
    print("\nСтатистика после обработки:")
    print(f"Обработанных статей: {stats_after['processed_articles']}")
    print(f"Всего тезисов: {stats_after['total_theses']}")
    print(f"Статей с 3+ тезисами: {stats_after['articles_with_multiple_theses']}")

if __name__ == "__main__":
    main()