spb/uqo-research-db

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 argparse20import sys21from pathlib import Path2223sys.path.insert(0, str(Path(__file__).resolve().parent.parent))2425from src.embeddings import EmbeddingManager, EmbeddingModelRegistry26from src.storage.database import DatabaseManager27from src.storage.parquet_manager import ParquetManager28from src.utils.logging import setup_logging2930logger = setup_logging("build_embeddings")313233def parse_args() -> argparse.Namespace:34    parser = argparse.ArgumentParser(35        description="Build embedding vectors for filing chunks or sections"36    )37    parser.add_argument(38        "--model",39        type=str,40        default="all-MiniLM-L6-v2",41        help="Embedding model name (default: all-MiniLM-L6-v2)",42    )43    parser.add_argument(44        "--batch-size",45        type=int,46        default=256,47        help="Batch size for encoding (default: 256)",48    )49    parser.add_argument(50        "--source",51        type=str,52        choices=["chunk", "section"],53        default="chunk",54        help="Source type to embed: chunk or section (default: chunk)",55    )56    parser.add_argument(57        "--accession",58        type=str,59        default=None,60        help="Comma-separated accession numbers to embed (default: all)",61    )62    return parser.parse_args()636465def main() -> None:66    args = parse_args()6768    registry = EmbeddingModelRegistry()69    logger.info("Available models:")70    for model in registry.list_models():71        default_marker = " [default]" if model.get("is_default") else ""72        logger.info(73            "  %s (dim=%d, provider=%s)%s",74            model["name"], model["dimension"], model["provider"], default_marker,75        )7677    db = DatabaseManager()78    db.initialize()79    parquet = ParquetManager()8081    manager = EmbeddingManager(db, parquet, model_name=args.model)8283    accession_numbers = None84    if args.accession:85        accession_numbers = [a.strip() for a in args.accession.split(",")]86        logger.info("Embedding specific accessions: %s", accession_numbers)8788    try:89        if args.source == "chunk":90            count = manager.embed_chunks(91                accession_numbers=accession_numbers,92                model_name=args.model,93                batch_size=args.batch_size,94            )95        else:96            count = manager.embed_sections(97                accession_numbers=accession_numbers,98                model_name=args.model,99            )100101        logger.info("Embedded %d %ss with model '%s'", count, args.source, args.model)102103    except NotImplementedError as e:104        logger.error("Model not supported: %s", e)105        sys.exit(1)106    except Exception:107        logger.error("Embedding failed", exc_info=True)108        sys.exit(1)109110    stats = manager.get_embedding_stats()111    logger.info("Embedding statistics:")112    logger.info("  Total embeddings: %d", stats["total_embeddings"])113    for model_name, sources in stats.get("by_model", {}).items():114        for source_type, info in sources.items():115            logger.info(116                "  %s / %s: %d vectors",117                model_name, source_type, info["count"],118            )119120    db.close()121    logger.info("Embedding build complete")122123124if __name__ == "__main__":125    main()126