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 annotations1819import argparse20import sys21import time22from pathlib import Path2324sys.path.insert(0, str(Path(__file__).resolve().parent.parent))2526from src.embeddings.models import EmbeddingModelRegistry27from src.factors.anchors import AnchorRegistry28from src.factors.scoring import ScoringEngine29from src.factors.aggregation import FactorAggregator30from src.factors.longitudinal import LongitudinalBuilder31from src.factors.panels import PanelExporter32from src.factors.diagnostics import FactorDiagnostics33from src.storage.database import DatabaseManager34from src.utils.logging import setup_logging3536logger = setup_logging("build_factors")3738STEPS = ["all", "anchors", "scores", "aggregate", "longitudinal", "panels", "diagnostics"]394041def parse_args() -> argparse.Namespace:42 parser = argparse.ArgumentParser(43 description="Build thematic factor scores from SEC filing embeddings"44 )45 parser.add_argument(46 "--step",47 type=str,48 default="all",49 choices=STEPS,50 help="Pipeline step to run (default: all)",51 )52 parser.add_argument(53 "--batch-size",54 type=int,55 default=10_000,56 help="Batch size for scoring (default: 10000)",57 )58 parser.add_argument(59 "--skip-drift",60 action="store_true",61 default=False,62 help="Skip embedding drift computation (slower step)",63 )64 parser.add_argument(65 "--skip-novelty",66 action="store_true",67 default=False,68 help="Skip novelty score computation",69 )70 return parser.parse_args()717273def run_step(name: str) -> None:74 logger.info("=" * 60)75 logger.info("STEP: %s", name.upper())76 logger.info("=" * 60)777879def main() -> None:80 args = parse_args()81 t0 = time.time()8283 db = DatabaseManager()84 db.initialize()85 registry = EmbeddingModelRegistry()8687 anchor_reg = AnchorRegistry(registry)88 anchor_ids, anchor_matrix = anchor_reg.get_anchor_matrix()89 logger.info("Anchor IDs: %s", anchor_ids)90 logger.info("Anchor matrix shape: %s", anchor_matrix.shape)9192 # ---- Step 1: Anchors ----93 if args.step in ("all", "anchors"):94 run_step("anchors")95 anchor_reg.save_to_duckdb(db)96 anchor_reg.save_to_parquet()97 logger.info("Anchors embedded and saved")9899 # ---- Step 2: Scoring ----100 if args.step in ("all", "scores"):101 run_step("scoring")102 engine = ScoringEngine(db, batch_size=args.batch_size)103 scores_df = engine.compute_scores(anchor_ids, anchor_matrix)104 engine.save_chunk_scores(scores_df)105 del scores_df106 logger.info("Chunk scoring complete")107108 # ---- Step 3: Aggregation ----109 if args.step in ("all", "aggregate"):110 run_step("aggregation")111 agg = FactorAggregator(db)112113 filing_df = agg.aggregate_to_filing(anchor_ids)114 agg.save_filing_scores(filing_df)115116 quarterly_df = agg.aggregate_to_quarterly(anchor_ids)117 agg.save_quarterly_panel(quarterly_df)118119 annual_df = agg.aggregate_to_annual(anchor_ids)120 agg.save_annual_panel(annual_df)121 logger.info("Aggregation complete")122123 # ---- Step 4: Longitudinal ----124 drift_df = None125 novelty_df = None126 if args.step in ("all", "longitudinal"):127 run_step("longitudinal")128 longi = LongitudinalBuilder(db)129130 if not args.skip_drift:131 drift_df = longi.compute_embedding_drift()132 logger.info("Embedding drift computed")133 else:134 logger.info("Skipping embedding drift (--skip-drift)")135136 if not args.skip_novelty:137 novelty_df = longi.compute_novelty()138 logger.info("Novelty scores computed")139 else:140 logger.info("Skipping novelty (--skip-novelty)")141142 # ---- Step 5: Panels ----143 if args.step in ("all", "panels"):144 run_step("panels")145 exporter = PanelExporter(db)146147 annual = exporter.build_annual_panel(anchor_ids, drift_df, novelty_df)148 quarterly = exporter.build_quarterly_panel(anchor_ids)149 exporter.export_all(annual, quarterly)150 logger.info("Panel export complete")151152 # ---- Step 6: Diagnostics ----153 if args.step in ("all", "diagnostics"):154 run_step("diagnostics")155 try:156 annual = db.fetch_df("SELECT * FROM annual_final")157 except Exception:158 annual = db.fetch_df("SELECT * FROM annual_panel")159 companies = db.fetch_df("SELECT ticker, sector, sub_industry FROM companies")160 annual = annual.merge(companies, on="ticker", how="left")161162 diag = FactorDiagnostics(db)163 diag.save_diagnostics(annual, anchor_ids)164 logger.info("Diagnostics complete")165166 elapsed = time.time() - t0167 logger.info("=" * 60)168 logger.info("PIPELINE COMPLETE in %.1f seconds (%.1f minutes)", elapsed, elapsed / 60)169 logger.info("=" * 60)170171 stats = {}172 for table in ["anchor_embeddings", "chunk_scores", "filing_scores", "quarterly_panel", "annual_panel"]:173 try:174 count = db.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0]175 stats[table] = count176 except Exception:177 stats[table] = "N/A"178 logger.info("Final table counts:")179 for table, count in stats.items():180 logger.info(" %-25s %s rows", table, f"{count:,}" if isinstance(count, int) else count)181182 db.close()183184185if __name__ == "__main__":186 main()187