spb/uqo-research-db
Public
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