#!/usr/bin/env python3 from __future__ import annotations """ IRIS — IR Document Training Script ===================================== Trains the system on any IR PDF to reduce retrieval errors and hallucination. This script does TWO distinct kinds of "training": ───────────────────────────────────────────────────────────────── A. RETRIEVAL TRAINING (always runs) ───────────────────────────────────────────────────────────────── 1. Parse PDF → extract text, tables, section headings 2. Auto-generate page → section → KPI ground-truth mapping 3. Build (query, correct_page) training pairs from the mapping 4. Calibrate retrieval scoring weights via grid search (MRR metric) - metric_match_weight - section_heading_boost - caption_boost - min_confidence_threshold 5. Save calibrated weights to data/retrieval_config.json 6. Re-embed text chunks with updated config in ChromaDB 7. Re-index tables with enriched metadata ───────────────────────────────────────────────────────────────── B. EMBEDDING FINE-TUNING (optional flag --fine-tune) ───────────────────────────────────────────────────────────────── 8. Build contrastive training pairs: Positive : (KPI query, correct page text chunk) Hard negatives: (KPI query, wrong-section chunks) 9. Fine-tune BAAI/bge-small-en-v1.5 using MultipleNegativesRankingLoss (sentence-transformers) 10. Save fine-tuned model to data/models/_embed/ 11. Re-embed all chunks with the fine-tuned model 12. Re-index ChromaDB with new embeddings ───────────────────────────────────────────────────────────────── C. ANTI-HALLUCINATION HARDENING (always runs, step 13) ───────────────────────────────────────────────────────────────── 13. Extract EXACT numeric KPI values from verified tables 14. Save to data/kpi_ground_truth.json 15. Update generation agent to cross-check produced numbers against ground truth before returning a response 16. Apply confidence threshold — below threshold returns "Insufficient evidence" instead of a hallucinated answer ───────────────────────────────────────────────────────────────── D. INGESTION PIPELINE (always runs) ───────────────────────────────────────────────────────────────── 17. Render PDF pages to PNG images 18. Generate ColPali visual embeddings (batch_size=1, MPS safe) Usage: python train_document.py --pdf ../documents/enbd_q1_2026.pdf python train_document.py --pdf ../documents/enbd_q1_2026.pdf --fine-tune python train_document.py --pdf ../documents/enbd_q1_2026.pdf --skip-colpali """ import argparse import json import logging import math import os import re import sys import time from dataclasses import dataclass from pathlib import Path from typing import Optional sys.path.insert(0, str(Path(__file__).parent)) logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)-8s %(name)s %(message)s", datefmt="%H:%M:%S", ) logger = logging.getLogger("iris.train") BASE_DIR = Path(__file__).parent DATA_DIR = BASE_DIR / "data" PAGES_DIR = DATA_DIR / "pages" COLPALI_DIR = DATA_DIR / "colpali_index" CHROMA_DIR = DATA_DIR / "chroma" TABLES_FILE = DATA_DIR / "tables.json" DOCS_FILE = DATA_DIR / "documents.json" PROC_DIR = DATA_DIR / "processed" PAGEMAP_DIR = DATA_DIR / "page_maps" MODELS_DIR = DATA_DIR / "models" CONFIG_FILE = DATA_DIR / "retrieval_config.json" KPI_GT_FILE = DATA_DIR / "kpi_ground_truth.json" # ── Default retrieval weights (overridden by calibration) ──────────────────── DEFAULT_CONFIG = { "embed_model": "BAAI/bge-small-en-v1.5", "chunk_size": 300, "chunk_overlap": 60, "top_k_text": 6, "top_k_tables": 4, "top_k_visual": 3, "metric_match_weight": 3.0, "section_heading_boost": 4.0, "caption_boost": 2.0, "term_overlap_weight": 0.3, "min_confidence_threshold": 0.20, "ollama_temperature": 0.05, "ollama_num_predict": 2048, "generation_grounding": True, "reject_below_threshold": True, } # ── Section heading detection patterns ─────────────────────────────────────── SECTION_PATTERNS = [ (r"funding.*liquidity", "Liquidity"), (r"liquidity coverage ratio", "Liquidity"), (r"advances.*deposit.*ratio", "Liquidity"), (r"liquid assets.*aed", "Liquidity"), (r"\blcr\b.*\badr\b", "Liquidity"), (r"income statement", "Income Statement"), (r"profit before tax", "Income Statement"), (r"aed\s+[\d.]+bn.*profit", "Income Statement"), (r"net interest margin", "Net Interest Margin"), (r"margins remain", "Net Interest Margin"), (r"non[- ]funded income", "Non-Funded Income"), (r"loan growth", "Loans & Deposits"), (r"gross loan", "Loans & Deposits"), (r"deposit growth", "Loans & Deposits"), (r"robust credit quality", "Asset Quality"), (r"npl ratio", "Asset Quality"), (r"coverage ratio", "Asset Quality"), (r"cost of risk", "Asset Quality"), (r"cost[- ]to[- ]income", "Cost to Income"), (r"operating expense", "Cost to Income"), (r"common equity tier", "Capital Adequacy"), (r"cet[- ]?1.*ratio", "Capital Adequacy"), (r"capital adequacy", "Capital Adequacy"), (r"divisional performance", "Divisional Performance"), (r"\besg\b", "ESG"), (r"sustainability", "ESG"), (r"gdp.*growth", "Economic Environment"), (r"denizbank", "DenizBank / Türkiye"), (r"hyperinflation", "DenizBank / Türkiye"), (r"credit rating", "Credit Ratings"), (r"investment case", "Investment Case"), (r"financial results.*q", "Financial Appendix"), ] SECTION_BY_PAGE = { 1: "Cover", 2: "Important Information", 3: "Economic Environment", 4: "Economic Environment", 5: "Economic Environment", 6: "Group Overview", 7: "Group Overview", 8: "International Presence", 9: "Credit Ratings", 10: "Shareholder Base", 11: "Investment Case", 12: "Peer Comparison", 13: "Profitability", 14: "Financial & Operating Performance", 15: "Executive Summary", 16: "Income Statement", 17: "Net Interest Margin", 18: "Non-Funded Income", 19: "Loans & Deposits", 20: "Asset Quality", 21: "Cost to Income", 22: "Liquidity", 23: "Capital Adequacy", 24: "Divisional Performance", 25: "ESG", 26: "ESG", 27: "ESG", 28: "ESG", 29: "Appendix", 30: "Financial Results", 31: "USD Translation", 32: "Hyperinflation", 33: "Turkey Macro", 34: "Egypt Macro", 35: "KSA Macro", 36: "Contact", } SECTION_KPI_TAGS = { "Income Statement": ["net profit", "revenue", "nim", "cor", "cost_income"], "Net Interest Margin": ["nim"], "Non-Funded Income": ["revenue"], "Loans & Deposits": ["loans", "deposits"], "Asset Quality": ["npl", "cor"], "Cost to Income": ["cost_income"], "Liquidity": ["lcr", "deposits"], "Capital Adequacy": ["car"], "Divisional Performance": ["net profit", "revenue", "nim", "cor", "npl"], "Group Overview": ["net profit", "revenue", "deposits", "loans", "car"], "Financial Appendix": ["net profit", "revenue", "nim", "cor"], } # ── Canonical queries per KPI for training pair generation ─────────────────── KPI_CANONICAL_QUERIES = { "Income Statement": [ "net profit", "profit before tax", "total income", "operating profit", "group earnings", "net earnings", "profit after tax", "PAT", "profit growth year on year", "Q1 2026 profitability", ], "Net Interest Margin": [ "net interest margin", "NIM", "NIM trend", "interest margin performance", "margin compression", "NII growth", "net interest income", ], "Non-Funded Income": [ "non-funded income", "NFI", "fee income", "non-interest income", "fee and commission income", "trading income", ], "Loans & Deposits": [ "loan growth", "deposit growth", "advances growth", "loan book", "credit growth", "total loans", "customer deposits", ], "Asset Quality": [ "NPL ratio", "non-performing loans", "cost of risk", "credit quality", "coverage ratio", "impairment charges", "provisions", ], "Cost to Income": [ "cost to income", "cost income ratio", "operating efficiency", "operating expenses", "CIR", ], "Liquidity": [ "liquidity coverage ratio", "LCR", "advances to deposit ratio", "ADR", "liquid assets", "funding structure", "liquidity position", "what is the lcr", "lcr in q1 2026", "liquidity coverage ratio in q1 2026", "how strong is the lcr", "lcr performance", ], "Capital Adequacy": [ "capital adequacy ratio", "CAR", "CET1", "CET-1", "tier 1 ratio", "capital position", "common equity tier 1", "regulatory capital", ], "Divisional Performance": [ "divisional performance", "business segments", "retail banking", "corporate banking", "DenizBank performance", "segment results", ], } # ═══════════════════════════════════════════════════════════════════════════════ # A. RETRIEVAL TRAINING # ═══════════════════════════════════════════════════════════════════════════════ def detect_section(page_text: str) -> str: text_lower = page_text.lower() for pattern, section in SECTION_PATTERNS: if re.search(pattern, text_lower): return section return "" def detect_slide_number(page_text: str) -> Optional[int]: lines = [l.strip() for l in page_text.split("\n") if l.strip()] for line in lines[:4]: if re.match(r"^\d{1,3}$", line): try: return int(line) except ValueError: pass return None def load_slide_directory() -> dict[int, dict]: slide_dir_file = DATA_DIR / "slide_directory_index.json" if not slide_dir_file.exists(): return {} with open(slide_dir_file, "r", encoding="utf-8") as f: rows = json.load(f) return { int(row["slide"]): row for row in rows if str(row.get("slide", "")).isdigit() } def split_slide_list(value: str) -> list[str]: return [item.strip() for item in str(value or "").split(",") if item.strip()] def _unique_terms(terms: list[str]) -> list[str]: seen = set() result = [] for term in terms: cleaned = re.sub(r"\s+", " ", str(term or "")).strip() if len(cleaned) < 3: continue key = cleaned.lower() if key not in seen: seen.add(key) result.append(cleaned) return result def metadata_query_terms(info: dict) -> list[str]: """ Convert every slide's curated KPI, mapping item, description, and finance synonym metadata into retrieval training queries. This keeps training aligned with the slide directory instead of relying only on broad sections. """ terms: list[str] = [] for key in ("kpis", "mapping_items", "synonyms", "kpi_tags"): values = info.get(key, []) or [] if isinstance(values, str): values = split_slide_list(values) terms.extend(values) description = str(info.get("description", "") or "").strip() if description: terms.append(description) for fragment in re.split(r"[.;:]", description): fragment = fragment.strip() if 12 <= len(fragment) <= 140: terms.append(fragment) title = str(info.get("title", "") or "").strip() if title and not re.match(r"^\d+$", title): terms.append(title) return _unique_terms(terms) def build_page_map(extracted) -> dict: page_map = {} slide_directory = load_slide_directory() for page in extracted.pages: section = SECTION_BY_PAGE.get(page.page_number) or getattr(page, "section_heading", "") or detect_section(page.text) slide = getattr(page, "slide_number", None) or detect_slide_number(page.text) or page.page_number slide_meta = slide_directory.get(page.page_number) or slide_directory.get(slide) lines = [l.strip() for l in page.text.split("\n") if l.strip() and len(l.strip()) > 8 and not re.match(r"^\d+$", l.strip())] title = lines[0][:120] if lines else "" kpis = list(SECTION_KPI_TAGS.get(section, [])) if slide_meta: slide = int(slide_meta.get("slide") or page.page_number) kpis = split_slide_list(slide_meta.get("kpis", "")) page_map[page.page_number] = { "section": section, "slide": slide, "kpis": kpis, "title": title, "text_preview": page.text[:200].replace("\n", " ").strip(), } if slide_meta: page_map[page.page_number].update({ "period": slide_meta.get("period", ""), "mapping_items": split_slide_list(slide_meta.get("topics", "")), "synonyms": split_slide_list(slide_meta.get("synonyms", "")), "description": slide_meta.get("description", ""), "visual_layout": slide_meta.get("visual_layout", ""), }) return page_map @dataclass class TrainingPair: query: str positive_page: int positive_text: str negative_pages: list negative_texts: list section: str kpi_tag: str def generate_training_pairs(page_map: dict, chunks: list) -> list[TrainingPair]: """ Generate (query, positive_chunk, hard_negative_chunks) training pairs. Positive = chunk from the correct section page. Hard negative = chunk from a different section with overlapping keywords. """ # Group chunks by page chunks_by_page: dict[int, list] = {} for c in chunks: chunks_by_page.setdefault(c.page_number, []).append(c) # Group pages by section pages_by_section: dict[str, list[int]] = {} for pg, info in page_map.items(): sec = info["section"] if sec: pages_by_section.setdefault(sec, []).append(int(pg)) def page_info_for(page_number: int) -> dict: return page_map.get(page_number) or page_map.get(str(page_number)) or {} # Shared hard negatives are selected from other sections with financial data. all_financial_negatives = [ c for c in chunks if getattr(c, "has_financial_data", False) ] pairs = [] for section, queries in KPI_CANONICAL_QUERIES.items(): correct_pages = pages_by_section.get(section, []) if not correct_pages: continue # Positive chunks = all chunks from pages in this section pos_chunks = [] for pg in correct_pages: pos_chunks.extend(chunks_by_page.get(pg, [])) if not pos_chunks: continue # Hard negative chunks = chunks from other sections with financial keywords neg_chunks = [] for other_sec, other_pages in pages_by_section.items(): if other_sec == section: continue for pg in other_pages: neg_chunks.extend(chunks_by_page.get(pg, [])) # Keep only negatives that have financial content (hard negatives) neg_chunks = [c for c in neg_chunks if c.has_financial_data][:10] for query in queries: # Use the best positive chunk (most financial keywords) best_pos = max(pos_chunks, key=lambda c: len(c.financial_keywords_found)) pairs.append(TrainingPair( query=query, positive_page=best_pos.page_number, positive_text=best_pos.text, negative_pages=[c.page_number for c in neg_chunks[:5]], negative_texts=[c.text for c in neg_chunks[:5]], section=section, kpi_tag=section.lower().replace(" ", "_"), )) # Add page-specific slide training from curated slide metadata. These pairs # are intentionally tied to the exact slide, so queries such as ESG ratings, # green bond framework, or business segment performance cannot drift to NIM. existing = {(p.query.lower(), p.positive_page) for p in pairs} for pg, info in page_map.items(): page_chunks = chunks_by_page.get(int(pg), []) if not page_chunks: continue terms = metadata_query_terms(info) if not terms: continue best_pos = max( page_chunks, key=lambda c: ( len(getattr(c, "financial_keywords_found", []) or []), len(getattr(c, "text", "") or ""), ), ) section = info.get("section", "") neg_chunks = [ c for c in all_financial_negatives if c.page_number != best_pos.page_number and page_info_for(c.page_number).get("section") != section ][:10] for query in terms: key = (query.lower(), best_pos.page_number) if key in existing: continue existing.add(key) pairs.append(TrainingPair( query=query, positive_page=best_pos.page_number, positive_text=best_pos.text, negative_pages=[c.page_number for c in neg_chunks[:5]], negative_texts=[c.text for c in neg_chunks[:5]], section=section, kpi_tag=section.lower().replace(" ", "_"), )) logger.info(f" Generated {len(pairs)} training pairs from {len(pages_by_section)} sections") return pairs def calibrate_weights(pairs: list[TrainingPair], chunks: list) -> dict: """ Grid search over retrieval scoring weights. Evaluates each weight combination using Mean Reciprocal Rank (MRR). MRR = average of 1/rank for each query where rank = position of correct page. Returns the weight dict that maximises MRR. """ from sentence_transformers import SentenceTransformer logger.info("\n [Calibration] Loading embedding model for weight calibration…") embed_model = SentenceTransformer("BAAI/bge-small-en-v1.5") # Pre-compute chunk embeddings chunk_texts = [c.text for c in chunks] chunk_embs = embed_model.encode(chunk_texts, normalize_embeddings=True, show_progress_bar=False) # Pre-compute query embeddings once. The grid search reuses the same # training queries for every weight combination, so re-encoding them inside # the scoring loop makes training unnecessarily slow. pair_query_embs = embed_model.encode( [pair.query for pair in pairs], batch_size=64, normalize_embeddings=True, show_progress_bar=False, ) # Build quick lookup: page → chunks index page_to_idx: dict[int, list[int]] = {} for i, c in enumerate(chunks): page_to_idx.setdefault(c.page_number, []).append(i) def score_retrieval(pair: TrainingPair, q_emb, metric_w: float, section_boost: float) -> int: """ Score a single query and return the rank of the correct page. Lower rank = better. Returns 99 if correct page not in top-10. """ import numpy as np # Score each chunk: cosine_sim + section_boost if section matches scored: list[tuple[float, int]] = [] # (score, page_number) for i, c in enumerate(chunks): cos_sim = float(np.dot(q_emb, chunk_embs[i])) # Section heading boost section_match = ( pair.section.lower() in (c.section_heading or "").lower() or (c.section_heading or "").lower() in pair.section.lower() ) score = cos_sim + (section_boost * 0.1 if section_match else 0.0) scored.append((score, c.page_number)) # Sort by score descending, get unique page order scored.sort(key=lambda x: -x[0]) seen_pages = [] for _, pg in scored: if pg not in seen_pages: seen_pages.append(pg) if len(seen_pages) >= 10: break if pair.positive_page in seen_pages: return seen_pages.index(pair.positive_page) + 1 return 99 def compute_mrr(metric_w: float, section_boost: float) -> float: rr_sum = 0.0 for pair, q_emb in zip(pairs, pair_query_embs): rank = score_retrieval(pair, q_emb, metric_w, section_boost) rr_sum += 1.0 / rank return rr_sum / len(pairs) # Grid search logger.info(" [Calibration] Running grid search over weight combinations…") best_mrr = 0.0 best_config = {"metric_match_weight": 3.0, "section_heading_boost": 4.0} metric_weights = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] section_boosts = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 8.0] caption_boosts = [1.0, 2.0, 3.0] confidence_thrs = [0.10, 0.15, 0.20, 0.25, 0.30] total_combinations = len(metric_weights) * len(section_boosts) evaluated = 0 for mw in metric_weights: for sb in section_boosts: mrr = compute_mrr(mw, sb) evaluated += 1 if mrr > best_mrr: best_mrr = mrr best_config["metric_match_weight"] = mw best_config["section_heading_boost"] = sb logger.info( f" [Calibration] New best MRR={mrr:.4f} " f"(metric_w={mw}, section_boost={sb}) " f"[{evaluated}/{total_combinations}]" ) # Find best confidence threshold (based on max score distribution) logger.info(" [Calibration] Calibrating confidence threshold…") import numpy as np q_embs = embed_model.encode([p.query for p in pairs], normalize_embeddings=True) max_scores = [] for qi, pair in enumerate(pairs): sims = chunk_embs @ q_embs[qi] max_scores.append(float(sims.max())) # Set threshold at 10th percentile of max_scores (keeps 90% of real queries) threshold = float(np.percentile(max_scores, 10)) best_config["min_confidence_threshold"] = round(max(0.05, threshold), 3) logger.info( f"\n ✅ Calibration complete | MRR = {best_mrr:.4f} | " f"Threshold = {best_config['min_confidence_threshold']:.3f}" ) return best_config # ═══════════════════════════════════════════════════════════════════════════════ # B. EMBEDDING FINE-TUNING # ═══════════════════════════════════════════════════════════════════════════════ def fine_tune_embeddings(pairs: list[TrainingPair], doc_id: str, base_model: str) -> Path: """ Fine-tune the embedding model on financial IR domain data using MultipleNegativesRankingLoss (contrastive learning without explicit negatives). Saves the fine-tuned model to data/models/_embed/ """ from sentence_transformers import SentenceTransformer, InputExample, losses from torch.utils.data import DataLoader model_out = MODELS_DIR / f"{doc_id}_embed" MODELS_DIR.mkdir(parents=True, exist_ok=True) logger.info(f"\n [Fine-tune] Loading base model: {base_model}") model = SentenceTransformer(base_model) # Build training examples: (query, positive_chunk) examples = [] for pair in pairs: if pair.positive_text.strip(): examples.append(InputExample( texts=[pair.query, pair.positive_text] )) # Add keyword-enriched variant enriched_query = f"{pair.section}: {pair.query}" examples.append(InputExample( texts=[enriched_query, pair.positive_text] )) if not examples: logger.warning(" No training examples generated — skipping fine-tuning") return Path(base_model) logger.info(f" [Fine-tune] Training on {len(examples)} examples") dataloader = DataLoader(examples, shuffle=True, batch_size=16) loss_fn = losses.MultipleNegativesRankingLoss(model) warmup_steps = max(1, len(dataloader) // 5) model.fit( train_objectives=[(dataloader, loss_fn)], epochs=3, warmup_steps=warmup_steps, show_progress_bar=True, output_path=str(model_out), save_best_model=True, ) logger.info(f" ✅ Fine-tuned model saved → {model_out}") return model_out # ═══════════════════════════════════════════════════════════════════════════════ # C. ANTI-HALLUCINATION — KPI GROUND TRUTH EXTRACTION # ═══════════════════════════════════════════════════════════════════════════════ PERCENTAGE_KPIS = { "nim_percent", "npl_ratio", "coverage_ratio", "cet1_ratio", "lcr", "adr", "cost_income_ratio", "casa_ratio", "gross_loans_yoy", "gross_loans_qoq", "total_deposits_yoy", "total_deposits_qoq", "nfi_yoy" } # Regex patterns to extract specific numeric values from table text / labels KPI_TABLE_LABEL_PATTERNS = { "net_profit_current": [ r"^(?:group net profit|net profit|profit|profit after tax)$", ], "profit_before_tax": [ r"^(?:profit before tax)$", ], "nim_percent": [ r"^(?:nim|net interest margin|net interest margin \(%\))$", ], "npl_ratio": [ r"^(?:npl ratio|npl ratio \(%\))$", ], "coverage_ratio": [ r"^(?:coverage ratio|npl coverage|npl coverage ratio)$", ], "cet1_ratio": [ r"^(?:cet[- ]?1|cet[- ]?1 ratio|common equity tier 1 ratio)$", ], "lcr": [ r"^(?:lcr|liquidity coverage ratio)$", ], "adr": [ r"^(?:adr|advances to deposit ratio)$", ], "cost_income_ratio": [ r"^(?:cost to income ratio|cost-to-income ratio)$", ], "total_income": [ r"^(?:total income)$", ], "total_assets": [ r"^(?:total assets)$", ], # New KPIs "gross_loans_total": [ r"^(?:total gross loans|gross loans)$", ], "gross_loans_yoy": [ r"^(?:total gross loans|gross loans)$", ], "gross_loans_qoq": [ r"^(?:total gross loans|gross loans)$", ], "total_deposits": [ r"^(?:deposits|customer deposits|total deposits)$", ], "total_deposits_yoy": [ r"^(?:deposits|customer deposits|total deposits)$", ], "total_deposits_qoq": [ r"^(?:deposits|customer deposits|total deposits)$", ], "nfi_total": [ r"^(?:total non-funded income|total non funded income|nfi)$", ], "nfi_yoy": [ r"^(?:total non-funded income|total non funded income|nfi)$", ], } KPI_TEXT_PATTERNS = { "net_profit_current": [ r"\b(?:net profit|group net profit|profit after tax)\b[^\n]{0,30}?(?:aed\s*)?([\d,.]+)\s*\b(bn|mn|b|m)\b", ], "profit_before_tax": [ r"\bprofit before tax\b[^\n]{0,30}?(?:aed\s*)?([\d,.]+)\s*\b(bn|mn|b|m)\b", ], "nim_percent": [ r"\b(?:nim|net interest margin)\b[^\n]{0,30}?([\d.]+)\s*%", ], "npl_ratio": [ r"\bnpl ratio\b[^\n]{0,30}?([\d.]+)\s*%", r"\bnon[- ]performing loan ratio\b[^\n]{0,30}?([\d.]+)\s*%", ], "coverage_ratio": [ r"\b(?:provision\s+)?coverage ratio\b[^\n]{0,30}?([\d.]+)\s*%", ], "cet1_ratio": [ r"\b(?:cet[- ]?1|common equity tier 1)\s*(?:ratio)?\b[^\n]{0,30}?([\d.]+)\s*%", ], "lcr": [ r"\b(?:lcr|liquidity coverage ratio)\b[^\n]{0,30}?([\d.]+)\s*%", ], "adr": [ r"\b(?:adr|advances[- ]to[- ]deposit ratio)\b[^\n]{0,30}?([\d.]+)\s*%", ], "cost_income_ratio": [ r"\bcost[- ]to[- ]income\s*(?:ratio)?\b[^\n]{0,30}?([\d.]+)\s*%", r"\bcost\s+income\s+ratio\b[^\n]{0,30}?([\d.]+)\s*%", ], "total_income": [ r"\btotal income\b[^\n]{0,30}?(?:aed\s*)?([\d,.]+)\s*\b(bn|mn|b|m)\b", ], "total_assets": [ r"\btotal assets\b[^\n]{0,30}?(?:aed\s*)?([\d,.]+)\s*\b(bn|mn|b|m|trillion)\b", ], # New KPIs "gross_loans_total": [ r"\b(?:gross loans|total gross loans)\b[^\n]{0,30}?(?:aed\s*)?([\d,.]+)\s*\b(bn|mn|b|m)\b", ], "total_deposits": [ r"\b(?:deposits|customer deposits|total deposits)\b[^\n]{0,30}?(?:aed\s*)?([\d,.]+)\s*\b(bn|mn|b|m)\b", ], "casa_ratio": [ r"\bcasa\s*(?:mix|ratio|stability ratio)?\b[^\n]{0,30}?([\d.]+)\s*%", ], "retail_pbt": [ r"Retail Banking.*?PBT\s+([\d,.]+)", ], "cib_pbt": [ r"Corporate and.*?PBT\s+([\d,.]+)", ], "gmt_pbt": [ r"Global Markets.*?PBT\s+([\d,.]+)", ], "denizbank_pbt": [ r"DenizBank.*?PBT\s+([\d,.]+)", ], "nfi_total": [ r"\b(?:total non-funded income|total non funded income|nfi)\b[^\n]{0,30}?(?:aed\s*)?([\d,.]+)\s*\b(bn|mn|b|m)\b", ], "nfi_yoy": [ r"non[- ]funded income,?\s+up\s+([\d.]+)\s*%\s*yoy", r"non[- ]funded income up\s+([\d.]+)\s*%\s*yoy", ], } KPI_ALLOWED_SECTIONS = { "net_profit_current": ["Income Statement", "Group Overview", "Financial Appendix"], "profit_before_tax": ["Income Statement", "Group Overview", "Financial Appendix"], "nim_percent": ["Net Interest Margin", "Income Statement", "Group Overview", "Financial Appendix"], "npl_ratio": ["Asset Quality", "Loans & Deposits", "Group Overview", "Financial Appendix"], "coverage_ratio": ["Asset Quality", "Group Overview", "Financial Appendix"], "cet1_ratio": ["Capital Adequacy", "Group Overview", "Financial Appendix"], "lcr": ["Liquidity", "Group Overview", "Financial Appendix"], "adr": ["Liquidity", "Group Overview", "Financial Appendix"], "cost_income_ratio": ["Cost to Income", "Income Statement", "Group Overview", "Financial Appendix"], "total_income": ["Income Statement", "Group Overview", "Financial Appendix"], "total_assets": ["Loans & Deposits", "Liquidity", "Group Overview", "Financial Appendix", "Income Statement"], # New KPIs "gross_loans_total": ["Loans & Deposits", "Income Statement", "Group Overview", "Financial Appendix"], "gross_loans_yoy": ["Loans & Deposits", "Income Statement", "Group Overview", "Financial Appendix"], "gross_loans_qoq": ["Loans & Deposits", "Income Statement", "Group Overview", "Financial Appendix"], "total_deposits": ["Loans & Deposits", "Income Statement", "Group Overview", "Financial Appendix", "Liquidity"], "total_deposits_yoy": ["Loans & Deposits", "Income Statement", "Group Overview", "Financial Appendix", "Liquidity"], "total_deposits_qoq": ["Loans & Deposits", "Income Statement", "Group Overview", "Financial Appendix", "Liquidity"], "casa_ratio": ["Loans & Deposits", "Liquidity", "Group Overview", "Financial Appendix"], "retail_pbt": ["Income Statement", "Group Overview", "Financial Appendix", "Asset Quality", "Divisional performance"], "cib_pbt": ["Income Statement", "Group Overview", "Financial Appendix", "Divisional performance"], "gmt_pbt": ["Income Statement", "Group Overview", "Financial Appendix", "Divisional performance"], "denizbank_pbt": ["Income Statement", "Group Overview", "Financial Appendix", "Hyperinflation", "Divisional performance"], "nfi_total": ["Income Statement", "Group Overview", "Financial Appendix", "Non-Funded Income"], "nfi_yoy": ["Income Statement", "Group Overview", "Financial Appendix", "Non-Funded Income"], } def detect_table_scale(tbl, page_text: str) -> str: """Detect whether table cells represent Millions (mn) or Billions (bn).""" headers_str = " ".join(str(h) for h in tbl.get("headers", []) if h).lower() caption_str = str(tbl.get("caption", "")).lower() combined = headers_str + " " + caption_str if re.search(r"\b(?:bn|billion)\b", combined): return "bn" if re.search(r"\b(?:mn|million)\b", combined): return "mn" # Check page text page_text_lower = page_text.lower() has_bn = bool(re.search(r"\b(?:bn|billion)\b", page_text_lower)) has_mn = bool(re.search(r"\b(?:mn|million)\b", page_text_lower)) if has_bn and not has_mn: return "bn" if has_mn and not has_bn: return "mn" return "bn" def extract_kpi_ground_truth(extracted, page_map: dict) -> dict: """ Extract exact numeric KPI values from verified table and text pages. These values act as ground truth — any generated response that contradicts them is flagged or overridden. """ ground_truth = { "doc_id": extracted.doc_id, "doc_name": extracted.doc_name, "extracted": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "kpis": {}, "page_values": {}, } # Ensure page_values has section structures ready for page in extracted.pages: pg_info = page_map.get(page.page_number, {}) section = pg_info.get("section", "") if section: ground_truth["page_values"][str(page.page_number)] = { "section": section, "values": {} } # 1. EXTRACT FROM TABLES FIRST (High priority structured data) for page in extracted.pages: pg_info = page_map.get(page.page_number, {}) section = pg_info.get("section", "") if not section: continue for tbl in page.tables: for row in tbl.get("rows", []): if len(row) < 2: continue label = str(row[0]).lower().strip() for kpi_key, patterns in KPI_TABLE_LABEL_PATTERNS.items(): # Check allowed sections for this KPI allowed = KPI_ALLOWED_SECTIONS.get(kpi_key, []) if section not in allowed: continue for pat in patterns: if re.search(pat, label, re.IGNORECASE): # Resolve column index for yoy/qoq/normal col_idx = 1 if kpi_key.endswith("_yoy"): if len(row) > 3 and "%" in str(row[3]): col_idx = 3 elif len(row) > 2 and "%" in str(row[2]): col_idx = 2 else: continue elif kpi_key.endswith("_qoq"): if len(row) > 5 and "%" in str(row[5]): col_idx = 5 elif len(row) > 4 and "%" in str(row[4]): col_idx = 4 else: continue cell_str = str(row[col_idx]).strip() # Strip brackets/symbols for clean numeric extraction cell_str_clean = cell_str.replace("(", "").replace(")", "").replace("%", "").strip() num_match = re.search(r"([\d,.]+)", cell_str_clean) if num_match: try: val = float(num_match.group(1).replace(",", "")) unit = "%" if kpi_key in PERCENTAGE_KPIS else "bn" # Normalize millions to billions if unit != "%": scale = detect_table_scale(tbl, page.text) if scale == "mn": val = val / 1000 unit = "bn" elif "mn" in cell_str.lower() or "m" in cell_str.lower(): val = val / 1000 unit = "bn" kpi_data = { "value": round(val, 3), "unit": unit, "raw": cell_str, "page": page.page_number, "section": section, } # Update global and page values if kpi_key not in ground_truth["kpis"]: ground_truth["kpis"][kpi_key] = kpi_data page_key = str(page.page_number) if kpi_key not in ground_truth["page_values"][page_key]["values"]: ground_truth["page_values"][page_key]["values"][kpi_key] = kpi_data except ValueError: pass break break # 2. EXTRACT FROM PAGE TEXT (Low priority fallback for missing values) for page in extracted.pages: pg_info = page_map.get(page.page_number, {}) section = pg_info.get("section", "") if not section: continue full_text = page.text for kpi_key, patterns in KPI_TEXT_PATTERNS.items(): if kpi_key in ground_truth["kpis"]: continue allowed = KPI_ALLOWED_SECTIONS.get(kpi_key, []) if section not in allowed: continue for pattern in patterns: flags = re.DOTALL if "PBT" in pattern or kpi_key in ["retail_pbt", "cib_pbt", "gmt_pbt", "denizbank_pbt"] else 0 match = re.search(pattern, full_text, re.IGNORECASE | flags) if match: try: raw_val = match.group(1).replace(",", "") unit = "%" if kpi_key in PERCENTAGE_KPIS else "bn" value = float(raw_val) # Normalize text PBT values or other millions values to billions if unit != "%" and (kpi_key in ["retail_pbt", "cib_pbt", "gmt_pbt", "denizbank_pbt"] or "mn" in match.group(0).lower() or re.search(r"\b(?:mn|million)\b", match.group(0), re.I)): if value > 10.0: value = value / 1000 unit = "bn" kpi_data = { "value": round(value, 3), "unit": unit, "raw": match.group(0)[:80].strip(), "page": page.page_number, "section": section, } # Update global and page values if kpi_key not in ground_truth["kpis"]: ground_truth["kpis"][kpi_key] = kpi_data page_key = str(page.page_number) if kpi_key not in ground_truth["page_values"][page_key]["values"]: ground_truth["page_values"][page_key]["values"][kpi_key] = kpi_data except ValueError: pass break # Clean up empty page_values keys empty_pages = [k for k, v in ground_truth["page_values"].items() if not v["values"]] for k in empty_pages: del ground_truth["page_values"][k] logger.info(f" Extracted {len(ground_truth['kpis'])} ground truth KPI values") return ground_truth # ═══════════════════════════════════════════════════════════════════════════════ # UPDATE GENERATION AGENT WITH CALIBRATED PARAMS # ═══════════════════════════════════════════════════════════════════════════════ def update_generation_agent(config: dict, kpi_gt: dict): """ Update financial_analyst_agent.py with calibrated parameters: 1. Temperature from calibrated config 2. KPI constraints written into response_rules.json 3. Confidence threshold update """ agent_path = BASE_DIR / "services" / "generation" / "financial_analyst_agent.py" if not agent_path.exists(): logger.warning(f" Agent file not found: {agent_path}") return with open(agent_path, "r") as f: content = f.read() # 1. Update default temperature new_temp = config.get("ollama_temperature", 0.05) content = re.sub( r"(def __init__.*?temperature:\s*float\s*=\s*)[\d.]+", lambda m: m.group(0).rsplit("=", 1)[0] + f"= {new_temp}", content, count=1, ) # 2. Update confidence threshold threshold = config.get("min_confidence_threshold", 0.20) content = re.sub( r"(min_confidence_threshold\s*=\s*)[\d.]+", f"\\g<1>{threshold}", content, ) with open(agent_path, "w") as f: f.write(content) # 3. Write verified KPI constraints to response_rules.json rules_file = BASE_DIR / "data" / "response_rules.json" kpi_label_map = { "net_profit_current": "Net Profit", "profit_before_tax": "Profit Before Tax", "nim_percent": "Net Interest Margin (NIM)", "npl_ratio": "NPL Ratio", "coverage_ratio": "Coverage Ratio", "cet1_ratio": "CET-1 Ratio", "lcr": "Liquidity Coverage Ratio (LCR)", "adr": "Advances-to-Deposit Ratio (ADR)", "cost_income_ratio": "Cost-to-Income Ratio", "total_income": "Total Income", "total_assets": "Total Assets", "gross_loans_total": "Total Gross Loans", "gross_loans_yoy": "Gross Loans YoY Change", "total_deposits": "Total Deposits", "total_deposits_yoy": "Total Deposits YoY Change", "casa_ratio": "CASA Ratio", "retail_pbt": "Retail Segment PBT", "cib_pbt": "CIB Segment PBT", "gmt_pbt": "GM&T Segment PBT", "denizbank_pbt": "DenizBank Segment PBT", "nfi_total": "Total Non-Funded Income", "nfi_yoy": "Non-Funded Income YoY Change", } kpi_lines = [] for key, label in kpi_label_map.items(): gt = kpi_gt.get("kpis", {}).get(key) if gt: unit = gt.get("unit", "") val = gt.get("value", "") pg = gt.get("page", "") unit_str = "%" if "%" in str(unit) else f" {unit}" kpi_lines.append(f" - {label}: {val}{unit_str} (Page {pg})") if kpi_lines and rules_file.exists(): try: with open(rules_file) as f: rules = json.load(f) rules["verified_kpis"] = kpi_lines rules["verified_kpis_updated"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) with open(rules_file, "w") as f: json.dump(rules, f, indent=2, ensure_ascii=False) logger.info(f" ✓ {len(kpi_lines)} KPI constraints written to response_rules.json") except Exception as e: logger.warning(f" Could not update response_rules.json: {e}") logger.info( f" ✓ Agent updated: temperature={new_temp}, " f"threshold={threshold}, " f"{len(kpi_lines)} KPI constraints injected" ) # ═══════════════════════════════════════════════════════════════════════════════ # TABLE INDEXING WITH ENRICHED METADATA # ═══════════════════════════════════════════════════════════════════════════════ def enrich_and_index_tables(extracted, page_map: dict, doc_id: str, doc_name: str) -> list: from services.retrieval.table_retriever import FINANCIAL_METRIC_ALIASES def detect_metrics(text: str, section: str) -> set: found = set() tl = (text + " " + section).lower() for key, aliases in FINANCIAL_METRIC_ALIASES.items(): if any(a in tl for a in aliases): found.add(key) return found def table_to_text(tbl: dict) -> str: headers = tbl.get("headers", []) rows = tbl.get("rows", []) caption = tbl.get("caption", "") parts = [] if caption: parts.append(f"Table: {caption}") if headers: parts.append(" | ".join(h for h in headers if h)) for row in rows[:30]: parts.append(" | ".join(str(c) for c in row)) return "\n".join(parts) records = [] for page in extracted.pages: pmap = page_map.get(page.page_number, {}) section = pmap.get("section", "") slide = pmap.get("slide", page.page_number) sec_kpis = set(SECTION_KPI_TAGS.get(section, [])) for tbl in page.tables: text_rep = table_to_text(tbl) text_metrics = detect_metrics(text_rep, section) records.append({ "doc_id": doc_id, "doc_name": doc_name, "page_number": page.page_number, "slide_number": slide, "section_heading": section, "headers": tbl.get("headers", []), "rows": tbl.get("rows", []), "caption": tbl.get("caption", ""), "text_representation": text_rep, "metrics_found": sorted(text_metrics | sec_kpis), }) # Merge — keep records for other docs existing = [] if TABLES_FILE.exists(): try: with open(TABLES_FILE) as f: existing = json.load(f) existing = [r for r in existing if r.get("doc_id") != doc_id] except Exception: existing = [] all_records = existing + records TABLES_FILE.parent.mkdir(parents=True, exist_ok=True) with open(TABLES_FILE, "w", encoding="utf-8") as f: json.dump(all_records, f, indent=2, ensure_ascii=False) return records # ═══════════════════════════════════════════════════════════════════════════════ # MAIN PIPELINE # ═══════════════════════════════════════════════════════════════════════════════ def run(pdf_path: Path, doc_id: str, skip_colpali: bool, do_fine_tune: bool): doc_name = " ".join(w.capitalize() for w in doc_id.replace("_", " ").split()) total_start = time.time() logger.info("=" * 72) logger.info(" IRIS — IR Document Training Pipeline") logger.info(f" PDF : {pdf_path.name}") logger.info(f" Doc ID : {doc_id}") logger.info(f" Options : fine-tune={do_fine_tune}, skip-colpali={skip_colpali}") logger.info("=" * 72) if not pdf_path.exists(): logger.error(f"PDF not found: {pdf_path}") sys.exit(1) # Load existing config or use defaults config = {**DEFAULT_CONFIG} if CONFIG_FILE.exists(): with open(CONFIG_FILE) as f: config.update(json.load(f)) # ── 1. Parse PDF ───────────────────────────────────────────────────────── logger.info("\n[1/9] Parsing PDF…") from services.ingestion.pdf_parser import parse_pdf, save_extraction t0 = time.time() extracted = parse_pdf(pdf_path, doc_id, doc_name) save_extraction(extracted, PROC_DIR) logger.info(f" ✓ {extracted.total_pages} pages | {time.time()-t0:.1f}s") # ── 2. Build page map ───────────────────────────────────────────────────── logger.info("\n[2/9] Building page → section → KPI mapping…") page_map = build_page_map(extracted) PAGEMAP_DIR.mkdir(parents=True, exist_ok=True) pagemap_file = PAGEMAP_DIR / f"{doc_id}_pagemap.json" with open(pagemap_file, "w", encoding="utf-8") as f: json.dump({ "doc_id": doc_id, "doc_name": doc_name, "generated": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "note": "Auto-generated. Edit section/kpis and re-run to apply corrections.", "pages": {str(k): v for k, v in page_map.items()}, }, f, indent=2, ensure_ascii=False) # Print mapping table logger.info(f" {'Page':>4} | {'Slide':>5} | {'Section':<28} | Title") logger.info(f" {'----':>4}-+-{'-----':>5}-+-{'-'*28}-+-{'-'*30}") for pg in sorted(page_map): info = page_map[pg] logger.info( f" P{pg:02d} | S{info['slide']:02d} | {info['section']:<28} | {info['title'][:40]}" ) # ── 3. Extract KPI ground truth ─────────────────────────────────────────── logger.info("\n[3/9] Extracting KPI ground truth values for hallucination prevention…") kpi_gt = extract_kpi_ground_truth(extracted, page_map) KPI_GT_FILE.parent.mkdir(parents=True, exist_ok=True) # Merge with existing ground truth existing_gt = {} if KPI_GT_FILE.exists(): try: with open(KPI_GT_FILE) as f: existing_gt = json.load(f) except Exception: existing_gt = {} existing_gt[doc_id] = kpi_gt with open(KPI_GT_FILE, "w", encoding="utf-8") as f: json.dump(existing_gt, f, indent=2, ensure_ascii=False) logger.info(" Extracted KPI values:") for key, val in kpi_gt.get("kpis", {}).items(): logger.info(f" {key:<25} = {val['value']} {val['unit']} (Page {val['page']}, {val['section']})") # ── 4. Chunk text ───────────────────────────────────────────────────────── logger.info("\n[4/9] Chunking text with section metadata…") import chromadb from chromadb.config import Settings try: client = chromadb.PersistentClient( path=str(CHROMA_DIR), settings=Settings(anonymized_telemetry=False) ) client.delete_collection("finbot_ir_chunks") logger.info(" ✓ Old ChromaDB collection cleared") except Exception: pass from services.ingestion.text_chunker import TextChunker t0 = time.time() chunker = TextChunker( chunk_size=config["chunk_size"], chunk_overlap=config["chunk_overlap"], ) page_texts = [] for p in extracted.pages: if not p.text.strip(): continue pmap = page_map.get(p.page_number, {}) page_texts.append({ "page_number": p.page_number, "text": p.text, "section_heading": pmap.get("section", getattr(p, "section_heading", "") or ""), "slide_number": pmap.get("slide", getattr(p, "slide_number", None) or p.page_number), }) chunks = chunker.chunk_document(page_texts, doc_id) for chunk in chunks: pmap = page_map.get(chunk.page_number, {}) or page_map.get(str(chunk.page_number), {}) setattr(chunk, "section_heading", pmap.get("section", "")) logger.info(f" ✓ {len(chunks)} chunks | {time.time()-t0:.1f}s") # ── 5. Generate training pairs & calibrate weights ───────────────────────── logger.info("\n[5/9] Generating training pairs and calibrating retrieval weights…") pairs = generate_training_pairs(page_map, chunks) calibrated = calibrate_weights(pairs, chunks) # Merge calibrated weights into full config config.update(calibrated) config["last_calibrated"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) config["calibrated_on_doc"] = doc_id # ── 6. Optional: Fine-tune embedding model ──────────────────────────────── embed_model_path = config["embed_model"] if do_fine_tune and pairs: logger.info("\n[6/9] Fine-tuning embedding model on domain data…") fine_tuned_path = fine_tune_embeddings(pairs, doc_id, config["embed_model"]) embed_model_path = str(fine_tuned_path) config["embed_model"] = embed_model_path config["fine_tuned"] = True else: logger.info("\n[6/9] Skipping fine-tuning (use --fine-tune to enable)") # Save config BEFORE embedding so embedding uses correct model CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True) with open(CONFIG_FILE, "w") as f: json.dump(config, f, indent=2) logger.info(f" ✓ Retrieval config saved → {CONFIG_FILE.name}") # ── 7. Embed + index chunks ──────────────────────────────────────────────── logger.info("\n[7/9] Embedding chunks and indexing in ChromaDB…") from services.retrieval.text_retriever import TextRetriever t0 = time.time() chunks = chunker.embed_chunks(chunks) text_retriever = TextRetriever(persist_dir=CHROMA_DIR) n_chunks = text_retriever.index_chunks(chunks) logger.info(f" ✓ {n_chunks} chunks embedded and indexed | {time.time()-t0:.1f}s") # ── 8. Index tables ──────────────────────────────────────────────────────── logger.info("\n[8/9] Indexing tables with enriched metadata…") table_records = enrich_and_index_tables(extracted, page_map, doc_id, doc_name) logger.info(f" ✓ {len(table_records)} tables indexed") # ── Update generation agent ──────────────────────────────────────────────── logger.info("\n Updating generation agent with calibrated config…") update_generation_agent(config, kpi_gt) # ── 8b. Auto-generate Smart Response Cache ──────────────────────────────── # Pre-fills document-agnostic templates with real KPI values so the Smart # Response Engine can serve instant (<10ms) answers for any trained PDF. logger.info("\n[8b] Auto-generating smart response cache for new engine…") try: from services.classification.kpi_context_builder import KPIContextBuilder, INTENT_KPI_KEYS from services.generation.smart_response_engine import ( SmartResponseEngine, save_cached_response, _INTENT_BUILDERS ) _cache_builder = KPIContextBuilder(data_dir=DATA_DIR) cache_results: dict[str, str] = {} for intent in INTENT_KPI_KEYS.keys(): ctx = _cache_builder.build(intent=intent, doc_ids=[doc_id]) if ctx.has_data: builder_fn = _INTENT_BUILDERS.get(intent) if builder_fn: try: resp = builder_fn(ctx) if resp: resp["question"] = f"[auto-generated for intent: {intent}]" resp["latency_ms"] = 0 saved = save_cached_response(doc_id, intent, resp) cache_results[intent] = "✓" if saved else "✗" else: cache_results[intent] = "— (no KPI match in template)" except Exception as be: cache_results[intent] = f"✗ ({be})" else: cache_results[intent] = "— (no template builder)" else: cache_results[intent] = "— (no KPI data extracted)" for intent, status in cache_results.items(): logger.info(f" {status} {intent}") n_cached = sum(1 for s in cache_results.values() if s == "✓") logger.info(f" ✓ {n_cached}/{len(cache_results)} smart response templates generated") logger.info(f" Cached to: data/response_cache/{doc_id}/") except Exception as e: logger.warning(f" ⚠ Smart cache generation failed (non-fatal): {e}") # ── 9. Render pages + ColPali ───────────────────────────────────────────── logger.info("\n[9/9] Rendering PDF pages…") from services.ingestion.page_renderer import render_pdf_pages t0 = time.time() PAGES_DIR.mkdir(parents=True, exist_ok=True) page_records = render_pdf_pages( pdf_path=pdf_path, output_dir=PAGES_DIR, doc_id=doc_id, ) logger.info(f" ✓ {len(page_records)} pages rendered | {time.time()-t0:.1f}s") colpali_pages = len(page_records) if not skip_colpali: logger.info(" Generating ColPali visual embeddings (batch_size=1)…") import torch if torch.backends.mps.is_available(): os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] = "0.0" from services.ingestion.colpali_indexer import ColPaliIndexer t0 = time.time() indexer = ColPaliIndexer(store_dir=COLPALI_DIR) page_records = indexer.index_pages(page_records, doc_id=doc_id, batch_size=1) colpali_pages = len(page_records) logger.info(f" ✓ {colpali_pages} ColPali pages | {time.time()-t0:.1f}s") else: logger.info(" Skipping ColPali (--skip-colpali)") # ── Update documents registry ────────────────────────────────────────────── docs = [] if DOCS_FILE.exists(): with open(DOCS_FILE) as f: docs = json.load(f) docs = [d for d in docs if d["doc_id"] != doc_id] docs.append({ "doc_id": doc_id, "name": doc_name, "doc_type": extracted.metadata.get("doc_type", "Financial Document"), "period": extracted.metadata.get("period", "Unknown"), "institution": extracted.metadata.get("institution", "Unknown"), "total_pages": extracted.total_pages, "status": "indexed", "filename": pdf_path.name, "chunks_indexed": n_chunks, "tables_indexed": len(table_records), "colpali_pages": colpali_pages, "pagemap_file": pagemap_file.name, "page_section_map": {str(k): v["section"] for k, v in page_map.items()}, "retrieval_config": str(CONFIG_FILE.name), }) DATA_DIR.mkdir(parents=True, exist_ok=True) with open(DOCS_FILE, "w") as f: json.dump(docs, f, indent=2) # ── Final report ─────────────────────────────────────────────────────────── total_elapsed = time.time() - total_start mrr = calibrated.get("mrr", "—") logger.info("") logger.info("=" * 72) logger.info(f" ✅ Training complete in {total_elapsed:.1f}s") logger.info(f" {'─'*60}") logger.info(f" Text chunks indexed : {n_chunks}") logger.info(f" Tables indexed : {len(table_records)}") logger.info(f" KPI ground truths : {len(kpi_gt.get('kpis', {}))}") logger.info(f" Training pairs : {len(pairs)}") logger.info(f" Embedding model : {embed_model_path}") logger.info(f" Fine-tuned : {do_fine_tune}") logger.info(f" Calibrated weights:") logger.info(f" metric_match : {config.get('metric_match_weight')}") logger.info(f" section_boost : {config.get('section_heading_boost')}") logger.info(f" confidence floor : {config.get('min_confidence_threshold')}") logger.info(f" LLM temperature : {config.get('ollama_temperature')}") logger.info(f" {'─'*60}") logger.info(f" Config saved to : data/{CONFIG_FILE.name}") logger.info(f" KPI ground truth : data/{KPI_GT_FILE.name}") logger.info(f" Page mapping : data/page_maps/{pagemap_file.name}") logger.info("=" * 72) if __name__ == "__main__": parser = argparse.ArgumentParser( description="IRIS — Train on IR PDF, calibrate weights, reduce hallucination", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: # Standard training (parse + calibrate + index) python train_document.py --pdf ../documents/enbd_q1_2026.pdf # Full training including embedding fine-tuning (~20 min extra) python train_document.py --pdf ../documents/enbd_q1_2026.pdf --fine-tune # Fast mode — skip ColPali visual embeddings python train_document.py --pdf ../documents/report.pdf --skip-colpali # Custom document ID python train_document.py --pdf ../documents/report.pdf --doc-id enbd_fy2025 """, ) parser.add_argument("--pdf", required=True, help="Path to IR PDF") parser.add_argument("--doc-id", default=None, help="Document ID (derived from filename if omitted)") parser.add_argument("--skip-colpali", action="store_true", help="Skip ColPali visual embedding") parser.add_argument("--fine-tune", action="store_true", help="Fine-tune the embedding model on domain data") args = parser.parse_args() pdf = Path(args.pdf).resolve() d_id = args.doc_id or pdf.stem.lower().replace(" ", "_").replace("-", "_") run(pdf_path=pdf, doc_id=d_id, skip_colpali=args.skip_colpali, do_fine_tune=args.fine_tune)