# ============================================================ # Project: EDGAR Research Database # Description: SEC EDGAR Research Infrastructure # # Authors: # - Manel Kammoun — manel.kammoun@uqo.ca # - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca # - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca # # Institution: # Université du Québec en Outaouais (UQO) # # License: # Research / Academic Use # ============================================================ from __future__ import annotations import time import uuid from typing import TYPE_CHECKING import numpy as np import pyarrow as pa import pandas as pd from tqdm import tqdm from src.embeddings.models import EmbeddingModelRegistry from src.utils.logging import get_logger if TYPE_CHECKING: from src.storage.database import DatabaseManager from src.storage.parquet_manager import ParquetManager logger = get_logger("embeddings.manager") class EmbeddingManager: def __init__( self, db: DatabaseManager, parquet: ParquetManager, model_name: str | None = None, ) -> None: self.db = db self.parquet = parquet self.registry = EmbeddingModelRegistry() self.model_name = model_name or self.registry.default_model def _get_already_embedded_ids(self, model_name: str, source_type: str) -> set[str]: sql = ( "SELECT source_id FROM embeddings_metadata " "WHERE model_name = ? AND source_type = ?" ) df = self.db.fetch_df(sql, [model_name, source_type]) return set(df["source_id"].tolist()) def embed_chunks( self, accession_numbers: list[str] | None = None, model_name: str | None = None, batch_size: int = 256, fetch_size: int = 50_000, ) -> int: """Stream chunks that have no embedding yet, ``fetch_size`` rows at a time. Memory stays bounded whatever the corpus size; each fetch produces one parquet file under ``parquet/embeddings/model_name=/``. """ model_name = model_name or self.model_name model_info = self.registry.get_model_info(model_name) 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')" params: list = [model_name] if accession_numbers: placeholders = ", ".join(["?"] * len(accession_numbers)) where += f" AND c.accession_number IN ({placeholders})" params += list(accession_numbers) total = self.db.execute(f"SELECT COUNT(*) FROM filing_chunks c WHERE {where}", params).fetchone()[0] if total == 0: logger.info("All chunks already embedded with model '%s'", model_name) return 0 logger.info("Embedding %d chunks with model '%s' (fetch=%d, batch=%d)", total, model_name, fetch_size, batch_size) embedded_count = 0 t0 = time.time() while True: chunks_df = self.db.fetch_df( f"SELECT c.chunk_id, c.chunk_text FROM filing_chunks c WHERE {where} LIMIT {int(fetch_size)}", params, ) if chunks_df.empty: break texts = chunks_df["chunk_text"].fillna("").tolist() vectors = self.registry.encode(texts, model_name=model_name, batch_size=batch_size, show_progress=False) vectors = np.asarray(vectors, dtype=np.float32) ids = chunks_df["chunk_id"].tolist() emb_ids = [str(uuid.uuid4()) for _ in ids] table = pa.table({ "embedding_id": pa.array(emb_ids, pa.string()), "source_id": pa.array(ids, pa.string()), "source_type": pa.array(["chunk"] * len(ids), pa.string()), "vector": pa.FixedSizeListArray.from_arrays(pa.array(vectors.ravel(), pa.float32()), vectors.shape[1]).cast(pa.list_(pa.float32())), "model_name": pa.array([model_name] * len(ids), pa.string()), }) self.parquet.write_embeddings_table(table, model_name) meta_df = pd.DataFrame({ "embedding_id": emb_ids, "source_type": "chunk", "source_id": ids, "model_name": model_name, "model_version": model_info.get("provider", ""), "embedding_dimension": int(model_info.get("dimension", vectors.shape[1])), }) self.db.insert_df("embeddings_metadata", meta_df) embedded_count += len(ids) el = time.time() - t0 logger.info("embedded %d/%d chunks (%.0f chunks/s)", embedded_count, total, embedded_count / max(el, 1)) if len(ids) < fetch_size: break logger.info("Embedded %d chunks with model '%s'", embedded_count, model_name) return embedded_count def embed_sections( self, accession_numbers: list[str] | None = None, model_name: str | None = None, ) -> int: model_name = model_name or self.model_name model_info = self.registry.get_model_info(model_name) batch_size = 64 if accession_numbers: placeholders = ", ".join(["?"] * len(accession_numbers)) sql = f"SELECT section_id, text, accession_number, section_name FROM filing_sections WHERE accession_number IN ({placeholders})" sections_df = self.db.fetch_df(sql, accession_numbers) else: sql = "SELECT section_id, text, accession_number, section_name FROM filing_sections" sections_df = self.db.fetch_df(sql) if sections_df.empty: logger.info("No sections found to embed") return 0 already_embedded = self._get_already_embedded_ids(model_name, "section") mask = ~sections_df["section_id"].isin(already_embedded) sections_df = sections_df[mask].reset_index(drop=True) if sections_df.empty: logger.info("All sections already embedded with model '%s'", model_name) return 0 total = len(sections_df) logger.info("Embedding %d sections with model '%s'", total, model_name) embedded_count = 0 for start in tqdm(range(0, total, batch_size), desc="Embedding sections", unit="batch"): end = min(start + batch_size, total) batch_df = sections_df.iloc[start:end] texts = batch_df["text"].fillna("").tolist() max_tokens = model_info.get("max_tokens", 256) truncated = [" ".join(t.split()[:max_tokens]) for t in texts] vectors = self.registry.encode(truncated, model_name=model_name, batch_size=batch_size, show_progress=False) embeddings_records: list[dict] = [] metadata_records: list[list] = [] for i, row in enumerate(batch_df.itertuples(index=False)): embedding_id = str(uuid.uuid4()) embeddings_records.append({ "embedding_id": embedding_id, "source_id": row.section_id, "source_type": "section", "vector": vectors[i].tolist(), }) metadata_records.append([ embedding_id, "section", row.section_id, model_name, model_info.get("provider", ""), model_info.get("dimension", 0), ]) self.parquet.write_embeddings(embeddings_records, model_name) meta_sql = ( "INSERT OR REPLACE INTO embeddings_metadata " "(embedding_id, source_type, source_id, model_name, model_version, embedding_dimension) " "VALUES (?, ?, ?, ?, ?, ?)" ) self.db.execute_many(meta_sql, metadata_records) embedded_count += len(batch_df) logger.info("Embedded %d sections with model '%s'", embedded_count, model_name) return embedded_count def get_vectors( self, source_ids: list[str], model_name: str | None = None, ) -> np.ndarray: model_name = model_name or self.model_name try: emb_df = self.parquet.read_df("embeddings") except Exception: logger.warning("No embeddings found in parquet store") return np.array([]) filtered = emb_df[ (emb_df["source_id"].isin(source_ids)) & (emb_df["model_name"] == model_name) ] if filtered.empty: return np.array([]) vectors = np.stack(filtered["vector"].values) return vectors def get_embedding_stats(self) -> dict: sql = """ SELECT model_name, source_type, COUNT(*) as count, MIN(created_at) as earliest, MAX(created_at) as latest FROM embeddings_metadata GROUP BY model_name, source_type ORDER BY model_name, source_type """ df = self.db.fetch_df(sql) stats: dict = { "total_embeddings": int(df["count"].sum()) if not df.empty else 0, "by_model": {}, } for _, row in df.iterrows(): model = row["model_name"] if model not in stats["by_model"]: stats["by_model"][model] = {} stats["by_model"][model][row["source_type"]] = { "count": int(row["count"]), "earliest": str(row["earliest"]), "latest": str(row["latest"]), } return stats