scripts/pubmed_load.py - common_agents_notes - IvanMysin/Topics

scripts/pubmed_load.py
712 строк25 KB
ivan
Working on downloading scripts
17 фев 2026, 15:29
17 фев 2026, 15:297f61ba9
#!/usr/bin/env python
"""
Скрипт для загрузки коллекции документов из PubMed.

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

import argparse
import os
import sqlite3
from pathlib import Path
from typing import List, Dict, Optional, Iterator, Any, Tuple
import time
from urllib.error import URLError, HTTPError
from http.client import RemoteDisconnected
from functools import wraps
import xml.etree.ElementTree as ET
import sys
from Bio import Entrez
from tqdm import tqdm

PUBMED_API_KEY = os.getenv('PUBMED_API_KEY')
email = 'imysin@mail.ru'

# Добавляем корневую директорию проекта в sys.path
PROJECT_ROOT = Path(__file__).parent.parent
sys.path.append(str(PROJECT_ROOT))

# Настройка логирования
import logging
from datetime import datetime

LOG_DIR = PROJECT_ROOT / "logs"
LOG_DIR.mkdir(exist_ok=True)
LOG_FILE = LOG_DIR / f"pubmed_load_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(LOG_FILE, encoding='utf-8'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# habanero для Crossref API
try:
    from habanero import Crossref as HabaneroCrossref
    habanero_available = True
except ImportError:
    HabaneroCrossref = None
    habanero_available = False
    print("⚠️ Пакет habanero не установлен. pip install habanero")

# Импорты из проекта
from src.database_manager import (
    create_articles_table, document_exists, insert_article,
    validate_doi, validate_title, validate_abstract, generate_doi_from_title_abstract
)


def validate_abstract(abstract: str) -> Tuple[bool, str]:
    """
    Проверяет корректность абстракта.

    Args:
        abstract: Абстракт для проверки

    Returns:
        Tuple[is_valid, message] - валиден ли абстракт и сообщение об ошибке
    """
    if not abstract:
        return False, "Абстракт отсутствует"

    abstract = abstract.strip()
    if len(abstract) < 150:
        return False, f"Абстракт слишком короткий ({len(abstract)} символов): {abstract}"

    return True, ""


def log_invalid_document(doc_type: str, message: str, metadata: Dict[str, Any]) -> None:
    """
    Логирует информацию о некорректном документе.

    Args:
        doc_type: Тип документа (например, "no_doi", "no_title", "no_abstract")
        message: Сообщение об ошибке
        metadata: Метаданные документа для логирования
    """
    log_entry = {
        'type': doc_type,
        'pmid': metadata.get('uid', ''),
        'doi': metadata.get('doi', ''),
        'title': str(metadata.get('Title', ''))[:200],
        'message': message
    }

    logger.warning(
        f"🚨 INVALID DOCUMENT [{doc_type}]: {message}\n"
        f"   PMID: {log_entry['pmid']}\n"
        f"   DOI: {log_entry['doi']}\n"
        f"   Title: {log_entry['title']}\n"
        f"   ---"
    )


def query_crossref_by_doi(doi: str) -> Optional[Dict[str, Any]]:
    """
    Запрашивает метаданные из Crossref по DOI.

    Args:
        doi: DOI статьи

    Returns:
        Словарь с метаданными или None
    """
    if not habanero_available:
        return None

    try:
        cr = HabaneroCrossref()
        work = cr.works(ids=doi)
        return work.get('message', {})
    except Exception as e:
        logger.warning(f"Ошибка при запросе к Crossref по DOI {doi}: {e}")
        return None


def query_crossref_by_title(title: str) -> Optional[Dict[str, Any]]:
    """
    Запрашивает метаданные из Crossref по названию статьи.

    Args:
        title: Название статьи

    Returns:
        Словарь с метаданными или None
    """
    if not habanero_available:
        return None

    try:
        cr = HabaneroCrossref()
        results = cr.works(query=title, limit=1)
        items = results.get('message', {}).get('items', [])
        if items:
            return items[0]
        return None
    except Exception as e:
        logger.warning(f"Ошибка при запросе к Crossref по названию '{title[:50]}...': {e}")
        return None


def extract_metadata_from_crossref(crossref_data: Dict[str, Any]) -> Dict[str, Any]:
    """
    Извлекает метаданные из ответа Crossref.

    Args:
        crossref_data: Данные от Crossref

    Returns:
        Словарь с метаданными
    """
    metadata = {}

    publisher = crossref_data.get('publisher', '')
    short_container_title = crossref_data.get('short-container-title', [])
    container_title = crossref_data.get('container-title', [])
    journal = short_container_title[0] if short_container_title else (container_title[0] if container_title else publisher)
    metadata['journal'] = journal

    metadata['volume'] = crossref_data.get('volume', '')
    metadata['issue'] = crossref_data.get('issue', '')
    metadata['pages'] = crossref_data.get('page', '')
    metadata['pmid'] = crossref_data.get('pmid', '')

    authors = []
    author_list = crossref_data.get('author', [])
    for author in author_list:
        given_name = author.get('given', '')
        family_name = author.get('family', '')
        if family_name:
            author_name = f"{family_name} {given_name}".strip()
            authors.append(author_name)
    metadata['authors'] = ", ".join(authors) if authors else ""

    return metadata


def validate_metadata_completeness(metadata: Dict[str, Any]) -> bool:
    """
    Проверяет полноту метаданных.

    Args:
        metadata: Метаданные для проверки

    Returns:
        True если метаданные достаточно полные
    """
    required_fields = ['journal']
    for field in required_fields:
        if not metadata.get(field):
            return False
    return True


def enrich_metadata_from_crossref(metadata: Dict[str, Any], use_crossref: bool = True) -> Tuple[Dict[str, Any], bool]:
    """
    Обогащает метаданные данными из Crossref, если PubMed данные неполные.

    Args:
        metadata: Метаданные из PubMed
        use_crossref: Использовать ли Crossref

    Returns:
        Tuple(обогащенные метаданные, были ли обновлены данные)
    """
    updated = False
    doi = metadata.get('doi', '')
    title = metadata.get('Title', '')

    if not use_crossref:
        return metadata, updated

    if validate_metadata_completeness(metadata):
        return metadata, updated

    if doi and not doi.startswith('PMID:') and not doi.startswith('GEN:'):
        crossref_data = query_crossref_by_doi(doi)
        if crossref_data:
            crossref_metadata = extract_metadata_from_crossref(crossref_data)
            for key, value in crossref_metadata.items():
                if not metadata.get(key) and value:
                    metadata[key] = value
                    updated = True

    if not validate_metadata_completeness(metadata) and title:
        crossref_data = query_crossref_by_title(title)
        if crossref_data:
            crossref_metadata = extract_metadata_from_crossref(crossref_data)
            for key, value in crossref_metadata.items():
                if not metadata.get(key) and value:
                    metadata[key] = value
                    updated = True

    return metadata, updated


def fetch_pubmed_ids(query: str, load_max_docs: int, email: str, api_key: Optional[str] = "", from_date: str = None, to_date: str = None) -> List[str]:
    """Fetch PubMed IDs matching the query."""
    Entrez.email = email
    if api_key:
        Entrez.api_key = api_key

    # Формируем расширенный запрос с фильтрацией по дате
    full_query = query
    if from_date or to_date:
        date_filter = ""
        if from_date:
            date_filter += f"{from_date}[Date - Publication] : "
        if to_date:
            date_filter += f"{to_date}[Date - Publication]"
        else:
            # Установим дату по умолчанию
            date_filter += "2023/12/31[Date - Publication]"
            
        full_query = f"{query} AND ({date_filter})"
        
    handle = Entrez.esearch(
        db="pubmed",
        term=full_query,
        retmax=load_max_docs,
        retmode="xml"
    )
    record = Entrez.read(handle)
    handle.close()
    return record.get("IdList", [])


def fetch_article_details(pubmed_ids: List[str]) -> List[Dict[str, Any]]:
    """Fetch detailed information for PubMed IDs."""
    if not pubmed_ids:
        return []

    batch_size = 200
    all_articles = []

    for i in range(0, len(pubmed_ids), batch_size):
        batch_ids = pubmed_ids[i:i + batch_size]
        handle = Entrez.efetch(
            db="pubmed",
            id=batch_ids,
            rettype="medline",
            retmode="xml"
        )
        xml_data = handle.read()
        handle.close()

        articles = parse_medline_xml(xml_data)
        all_articles.extend(articles)

    return all_articles


def parse_medline_xml(xml_data: str) -> List[Dict[str, Any]]:
    """Parse MEDLINE XML data and extract article information."""
    articles = []
    root = ET.fromstring(xml_data)

    for article in root.findall(".//PubmedArticle"):
        article_data = {}

        medline_citation = article.find("MedlineCitation")
        if medline_citation is None:
            continue

        pubmed_data = article.find("PubmedData")
        article_data["uid"] = medline_citation.get("PMID", "")

        article_elem = medline_citation.find("Article")
        if article_elem is None:
            continue

        title_elem = article_elem.find("ArticleTitle")
        article_data["Title"] = extract_text(title_elem) if title_elem is not None else ""

        abstract_elem = article_elem.find("Abstract")
        abstract_texts = []
        if abstract_elem is not None:
            for abstract_text in abstract_elem.findall("AbstractText"):
                label = abstract_text.get("Label", "")
                text = extract_text(abstract_text)
                if label:
                    abstract_texts.append(f"{label}: {text}")
                else:
                    abstract_texts.append(text)
        article_data["Abstract"] = " ".join(abstract_texts)

        journal_elem = article_elem.find("Journal")
        if journal_elem is not None:
            journal_title_elem = journal_elem.find("Title")
            article_data["journal"] = extract_text(journal_title_elem) if journal_title_elem is not None else ""

            journal_issue_elem = journal_elem.find("JournalIssue")
            if journal_issue_elem is not None:
                volume_elem = journal_issue_elem.find("Volume")
                article_data["volume"] = extract_text(volume_elem) if volume_elem is not None else ""

                issue_elem = journal_issue_elem.find("Issue")
                article_data["issue"] = extract_text(issue_elem) if issue_elem is not None else ""

                pub_date_elem = journal_issue_elem.find("PubDate")
                if pub_date_elem is not None:
                    year_elem = pub_date_elem.find("Year")
                    month_elem = pub_date_elem.find("Month")
                    day_elem = pub_date_elem.find("Day")

                    year = extract_text(year_elem) if year_elem is not None else ""
                    month = extract_text(month_elem) if month_elem is not None else ""
                    day = extract_text(day_elem) if day_elem is not None else ""

                    article_data["Published"] = f"{year}-{month}-{day}".strip("-")

        if "Published" not in article_data:
            article_data["Published"] = ""

        if "journal" not in article_data:
            article_data["journal"] = ""

        if "volume" not in article_data:
            article_data["volume"] = ""

        if "issue" not in article_data:
            article_data["issue"] = ""

        pagination_elem = article_elem.find("Pagination")
        if pagination_elem is not None:
            medline_pgn_elem = pagination_elem.find("MedlinePgn")
            article_data["pages"] = extract_text(medline_pgn_elem) if medline_pgn_elem is not None else ""
        else:
            article_data["pages"] = ""

        author_list_elem = article_elem.find("AuthorList")
        authors = []
        if author_list_elem is not None:
            for author in author_list_elem.findall("Author"):
                last_name = extract_text(author.find("LastName"))
                initials = extract_text(author.find("Initials"))
                if last_name:
                    author_name = f"{last_name} {initials}".strip()
                    authors.append(author_name)
        article_data["authors"] = ", ".join(authors) if authors else ""

        if pubmed_data is not None:
            article_id_list = pubmed_data.find("ArticleIdList")
            if article_id_list is not None:
                for article_id in article_id_list.findall("ArticleId"):
                    if article_id.get("IdType") == "doi":
                        article_data["doi"] = article_id.text
                        break

        if "doi" not in article_data:
            article_data["doi"] = article_data.get("uid", "")

        articles.append(article_data)

    return articles


def extract_text(element: Any) -> str:
    """Extract text from XML element, handling nested elements."""
    if element is None:
        return ""
    if element.text:
        return element.text.strip()
    return "".join(extract_text(child) + (child.tail or "") for child in element)


def retry_on_failure(max_retries: int = 3, delay: float = 5, backoff: float = 2):
    """Декоратор для повторных попыток"""
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_exception: Optional[Exception] = None
            current_delay = delay

            for attempt in range(max_retries):
                try:
                    return func(*args, **kwargs)
                except (RemoteDisconnected, URLError, HTTPError, TimeoutError) as e:
                    last_exception = e
                    if attempt < max_retries - 1:
                        print(f"Попытка {attempt + 1} не удалась: {e}. "
                              f"Повтор через {current_delay} секунд...")
                        time.sleep(current_delay)
                        current_delay *= backoff
                    else:
                        print(f"Все {max_retries} попыток не удались")

            if last_exception is None:
                raise Exception("Unknown error occurred")
            raise last_exception
        return wrapper
    return decorator


@retry_on_failure(max_retries=5, delay=3, backoff=1.5)
def load_pubmed_with_retry(query: str, **kwargs) -> Iterator[Dict[str, Any]]:
    """Load documents from PubMed with retry logic."""
    pubmed_ids = fetch_pubmed_ids(query, kwargs.get('load_max_docs', 1000),
                                    kwargs.get('email', email),
                                    kwargs.get('api_key', ''),
                                    kwargs.get('from_date', None),
                                    kwargs.get('to_date', None))
    print(f"Найдено {len(pubmed_ids)} документов по запросу")

    articles = fetch_article_details(pubmed_ids)

    for article in articles:
        yield {
            "metadata": article,
            "page_content": article.get("Abstract", "")
        }


def clear_text(text: str) -> str:
    lines = text.split('\n')

    cleaned_lines = []
    for line in lines:
        line = line.strip()
        if len(line) == 0:
            continue
        if line[0] == '[' and line[-1] == ']':
            continue

        if len(line) < 20:
            continue

        cleaned_lines.append(line)

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


def read_queries(query_file: str) -> List[str]:
    """
    Читает запросы из файла.

    Args:
        query_file: Путь к файлу с запросами

    Returns:
        Список запросов
    """
    queries = []
    with open(query_file, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line and (not line.startswith('#')):
                queries.append(line)
    return queries





def save_documents(documents: List[Any], db_path: str, use_crossref: bool = True) -> Dict[str, int]:
    """
    Сохраняет документы в базу данных.

    Args:
        documents: Список документов
        db_path: Путь к файлу базы данных
        use_crossref: Использовать Crossref для обогащения метаданных

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

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

    for doc in tqdm(documents, desc="Обработка документов", unit="статья"):
        try:
            metadata = doc["metadata"]
            doi = metadata.get('doi', '')
            title = metadata.get('Title', '')
            publish_time = metadata.get('Published', '')

            if isinstance(title, dict):
                try:
                    title = str(title.get('#text', ''))
                except:
                    title = str(title)

            if not title:
                log_invalid_document("no_title", "Название отсутствует", metadata)
                stats['errors'] += 1
                continue

            title_valid, title_msg = validate_title(title)
            if not title_valid:
                log_invalid_document("invalid_title", title_msg, metadata)
                stats['errors'] += 1
                continue

            if document_exists(cursor, doi, title):
                stats['skipped'] += 1
                continue

            abstract = doc.get('page_content', '')
            abstract = clear_text(abstract)

            abstract_valid, abstract_msg = validate_abstract(abstract)
            if not abstract_valid:
                log_invalid_document("invalid_abstract", abstract_msg, metadata)
                stats['errors'] += 1
                continue

            metadata['Title'] = title
            metadata, crossref_updated = enrich_metadata_from_crossref(metadata, use_crossref)
            if crossref_updated:
                stats['crossref_updated'] += 1
                logger.info(f"📡 Данные из Crossref обновлены для: {title[:50]}...")

            doi_valid, doi_msg = validate_doi(doi)
            if not doi_valid:
                logger.warning(f"⚠️ {doi_msg}. Генерация DOI из названия и абстракта.")
                doi = generate_doi_from_title_abstract(title, abstract)
                metadata['doi'] = doi

            article_data = {
                'doi': doi,
                'title': title,
                'abstract': abstract,
                'date': publish_time,
                'full_text': None,
                'journal': metadata.get('journal', ''),
                'authors': metadata.get('authors', ''),
                'volume': metadata.get('volume', ''),
                'issue': metadata.get('issue', ''),
                'pages': metadata.get('pages', ''),
                'pmid': metadata.get('uid', ''),
                'filepath': None
            }

            insert_article(cursor, article_data)
            conn.commit()

            stats['saved'] += 1

            time.sleep(0.5)

        except Exception as e:
            logger.error(f"Ошибка при обработке документа: {e}")
            stats['errors'] += 1
            time.sleep(0.5)
            continue

        stats['processed'] += 1

    conn.close()
    return stats


def main():
    parser = argparse.ArgumentParser(description='Загрузка документов из PubMed')
    parser.add_argument('--query_file', type=str, required=True, help='Путь к файлу с запросами')
    parser.add_argument('--db_name', type=str, required=True, help='Название базы данных')
    parser.add_argument('--no-crossref', action='store_true', help='Не использовать Crossref для обогащения метаданных')
    parser.add_argument('--from-date', type=str, help='Дата начала (YYYY-MM-DD)')
    parser.add_argument('--to-date', type=str, help='Дата окончания (YYYY-MM-DD)')
    parser.add_argument('--days-back', type=int, help='Количество дней назад для загрузки новых статей')
    parser.add_argument('--update-only', action='store_true', help='Обновлять только существующие статьи (без проверки дубликатов)')
    args = parser.parse_args()

    db_path = Path(args.db_name)
    if not db_path.suffix:
        db_path = db_path.with_suffix('.db')

    if not db_path.is_absolute():
        db_path = PROJECT_ROOT / "data" / db_path

    db_path.parent.mkdir(parents=True, exist_ok=True)
    db_path_str = str(db_path)

    if not os.path.exists(args.query_file):
        raise FileNotFoundError(f"Файл запросов не найден: {args.query_file}")

    queries = read_queries(args.query_file)
    if not queries:
        print("Нет действительных запросов для выполнения")
        return

    print(f"Найдено {len(queries)} запросов")

    # Обработка параметров даты
    from_date = args.from_date
    to_date = args.to_date
    
    if args.days_back:
        from datetime import datetime, timedelta
        to_date = datetime.now().strftime('%Y/%m/%d')
        from_date = (datetime.now() - timedelta(days=args.days_back)).strftime('%Y/%m/%d')
    
    # Используем параметры даты при выполнении запросов
    create_articles_table(db_path_str)
    print(f"База данных создана/проверена: {db_path}")

    use_crossref = not args.no_crossref
    total_stats = {'processed': 0, 'saved': 0, 'skipped': 0, 'errors': 0, 'crossref_updated': 0}

    for i, query in enumerate(queries, 1):
        print(f"\nОбработка запроса {i}/{len(queries)}: {query}")
        
        try:
            documents_iter = load_pubmed_with_retry(
                query=query, 
                load_max_docs=5000, 
                email=email, 
                api_key=PUBMED_API_KEY or "",
                from_date=from_date,
                to_date=to_date
            )
            documents = list(documents_iter)
        except Exception as e:
            documents = []
            print(f"Все попытки не удались: {e}")

        if documents:
            stats = save_documents(documents, db_path_str, use_crossref)

            print(f"\nСтатистика сохранения:")
            print(f"  Обработано: {stats['processed']}")
            print(f"  Сохранено: {stats['saved']}")
            print(f"  Пропущено (дубликаты): {stats['skipped']}")
            print(f"  Ошибок: {stats['errors']}")
            print(f"  Обновлено через Crossref: {stats['crossref_updated']}")

            for key in total_stats:
                total_stats[key] += stats[key]

        if i < len(queries):
            print("\nОжидание 2 секунды перед следующим запросом...")
            time.sleep(2)


    print("\n" + "="*50)
    print("ОБЩАЯ СТАТИСТИКА")
    print("="*50)
    print(f"Всего обработано документов: {total_stats['processed']}")
    print(f"Всего сохранено документов: {total_stats['saved']}")
    print(f"Всего пропущено документов: {total_stats['skipped']}")
    print(f"Всего ошибок: {total_stats['errors']}")
    print(f"Всего обновлено через Crossref: {total_stats['crossref_updated']}")
    print("="*50)


if __name__ == "__main__":
    main()