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 json20import traceback21import uuid22from concurrent.futures import ProcessPoolExecutor, as_completed23from datetime import datetime, timezone24from pathlib import Path2526from tqdm import tqdm2728from src.parsing.chunker import TextChunker29from src.parsing.html_parser import HTMLCleaner30from src.parsing.section_extractor import SectionExtractor31from src.utils.config import get_parsing_config32from src.utils.constants import DATA_DIRS33from src.utils.logging import get_logger3435logger = get_logger("parsing.pipeline")363738def _parse_single(args: tuple[str, str, str, dict]) -> dict:39    """Top-level function so ProcessPoolExecutor can pickle it."""40    file_path_str, form_type, accession_number, config = args41    parser = FilingParser(config=config)42    return parser.parse_filing(Path(file_path_str), form_type, accession_number)434445class FilingParser:46    def __init__(self, config: dict | None = None) -> None:47        self._config = config or get_parsing_config()48        chunking_cfg = self._config.get("chunking", {})49        self._cleaner = HTMLCleaner(config=self._config)50        self._extractor = SectionExtractor()51        self._chunker = TextChunker(52            chunk_size=chunking_cfg.get("chunk_size", 512),53            overlap=chunking_cfg.get("chunk_overlap", 64),54            min_chunk_size=chunking_cfg.get("min_chunk_size", 50),55        )56        parallel_cfg = self._config.get("parallel", {})57        self._max_workers: int = parallel_cfg.get("max_workers", 4)58        self._output_dir: Path = DATA_DIRS["processed"]59        self._save_json: bool = bool(self._config.get("save_json", True))6061    def parse_filing(62        self, file_path: Path, form_type: str, accession_number: str63    ) -> dict:64        file_path = Path(file_path)65        logger.debug(66            "Parsing filing %s (%s) from %s", accession_number, form_type, file_path67        )6869        try:70            clean_text = self._cleaner.clean_file(file_path)71        except Exception:72            logger.error(73                "Failed to clean %s: %s", accession_number, traceback.format_exc()74            )75            return self._error_result(accession_number, form_type, "clean_failed")7677        if not clean_text:78            logger.warning("Empty text after cleaning: %s", accession_number)79            return self._error_result(accession_number, form_type, "empty_after_clean")8081        sections = self._extractor.extract_sections(clean_text, form_type)82        if not sections and clean_text:83            # no Item markers found (e.g. exhibit-only 10-Q/10-K): keep the whole text as one84            # real section so the chunks satisfy the filing_chunks -> filing_sections foreign key85            sections = [{86                "section_id": str(uuid.uuid4()),87                "section_name": "full_document",88                "section_order": 0,89                "item_number": "N/A",90                "text": clean_text.strip(),91                "word_count": len(clean_text.split()),92                "char_count": len(clean_text),93            }]94        for section in sections:95            section["accession_number"] = accession_number9697        chunks: list[dict] = []98        for section in sections:99            section_chunks = self._chunker.chunk_section(section)100            for chunk in section_chunks:101                chunk["accession_number"] = accession_number102            chunks.extend(section_chunks)103104        result = {105            "accession_number": accession_number,106            "form_type": form_type,107            "clean_text": clean_text,108            "sections": sections,109            "chunks": chunks,110            "metadata": {111                "word_count": len(clean_text.split()),112                "char_count": len(clean_text),113                "section_count": len(sections),114                "chunk_count": len(chunks),115                "parse_timestamp": datetime.now(timezone.utc).isoformat(),116            },117            "status": "success",118        }119120        if self._save_json:121            self._save_result(result)122        logger.debug(123            "Parsed %s: %d sections, %d chunks",124            accession_number,125            len(sections),126            len(chunks),127        )128        return result129130    def parse_batch(131        self, file_list: list[dict], max_workers: int | None = None132    ) -> list[dict]:133        max_workers = max_workers or self._max_workers134        results: list[dict] = []135        failed = 0136        succeeded = 0137138        tasks = [139            (140                str(item["file_path"]),141                item["form_type"],142                item["accession_number"],143                self._config,144            )145            for item in file_list146        ]147148        with ProcessPoolExecutor(max_workers=max_workers) as executor:149            futures = {150                executor.submit(_parse_single, task): task[2] for task in tasks151            }152153            with tqdm(total=len(futures), desc="Parsing filings", unit="filing") as pbar:154                for future in as_completed(futures):155                    accession = futures[future]156                    try:157                        result = future.result()158                        results.append(result)159                        if result.get("status") == "success":160                            succeeded += 1161                        else:162                            failed += 1163                    except Exception:164                        logger.error(165                            "Batch parse failed for %s: %s",166                            accession,167                            traceback.format_exc(),168                        )169                        results.append(170                            self._error_result(accession, "unknown", "exception")171                        )172                        failed += 1173                    pbar.update(1)174175        logger.info(176            "Batch complete: %d succeeded, %d failed out of %d total",177            succeeded,178            failed,179            len(file_list),180        )181        return results182183    def _save_result(self, result: dict) -> None:184        self._output_dir.mkdir(parents=True, exist_ok=True)185        output_path = self._output_dir / f"{result['accession_number']}.json"186187        serializable = {188            "accession_number": result["accession_number"],189            "form_type": result["form_type"],190            "sections": result["sections"],191            "chunks": result["chunks"],192            "metadata": result["metadata"],193            "status": result.get("status", "success"),194        }195196        with open(output_path, "w", encoding="utf-8") as f:197            json.dump(serializable, f, indent=2, ensure_ascii=False)198199    def _error_result(200        self, accession_number: str, form_type: str, error_type: str201    ) -> dict:202        return {203            "accession_number": accession_number,204            "form_type": form_type,205            "clean_text": "",206            "sections": [],207            "chunks": [],208            "metadata": {209                "word_count": 0,210                "char_count": 0,211                "section_count": 0,212                "chunk_count": 0,213                "parse_timestamp": datetime.now(timezone.utc).isoformat(),214            },215            "status": f"error:{error_type}",216        }217