scripts/scihub_loader.py - master - IvanMysin/Topics

scripts/scihub_loader.py
652 строки21 KB
ivan
Working on preport scripts
19 фев 2026, 15:57
19 фев 2026, 15:57391c370
#!/usr/bin/env python
"""
Скрипт для скачивания PDF версий документов с Sci-Hub.

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

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

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

import requests
from bs4 import BeautifulSoup


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

from config import SCIHUB_CONFIG, LLM_CONFIG, DB_CONFIG
from src.database_manager import get_articles_config, validate_doi, update_article_filepath

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)


def extract_year_from_date(date_str: str) -> Optional[int]:
    """
    Извлекает год из строки даты.

    Args:
        date_str: Строка с датой

    Returns:
        Год или None
    """
    if not date_str:
        return None

    date_str = str(date_str).strip()

    year_match = re.search(r'\d{4}', date_str)
    if year_match:
        try:
            return int(year_match.group())
        except (ValueError, IndexError):
            pass

    return None


def get_pdf_filename(title: str, doi: str = None) -> str:
    """
    Генерирует имя файла PDF на основе названия статьи или DOI.

    Args:
        title: Название статьи
        doi: DOI статьи (опционально)

    Returns:
        Имя PDF файла
    """
    if doi and doi.startswith("GEN:"):
        filename = f"{doi}.pdf"
    elif title:
        filename = f"{re.sub(r'[^\w\-_.]', '_', title[:50])}.pdf"
    else:
        filename = "unknown.pdf"

    return filename


def pdf_exists(output_dir: Path, article: Dict[str, any]) -> Optional[Path]:
    """
    Проверяет, существует ли PDF файл для статьи в указанной директории.

    Args:
        output_dir: Директория с PDF файлами
        article: Данные статьи

    Returns:
        Путь к файлу или None
    """
    title = article.get('title')
    doi = article.get('doi')

    pdf_filename = get_pdf_filename(title, doi)
    pdf_path = output_dir / pdf_filename

    if pdf_path.exists():
        return pdf_path

    return None


def get_articles_without_full_text(db_path: str, max_year: int = None) -> List[Dict[str, any]]:
    """
    Получает список статей без полного текста с валидным DOI.

    Args:
        db_path: Путь к базе данных
        max_year: Максимальный год публикации (по умолчанию 2022 для Sci-Hub)

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

    articles_config = get_articles_config()

    where_conditions = [
        f"{articles_config['doi_column']} IS NOT NULL",
        f"{articles_config['doi_column']} != ''",
        f"({articles_config['text_column']} IS NULL OR {articles_config['text_column']} = '')"
    ]

    if max_year:
        where_conditions.append(f"CAST(SUBSTR({articles_config['date_column']}, 1, 4) AS INTEGER) <= {max_year}")

    query = f'''
    SELECT
        {articles_config["id_column"]},
        {articles_config["doi_column"]},
        {articles_config["title_column"]},
        {articles_config["text_column"]},
        {articles_config["date_column"]}
    FROM Articles
    WHERE {' AND '.join(where_conditions)}
    ORDER BY {articles_config["id_column"]}
    '''

    cursor.execute(query)
    rows = cursor.fetchall()
    conn.close()

    articles = []
    for row in rows:
        year = extract_year_from_date(row[4])

        articles.append({
            'id': row[0],
            'doi': row[1],
            'title': row[2],
            'has_text': row[3] is not None and row[3] != '',
            'date': row[4],
            'year': year
        })

    return articles


def fetch_scihub_page(doi: str, max_retries: int = None) -> Optional[str]:
    """
    Загружает страницу Sci-Hub для заданного DOI.

    Args:
        doi: DOI статьи
        max_retries: Максимальное количество попыток

    Returns:
        HTML страницы или None
    """
    max_retries = max_retries or SCIHUB_CONFIG["max_retries"]

    url = f"{SCIHUB_CONFIG['base_url'].rstrip('/')}/{doi}"

    for attempt in range(max_retries):
        try:
            response = requests.get(
                url,
                timeout=SCIHUB_CONFIG["timeout"],
                headers={
                    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
                }
            )

            if response.status_code == 200:
                return response.text
            else:
                print(f"  ⚠️ Status code: {response.status_code}")

        except requests.exceptions.RequestException as e:
            print(f"  ❌ Attempt {attempt + 1}/{max_retries} failed: {e}")

        if attempt < max_retries - 1:
            time.sleep(SCIHUB_CONFIG["retry_delay"])

    return None


def find_download_link(html: str) -> Optional[str]:
    """
    Находит ссылку для скачивания в HTML странице Sci-Hub.

    Args:
        html: HTML страница Sci-Hub

    Returns:
        URL для скачивания или None
    """
    soup = BeautifulSoup(html, 'html.parser')

    download_div = soup.find('div', class_='download')
    if download_div:
        link = download_div.find('a')
        if link and link.get('href'):
            return link.get('href')

    return None


def download_pdf(url: str, save_path: Path) -> bool:
    """
    Скачивает PDF файл по URL.

    Args:
        url: URL для скачивания
        save_path: Путь для сохранения файла

    Returns:
        True если скачивание успешно
    """
    full_url = urljoin(SCIHUB_CONFIG['base_url'], url)

    for attempt in range(SCIHUB_CONFIG["max_retries"]):
        try:
            response = requests.get(
                full_url,
                timeout=SCIHUB_CONFIG["timeout"],
                stream=True
            )

            if response.status_code == 200:
                with open(save_path, 'wb') as f:
                    for chunk in response.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
                return True
            else:
                print(f"  ❌ Status code: {response.status_code}")

        except requests.exceptions.RequestException as e:
            print(f"  ❌ Download attempt {attempt + 1} failed: {e}")

        if attempt < SCIHUB_CONFIG["max_retries"] - 1:
            time.sleep(SCIHUB_CONFIG["retry_delay"])

    return False


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

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

    Returns:
        Кортеж (abstract, full_text)
    """
    try:
        # Создаем клиент GROBID
        client = GrobidClient(grobid_server=grobid_url)
        
        # Обрабатываем PDF через GROBID
        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
        )
        
        if not response or response[0] is None:
            print(f"  ⚠️ GROBID вернул пустой ответ для {pdf_path}")
            return '', ''
        
        # Парсим XML ответ
        xml_content = response[0]
        root = ET.fromstring(xml_content)
        
        # Извлекаем абстракт
        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
    
    except Exception as e:
        print(f"  ❌ Error reading PDF {pdf_path} via 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 clean_text_with_llm(text: str, model: str) -> str:
    """
    Очищает текст с помощью LLM модели.

    Args:
        text: Исходный текст
        model: Название модели Ollama

    Returns:
        Очищенный текст
    """
    try:
        base_url = LLM_CONFIG["ollama"]["base_url"]
        response = requests.post(
            f"{base_url}/api/generate",
            json={
                "model": model,
                "prompt": f"""Clean and format following scientific text. Remove:
- Special characters and artifacts from PDF parsing
- Page numbers and headers
- Broken words from line breaks
- Duplicate sentences

Keep text meaningful and well-formatted:

{text[:10000]}""",
                "stream": False,
                "options": {
                    "temperature": 0.3,
                    "top_p": 0.9,
                }
            },
            timeout=LLM_CONFIG["ollama"]["timeout"]
        )

        if response.status_code == 200:
            result = response.json()
            return result.get('response', text)

    except Exception as e:
        print(f"  ⚠️ LLM cleaning failed: {e}")

    return text


def update_article_text(db_path: str, article_id: int, full_text: str, abstract: str = None) -> bool:
    """
    Обновляет текст статьи в базе данных.

    Args:
        db_path: Путь к базе данных
        article_id: ID статьи
        full_text: Полный текст
        abstract: Абстракт (опционально)

    Returns:
        True если обновление успешно
    """
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    articles_config = get_articles_config()

    try:
        if abstract:
            cursor.execute(f'''
            UPDATE Articles
            SET {articles_config["text_column"]} = ?, {articles_config["abstract_column"]} = ?
            WHERE {articles_config["id_column"]} = ?
            ''', (full_text, abstract, article_id))
        else:
            cursor.execute(f'''
            UPDATE Articles
            SET {articles_config["text_column"]} = ?
            WHERE {articles_config["id_column"]} = ?
            ''', (full_text, article_id))

        conn.commit()
        conn.close()
        return True

    except Exception as e:
        print(f"  ❌ Failed to update database: {e}")
        conn.close()
        return False


def process_article(article: Dict[str, any], output_dir: Path, db_path: str, use_llm: bool = False, llm_model: str = None, max_year: int = 2022, grobid_url: str = "http://localhost:8070") -> Tuple[bool, str, bool, Optional[Path]]:
    """
    Обрабатывает одну статью: скачивает PDF, извлекает текст, обновляет базу.

    Args:
        article: Данные статьи
        output_dir: Директория для сохранения PDF
        db_path: Путь к базе данных
        use_llm: Использовать LLM для очистки текста
        llm_model: Название модели LLM
        max_year: Максимальный год для Sci-Hub
        grobid_url: URL GROBID сервера

    Returns:
        Tuple(success, message, was_downloaded, pdf_path)
    """
    doi = article['doi']
    title = article['title']
    year = article.get('year')

    if year and year > max_year:
        return False, f"Year {year} > {max_year} (Sci-Hub limit)", False, None

    print(f"\n📄 Обработка: {title[:60]}...")
    print(f"  DOI: {doi}")
    if year:
        print(f"  Год: {year}")

    doi_valid, doi_msg = validate_doi(doi)
    if not doi_valid:
        return False, f"Invalid DOI: {doi_msg}", False, None

    existing_pdf = pdf_exists(output_dir, article)
    if existing_pdf:
        print(f"  ✅ PDF уже существует: {existing_pdf.name}")
        print(f"  📝 Извлечение текста из существующего PDF...")

        abstract, full_text = extract_text_from_pdf(existing_pdf, grobid_url)

        if not full_text:
            return False, "Failed to extract text from existing PDF", False, None

        full_text = clean_text(full_text)

        if use_llm and llm_model:
            print(f"  🤖 Очистка текста с помощью LLM ({llm_model})...")
            full_text = clean_text_with_llm(full_text, llm_model)

        return True, "Text extracted from existing PDF", False, existing_pdf

    print(f"  🔍 Поиск на Sci-Hub...")
    html = fetch_scihub_page(doi)
    if not html:
        return False, "Failed to fetch Sci-Hub page", False, None

    print(f"  🔎 Поиск ссылки для скачивания...")
    download_url = find_download_link(html)
    if not download_url:
        return False, "Download link not found", False, None

    print(f"  📥 Скачивание PDF...")
    pdf_filename = get_pdf_filename(title, doi)
    pdf_path = output_dir / pdf_filename

    if not download_pdf(download_url, pdf_path):
        return False, "Failed to download PDF", False, None

    print(f"  ✅ PDF сохранен: {pdf_path.name}")
    print(f"  📝 Извлечение текста...")

    abstract, full_text = extract_text_from_pdf(pdf_path, grobid_url)

    if not full_text:
        pdf_path.unlink()
        return False, "Failed to extract text from PDF", False, None

    full_text = clean_text(full_text)

    if use_llm and llm_model:
        print(f"  🤖 Очистка текста с помощью LLM ({llm_model})...")
        full_text = clean_text_with_llm(full_text, llm_model)

    return True, "Success", True, pdf_path


def main():
    parser = argparse.ArgumentParser(description='Скачивание PDF документов с Sci-Hub')
    parser.add_argument('--db-name', type=str, required=True,
                        help='Название базы данных')
    parser.add_argument('--output-dir', type=str, required=True,
                        help='Директория для сохранения PDF файлов')
    parser.add_argument('--max-articles', type=int, default=None,
                        help='Максимальное количество статей для обработки')
    parser.add_argument('--max-year', type=int, default=SCIHUB_CONFIG["max_year"],
                        help=f'Максимальный год публикации (по умолчанию {SCIHUB_CONFIG["max_year"]} для Sci-Hub)')
    parser.add_argument('--use-llm', action='store_true',
                        help='Использовать LLM для очистки текста')
    parser.add_argument('--llm-model', type=str, default='llama3.2',
                        help='Модель LLM для очистки текста')
    parser.add_argument('--scihub-url', type=str, default=None,
                        help='URL Sci-Hub (по умолчанию из конфига)')
    parser.add_argument('--grobid-url', type=str, default='http://localhost:8070',
                        help='URL GROBID сервера (по умолчанию: http://localhost:8070')

    args = parser.parse_args()

    db_path = PROJECT_ROOT / "data" / args.db_name
    if not db_path.suffix:
        db_path = db_path.with_suffix('.db')

    if not db_path.exists():
        raise FileNotFoundError(f"База данных не найдена: {db_path}")

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    scihub_url = args.scihub_url or SCIHUB_CONFIG["base_url"]

    print("\n" + "="*60)
    print("📥 СКАЧИВАНИЕ PDF С SCI-HUB")
    print("="*60)
    print(f"💾 База данных: {db_path}")
    print(f"📂 Директория для PDF: {output_dir}")
    print(f"🌐 Sci-Hub URL: {scihub_url}")
    print(f"🔗 GROBID сервер: {args.grobid_url}")
    print(f"📅 Максимальный год: {args.max_year}")
    if args.use_llm:
        print(f"🤖 LLM очистка: {args.llm_model}")
    print("="*60)

    print(f"\n🔍 Поиск статей без полного текста (до {args.max_year} года)...")
    articles = get_articles_without_full_text(str(db_path), max_year=args.max_year)

    if not articles:
        print("✅ Все статьи имеют полный текст или не подходят по году")
        return

    print(f"📊 Найдено {len(articles)} статей без полного текста")

    if args.max_articles:
        articles = articles[:args.max_articles]
        print(f"🔢 Будет обработано: {len(articles)} статей")

    stats = {
        'total': len(articles),
        'success': 0,
        'failed': 0,
        'skipped': 0,
        'downloaded': 0,
        'existing': 0
    }

    for i, article in enumerate(articles, 1):
        print(f"\n[{i}/{stats['total']}]")

        success, message, was_downloaded, pdf_path = process_article(
            article,
            output_dir,
            str(db_path),
            use_llm=args.use_llm,
            llm_model=args.llm_model,
            max_year=args.max_year,
            grobid_url=args.grobid_url
        )

        if success:
            stats['success'] += 1
            print(f"  ✅ {message}")

            if was_downloaded:
                stats['downloaded'] += 1
            else:
                stats['existing'] += 1

            if pdf_path:
                abstract, full_text = extract_text_from_pdf(pdf_path, args.grobid_url)

                if full_text:
                    update_article_text(str(db_path), article['id'], full_text, abstract)

                    conn = sqlite3.connect(str(db_path))
                    cursor = conn.cursor()
                    update_article_filepath(cursor, article['id'], str(pdf_path))
                    conn.commit()
                    conn.close()

        else:
            stats['failed'] += 1
            print(f"  ❌ {message}")

        time.sleep(2)

    print("\n" + "="*60)
    print("📊 ОБЩАЯ СТАТИСТИКА")
    print("="*60)
    print(f"Всего статей: {stats['total']}")
    print(f"✅ Успешно: {stats['success']}")
    print(f"  📥 Скачано: {stats['downloaded']}")
    print(f"  📁 Из существующих PDF: {stats['existing']}")
    print(f"❌ Неудачно: {stats['failed']}")
    print(f"⏭️ Пропущено: {stats['skipped']}")
    print(f"Успешность: {stats['success']/stats['total']*100:.1f}%" if stats['total'] > 0 else "")
    print("="*60)


if __name__ == "__main__":
    main()