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# Université du Québec en Outaouais (UQO)12#13# License:14# Research / Academic Use15# ============================================================1617from __future__ import annotations1819import json20import traceback21from concurrent.futures import ProcessPoolExecutor, as_completed22from datetime import datetime, timezone23from pathlib import Path2425from tqdm import tqdm2627from src.parsing.chunker import TextChunker28from src.parsing.html_parser import HTMLCleaner29from src.parsing.section_extractor import SectionExtractor30from src.utils.config import get_parsing_config31from src.utils.constants import DATA_DIRS32from src.utils.logging import get_logger3334logger = get_logger("parsing.pipeline")353637def _parse_single(args: tuple[str, str, str, dict]) -> dict:38 """Top-level function so ProcessPoolExecutor can pickle it."""39 file_path_str, form_type, accession_number, config = args40 parser = FilingParser(config=config)41 return parser.parse_filing(Path(file_path_str), form_type, accession_number)424344class FilingParser:45 def __init__(self, config: dict | None = None) -> None:46 self._config = config or get_parsing_config()47 chunking_cfg = self._config.get("chunking", {})48 self._cleaner = HTMLCleaner(config=self._config)49 self._extractor = SectionExtractor()50 self._chunker = TextChunker(51 chunk_size=chunking_cfg.get("chunk_size", 512),52 overlap=chunking_cfg.get("chunk_overlap", 64),53 min_chunk_size=chunking_cfg.get("min_chunk_size", 50),54 )55 parallel_cfg = self._config.get("parallel", {})56 self._max_workers: int = parallel_cfg.get("max_workers", 4)57 self._output_dir: Path = DATA_DIRS["processed"]58 self._save_json: bool = bool(self._config.get("save_json", True))5960 def parse_filing(61 self, file_path: Path, form_type: str, accession_number: str62 ) -> dict:63 file_path = Path(file_path)64 logger.debug(65 "Parsing filing %s (%s) from %s", accession_number, form_type, file_path66 )6768 try:69 clean_text = self._cleaner.clean_file(file_path)70 except Exception:71 logger.error(72 "Failed to clean %s: %s", accession_number, traceback.format_exc()73 )74 return self._error_result(accession_number, form_type, "clean_failed")7576 if not clean_text:77 logger.warning("Empty text after cleaning: %s", accession_number)78 return self._error_result(accession_number, form_type, "empty_after_clean")7980 sections = self._extractor.extract_sections(clean_text, form_type)81 for section in sections:82 section["accession_number"] = accession_number8384 chunks: list[dict] = []85 for section in sections:86 section_chunks = self._chunker.chunk_section(section)87 for chunk in section_chunks:88 chunk["accession_number"] = accession_number89 chunks.extend(section_chunks)9091 if not chunks and clean_text:92 fallback_chunks = self._chunker.chunk_text(clean_text)93 for chunk in fallback_chunks:94 chunk["section_id"] = "full_document"95 chunk["accession_number"] = accession_number96 chunks.extend(fallback_chunks)9798 result = {99 "accession_number": accession_number,100 "form_type": form_type,101 "clean_text": clean_text,102 "sections": sections,103 "chunks": chunks,104 "metadata": {105 "word_count": len(clean_text.split()),106 "char_count": len(clean_text),107 "section_count": len(sections),108 "chunk_count": len(chunks),109 "parse_timestamp": datetime.now(timezone.utc).isoformat(),110 },111 "status": "success",112 }113114 if self._save_json:115 self._save_result(result)116 logger.debug(117 "Parsed %s: %d sections, %d chunks",118 accession_number,119 len(sections),120 len(chunks),121 )122 return result123124 def parse_batch(125 self, file_list: list[dict], max_workers: int | None = None126 ) -> list[dict]:127 max_workers = max_workers or self._max_workers128 results: list[dict] = []129 failed = 0130 succeeded = 0131132 tasks = [133 (134 str(item["file_path"]),135 item["form_type"],136 item["accession_number"],137 self._config,138 )139 for item in file_list140 ]141142 with ProcessPoolExecutor(max_workers=max_workers) as executor:143 futures = {144 executor.submit(_parse_single, task): task[2] for task in tasks145 }146147 with tqdm(total=len(futures), desc="Parsing filings", unit="filing") as pbar:148 for future in as_completed(futures):149 accession = futures[future]150 try:151 result = future.result()152 results.append(result)153 if result.get("status") == "success":154 succeeded += 1155 else:156 failed += 1157 except Exception:158 logger.error(159 "Batch parse failed for %s: %s",160 accession,161 traceback.format_exc(),162 )163 results.append(164 self._error_result(accession, "unknown", "exception")165 )166 failed += 1167 pbar.update(1)168169 logger.info(170 "Batch complete: %d succeeded, %d failed out of %d total",171 succeeded,172 failed,173 len(file_list),174 )175 return results176177 def _save_result(self, result: dict) -> None:178 self._output_dir.mkdir(parents=True, exist_ok=True)179 output_path = self._output_dir / f"{result['accession_number']}.json"180181 serializable = {182 "accession_number": result["accession_number"],183 "form_type": result["form_type"],184 "sections": result["sections"],185 "chunks": result["chunks"],186 "metadata": result["metadata"],187 "status": result.get("status", "success"),188 }189190 with open(output_path, "w", encoding="utf-8") as f:191 json.dump(serializable, f, indent=2, ensure_ascii=False)192193 def _error_result(194 self, accession_number: str, form_type: str, error_type: str195 ) -> dict:196 return {197 "accession_number": accession_number,198 "form_type": form_type,199 "clean_text": "",200 "sections": [],201 "chunks": [],202 "metadata": {203 "word_count": 0,204 "char_count": 0,205 "section_count": 0,206 "chunk_count": 0,207 "parse_timestamp": datetime.now(timezone.utc).isoformat(),208 },209 "status": f"error:{error_type}",210 }211