scripts/load_pdf.py - common_agents_notes - IvanMysin/Topics

scripts/load_pdf.py
617 строк22 KB
ivan
Working on load scripts
26 фев 2026, 16:11
26 фев 2026, 16:11123bd6d
#!/usr/bin/env python
"""
Скрипт для загрузки коллекции документов из PDF файлов.

Требования:
- grobid_client_python
- sqlite3

Примечание: Для работы требуется запущенный GROBID сервер.
Запустите GROBID: docker run -t --rm -p 8070:8070 grobid/grobid:0.8.0
"""

import argparse
import os
import sqlite3
from pathlib import Path
from typing import List, Dict, Optional, Tuple
import re
import time
import sys

PROJECT_ROOT = Path(__file__).parent.parent
sys.path.append(str(PROJECT_ROOT))

try:
    from grobid_client.grobid_client import GrobidClient
    import xml.etree.ElementTree as ET
except ImportError:
    print("Ошибка: grobid_client_python не установлен")
    print("Установите его: pip install grobid_client_python")
    exit(1)

from src.database_manager import create_articles_table, document_exists, generate_doi_from_title_abstract, insert_article


GARBAGE_PATTERNS = [
    'jstor',
    'jstor.org',
    'nomes de lugar',
    'revista de letras',
    'stable url',
    'proquest',
    'gale academic',
    'this is a pdf file',
    'click here to view',
    'unesp universidade',
    'massimo cacciari',
    'not-for-profit service',
    'digital archive',
]

MIN_ABSTRACT_LENGTH = 100


def is_abstract_valid(abstract: Optional[str]) -> Tuple[bool, str]:
    """
    Проверяет корректность абстракта.
    
    Args:
        abstract: Абстракт для проверки
        
    Returns:
        Tuple (is_valid, reason) - валиден ли абстракт и причина если нет
    """
    if not abstract:
        return False, "Абстракт отсутствует"
    
    abstract_clean = abstract.strip()
    abstract_lower = abstract_clean.lower()
    
    if len(abstract_clean) < MIN_ABSTRACT_LENGTH:
        return False, f"Абстракт слишком короткий ({len(abstract_clean)} символов)"
    
    for pattern in GARBAGE_PATTERNS:
        if pattern in abstract_lower:
            return False, f"Обнаружен мусорный паттерн: '{pattern}'"
    
    return True, ""


def get_existing_abstract(cursor, article_id: int) -> Optional[str]:
    """
    Получает текущий абстракт статьи из базы.
    
    Args:
        cursor: Курсор базы данных
        article_id: ID статьи
        
    Returns:
        Текущий абстракт или None
    """
    cursor.execute("SELECT abstract FROM Articles WHERE id = ?", (article_id,))
    result = cursor.fetchone()
    if result:
        return result[0]
    return None


def find_pdf_files(directory: Path) -> List[Path]:
    """
    Находит все PDF файлы в директории и поддиректориях.

    Args:
        directory: Путь к директории для поиска PDF файлов

    Returns:
        Список путей к PDF файлам
    """
    pdf_files = []

    print(f"🔍 Поиск PDF файлов в директории: {directory}")

    # Рекурсивно ищем PDF файлы
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.lower().endswith('.pdf'):
                full_path = Path(root) / file
                pdf_files.append(full_path)

    print(f"✅ Найдено {len(pdf_files)} PDF файлов")
    return pdf_files


def extract_doi_from_filename(filename: str) -> Optional[str]:
    """
    Пытается извлечь DOI из имени файла.

    Args:
        filename: Имя файла

    Returns:
        DOI или None
    """
    # Удаляем расширение и лишние символы
    name = Path(filename).stem

    # Ищем DOI в формате xxxx.xxxx/xxxxx.xxxx
    doi_pattern = r'(\d{4}\.\d{4}/\d{4}\.\d{4})'
    match = re.search(doi_pattern, name)

    if match:
        return match.group(1)

    # Ищем DOI в формате 10.xxxx/xxxx
    doi_pattern2 = r'10\.\d{4}/[a-z]+\.\d{4,6}'
    match2 = re.search(doi_pattern2, name.lower())

    if match2:
        return match2.group(0)

    return None


def extract_metadata_from_filename(filename: str) -> Dict[str, str]:
    """
    Извлекает метаданные из имени файла.

    Args:
        filename: Имя файла

    Returns:
        Словарь с метаданными
    """
    metadata = {
        'title': '',
        'date': '',
        'doi': None
    }

    # Попытка извлечь DOI
    doi = extract_doi_from_filename(filename)
    if doi:
        metadata['doi'] = doi
        # Используем DOI как заголовок
        metadata['title'] = f"Paper {doi}"

    # Если DOI не найден, используем имя файла
    if not metadata['title']:
        name = Path(filename).stem
        # Заменяем подчеркивания и дефисы на пробелы
        title = re.sub(r'[_-]', ' ', name)
        # Делаем первую букву заглавной
        title = title.capitalize()
        metadata['title'] = title

    return metadata


def extract_text_and_metadata_from_pdf(pdf_path: Path, grobid_url: str = "http://localhost:8070") -> Tuple[str, str, Dict]:
    """
    Извлекает текст и метаданные из PDF файла используя GROBID.

    Args:
        pdf_path: Путь к PDF файлу
        grobid_url: URL GROBID сервера

    Returns:
        Кортеж (abstract, full_text, metadata)
    """
    # try:
    if True:
        client = GrobidClient(grobid_server=grobid_url)
        
        response = client.process_pdf(
            service="processFulltextDocument",
            pdf_file=str(pdf_path),
            generateIDs=True,
            consolidate_header=True,
            consolidate_citations=False,
            include_raw_citations=False,
            include_raw_affiliations=False,
            tei_coordinates=False,
            segment_sentences=False
        )
        xml_content = response[2]
        # print(xml_content)

        if not response or xml_content is None:
            print(f"⚠️ GROBID вернул пустой ответ для {pdf_path}")
            return '', '', {}



        try:
            root = ET.fromstring(xml_content)
        except ET.ParseError as e:
            print(f"⚠️ GROBID вернул некорректный XML для {pdf_path}: {e}")
            if len(xml_content) < 500:
                print(f"   Ответ GROBID: {xml_content[:500]}")
            return '', '', {}
        
        # Извлекаем метаданные
        metadata = {}
        
        # Заголовок
        title_elem = root.find('.//{http://www.tei-c.org/ns/1.0}titleStmt/{http://www.tei-c.org/ns/1.0}title')
        if title_elem is not None:
            metadata['title'] = ''.join(title_elem.itertext()).strip()
        
        # DOI
        doi_elem = root.find('.//{http://www.tei-c.org/ns/1.0}idno[@type="DOI"]')
        if doi_elem is not None:
            metadata['doi'] = doi_elem.text.strip() if doi_elem.text else ''
        
        # Авторы
        authors = []
        for author in root.findall('.//{http://www.tei-c.org/ns/1.0}author'):
            pers_name = author.find('.//{http://www.tei-c.org/ns/1.0}persName')
            if pers_name is not None:
                author_parts = []
                for name_part in pers_name:
                    if name_part.text:
                        author_parts.append(name_part.text)
                if author_parts:
                    authors.append(' '.join(author_parts))
        metadata['authors'] = ', '.join(authors) if authors else None
        
        # Журнал
        journal_elem = root.find('.//{http://www.tei-c.org/ns/1.0}title[@level="j"]')
        if journal_elem is not None:
            metadata['journal'] = ''.join(journal_elem.itertext()).strip()
        
        # Дата
        date_elem = root.find('.//{http://www.tei-c.org/ns/1.0}publicationStmt/{http://www.tei-c.org/ns/1.0}date')
        if date_elem is not None:
            metadata['date'] = date_elem.get('when', '')
        
        # Том, выпуск, страницы
        bibl_scope = root.findall('.//{http://www.tei-c.org/ns/1.0}biblScope')
        for scope in bibl_scope:
            unit = scope.get('unit', '')
            if unit == 'volume':
                metadata['volume'] = scope.text
            elif unit == 'issue':
                metadata['issue'] = scope.text
            elif unit == 'page':
                from_page = scope.get('from', '')
                to_page = scope.get('to', '')
                if from_page and to_page:
                    metadata['pages'] = f"{from_page}-{to_page}"
                elif from_page:
                    metadata['pages'] = from_page
        
        # Извлекаем абстракт
        abstract = ""
        abstract_elem = root.find('.//{http://www.tei-c.org/ns/1.0}abstract')
        if abstract_elem is not None:
            abstract = ''.join(abstract_elem.itertext()).strip()
        
        # Извлекаем полный текст из body
        full_text_parts = []
        body = root.find('.//{http://www.tei-c.org/ns/1.0}body')
        if body is not None:
            for div in body.findall('.//{http://www.tei-c.org/ns/1.0}div'):
                for p in div.findall('.//{http://www.tei-c.org/ns/1.0}p'):
                    text = ''.join(p.itertext()).strip()
                    if text:
                        full_text_parts.append(text)
        
        full_text = '\n\n'.join(full_text_parts)
        
        # Если абстракт не найден, берем первые 500 слов из полного текста
        if not abstract and full_text:
            words = full_text.split()
            abstract_words = words[:500]
            abstract = ' '.join(abstract_words)
        
        return abstract, full_text, metadata
    
    # except ET.ParseError as e:
    #     print(f"⚠️ Ошибка парсинга XML для {pdf_path}: {e}")
    #     return '', '', {}
    # except Exception as e:
    #     print(f"⚠️ Ошибка при чтении PDF {pdf_path} через GROBID: {e}")
    #     return '', '', {}





def clean_text(text: str) -> str:
    """
    Очищает текст от лишних пробелов и пустых строк.

    Args:
        text: Исходный текст

    Returns:
        Очищенный текст
    """
    lines = text.split('\n')

    cleaned_lines = []
    for line in lines:
        line = line.strip()
        if len(line) < 20:
            continue
        cleaned_lines.append(line)

    cleaned_text = '\n'.join(cleaned_lines)
    return cleaned_text


def find_article_by_metadata(cursor, title: str, doi: str = None) -> Optional[int]:
    """
    Ищет статью в базе по заголовку и/или DOI.

    Args:
        cursor: Курсор базы данных
        title: Заголовок статьи
        doi: DOI статьи (опционально)

    Returns:
        ID статьи или None если не найдена
    """
    # Сначала ищем по DOI (если есть)
    if doi:
        cursor.execute("SELECT id FROM Articles WHERE doi = ?", (doi,))
        result = cursor.fetchone()
        if result:
            return result[0]
    
    # Затем ищем по заголовку
    if title:
        cursor.execute("SELECT id FROM Articles WHERE LOWER(title) = LOWER(?)", (title,))
        result = cursor.fetchone()
        if result:
            return result[0]
    
    return None


def update_article_full_text(cursor, article_id: int, full_text: str, abstract: Optional[str] = None, filepath: Optional[str] = None) -> bool:
    """
    Обновляет поля full_text, abstract и filepath для существующей статьи.

    Args:
        cursor: Курсор базы данных
        article_id: ID статьи
        full_text: Полный текст статьи
        abstract: Новый абстракт (опционально)
        filepath: Путь к PDF файлу (опционально)

    Returns:
        True если обновление успешно
    """
    try:
        if abstract and filepath:
            cursor.execute(
                "UPDATE Articles SET full_text = ?, abstract = ?, filepath = ? WHERE id = ?",
                (full_text, abstract, filepath, article_id)
            )
        elif abstract:
            cursor.execute(
                "UPDATE Articles SET full_text = ?, abstract = ? WHERE id = ?",
                (full_text, abstract, article_id)
            )
        elif filepath:
            cursor.execute(
                "UPDATE Articles SET full_text = ?, filepath = ? WHERE id = ?",
                (full_text, filepath, article_id)
            )
        else:
            cursor.execute(
                "UPDATE Articles SET full_text = ? WHERE id = ?",
                (full_text, article_id)
            )
        return cursor.rowcount > 0
    except Exception as e:
        print(f"❌ Ошибка при обновлении статьи: {e}")
        return False


def save_documents(documents: List[Dict], db_path: str, grobid_url: str = "http://localhost:8070") -> Dict[str, int]:
    """
    Сохраняет документы в базу данных.
    Если статья существует - обновляет full_text.
    Если абстракт некорректен - обновляет и его.
    Если статьи нет - создает новую запись с метаданными.

    Args:
        documents: Список документов для сохранения
        db_path: Путь к файлу базы данных
        grobid_url: URL GROBID сервера

    Returns:
        Статистика по сохраненным, обновленным и пропущенным документам
    """
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    stats = {
        'processed': 0,
        'saved': 0,
        'updated_fulltext': 0,
        'updated_abstract': 0,
        'skipped': 0,
        'errors': 0
    }

    for doc in documents:
        try:
            stats['processed'] += 1

            filename = doc['filename']
            pdf_path = doc['path']

            print(f"\n📄 Обработка: {filename}")

            abstract, full_text, metadata = extract_text_and_metadata_from_pdf(pdf_path, grobid_url)

            if not full_text:
                stats['errors'] += 1
                print(f"❌ Не удалось извлечь текст из: {filename}")
                continue

            abstract = clean_text(abstract)

            title = metadata.get('title', '') or extract_metadata_from_filename(filename)['title']
            doi = metadata.get('doi', '') or extract_metadata_from_filename(filename)['doi']
            date = metadata.get('date', '')
            journal = metadata.get('journal')
            authors = metadata.get('authors')
            volume = metadata.get('volume')
            issue = metadata.get('issue')
            pages = metadata.get('pages')

            existing_id = find_article_by_metadata(cursor, title, doi)

            if existing_id:
                existing_abstract = get_existing_abstract(cursor, existing_id)
                is_valid, reason = is_abstract_valid(existing_abstract)
                new_abstract = None
                
                if not is_valid:
                    print(f"  ⚠️ Текущий абстракт некорректен: {reason}")
                    if abstract:
                        new_abstract = abstract
                        print(f"  📝 Абстракт будет обновлен из PDF")
                    stats['updated_abstract'] += 1
                
                if update_article_full_text(cursor, existing_id, full_text, new_abstract, str(pdf_path)):
                    conn.commit()
                    if new_abstract:
                        print(f"🔄 Обновлен full_text, abstract и filepath для: {title[:50]}... (ID: {existing_id})")
                    else:
                        stats['updated_fulltext'] += 1
                        stats['updated_abstract'] -= 1  # Компенсируем, т.к. абстракт не обновлялся
                        print(f"🔄 Обновлен full_text и filepath для: {title[:50]}... (ID: {existing_id})")
                else:
                    stats['errors'] += 1
                    print(f"❌ Не удалось обновить данные для: {title[:50]}...")
            else:
                article_data = {
                    'doi': doi,
                    'title': title,
                    'abstract': abstract,
                    'date': date,
                    'full_text': full_text,
                    'journal': journal,
                    'authors': authors,
                    'volume': volume,
                    'issue': issue,
                    'pages': pages,
                    'pmid': None,
                    'filepath': str(pdf_path)
                }
                insert_article(cursor, article_data)
                conn.commit()
                stats['saved'] += 1
                print(f"✅ Сохранена новая статья: {title[:50]}...")

        except Exception as e:
            stats['errors'] += 1
            print(f"❌ Ошибка при обработке документа {doc.get('filename', 'unknown')}: {e}")
            import traceback
            traceback.print_exc()

    conn.close()
    return stats


def process_directory(pdf_dir: Path, db_path: Path, grobid_url: str = "http://localhost:8070") -> Dict[str, int]:
    """
    Обрабатывает директорию с PDF файлами.

    Args:
        pdf_dir: Директория с PDF файлами
        db_path: Путь к базе данных
        grobid_url: URL GROBID сервера

    Returns:
        Общая статистика обработки
    """
    print("\n" + "="*60)
    print("📁 ЗАГРУЗКА PDF ДОКУМЕНТОВ")
    print("="*60)
    print(f"🔗 GROBID сервер: {grobid_url}")

    create_articles_table(str(db_path))
    print(f"✅ База данных создана/проверена: {db_path}")

    # Находим все PDF файлы
    pdf_files = find_pdf_files(pdf_dir)

    if not pdf_files:
        print("⚠️ PDF файлы не найдены")
        return {'processed': 0, 'saved': 0, 'skipped': 0, 'errors': 0}

    # Подготавливаем данные для сохранения
    documents = []
    for pdf_path in pdf_files:
        documents.append({
            'path': pdf_path,
            'filename': pdf_path.name
        })

    # Сохраняем документы
    stats = save_documents(documents, str(db_path), grobid_url)

    return stats


def main():
    parser = argparse.ArgumentParser(description='Загрузка PDF документов в базу данных')
    parser.add_argument('--pdf-dir', type=str, required=True,
                        help='Путь к директории с PDF файлами')
    parser.add_argument('--db-name', type=str, required=True,
                        help='Название базы данных')
    parser.add_argument('--grobid-url', type=str, default='http://localhost:8070',
                        help='URL GROBID сервера (по умолчанию: http://localhost:8070)')

    args = parser.parse_args()

    # Преобразуем имя базы данных в путь
    db_path = Path(args.db_name)
    if not db_path.suffix:  # Если нет расширения
        db_path = db_path.with_suffix('.db')

    # Если путь не абсолютный, делаем его относительным к '../data'
    if not db_path.is_absolute():
        db_path = PROJECT_ROOT / 'data' / db_path

    # Проверяем существование директории с PDF файлами
    pdf_dir = Path(args.pdf_dir)
    if not pdf_dir.exists():
        raise FileNotFoundError(f"Директория не найдена: {args.pdf_dir}")

    if not pdf_dir.is_dir():
        raise NotADirectoryError(f"Указанный путь не является директорией: {args.pdf_dir}")

    # Создаем родительскую директорию, если нужно
    if not db_path.parent.exists():
        db_path.parent.mkdir(parents=True, exist_ok=True)

    print(f"\n📂 Директория с PDF: {pdf_dir}")
    print(f"💾 База данных: {db_path}")
    print(f"🔗 GROBID сервер: {args.grobid_url}")

    # Обрабатываем директорию
    stats = process_directory(pdf_dir, db_path, args.grobid_url)

    # Выводим финальную статистику
    print("\n" + "="*60)
    print("📊 ОБЩАЯ СТАТИСТИКА")
    print("="*60)
    print(f"📄 Обработано файлов: {stats['processed']}")
    print(f"✅ Новых статей сохранено: {stats['saved']}")
    print(f"🔄 Обновлен full_text: {stats['updated_fulltext']}")
    print(f"📝 Обновлен abstract (был некорректен): {stats['updated_abstract']}")
    print(f"⏭️ Пропущено: {stats['skipped']}")
    print(f"❌ Ошибок обработки: {stats['errors']}")
    print("="*60)


if __name__ == "__main__":
    main()