# ============================================================ # 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: # Universite du Quebec en Outaouais (UQO) # # License: # Research / Academic Use # ============================================================ from __future__ import annotations from typing import TYPE_CHECKING import numpy as np import pandas as pd import pyarrow as pa import pyarrow.dataset as ds import pyarrow.parquet as pq from tqdm import tqdm from src.utils.config import get_config from src.utils.constants import DATA_DIRS, PARQUET_COMPRESSION, PARQUET_COMPRESSION_LEVEL from src.utils.logging import get_logger if TYPE_CHECKING: from src.storage.database import DatabaseManager logger = get_logger("factors.scoring") def _vectors_from_column(col) -> np.ndarray: """Decode a (chunked) list Arrow column into a 2-D float32 array without Python loops. ``flatten()`` (not ``.values``) is used so that sliced batches with a non-zero offset decode to exactly ``len(col)`` rows. """ import pyarrow as pa if isinstance(col, pa.ChunkedArray): col = col.combine_chunks() n = len(col) if n == 0: return np.zeros((0, 0), dtype=np.float32) flat = col.flatten() arr = np.asarray(flat.to_numpy(zero_copy_only=False), dtype=np.float32) return arr.reshape(n, -1) class ScoringEngine: def __init__( self, db: DatabaseManager, batch_size: int | None = None, ) -> None: self.db = db cfg = get_config("factors").get("factors", {}).get("scoring", {}) self._batch_size = batch_size or cfg.get("batch_size", 10_000) def _load_chunk_metadata(self) -> pd.DataFrame: sql = """ SELECT c.chunk_id, c.accession_number, f.cik, f.ticker, f.form_type, f.filing_date, EXTRACT(YEAR FROM f.filing_date)::INTEGER AS fiscal_year, CEIL(EXTRACT(MONTH FROM f.filing_date) / 3.0)::INTEGER AS fiscal_quarter, COALESCE(s.section_name, 'full_body') AS section_name FROM filing_chunks c JOIN filings f ON c.accession_number = f.accession_number LEFT JOIN filing_sections s ON c.section_id = s.section_id """ logger.info("Loading chunk metadata from DuckDB...") meta = self.db.fetch_df(sql) logger.info("Loaded metadata for %d chunks", len(meta)) return meta def _check_normalization(self, sample_vectors: np.ndarray) -> bool: norms = np.linalg.norm(sample_vectors, axis=1) mean_norm = norms.mean() is_normalized = 0.99 <= mean_norm <= 1.01 logger.info( "Normalization check: mean norm=%.4f, pre-normalized=%s", mean_norm, is_normalized, ) return is_normalized def compute_scores( self, anchor_ids: list[str], anchor_matrix: np.ndarray, ) -> pd.DataFrame: anchor_matrix = anchor_matrix.astype(np.float32) chunk_meta = self._load_chunk_metadata() chunk_meta_indexed = chunk_meta.set_index("chunk_id") emb_dir = DATA_DIRS["parquet"] / "embeddings" dataset = ds.dataset(str(emb_dir), format="parquet", partitioning="hive") scanner = dataset.scanner( columns=["source_id", "vector"], batch_size=self._batch_size, ) all_ids: list[list[str]] = [] all_scores: list[np.ndarray] = [] normalization_checked = False is_prenormalized = True total_processed = 0 for batch in tqdm(scanner.to_batches(), desc="Scoring chunks", unit="batch"): if batch.num_rows == 0: continue source_ids = batch.column("source_id").to_pylist() vectors = _vectors_from_column(batch.column("vector")) if not normalization_checked: is_prenormalized = self._check_normalization(vectors[:min(100, len(vectors))]) normalization_checked = True if not is_prenormalized: norms = np.linalg.norm(vectors, axis=1, keepdims=True) norms[norms == 0] = 1.0 vectors = vectors / norms scores = vectors @ anchor_matrix.T all_ids.append(source_ids) all_scores.append(scores.astype(np.float32)) total_processed += len(source_ids) logger.info("Computed scores for %d chunks", total_processed) flat_ids = [] for id_list in all_ids: flat_ids.extend(id_list) score_matrix = np.vstack(all_scores) score_cols = {f"score_{aid}": score_matrix[:, i] for i, aid in enumerate(anchor_ids)} scores_df = pd.DataFrame({"chunk_id": flat_ids, **score_cols}) scores_df = scores_df.merge(chunk_meta, on="chunk_id", how="inner") dropped = total_processed - len(scores_df) if dropped > 0: logger.warning("%d chunks had no metadata match (dropped)", dropped) logger.info("Final chunk_scores: %d rows x %d columns", len(scores_df), len(scores_df.columns)) return scores_df def save_chunk_scores(self, scores_df: pd.DataFrame) -> None: anchor_ids = [c.replace("score_", "") for c in scores_df.columns if c.startswith("score_")] self.db.execute("DROP TABLE IF EXISTS chunk_scores") score_cols_ddl = ",\n ".join( f"score_{aid} DOUBLE" for aid in anchor_ids ) ddl = f""" CREATE TABLE chunk_scores ( chunk_id VARCHAR PRIMARY KEY, accession_number VARCHAR, cik VARCHAR, ticker VARCHAR, form_type VARCHAR, filing_date DATE, fiscal_year INTEGER, fiscal_quarter INTEGER, section_name VARCHAR, {score_cols_ddl}, created_at TIMESTAMP DEFAULT current_timestamp ) """ self.db.execute(ddl) cols = [ "chunk_id", "accession_number", "cik", "ticker", "form_type", "filing_date", "fiscal_year", "fiscal_quarter", "section_name", ] + [f"score_{aid}" for aid in anchor_ids] insert_df = scores_df[cols].copy() col_list = ", ".join(cols) self.db.conn.execute(f"INSERT INTO chunk_scores ({col_list}) SELECT * FROM insert_df") self.db.execute("CREATE INDEX IF NOT EXISTS idx_cs_accession ON chunk_scores(accession_number)") self.db.execute("CREATE INDEX IF NOT EXISTS idx_cs_ticker_year ON chunk_scores(ticker, fiscal_year)") logger.info("Saved %d rows to DuckDB chunk_scores", len(insert_df)) out_dir = DATA_DIRS["parquet"] / "chunk_scores" out_dir.mkdir(parents=True, exist_ok=True) table = pa.Table.from_pandas(insert_df, preserve_index=False) pq.write_to_dataset( table, root_path=str(out_dir), partition_cols=["form_type"], compression=PARQUET_COMPRESSION, compression_level=PARQUET_COMPRESSION_LEVEL, existing_data_behavior="overwrite_or_ignore", ) logger.info("Exported chunk_scores to parquet")