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# Universite du Quebec en Outaouais (UQO)12#13# License:14# Research / Academic Use15# ============================================================1617from __future__ import annotations1819from typing import TYPE_CHECKING2021import numpy as np22import pandas as pd23import pyarrow as pa24import pyarrow.dataset as ds25import pyarrow.parquet as pq26from tqdm import tqdm2728from src.utils.config import get_config29from src.utils.constants import DATA_DIRS, PARQUET_COMPRESSION, PARQUET_COMPRESSION_LEVEL30from src.utils.logging import get_logger3132if TYPE_CHECKING:33 from src.storage.database import DatabaseManager3435logger = get_logger("factors.scoring")363738def _vectors_from_column(col) -> np.ndarray:39 """Decode a (chunked) list<float|double> Arrow column into a 2-D float32 array without Python loops.4041 ``flatten()`` (not ``.values``) is used so that sliced batches with a non-zero offset42 decode to exactly ``len(col)`` rows.43 """44 import pyarrow as pa45 if isinstance(col, pa.ChunkedArray):46 col = col.combine_chunks()47 n = len(col)48 if n == 0:49 return np.zeros((0, 0), dtype=np.float32)50 flat = col.flatten()51 arr = np.asarray(flat.to_numpy(zero_copy_only=False), dtype=np.float32)52 return arr.reshape(n, -1)535455class ScoringEngine:5657 def __init__(58 self,59 db: DatabaseManager,60 batch_size: int | None = None,61 ) -> None:62 self.db = db63 cfg = get_config("factors").get("factors", {}).get("scoring", {})64 self._batch_size = batch_size or cfg.get("batch_size", 10_000)6566 def _load_chunk_metadata(self) -> pd.DataFrame:67 sql = """68 SELECT69 c.chunk_id,70 c.accession_number,71 f.cik,72 f.ticker,73 f.form_type,74 f.filing_date,75 EXTRACT(YEAR FROM f.filing_date)::INTEGER AS fiscal_year,76 CEIL(EXTRACT(MONTH FROM f.filing_date) / 3.0)::INTEGER AS fiscal_quarter,77 COALESCE(s.section_name, 'full_body') AS section_name78 FROM filing_chunks c79 JOIN filings f ON c.accession_number = f.accession_number80 LEFT JOIN filing_sections s ON c.section_id = s.section_id81 """82 logger.info("Loading chunk metadata from DuckDB...")83 meta = self.db.fetch_df(sql)84 logger.info("Loaded metadata for %d chunks", len(meta))85 return meta8687 def _check_normalization(self, sample_vectors: np.ndarray) -> bool:88 norms = np.linalg.norm(sample_vectors, axis=1)89 mean_norm = norms.mean()90 is_normalized = 0.99 <= mean_norm <= 1.0191 logger.info(92 "Normalization check: mean norm=%.4f, pre-normalized=%s",93 mean_norm, is_normalized,94 )95 return is_normalized9697 def compute_scores(98 self,99 anchor_ids: list[str],100 anchor_matrix: np.ndarray,101 ) -> pd.DataFrame:102 anchor_matrix = anchor_matrix.astype(np.float32)103 chunk_meta = self._load_chunk_metadata()104 chunk_meta_indexed = chunk_meta.set_index("chunk_id")105106 emb_dir = DATA_DIRS["parquet"] / "embeddings"107 dataset = ds.dataset(str(emb_dir), format="parquet", partitioning="hive")108109 scanner = dataset.scanner(110 columns=["source_id", "vector"],111 batch_size=self._batch_size,112 )113114 all_ids: list[list[str]] = []115 all_scores: list[np.ndarray] = []116 normalization_checked = False117 is_prenormalized = True118 total_processed = 0119120 for batch in tqdm(scanner.to_batches(), desc="Scoring chunks", unit="batch"):121 if batch.num_rows == 0:122 continue123 source_ids = batch.column("source_id").to_pylist()124 vectors = _vectors_from_column(batch.column("vector"))125126 if not normalization_checked:127 is_prenormalized = self._check_normalization(vectors[:min(100, len(vectors))])128 normalization_checked = True129130 if not is_prenormalized:131 norms = np.linalg.norm(vectors, axis=1, keepdims=True)132 norms[norms == 0] = 1.0133 vectors = vectors / norms134135 scores = vectors @ anchor_matrix.T136137 all_ids.append(source_ids)138 all_scores.append(scores.astype(np.float32))139 total_processed += len(source_ids)140141 logger.info("Computed scores for %d chunks", total_processed)142143 flat_ids = []144 for id_list in all_ids:145 flat_ids.extend(id_list)146 score_matrix = np.vstack(all_scores)147148 score_cols = {f"score_{aid}": score_matrix[:, i] for i, aid in enumerate(anchor_ids)}149 scores_df = pd.DataFrame({"chunk_id": flat_ids, **score_cols})150151 scores_df = scores_df.merge(chunk_meta, on="chunk_id", how="inner")152 dropped = total_processed - len(scores_df)153 if dropped > 0:154 logger.warning("%d chunks had no metadata match (dropped)", dropped)155156 logger.info("Final chunk_scores: %d rows x %d columns", len(scores_df), len(scores_df.columns))157 return scores_df158159 def save_chunk_scores(self, scores_df: pd.DataFrame) -> None:160 anchor_ids = [c.replace("score_", "") for c in scores_df.columns if c.startswith("score_")]161162 self.db.execute("DROP TABLE IF EXISTS chunk_scores")163164 score_cols_ddl = ",\n ".join(165 f"score_{aid} DOUBLE" for aid in anchor_ids166 )167 ddl = f"""168 CREATE TABLE chunk_scores (169 chunk_id VARCHAR PRIMARY KEY,170 accession_number VARCHAR,171 cik VARCHAR,172 ticker VARCHAR,173 form_type VARCHAR,174 filing_date DATE,175 fiscal_year INTEGER,176 fiscal_quarter INTEGER,177 section_name VARCHAR,178 {score_cols_ddl},179 created_at TIMESTAMP DEFAULT current_timestamp180 )181 """182 self.db.execute(ddl)183184 cols = [185 "chunk_id", "accession_number", "cik", "ticker", "form_type",186 "filing_date", "fiscal_year", "fiscal_quarter", "section_name",187 ] + [f"score_{aid}" for aid in anchor_ids]188 insert_df = scores_df[cols].copy()189190 col_list = ", ".join(cols)191 self.db.conn.execute(f"INSERT INTO chunk_scores ({col_list}) SELECT * FROM insert_df")192193 self.db.execute("CREATE INDEX IF NOT EXISTS idx_cs_accession ON chunk_scores(accession_number)")194 self.db.execute("CREATE INDEX IF NOT EXISTS idx_cs_ticker_year ON chunk_scores(ticker, fiscal_year)")195196 logger.info("Saved %d rows to DuckDB chunk_scores", len(insert_df))197198 out_dir = DATA_DIRS["parquet"] / "chunk_scores"199 out_dir.mkdir(parents=True, exist_ok=True)200 table = pa.Table.from_pandas(insert_df, preserve_index=False)201 pq.write_to_dataset(202 table,203 root_path=str(out_dir),204 partition_cols=["form_type"],205 compression=PARQUET_COMPRESSION,206 compression_level=PARQUET_COMPRESSION_LEVEL,207 existing_data_behavior="overwrite_or_ignore",208 )209 logger.info("Exported chunk_scores to parquet")210