spb/uqo-research-db

1# ============================================================2# Project: EDGAR Research Database3# Description: SEC EDGAR Research Infrastructure4#5# Authors:6# - Manel Kammoun — manel.kammoun@uqo.ca7# - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca8# - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca9#10# Institution:11# Université du Québec en Outaouais (UQO)12#13# License:14# Research / Academic Use15# ============================================================1617from __future__ import annotations1819import time20import uuid21from typing import TYPE_CHECKING2223import numpy as np24import pyarrow as pa25import pandas as pd26from tqdm import tqdm2728from src.embeddings.models import EmbeddingModelRegistry29from src.utils.logging import get_logger3031if TYPE_CHECKING:32    from src.storage.database import DatabaseManager33    from src.storage.parquet_manager import ParquetManager3435logger = get_logger("embeddings.manager")363738class EmbeddingManager:3940    def __init__(41        self,42        db: DatabaseManager,43        parquet: ParquetManager,44        model_name: str | None = None,45    ) -> None:46        self.db = db47        self.parquet = parquet48        self.registry = EmbeddingModelRegistry()49        self.model_name = model_name or self.registry.default_model5051    def _get_already_embedded_ids(self, model_name: str, source_type: str) -> set[str]:52        sql = (53            "SELECT source_id FROM embeddings_metadata "54            "WHERE model_name = ? AND source_type = ?"55        )56        df = self.db.fetch_df(sql, [model_name, source_type])57        return set(df["source_id"].tolist())5859    def embed_chunks(60        self,61        accession_numbers: list[str] | None = None,62        model_name: str | None = None,63        batch_size: int = 256,64        fetch_size: int = 50_000,65    ) -> int:66        """Stream chunks that have no embedding yet, ``fetch_size`` rows at a time.6768        Memory stays bounded whatever the corpus size; each fetch produces one69        parquet file under ``parquet/embeddings/model_name=<model>/``.70        """71        model_name = model_name or self.model_name72        model_info = self.registry.get_model_info(model_name)7374        where = "NOT EXISTS (SELECT 1 FROM embeddings_metadata m WHERE m.source_id = c.chunk_id AND m.model_name = ? AND m.source_type = 'chunk')"75        params: list = [model_name]76        if accession_numbers:77            placeholders = ", ".join(["?"] * len(accession_numbers))78            where += f" AND c.accession_number IN ({placeholders})"79            params += list(accession_numbers)8081        total = self.db.execute(f"SELECT COUNT(*) FROM filing_chunks c WHERE {where}", params).fetchone()[0]82        if total == 0:83            logger.info("All chunks already embedded with model '%s'", model_name)84            return 085        logger.info("Embedding %d chunks with model '%s' (fetch=%d, batch=%d)", total, model_name, fetch_size, batch_size)8687        embedded_count = 088        t0 = time.time()89        while True:90            chunks_df = self.db.fetch_df(91                f"SELECT c.chunk_id, c.chunk_text FROM filing_chunks c WHERE {where} LIMIT {int(fetch_size)}",92                params,93            )94            if chunks_df.empty:95                break96            texts = chunks_df["chunk_text"].fillna("").tolist()97            vectors = self.registry.encode(texts, model_name=model_name, batch_size=batch_size, show_progress=False)98            vectors = np.asarray(vectors, dtype=np.float32)99100            ids = chunks_df["chunk_id"].tolist()101            emb_ids = [str(uuid.uuid4()) for _ in ids]102            table = pa.table({103                "embedding_id": pa.array(emb_ids, pa.string()),104                "source_id": pa.array(ids, pa.string()),105                "source_type": pa.array(["chunk"] * len(ids), pa.string()),106                "vector": pa.FixedSizeListArray.from_arrays(pa.array(vectors.ravel(), pa.float32()), vectors.shape[1]).cast(pa.list_(pa.float32())),107                "model_name": pa.array([model_name] * len(ids), pa.string()),108            })109            self.parquet.write_embeddings_table(table, model_name)110111            meta_df = pd.DataFrame({112                "embedding_id": emb_ids,113                "source_type": "chunk",114                "source_id": ids,115                "model_name": model_name,116                "model_version": model_info.get("provider", ""),117                "embedding_dimension": int(model_info.get("dimension", vectors.shape[1])),118            })119            self.db.insert_df("embeddings_metadata", meta_df)120            embedded_count += len(ids)121            el = time.time() - t0122            logger.info("embedded %d/%d chunks (%.0f chunks/s)", embedded_count, total, embedded_count / max(el, 1))123            if len(ids) < fetch_size:124                break125126        logger.info("Embedded %d chunks with model '%s'", embedded_count, model_name)127        return embedded_count128129    def embed_sections(130        self,131        accession_numbers: list[str] | None = None,132        model_name: str | None = None,133    ) -> int:134        model_name = model_name or self.model_name135        model_info = self.registry.get_model_info(model_name)136        batch_size = 64137138        if accession_numbers:139            placeholders = ", ".join(["?"] * len(accession_numbers))140            sql = f"SELECT section_id, text, accession_number, section_name FROM filing_sections WHERE accession_number IN ({placeholders})"141            sections_df = self.db.fetch_df(sql, accession_numbers)142        else:143            sql = "SELECT section_id, text, accession_number, section_name FROM filing_sections"144            sections_df = self.db.fetch_df(sql)145146        if sections_df.empty:147            logger.info("No sections found to embed")148            return 0149150        already_embedded = self._get_already_embedded_ids(model_name, "section")151        mask = ~sections_df["section_id"].isin(already_embedded)152        sections_df = sections_df[mask].reset_index(drop=True)153154        if sections_df.empty:155            logger.info("All sections already embedded with model '%s'", model_name)156            return 0157158        total = len(sections_df)159        logger.info("Embedding %d sections with model '%s'", total, model_name)160161        embedded_count = 0162        for start in tqdm(range(0, total, batch_size), desc="Embedding sections", unit="batch"):163            end = min(start + batch_size, total)164            batch_df = sections_df.iloc[start:end]165166            texts = batch_df["text"].fillna("").tolist()167            max_tokens = model_info.get("max_tokens", 256)168            truncated = [" ".join(t.split()[:max_tokens]) for t in texts]169170            vectors = self.registry.encode(truncated, model_name=model_name, batch_size=batch_size, show_progress=False)171172            embeddings_records: list[dict] = []173            metadata_records: list[list] = []174            for i, row in enumerate(batch_df.itertuples(index=False)):175                embedding_id = str(uuid.uuid4())176                embeddings_records.append({177                    "embedding_id": embedding_id,178                    "source_id": row.section_id,179                    "source_type": "section",180                    "vector": vectors[i].tolist(),181                })182                metadata_records.append([183                    embedding_id,184                    "section",185                    row.section_id,186                    model_name,187                    model_info.get("provider", ""),188                    model_info.get("dimension", 0),189                ])190191            self.parquet.write_embeddings(embeddings_records, model_name)192193            meta_sql = (194                "INSERT OR REPLACE INTO embeddings_metadata "195                "(embedding_id, source_type, source_id, model_name, model_version, embedding_dimension) "196                "VALUES (?, ?, ?, ?, ?, ?)"197            )198            self.db.execute_many(meta_sql, metadata_records)199            embedded_count += len(batch_df)200201        logger.info("Embedded %d sections with model '%s'", embedded_count, model_name)202        return embedded_count203204    def get_vectors(205        self,206        source_ids: list[str],207        model_name: str | None = None,208    ) -> np.ndarray:209        model_name = model_name or self.model_name210        try:211            emb_df = self.parquet.read_df("embeddings")212        except Exception:213            logger.warning("No embeddings found in parquet store")214            return np.array([])215216        filtered = emb_df[217            (emb_df["source_id"].isin(source_ids)) & (emb_df["model_name"] == model_name)218        ]219        if filtered.empty:220            return np.array([])221222        vectors = np.stack(filtered["vector"].values)223        return vectors224225    def get_embedding_stats(self) -> dict:226        sql = """227            SELECT228                model_name,229                source_type,230                COUNT(*) as count,231                MIN(created_at) as earliest,232                MAX(created_at) as latest233            FROM embeddings_metadata234            GROUP BY model_name, source_type235            ORDER BY model_name, source_type236        """237        df = self.db.fetch_df(sql)238        stats: dict = {239            "total_embeddings": int(df["count"].sum()) if not df.empty else 0,240            "by_model": {},241        }242        for _, row in df.iterrows():243            model = row["model_name"]244            if model not in stats["by_model"]:245                stats["by_model"][model] = {}246            stats["by_model"][model][row["source_type"]] = {247                "count": int(row["count"]),248                "earliest": str(row["earliest"]),249                "latest": str(row["latest"]),250            }251        return stats252