scripts/chunk_fulltext.py - common_agents_notes - IvanMysin/Topics

scripts/chunk_fulltext.py
377 строк14 KB
ivan
Add chunking
21 апр 2026, 14:47
21 апр 2026, 14:4749c216b
#!/usr/bin/env python3
"""
Скрипт для чанкования полных текстов научных статей.

Цель: создать эмбеддинги для чанков полного текста, чтобы находить конкретные параграфы.

Реализация:
- Разбить каждый full_text на чанки по ~512 токенов с перекрытием 64 токена
- Каждому чанку присвоить метаданные: article_id, chunk_index, section (эвристически)
- Закодировать чанки через SPECTER2
- Сохранить в отдельную коллекцию ChromaDB (fulltext_chunks)
"""

import os
import sqlite3
import hashlib
import re
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass

import chromadb
import numpy as np
import pandas as pd
from tqdm.auto import tqdm
import gc
import torch
from transformers import AutoTokenizer
from adapters import AutoAdapterModel
from pathlib import Path
import sys

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

from config import DB_CONFIG, MODEL_CONFIG, CHROMA_CONFIG


CHUNK_SIZE = 512
CHUNK_OVERLAP = 64
COLLECTION_NAME = "fulltext_chunks"


SECTION_PATTERNS = {
    "abstract": r"^\s*(?:abstract|summary|overview)\s*$",
    "introduction": r"^\s*(?:introduction|background|aim|purpose|objective)\s*$",
    "methods": r"^\s*(?:methods|materials?\s*(?:and|&)?\s*methods?|methodology|approach|experimental|procedure|protocol)\s*$",
    "results": r"^\s*(?:results|findings|outcome|observations?)\s*$",
    "discussion": r"^\s*(?:discussion|interpretation|limitations?)\s*$",
    "conclusion": r"^\s*(?:conclusion|conclusions|summaries?|concluding\s+remarks?)\s*$",
    "references": r"^\s*(?:references?|bibliography|literature\s+cited|literature)\s*$",
}


@dataclass
class TextChunk:
    """Представляет один чанк текста"""
    article_id: int
    chunk_index: int
    section: str
    text: str
    start_char: int
    end_char: int


class FulltextChunker:
    """Разбивает текст на чанки с перекрытием"""
    
    def __init__(self, tokenizer, chunk_size: int = CHUNK_SIZE, chunk_overlap: int = CHUNK_OVERLAP):
        self.tokenizer = tokenizer
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.section_patterns = {k: re.compile(v, re.IGNORECASE | re.MULTILINE) for k, v in SECTION_PATTERNS.items()}
    
    def split_into_sentences(self, text: str) -> List[str]:
        """Разбивает текст на предложения (упрощенно)"""
        text = re.sub(r'([.!?])\s+', r'\1<END>', text)
        text = text.replace('\n', ' ')
        sentences = text.split('<END>')
        return [s.strip() for s in sentences if s.strip()]
    
    def detect_section(self, text: str) -> str:
        """Определяет секцию по заголовку"""
        lines = text.split('\n')
        for line in lines[:10]:
            line_clean = line.strip()
            if len(line_clean) < 100 and line_clean:
                for section_name, pattern in self.section_patterns.items():
                    if pattern.match(line_clean):
                        return section_name
        return "body"
    
    def chunk_text(self, article_id: int, text: str) -> List[TextChunk]:
        """Разбивает текст на чанки"""
        if not text or not text.strip():
            return []
        
        tokens = self.tokenizer.encode(text, add_special_tokens=False)
        
        if len(tokens) < 50:
            return []
        
        chunks = []
        step = self.chunk_size - self.chunk_overlap
        
        for i in range(0, len(tokens), step):
            token_chunk = tokens[i:i + self.chunk_size]
            
            chunk_text = self.tokenizer.decode(token_chunk, skip_special_tokens=True)
            chunk_text = chunk_text.strip()
            
            if len(chunk_text) < 100:
                continue
            
            start_char = len(self.tokenizer.decode(tokens[:i], skip_special_tokens=True))
            end_char = start_char + len(chunk_text)
            
            section = self.detect_section(chunk_text)
            
            chunks.append(TextChunk(
                article_id=article_id,
                chunk_index=len(chunks),
                section=section,
                text=chunk_text,
                start_char=start_char,
                end_char=end_char
            ))
            
            if i + self.chunk_size >= len(tokens):
                break
        
        return chunks


class EmbeddingChunkManager:
    """Менеджер для работы с чанками и их эмбеддингами"""
    
    def __init__(self):
        self.model_name = MODEL_CONFIG["name"]
        self.embedding_dim = MODEL_CONFIG["embedding_dimension"]
        self.batch_size = MODEL_CONFIG["batch_size"]
        
        os.makedirs(CHROMA_CONFIG["path"], exist_ok=True)
        
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoAdapterModel.from_pretrained(self.model_name)
        self.model.load_adapter("allenai/specter2", source="hf", set_active=True)
        
        self.chroma_client = chromadb.PersistentClient(path=str(CHROMA_CONFIG["path"]))
        self.collection = self._get_or_create_collection()
    
    def _get_or_create_collection(self) -> chromadb.Collection:
        """Создает или загружает коллекцию Chunks"""
        existing = self.chroma_client.list_collections()
        names = [c.name for c in existing]
        
        if COLLECTION_NAME in names:
            return self.chroma_client.get_collection(COLLECTION_NAME)
        
        print(f"📝 Создаем коллекцию {COLLECTION_NAME}...")
        try:
            return self.chroma_client.create_collection(
                name=COLLECTION_NAME,
                metadata={"description": "Fulltext chunks embeddings", "model": self.model_name}
            )
        except TypeError:
            return self.chroma_client.create_collection(
                name=COLLECTION_NAME,
                metadata={"description": "Fulltext chunks embeddings", "model": self.model_name},
                get_or_create=True
            )
    
    def _compute_hash(self, text: str) -> str:
        return hashlib.md5(text.encode('utf-8')).hexdigest()
    
    def compute_embeddings_batch(self, texts: List[str]) -> np.ndarray:
        with torch.no_grad():
            inputs = self.tokenizer(
                texts, 
                padding=True, 
                truncation=True,
                return_tensors="pt", 
                return_token_type_ids=False, 
                max_length=512
            )
            outputs = self.model(**inputs)
            embeddings = outputs.last_hidden_state[:, 0, :].detach().numpy()
        return embeddings
    
    def load_articles_with_fulltext(self, limit: Optional[int] = None) -> pd.DataFrame:
        """Загружает статьи с полным текстом из БД"""
        conn = sqlite3.connect(DB_CONFIG["db_path"])
        
        query = f"""
            SELECT id, doi, title, abstract, full_text, journal, date 
            FROM {DB_CONFIG['table_name']}
            WHERE full_text IS NOT NULL 
            AND LENGTH(full_text) > 500
        """
        
        if limit:
            query += f" LIMIT {limit}"
        
        df = pd.read_sql_query(query, conn)
        conn.close()
        
        print(f"📥 Загружено {len(df)} статей с full_text")
        return df
    
    def process_articles(
        self, 
        df: pd.DataFrame, 
        force_update: bool = False,
        batch_size: int = 16
    ) -> int:
        """Обрабатывает статьи и создает чанки"""
        
        chunker = FulltextChunker(
            tokenizer=self.tokenizer,
            chunk_size=CHUNK_SIZE,
            chunk_overlap=CHUNK_OVERLAP
        )
        
        existing_chunks = self._get_existing_chunk_ids()
        
        total_chunks = 0
        
        articles_to_process = []
        
        for _, row in df.iterrows():
            article_id = int(row['id'])
            full_text = row.get('full_text', '')
            
            if not full_text or not isinstance(full_text, str) or len(full_text) < 500:
                continue
            
            chunks = chunker.chunk_text(article_id, full_text)
            
            if not chunks:
                continue
            
            needs_update = force_update
            if not force_update:
                first_chunk_id = f"{article_id}_0"
                needs_update = first_chunk_id not in existing_chunks
            
            if needs_update:
                articles_to_process.append((article_id, chunks))
        
        print(f"📋 Подготовлено к обработке {len(articles_to_process)} статей...")
        
        for article_id, chunks in tqdm(articles_to_process, desc="Обработка чанков"):
            chunk_texts = [c.text for c in chunks]
            
            embeddings = self.compute_embeddings_batch(chunk_texts)
            
            ids = [f"{c.article_id}_{c.chunk_index}" for c in chunks]
            metadatas = [
                {
                    "article_id": str(c.article_id),
                    "chunk_index": c.chunk_index,
                    "section": c.section,
                    "doi": df.loc[df['id'] == c.article_id, 'doi'].iloc[0] if len(df.loc[df['id'] == c.article_id]) > 0 else "",
                    "title": df.loc[df['id'] == c.article_id, 'title'].iloc[0] if len(df.loc[df['id'] == c.article_id]) > 0 else "",
                    "text_hash": self._compute_hash(c.text),
                }
                for c in chunks
            ]
            
            self.collection.upsert(
                ids=ids,
                embeddings=embeddings.tolist(),
                metadatas=metadatas,
                documents=chunk_texts
            )
            
            total_chunks += len(chunks)
            
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        
        print(f"✅ Создано {total_chunks} чанков")
        return total_chunks
    
    def _get_existing_chunk_ids(self) -> set:
        """Получает ID существующих чанков с пагинацией"""
        total = self.collection.count()
        if total == 0:
            return set()
        
        all_ids = set()
        batch_size = 1000
        offset = 0
        
        while offset < total:
            results = self.collection.get(
                include=["metadatas"],
                limit=batch_size,
                offset=offset
            )
            if not results["ids"]:
                break
            all_ids.update(results["ids"])
            offset += batch_size
        
        return all_ids
    
    def get_stats(self) -> Dict:
        """Возвращает статистику коллекции с пагинацией"""
        total = self.collection.count()
        if total == 0:
            return {"total_chunks": 0}
        
        article_ids = set()
        sections = {}
        batch_size = 1000
        offset = 0
        
        while offset < total:
            results = self.collection.get(
                include=["metadatas"],
                limit=batch_size,
                offset=offset
            )
            if not results["metadatas"]:
                break
            for meta in results["metadatas"]:
                article_ids.add(meta.get("article_id", ""))
                section = meta.get("section", "unknown")
                sections[section] = sections.get(section, 0) + 1
            offset += batch_size
        
        return {
            "total_chunks": total,
            "total_articles": len(article_ids),
            "sections": sections
        }


def main():
    import argparse
    
    parser = argparse.ArgumentParser(description="Чанкование полных текстов научных статей")
    parser.add_argument("--limit", type=int, default=None, help="Ограничение количества статей")
    parser.add_argument("--force-update", action="store_true", help="Принудительно обновить все чанки")
    parser.add_argument("--stats", action="store_true", help="Показать статистику")
    args = parser.parse_args()
    
    manager = EmbeddingChunkManager()
    
    if args.stats:
        stats = manager.get_stats()
        print("📊 Статистика коллекции чанков:")
        print(f"  Всего чанков: {stats.get('total_chunks', 0)}")
        print(f"  Всего статей: {stats.get('total_articles', 0)}")
        sections = stats.get("sections", {})
        print("  По секциям:")
        for section, count in sorted(sections.items(), key=lambda x: -x[1]):
            print(f"    {section}: {count}")
        return
    
    df = manager.load_articles_with_fulltext(limit=args.limit)
    
    if df.empty:
        print("⚠️ Нет статей с full_text для обработки")
        return
    
    manager.process_articles(df, force_update=args.force_update)
    
    stats = manager.get_stats()
    print(f"\n📊 Итоговая статистика:")
    print(f"  Всего чанков: {stats.get('total_chunks', 0)}")
    print(f"  Всего статей: {stats.get('total_articles', 0)}")


if __name__ == "__main__":
    main()