diff --git "a/canlex/index.py" "b/canlex/index.py" new file mode 100644--- /dev/null +++ "b/canlex/index.py" @@ -0,0 +1,2527 @@ +"""Hybrid retrieval (BM25 + semantic) with cross-encoder reranking.""" +import importlib.metadata +import json +import math +import os +import re +import sys +import unicodedata +from collections import Counter, defaultdict +from pathlib import Path + +import numpy as np +import snowballstemmer + +from . import bm25 +from .config import DATA_DIR, PROCESSED_DIR +from .airports import airport_terms +from .synonyms import expand_query + +K1 = 1.5 +B = 0.75 +RRF_K = 60 # reciprocal-rank-fusion damping constant +W_SEM = 2.0 # weight on the semantic retriever in the fusion (1.0 = equal; eval-tuned) +CANDIDATES = 80 # hits each retriever contributes to the fusion +RERANK_POOL = int(os.environ.get("CANLEX_RERANK_POOL", "16")) + # top fused candidates the cross-encoder rescores. Profiled + # at ~96% of query latency (the pool is cross-encoded on + # CPU), so this is THE latency knob; env-sweepable, and any + # change must be eval-gated (a smaller pool limits which + # low-fusion candidates the promote-only rerank can lift). + # Swept 2026-06: 50 -> 16 (with doc cap 3000 -> 1000) held + # the 159-Q eval exactly (MRR +0.01) at ~5.5x lower latency. +POOL_CAP = int(os.environ.get("CANLEX_POOL_CAP", "8")) + # max chunks one capped source (a decision, a memorandum) + # may hold in the rerank pool. The pool is the cross- + # encoder's whole world: with RECALL_CAP=20 a deep + # decision can take 14 of 16 slots (Fearon, 88 paragraphs + # on cell-phone searches, for a device-search query), so + # the governing border authorities sitting at fused ranks + # 19-32 (Canfield, Pike, ENF 12 s. 6.15) are never scored + # at all -- _diversify runs after the rerank and can only + # reorder what was scored. Capping while the pool fills + # backfills those slots with the next sources. Swept + # 2026-09-06 over {5,6,8,10,12,16,off}: 5 let two + # incidental FB articles into the top 5 of the + # discipline question, so _ensure_primary no longer + # pulled art. 17 in (absent from the top 20); 6-16 keep + # it at 3 with Canfield at 4 on the device query. 8 is + # the comfortable middle. Env-sweepable; eval-gate any + # change. +MN_WEIGHT = float(os.environ.get("CANLEX_MN_WEIGHT", "0.0024")) + # title-match boost per unit of idf-weighted overlap between + # the query and a candidate's marginal note (section title) +MN_CAP = float(os.environ.get("CANLEX_MN_CAP", "0.024")) + # ceiling on the title-match boost -- it nudges the ranking + # without overriding a strong base score +REG_PENALTY = float(os.environ.get("CANLEX_REG_PENALTY", "0.004")) + # small fusion penalty on regulation sections, so the Act + # that creates a duty outranks the regulation elaborating it + # (sweep-tuned 2026-05-23 from 0.008 -> 0.004; see sweep.log) +NIF_PENALTY = float(os.environ.get("CANLEX_NIF_PENALTY", "0.012")) + # fusion penalty on not-in-force amendment chunks, so + # pending law never outranks the live provision it amends +MEMO_PENALTY = float(os.environ.get("CANLEX_MEMO_PENALTY", "0.008")) + # the same hierarchy prior one rung further: the statute + # creating a duty outranks the guidance explaining it. + # Measured 2026-08-27: 11/54 statute-gold questions had a + # guidance chunk at rank 1, mostly ENF overview sections + # that quote the Act's own words -- keyword magnets, not + # better answers. + # Refit 0.004 -> 0.008 on 2026-09-10, after IRCC's + # rehabilitation guide and the IRB Legal Services papers + # joined the corpus: a plain-language guide PARAPHRASES + # its provision, so it out-matches the provision's own + # words on a natural-language question (Guide 5312 beat + # IRPR s. 17 by 0.0115, three times the old penalty). + # Swept 0.006 / 0.008 / 0.010: 0.008 restores ALL Hit@1 + # and adds 4 to Hit@3 (legacy 0.95 -> 0.96, holdout 0.93 + # -> 0.95); 0.010 costs the French slice (0.58 -> 0.42) + # and Hit@5. +REPEALED_PENALTY = float(os.environ.get("CANLEX_REPEALED_PENALTY", "0.008")) + # likewise for repealed provisions (kept so lookups still + # inform): a repealed section's on-topic title must not + # beat live law in a close contest -- the CDSA-to-Cannabis + # -Act migration is the live risk zone +INAPPLICABLE_PENALTY = float(os.environ.get("CANLEX_INAPPLICABLE_PENALTY", "0.004")) + # fusion penalty on provisions the corpus's own audience + # is excluded from (see _INAPPLICABLE): Canada Labour + # Code Part III does not apply to federal departments + # (s. 167(1)(d)), yet its hours-of-work sections out- + # score the FB agreement article that actually governs + # a shift-change grievance. The chunks stay retrievable + # and carry an applicability caveat when shown. +BACKMATTER_PENALTY = float(os.environ.get("CANLEX_BACKMATTER_PENALTY", "0.004")) + # likewise for a collective agreement's back-matter + # (memoranda, letters of understanding) vs its numbered articles + # (sweep-tuned 2026-05-23 from 0.008 -> 0.004) +SOURCE_CAP = 2 # max chunks one case or memorandum may contribute + # (and the 51 US-state commentary pages together, as one + # source -- see _source_key) +COMMENTARY_PENALTY = float(os.environ.get("CANLEX_COMMENTARY_PENALTY", "0")) + # the hierarchy prior's lowest tier: CanLex's own curated + # commentary (the US Dispositions Helper) is analysis, + # not a source, yet it took no penalty while the + # official guidance beside it takes MEMO_PENALTY + # (roadmap T5.4). Env-sweepable (canlex.sweep). +ATIP_PENALTY = float(os.environ.get("CANLEX_ATIP_PENALTY", "0.004")) + # extra, on top of MEMO_PENALTY, for guidance its ingester + # flagged as an unofficial ATIP copy of a withdrawn + # chapter (source_status == ATIP_STATUS: ENF 2, ENF 14). + # 0.004 measured in the 2026-09-12 eval diff (loss B: + # legacy +1, one other question better, no slice worse); + # maintainer-approved 2026-09-14. +APPLICANT_PENALTY = float(os.environ.get("CANLEX_APPLICANT_PENALTY", "0")) + # extra, on top of MEMO_PENALTY, for IRCC's applicant- + # facing pages (_APPLICANT_CODES): plain-language + # summaries for applicants, not officer instructions. +# Roadmap T5.4, swept 2026-09-15 (B10) on the eval set with the 30-item +# criminality slice and the criminality route as shipped (Hit@1 / MRR): +# ALL legacy holdout criminality +# all three at 0 .6985/.8012 .7812/.8685 .6600/.7451 .3000/.4410 +# COMMENTARY .002 .6954/.7989 .7760/.8646 .6600/.7451 .3000/.4410 +# COMMENTARY .004 .6985/.8005 .7760/.8646 .6600/.7451 .3000/.4410 +# COMMENTARY .006 .6954/.7984 .7708/.8612 .6400/.7317 .3000/.4410 +# MEMO .006 .6800/.7878 .7604/.8544 .6600/.7452 .3000/.4312 +# MEMO .010 .6954/.7994 .7812/.8675 .6600/.7454 .3000/.4438 +# APPLICANT .008 .7015/.8029 .7812/.8667 .6600/.7469 .3333/.4706 +# APPLICANT .014 .7015/.8030 .7812/.8667 .6600/.7475 .3333/.4717 +# APPLICANT .02 .7015/.8031 .7812/.8667 .6600/.7475 .3333/.4737 +# COMMENTARY never moves the criminality slice and costs legacy Hit@1 (and +# holdout at .006), with MEMO at .006, .008 or .010. MEMO .006 costs legacy +# .021; MEMO .010 lowers criminality guid@6 .345 -> .310. APPLICANT holds +# ALL and holdout within 0.01 and lifts the criminality slice at every +# value, but it turns the legacy 'How long after completing a sentence can +# someone apply for criminal rehabilitation?' from rank 2 into a miss (7), +# and on 'When is a person deemed rehabilitated after a single foreign +# conviction?' Guide 5312, not IRPR s. 18, takes the IRCC page's place at +# #1. So all three stay at their values: 0, 0.008 and 0. +# The value canlex/enf.py stamps on every chunk of an unofficial ATIP copy. +ATIP_STATUS = "unofficial-atip-copy" +# IRCC's applicant-facing pages (canlex/conduct.py marks them +# '(applicant-facing)' in act_name): the IRCC public pages and the IMM 1444 +# form page. Keyed on act_code, a structured field. +_APPLICANT_CODES = ("IRCC-PUB-", "IMM-1444") +READ_NEXT_CAP = 3 # max official guidance chunks co-surfaced from see_also +HIGHLIGHT_DEPTH = int(os.environ.get("CANLEX_HIGHLIGHT_DEPTH", "3")) + # the results the pinpoint pass (_highlight) scores: it + # cross-encodes one pair per provision unit of every + # legislation result it reaches (median 11 pairs a query, + # p95 45, max 53 at depth 8). 8 -> 3 adopted 2026-09-21 + # (audit follow-up, group A item 8b): search p50 951 -> + # 877 ms; results 4-8 lose their 'Most on point' + # pinpoint (159 of 339 eval answers); no rank can move + # (the pass runs after ranking). +# Feature switches for the two ranking changes below, so an eval dump can +# measure each in isolation (set the module attribute, or the env var to 0). +CRIMINALITY_ROUTE = os.environ.get("CANLEX_CRIMINALITY_ROUTE", "1") != "0" + # roadmap T5.2: on an s. 36 question, restrict what + # _ensure_primary may promote to criminality provisions +STATE_ONE_SOURCE = os.environ.get("CANLEX_STATE_ONE_SOURCE", "1") != "0" + # roadmap T5.3: the 51 US-state commentary pages count as + # ONE source for the diversity caps unless the query + # names the state + +# --- Criminality route (roadmap T5.2) ---------------------------------------- +# _ensure_primary's promotion is chosen by title overlap, so on criminality +# questions it pulled off-point statutes into ranks 2-3 (IRPA s. 64 'No +# appeal' on a foreign pardon, CC s. 748 on a state pardon, CC s. 57 on a +# sealed record) while IRPA s. 36 stayed out of the top 20. On an s. 36 +# question the route restricts WHICH statute may be promoted. It fires only +# when the query carries BOTH a criminality marker and an s. 36 context cue: +# a bare 'criminal' ('use force under the Criminal Code', 'criminal +# harassment of an employee') is not an inadmissibility question, and the +# verifier showed a 'crimin*' trigger injected IRPA s. 36 into 5 of 6 such +# queries. Words that are themselves criminality-and-inadmissibility terms +# (pardon, expungement, record suspension, deferred adjudication) are in +# both lists, so they fire alone. 'traveller' and 'american' are not cues on +# their own (2026-09-15 verifier: 'An American driver arrives impaired at +# the port of entry - what criminal offence applies under the Criminal +# Code?' fired the route and promoted CC s. 145 'Escape and being at +# large'); a traveller's inadmissibility question says 'inadmissible', +# 'admissible' or asks about entering Canada, which are cues. +# '(in)admissible' is also the law of evidence: 'Is evidence from a +# warrantless search admissible at a criminal trial?' fired the route and +# put IRPA s. 36 at rank 2 over CC s. 655 (the 2026-09-15 track B +# verifier: 7 of 7 evidence questions fired, 3 took s. 36). So the word +# counts only when the query carries no evidence vocabulary (named groups +# 'adm_ev' and 'ev' below) or ties it to Canada ('inadmissible to +# Canada'); every other cue counts as before (criminality_query). +_CRIM_MARKERS = re.compile( + r"(?(?:in)?admissib\w* (?:against|" + r"at (?:a |the |his |her |any )?(?:later |subsequent )?(?:criminal )?trial|" + r"in (?:a |the |his |her |any )?(?:later |subsequent )?(?:criminal )?" + r"(?:trial|prosecution|proceedings?)))|" + r"(?Pevidence|evidentiary|confessions?|hearsay|voir dire|testimony|" + r"breath (?:samples?|tests?|results?|demands?)|certificates? of analysis)|" + r"(?:in)?admissib\w* (?:to|into|in) canada|" + r"(?P(?:in)?admissib\w*)|" + r"irpa|a36|" + r"(?:s\.?|sec\.?|section|subsection|paragraph)\s*36|" + r"a?36\s*\(\s*(?:1|2|2\.1|3)\s*\)|" + r"rehabilit\w*|pardon\w*|record suspension\w*|expung\w*|seal(?:ed|ing)|" + r"deferred adjudication|withheld adjudication|" + r"adjudication (?:was |is |being )?withheld|in absentia|" + r"abroad|outside (?:of )?canada|overseas|foreign\w*|u\.s\.?|" + r"united states|usa|(?:in|from) the us|" + r"us (?:conviction|dui|dwi|felony|misdemeanou?r|record|court|state|" + r"charge|offen[cs]e|pardon|expungement|criminal)\w*|" + r"permanent residen\w*|" + r"(?:enter|entering|entry to|admitted to|admission to|come to|coming to|" + r"travel to|travelling to|traveling to) canada)(?!\w)") +# Words that point at IRPA s. 36(3) -- record suspensions and pardons, the +# youth-record rule -- when the route has to supply an s. 36 piece. +_S36_3_WORDS = re.compile( + r"(?' (equivalency +# pairings) -- see canlex/commentary.py -- so that is the lookup key. +SEE_ALSO_DRAFT = (DATA_DIR / "curated" / "drafts" + / "see-also-draft-2026-09-12.json") + + +def curated_see_also(curated_dir=None): + """{commentary section: [see_also item, ...]} from data/curated/*.json. + Empty when no curated entry carries see_also -- the approved baseline + predates the field, so a composed release runs with read-next off.""" + root = Path(curated_dir) if curated_dir else DATA_DIR / "curated" + out = {} + + def take(section, entry): + items = entry.get("see_also") if isinstance(entry, dict) else None + if section and items: + out[section] = list(items) + + disp = _read_json(root / "us_dispositions.json") or {} + for e in disp.get("methodology") or []: + take(e.get("id"), e) + for e in disp.get("dispositions") or []: + take(e.get("id"), e) + eq = _read_json(root / "us_equivalency.json") or {} + for e in eq.get("methodology") or []: + take(e.get("id"), e) + for p in eq.get("pairings") or []: + take(p.get("id") and f"equiv-{p['id']}", p) + return out + + +def draft_see_also(path): + """({commentary section: [item, ...]}, {key: key spec}) from a see-also + research draft (data/curated/drafts/see-also-draft-*.json), whose items + name stable keys in its 'keys' table. Used only when asked for + (CANLEX_SEE_ALSO_DRAFT): a draft is unreviewed.""" + data = _read_json(path) or {} + out = {} + for group in ("disposition_methodology", "dispositions", + "equivalency_methodology"): + for eid, items in (data.get(group) or {}).items(): + if items: + out[eid] = list(items) + for pid, items in (data.get("pairings") or {}).items(): + if items: + out[f"equiv-{pid}"] = list(items) + return out, dict(data.get("keys") or {}) + + +def _see_also_spec(item, keys): + """(key spec dict, pinned chunk id or None) for one see_also item. An + item is a stable key string ('ENF-3|ENF 3 s. 10.9', 'I-2.5|s. 36', + 'IRB-LS-ROA8|Effect of a Pardon', 'ENF-1|ENF 1 App. A|'), a + dict carrying such a 'key' (and optionally the chunk it quotes), or a + dict naming act_code with section and/or marginal_note.""" + pin = None + if isinstance(item, dict): + pin = item.get("chunk") or item.get("quote_chunk") + key = item.get("key") + if not key and item.get("act_code"): + spec = {k: item[k] for k in ("act_code", "section", + "marginal_note") if item.get(k)} + spec["match"] = "+".join(k for k in ("section", "marginal_note") + if k in spec) + return spec, pin + else: + key = item + if not isinstance(key, str) or "|" not in key: + return None, pin + if key in keys: + return keys[key], pin + parts = key.split("|") + if len(parts) == 3: + return {"act_code": parts[0], "section": parts[1], + "marginal_note": parts[2], + "match": "section+marginal_note"}, pin + rest = parts[1] + if rest.startswith("s. "): + return {"act_code": parts[0], "section": rest[3:], + "match": "section"}, pin + return {"act_code": parts[0], "section": rest, "marginal_note": rest, + "match": "section|marginal_note"}, pin + + +def _norm_text(text): + return " ".join((text or "").split()).lower() + + +def _see_also_match(chunk, spec): + """True when an English chunk satisfies a see_also key spec: act_code + equal and, per spec['match'], its section and/or marginal note equal + (marginal notes compared case- and whitespace-folded).""" + if not spec or chunk.get("lang", "en") == "fr": + return False + if chunk.get("act_code") != spec.get("act_code"): + return False + by_section = chunk.get("section", "") == spec.get("section") + by_note = _norm_text(chunk.get("marginal_note")) == _norm_text( + spec.get("marginal_note")) + match = spec.get("match") or "section" + if match == "section": + return by_section + if match == "marginal_note": + return by_note + if match == "section+marginal_note": + return by_section and by_note + return by_section or by_note # 'section|marginal_note' + +# Court hint: a query that names the deciding court ("what did the Ontario +# Court of Appeal decide about ...") is asking for THAT court's decision, so +# its case-law chunks get a small fusion bonus -- enough to win a close +# contest against a sister court's decision on the same subject, never +# enough to outrank a clearly better match. The patterns map the ways a +# court is written to a predicate over the chunk's 'part' (which carries the +# deciding court, see _filter_ok). +COURT_HINT = float(os.environ.get("CANLEX_COURT_HINT", "0.004")) +# Act hint: the court hint's twin for everything that is not case law. A +# query that names an instrument by its short or full name ('under the +# Income Tax Act', 'IRPA', 'the Customs Act') gets that instrument's chunks a +# small fusion bonus. Case law is excluded (its act_short is a party name -- +# 'Hill', 'Baker', 'Pike' -- which common words would trip) and has +# COURT_HINT instead. +ACT_HINT = float(os.environ.get("CANLEX_ACT_HINT", "0.012")) +# Swept 2026-09-10 on the 295-question set after the Income Tax Act joined +# the corpus: 0.004 / 0.008 / 0.012 / 0.016 all beat no hint on every slice. +# 0.012 is the conservative pick -- it puts ITA s. 239 at rank 1 for "penalty +# for tax evasion under the Income Tax Act" (the Luxury Tax Act and the +# Excise Act share its evasion boilerplate) and lifts ALL Hit@1 212 -> 215 +# over the pre-ITA baseline. 0.016 scored ~3 Hit@1 questions higher with no +# slice lower, so the curve had not turned: the ceiling is unmeasured. + + +def _part_is(*fragments, exclude=()): + fragments = tuple(f.lower() for f in fragments) + exclude = tuple(e.lower() for e in exclude) + def ok(part): + p = (part or "").lower() + return (any(f in p for f in fragments) + and not any(e in p for e in exclude)) + return ok + + +_COURT_HINTS = ( + (re.compile(r"\b(supreme court|scc)\b"), _part_is("supreme court of canada")), + (re.compile(r"\b(federal court of appeal|fca)\b"), _part_is("federal court of appeal")), + # "Federal Court" alone means the trial court, not the FCA + (re.compile(r"\bfederal court\b(?! of appeal)|\bfctd\b|\bfc\b"), + _part_is("federal court", exclude=("appeal",))), + (re.compile(r"\b(ontario court of appeal|court of appeal for ontario|onca)\b"), + _part_is("court of appeal for ontario")), + (re.compile(r"\b(alberta court of appeal|court of appeal of alberta|abca)\b"), + _part_is("court of appeal of alberta")), + (re.compile(r"\b(refugee appeal division|rad)\b"), _part_is("refugee appeal division")), + (re.compile(r"\b(fpslreb|pslreb|pslrb|pssrb|labour relations and employment board|" + r"labour relations board)\b"), + _part_is("labour relations and employment board", "labour relations board")), + (re.compile(r"\b(cirb|canada industrial relations board)\b"), + _part_is("canada industrial relations board")), +) + + +# The hint means "I want that court's DECISION": it needs a decision word +# beside the court's name. "What is the Federal Public Sector Labour +# Relations and Employment Board?" names the Board as an institution and +# wants the Act that creates it, not its rulings. +_DECISION_WORDS = re.compile( + r"\b(decid\w*|held|hold\w*|rul\w*|found|finding\w*|judg\w*|decision\w*|" + r"case\w*|jurisprudence|precedent\w*|reason\w*|dismiss\w*|allow\w*|quash\w*)\b") + + +# Employer policy instruments (the CBSA Code of Conduct, the TBS suite) govern +# the employment relationship and nothing else. _ensure_primary may promote +# one as "the governing instrument" only when the query is about that +# relationship; otherwise the Code's section titles ("Use of Personal +# Electronic Devices at Work") match traveller questions by accident. +_EMPLOYMENT_WORDS = re.compile( + r"\b(employee\w*|employer\w*|employment|manager\w*|supervisor\w*|disciplin\w*|reprimand\w*|" + r"grievance\w*|grievor|harass\w*|reliability status|security screening|security clearance|" + r"acting appointment\w*|rate of pay|overtime|union\w*|steward\w*|bargaining|workplace|" + r"off-duty|off duty|social media|code of conduct|values and ethics|conflict of interest|" + r"misconduct|public servant\w*|public service|fonctionnaire\w*|employé\w*|" + r"code de conduite|harcèlement|mesures? disciplinaires?|milieu de travail)\b") + + +def employment_query(query_lc): + """True when a lower-cased query is about the employment relationship.""" + return bool(_EMPLOYMENT_WORDS.search(query_lc)) + + +# Narrow-domain instruments _ensure_primary may pull in only when the query +# is about their domain. The Income Tax Act's marginal notes ('Search +# warrant', 'Inspections', 'Seizure', 'Employee', 'Benefit') match border +# and labour questions by accident: on 2026-09-10 the guarantee put +# ITA s. 231.3 (a CRA audit warrant) at rank 2 of "can an officer search a +# traveller's cell phone without a warrant", on a fused score well below the +# decisions around it. The Act still competes on its own merits through +# BM25 and the embedder; it is only never PROMOTED into a non-tax answer. +_PROMOTION_GATES = { + "I-3.3": re.compile( + r"\b(tax\w*|income|taxpayer\w*|cra|revenue agency|canada revenue|" + r"return\w*|evasion|evade\w*|deduct\w*|rrsp|tfsa|t4|t1|" + r"assessment\w*|reassess\w*|capital gain\w*|dividend\w*|" + r"withholding|source deduction\w*|impot\w*|impôt\w*|fiscal\w*|" + r"revenu|contribuable\w*|déclaration de revenus)\b"), +} + + +def promotion_allowed(act_code, query_lc): + """False when the instrument is domain-gated and the query is not about + its domain (see _PROMOTION_GATES).""" + gate = _PROMOTION_GATES.get(act_code) + return gate is None or bool(gate.search(query_lc or "")) + + +def court_hints(query_lc): + """The court predicates a lower-cased query names, if any -- only when the + query also asks what that court decided (see _DECISION_WORDS).""" + if not _DECISION_WORDS.search(query_lc): + return [] + return [pred for pat, pred in _COURT_HINTS if pat.search(query_lc)] +RECALL_CAP = int(os.environ.get("CANLEX_RECALL_CAP", "20")) + # max chunks one capped source (anything _source_key keys: + # a decision, a memorandum) may contribute to EACH + # retriever's CANDIDATES slice at the recall stage. Without + # it, a deep decision whose every paragraph matches a + # specific query (Canfield: 94 paragraphs on device border + # searches) fills the entire fusion pool, so the result- + # stage SOURCE_CAP defers its surplus with nothing left to + # promote and _ensure_primary finds no statute to pull in. + # Generous by design: SOURCE_CAP=2 governs display; this + # only has to leave room for other sources to survive + # recall. Uncapped sources (_source_key None) are unlimited. + # Swept 2026-07-10 over {12,16,20,24} on the 159-Q eval: + # 12 and 20 tie at 0.80/0.94/0.97/0.99/0.87 (4 misses), + # both above the uncapped baseline; 20 chosen as the + # lightest-touch winner (24 regresses Hit@3/@5). +APPENDIX_CAP = 3 # max referenced appendices co-surfaced into a result set + +# Primary instruments -- enacted law, collective agreements, the NJC directives +# incorporated into them, and the IRPA delegation instruments. Their sections or +# items are distinct provisions, so (like legislation) they are never collapsed +# under the diversity cap. +PRIMARY_DOC_TYPES = frozenset({"legislation", "agreement", "directive", + "delegation", "policy"}) +# Memorandum families whose act_code is ONE shared constant over many +# independent items, so the diversity cap keys on `section` instead (see +# _source_key): a D-memo number, an AMPS contravention code. +_FLAT_MEMO_CODES = frozenset({"D-Memo", "AMPS"}) + +_TOKEN = re.compile(r"[a-z0-9]+") +# A chunk id ending '-p' marks one embeddable piece of a long section that +# ingest.py split (see _split_text); same-section pieces are capped together. +_PIECE_ID = re.compile(r"-p\d+$") +# A section number the query cites. The trigger is 'section(s)', 'sec(s).', +# 's.'/'ss.' or a bare 's'/'ss' -- the bare form only when it is not the +# possessive of the word before it ("the officer's 5 minute break") and not +# preceded by a number ('5 s 6' names no section). A plural trigger carries a +# whole list ('ss. 17 and 18', 'sections 17, 18 and 18.1', 'sections 17 to +# 18.1'); a singular one carries only a range ('s. 17-18', 's. 17 to 18'), +# never an 'and' -- 's. 36 and 2 prior convictions' names one section. Code +# audit 2026-09-21: one number was captured per trigger, so 'ss. 17 and 18' +# pinned s. 17 alone. Audit follow-up 2026-09-21: 'sect. 36' and the +# article forms ('art. 36', 'article 36', 'articles 17 et 18' -- how the +# French twins cite a section) trigger as canlex_get_section accepts them, +# an article number followed by a letter ('Article 1E' of the Refugee +# Convention) is none, and an Oxford comma before a list's last item +# ('ss. 17, and 18') no longer ends the list. An article form (the named +# group 'article') is then read by _article_cites_section: in English an +# 'Article N' is usually a treaty's (verification 2026-09-22). +_SECTION_NUM = r"\d+(?:\.\d+)?" +_SECTION_REF = re.compile( + r"(?" + r"(?:articles|arts\.?)\s*" + _SECTION_NUM + r"(?![a-z\d])" + + r"(?:\s*(?:,\s*(?:and|or|&|et|ou)|,|and|or|&|et|ou|to|à|-|–)\s*" + + _SECTION_NUM + r"(?![a-z\d]))*" + r"|(?:article|art\.?)\s*" + _SECTION_NUM + r"(?![a-z\d])" + + r"(?:\s*(?:,|to|-|–)\s*" + _SECTION_NUM + r"(?![a-z\d]))*" + r")", re.IGNORECASE) +_SECTION_NUMS = re.compile(_SECTION_NUM) +# A treaty an English query names, which makes its 'Article N' the treaty's +# (_article_cites_section): convention, covenant, protocol, treaty (and the +# French words) -- but not 'Convention refugee', IRPA's own term (s. 96), nor +# 'convention collective', the French for the FB collective agreement -- the +# Safe Third Country Agreement (its articles are no numbered section of the +# corpus: 'ARTICLE 4'), the Rome Statute, the Universal Declaration, and the +# treaties' acronyms. 'Agreement' alone is none: the FB agreement's articles +# ARE its sections ('FB article 17.1'). +_TREATY = re.compile( + r"(? 3; unsplit -> 0).""" + m = _PIECE_ID.search(chunk.get("id", "")) + return int(m.group()[2:]) if m else 0 + + +_STEMMER = snowballstemmer.stemmer("english") +_STEM_CACHE = {} + + +# Stem pairs Snowball does not merge but that share a legal meaning, so a +# query naming the verb still matches a provision titled with the noun (or +# vice versa). Mapped to the verb form on both index and query sides, which +# is consistent and arbitrary -- the merge is what matters. +_STEM_NORMALIZE = { + "seizur": "seiz", # seizure -> seize + "forfeitur": "forfeit", # forfeiture -> forfeit + "appel": "appeal", # appellate/appellant -> appeal + "detent": "detain", # detention -> detain + "exclus": "exclud", # exclusion -> exclude + "admiss": "admit", # admission/admissibility -> admit + "applic": "appli", # application -> apply + "complianc": "compli", # compliance -> comply + "grievanc": "griev", # grievance -> grieve +} + + +def _stem(word): + """Snowball-stem a word, memoised -- legal text repeats terms heavily. + A small post-stem normalization merges a few verb/noun pairs Snowball + leaves apart ('seize'/'seizure', 'forfeit'/'forfeiture').""" + stemmed = _STEM_CACHE.get(word) + if stemmed is None: + stemmed = _STEMMER.stemWord(word) + stemmed = _STEM_NORMALIZE.get(stemmed, stemmed) + _STEM_CACHE[word] = stemmed + return stemmed + + +# Query-language detection for the bilingual corpus: two-plus French function +# words mark a French query, or an accented character together with one. +# French queries may draw on BOTH languages (guidance and case law exist only +# in English); English queries never surface French chunks (the English +# corpus is complete, and a French twin would only duplicate results). +_FR_HINT = re.compile(r"[àâçéèêëîïôùûüÿœæ]", re.IGNORECASE) +# ASCII-only entries: tokens are matched pre-folding with [a-zA-Z']+ (an +# accented token shatters and never matches -- the accent itself is read by +# _FR_HINT). +_FR_WORDS = frozenset( + "le la les de des une du et est dans pour que qui aux au sur par ne pas " + "peut doit selon lors quelle quel cette ces son ses leur d'un d'une " + "interdiction fouille saisie renvoi territoire".split()) +# The function words an accent may pair with: those that are not also +# English words, so 'travelling with her son from Québec' is not French on +# the strength of 'son' ('pour', 'par' and 'est' likewise). +_FR_ACCENT_COMPANIONS = _FR_WORDS - frozenset("son pour par est".split()) + + +def query_lang(query): + """'fr' when the query reads as French, else 'en'. + + Two French function words mark a French query; an accented character + marks one only together with at least one French function word. An + accent alone is not French evidence: English officer queries routinely + carry accented proper nouns and loanwords (Québec, Montréal, Côte + d'Ivoire, fiancé, résumé, cafés), and routing them as French drops the + English mask, so the French twins of the same provisions displaced the + English results ('traveller from Québec' answered with six French chunks; + code audit 2026-09-21). A lone accented word ('détention') therefore + reads as English; its folded form still matches the English corpus.""" + tokens = re.findall(r"[a-zA-Z']+", query.lower()) + if sum(1 for t in tokens if t in _FR_WORDS) >= 2: + return "fr" + if _FR_HINT.search(query) and any(t in _FR_ACCENT_COMPANIONS for t in tokens): + return "fr" + return "en" + + +def _fold_accents(text): + """ASCII-fold accented characters so a French or mixed query tokenizes + usefully instead of shattering: the [a-z0-9] token pattern used to split + 'détention' into ['d','tention']. NFKD + ASCII-encode drops the combining + marks ('détention' -> 'detention'), which both matches the English corpus + directly and lets the FR->EN synonym bridges fire.""" + return unicodedata.normalize("NFKD", text).encode( + "ascii", "ignore").decode("ascii") + + +def tokenize(text): + """Lower-case, accent-fold, split on word characters, and Snowball-stem + each token, so a query matches a provision even when their word forms + differ -- 'possession' vs 'possess', 'reporting' vs 'report', + 'importation' vs 'import', 'détention' vs 'detention'.""" + return [_stem(w) for w in _TOKEN.findall(_fold_accents(text.lower()))] + + +def _fold_title(text): + """Lower-cased, apostrophes straightened, accents folded: the form a + French Act title and a query are compared in ('Loi sur l’immigration' + and "loi sur l'immigration" meet as "loi sur l'immigration").""" + return _fold_accents(text.lower().replace("’", "'").replace("‘", "'")) + + +# The marginal note of the section that states a French Act's short title. +_FR_TITLE_NOTE = re.compile(r"titre abr[ée]g[ée]", re.IGNORECASE) + + +def _french_short_title(chunk): + """The folded French short title a French twin's 'Titre abrégé' section + states -- '1 Loi sur l’immigration et la protection des réfugiés.' -> + "loi sur l'immigration et la protection des refugies" -- or None for any + other chunk. A trailing parenthetical ('(2016)') is not part of it.""" + if (chunk.get("lang") != "fr" + or chunk.get("doc_type", "legislation") != "legislation" + or not _FR_TITLE_NOTE.match(chunk.get("marginal_note") or "")): + return None + m = re.match(r"\s*\d+(?:\.\d+)?\s+(.+)", chunk.get("text") or "", re.DOTALL) + if not m: + return None + title = re.sub(r"\s*\([^)]*\)$", "", m.group(1).strip().rstrip(".").strip()) + title = _fold_title(title).strip() + return title if len(title) >= 8 else None + + +def _bm25_blob(chunk): + """The text BM25 indexes for a chunk. The topical title is repeated to + weight it above body text; the Act name, code and section are indexed + too, so an Act's own terminology (e.g. "controlled substance") and its + codes/numbers are searchable even when a section's text omits them. The + title is doc_type-aware via topical_title -- for case-law it picks the + case proposition (heading), not the paragraph range (marginal_note), so + a leading case surfaces on a topical query.""" + title = topical_title(chunk) + return " ".join((title, title, chunk["heading"], chunk["part"], + chunk["division"], chunk["act_name"], chunk["act_code"], + chunk["section"], chunk["text"])) + + +# The typo corrector's vocabulary words: four letters or more, lower-cased +# and accent-folded (see LegislationIndex._word_index). +_CORRECTOR_WORD = re.compile(r"[a-z]{4,}") + + +def _corrector_words(text): + """The corrector's vocabulary words of one chunk's text.""" + return _CORRECTOR_WORD.findall(_fold_accents(text.lower())) + + +# The French corrector vocabulary keeps the accents: a French word, four +# letters or more, lower-cased (see LegislationIndex._correct_query). +_CORRECTOR_FR_WORD = re.compile(r"[a-zàâäçéèêëîïôöùûüÿœæ]{4,}") + + +def _corrector_fr_words(text): + """The French corrector's vocabulary words of one French chunk's text.""" + return _CORRECTOR_FR_WORD.findall(text.lower()) + + +def _bm25_code_key(): + """The code half of the BM25 cache key (bm25.code_fingerprint): every + function and table that shapes the cached postings, statistics and + corrector words, and the stemmer's version and class. A change to any + of them rebuilds the cache instead of serving postings a different + tokenizer produced. The class is named by module and qualname, never by + repr (an address, new on every run). Verification 2026-09-22 + (A-CACHE-KEY): swapping _STEMMER for the Porter stemmer left the key + unchanged, so the English stemmer's postings were served to it.""" + try: + stemmer = importlib.metadata.version("snowballstemmer") + except importlib.metadata.PackageNotFoundError: + stemmer = "unknown" + stemmer_class = type(_STEMMER).__module__ + "." + type(_STEMMER).__qualname__ + return bm25.code_fingerprint( + tokenize, _stem, _fold_accents, topical_title, _bm25_blob, + _corrector_words, _corrector_fr_words, sorted(_STEM_NORMALIZE.items()), + _TOKEN.pattern, _CORRECTOR_WORD.pattern, _CORRECTOR_FR_WORD.pattern, + stemmer, stemmer_class) + + +def _section_refs(query): + """Pull explicit section numbers from a query -- 'section 34', 's. 20.1', + 'sect. 36', 'art. 36', and every number of a list such as 'ss. 17 and + 18', 'ss. 17, and 18' or 'sections 17 to 18.1' (see _SECTION_REF). An + article form counts only where it cites a section, not a treaty's + article (_article_cites_section).""" + q = query.lower() + return {n for m in _SECTION_REF.finditer(q) + if m.group("article") is None or _article_cites_section(q, m) + for n in _SECTION_NUMS.findall(m.group(0))} + + +def _article_cites_section(query_lc, match): + """Whether an article form _SECTION_REF matched in the lower-cased query + cites a section of the corpus. + + In French syntax -- a French query (query_lang), "l'article 36", + 'article 36 de la LIPR', 'article 17 et 18' -- 'article' is how a statute + cites its sections, so it does, unless a treaty follows it ("l'article 33 + de la Convention"). In English an 'Article N' is a treaty's or a + collective agreement's, so it cites a section only in a query that names + no treaty (_TREATY), and only as 'article(s)' or a dotted 'art.'/'arts.' + ('martial arts 5 years' cites nothing). Verification 2026-09-22 + (A-PIN-ARTICLE): 'IRPA 98 exclusion under Article 1 of the Refugee + Convention' pinned IRPA s. 1 'Short title' to rank 1, 'Vienna Convention + on Consular Relations article 36 IRPA detention' IRPA s. 36 and 'ICCPR + article 14 fair trial IRPA' IRPA s. 14 'Regulations'.""" + start, end = match.span() + french = (query_lang(query_lc) == "fr" + or query_lc[max(0, start - 2):start] in ("l'", "l’") + or _FR_ARTICLE_TAIL.match(query_lc, end) is not None) + if french: + return _TREATY_AFTER.match(query_lc, end) is None + return (_DOTLESS_ART.match(match.group(0)) is None + and _TREATY.search(query_lc) is None) + + +def topical_title(chunk): + """Return the chunk's topic-bearing string, used wherever a section's + 'title' is weighted for retrieval -- BM25 indexing, the title-match boost, + and the semantic embedding. Differs by doc_type because the field that + carries the topic differs: legislation/agreement/directive/delegation use + the marginal_note (section heading); D-memoranda use 'part' because their + marginal_note is a generic banner; case-law uses 'heading' because its + marginal_note is just the paragraph range ('paras 11-13') and the case + proposition lives in heading.""" + doc_type = chunk.get("doc_type") + if doc_type == "memorandum": + return chunk.get("part") or chunk["marginal_note"] + if doc_type == "caselaw": + return chunk.get("heading") or chunk["marginal_note"] + return chunk["marginal_note"] + + +def _provision_units(text): + """Citable parts of a provision, for pinpoint scoring -- a list of + (citation_suffix, scoring_text, snippet). One entry per paragraph, with its + subsection chapeau prepended to scoring_text for context, plus one per + paragraph-less subsection. Returns [] when the provision is too flat to + pinpoint (fewer than two subdivisions).""" + marks = list(_MARKER.finditer(text)) + if len(marks) < 2: + return [] + spans = [] + for i, m in enumerate(marks): + end = marks[i + 1].start() if i + 1 < len(marks) else len(text) + spans.append((m.group(1), text[m.start():end].strip())) + units, cur_sub, cur_intro = [], "", "" + for j, (token, body) in enumerate(spans): + if "[Repealed" in body[:40]: + if token[0].isdigit(): + cur_sub, cur_intro = f"({token})", "" + continue + if token[0].isdigit(): + cur_sub, cur_intro = f"({token})", body + nxt = spans[j + 1][0] if j + 1 < len(spans) else "" + if not nxt or nxt[0].isdigit(): + units.append((cur_sub, body, body)) # subsection has no paragraphs + # otherwise the chapeau is emitted via its paragraphs below + else: + label = f"{cur_sub}({token})" if cur_sub else f"({token})" + units.append((label, f"{cur_intro} {body}".strip(), body)) + return units + + +class LegislationIndex: + # Chunks with no vector in embeddings.npz (_load_semantic): BM25-only + # for them. None until the store is read (a synthetic index). The + # store's staleness by content -- a vector embedded from older text -- + # is canlex.embed.currency's to count; tools/predeploy.py fails a deploy + # tree on either (audit follow-up 2026-09-21, Missing 15). None too when + # the store or the query embedder failed to load: semantic search is + # off, which `semantic` says. + embedding_missing = None + # Whether semantic search is on (_load_semantic): False when there is no + # store or it or the embedder failed to load -- keyword (BM25) search + # only -- and None until _load_semantic runs (a synthetic index). + semantic = None + + def __init__(self): + self.chunks = [] + paths = sorted(PROCESSED_DIR.glob("*.json")) + stamps = bm25.file_stamps(paths) # before reading: the cache key + for path in paths: + self.chunks.extend(json.loads(path.read_text(encoding="utf-8"))) + if not self.chunks: + raise RuntimeError( + f"No processed legislation in {PROCESSED_DIR}. Run 'canlex.ingest' first.") + self._tri_index = None # lazy; see _fuzzy_term + self._word_vocab = {} # lazy, per scope; see _word_index (typo repair) + self._word_tri = {} + # Cached recall mask for the common case (English query, no + # filters): everything but the French twins. + self._en_mask = [c.get("lang", "en") != "fr" for c in self.chunks] + self._build_bm25(cache_key=bm25.cache_key(stamps, _bm25_code_key())) + self._prebuild_vocabularies() + self._build_note_tokens() + self._build_xref() + self._build_appendix_index() + # Both read the curated JSON, and both are optional: the criminality + # route's promotion allowlist (_criminality_route) and the read-next + # map (_cosurface_guidance). + self._crim_allow = criminality_allowlist() + self._load_see_also() + self._load_semantic() + self._load_reranker() + + def _build_bm25(self, cache_key=None): + """BM25 postings plus TWO sets of collection statistics: one over the + whole corpus, one over the English scope alone. + + An English query never retrieves the French twins (see _en_mask), yet + a single corpus-wide idf/avgdl let those unreachable chunks set the + rarity of every shared token (numbers, act codes, proper nouns, folded + cognates) and drag the average document length. Ingesting the French + D-memos and NJC directives moved legacy Hit@1 0.767 -> 0.750 that way, + without a single French chunk ever being retrievable -- the statistics + must be scoped to what the query can actually reach. + + The postings are arrays (bm25.Postings). With a cache_key (the loaded + corpus; __init__ passes one) they are read from the BM25 cache when + it holds that key, else built and written through; a synthetic index + (the tests) builds in process. Audit follow-up 2026-09-21: the build + took 8.5-18 s of every start and its 7.04 M Python tuples 522 MB. + """ + data = bm25.load(cache_key, len(self.chunks)) if cache_key else None + if data is None: + data = bm25.build(self.chunks, _bm25_blob, tokenize, _corrector_words, + _corrector_fr_words) + if cache_key: + bm25.save(cache_key, data) + self._install_bm25(data) + + def _install_bm25(self, data): + """The index's BM25 fields from a bm25.Bm25Data: the postings, the + corpus-wide and English-scoped idf and average length, each chunk's + length normalisation for both scopes, and the corrector's word + counts. Arithmetic as it was in pure Python, so scores are + bit-identical: idf in Python floats, the per-chunk norm + K1 * (1 - B + B * dl / avgdl) in float64. A term the English corpus + never uses is absent from idf_en rather than zero-weighted: it can + only match chunks the mask already excludes.""" + n = len(self.chunks) + self.postings = data.postings + self.doc_len = data.doc_len + self.avgdl = int(data.doc_len.sum()) / n + self.idf = {t: math.log(1 + (n - d + 0.5) / (d + 0.5)) + for t, d in zip(data.postings.terms, data.df.tolist())} + en_docs = data.en_docs + self.avgdl_en = data.en_total_len / (en_docs or 1) + self.idf_en = {t: math.log(1 + (en_docs - d + 0.5) / (d + 0.5)) + for t, d in zip(data.postings.terms, data.df_en.tolist()) if d} + dl = data.doc_len.astype(np.float64) + self._bm25_norm = K1 * (1 - B + B * dl / self.avgdl) + self._bm25_norm_en = K1 * (1 - B + B * dl / self.avgdl_en) + self._corrector_counts = data.words + + def _build_note_tokens(self): + """Pre-tokenise each chunk's topical title (see topical_title) for the + title-match boost in search(). Each chunk is also flagged as a + regulation (act codes beginning SOR/C.R.C.) for the Act-over-regulation + preference, as collective-agreement back-matter (memoranda and letters + with no article number) for the back-matter penalty, and as a + not-in-force amendment for the pending-law penalty.""" + self._note_tokens = [] + self._is_regulation = [] + self._is_backmatter = [] + self._is_nif = [] + self._is_repealed = [] + self._inapplicable = [] + # Guidance sub-tiers (ATIP_PENALTY, APPLICANT_PENALTY): read from + # structured fields the ingesters set, never from the act_name text. + self._is_atip = [] + self._is_applicant = [] + for c in self.chunks: + self._note_tokens.append(set(tokenize(topical_title(c)))) + note = _applicability(c) + self._inapplicable.append(bool(note)) + if note: + c["applicability"] = note # rendered by server._format_section + self._is_regulation.append( + c.get("doc_type", "legislation") == "legislation" + and c["act_code"].startswith(("SOR", "C.R.C"))) + self._is_backmatter.append( + c.get("doc_type") == "agreement" + and not str(c["section"])[:1].isdigit()) + self._is_nif.append(c.get("status") == "not-in-force") + self._is_repealed.append(c.get("status") == "repealed") + memo = c.get("doc_type") == "memorandum" + self._is_atip.append(memo and c.get("source_status") == ATIP_STATUS) + self._is_applicant.append( + memo and str(c.get("act_code", "")).startswith(_APPLICANT_CODES)) + + def _load_see_also(self): + """Resolve the curated 'read next' keys (roadmap T3.6) to chunk + indices once, at load: self._see_also is {commentary section: + [chunk index, ...]} in curated order, one chunk per see_also item + (the item's pinned quote chunk when it resolves, else the first + English piece the key names). The curated data/curated/*.json + see_also fields are preferred; the research draft is read only when + CANLEX_SEE_ALSO_DRAFT is set ('1' for the default draft, or a path), + because a draft is unreviewed. With neither, read-next is off and + self.see_also_source is None -- the approved curated baseline + predates the field, so a composed release runs this way.""" + self._see_also = {} + self.see_also_source = None + self.see_also_unresolved = [] # (section, key) pairs that resolved to nothing + items, keys = curated_see_also(), {} + if items: + self.see_also_source = "curated" + else: + draft = os.environ.get("CANLEX_SEE_ALSO_DRAFT", "").strip() + if draft: + path = SEE_ALSO_DRAFT if draft == "1" else Path(draft) + items, keys = draft_see_also(path) + if items: + self.see_also_source = f"draft ({path.name})" + if not items: + return + by_id = {c.get("id"): i for i, c in enumerate(self.chunks)} + by_act = defaultdict(list) + for i, c in enumerate(self.chunks): + if c.get("lang", "en") != "fr": + by_act[c.get("act_code")].append(i) + for section, entries in items.items(): + out = [] + for item in entries: + spec, pin = _see_also_spec(item, keys) + idx = by_id.get(pin) if pin else None + if idx is not None and not _see_also_match(self.chunks[idx], spec): + idx = None # a pin that no longer fits its key + if idx is None and spec: + idx = next((i for i in by_act.get(spec.get("act_code"), ()) + if _see_also_match(self.chunks[i], spec)), None) + if idx is None: + key = item.get("key", spec) if isinstance(item, dict) else item + self.see_also_unresolved.append((section, key)) + elif idx not in out: + out.append(idx) + if out: + self._see_also[section] = out + + def _cosurface_guidance(self, top, keep=None): + """Append, after the results, the official guidance and statute the + curated commentary entries in `top` name as 'read next' (their + see_also keys, resolved by _load_see_also) -- like the directive + appendices, so a curated answer that ranks first no longer hides the + official passage on the same point. Never displaces a result: `top` + is returned as a prefix, extended by at most READ_NEXT_CAP chunks + not already shown. Items keep their curated order, first result + first, except that a chunk its ingester flagged as an unofficial + ATIP copy of a withdrawn chapter yields to official sources when the + cap binds. keep: the caller's filter as a predicate over chunk + indices (search() passes it on a filtered search); a chunk failing + it is never co-surfaced -- a doc_type='commentary' search must not + return statutes and memoranda (2026-09-15 verifier).""" + see_also = getattr(self, "_see_also", None) + if not see_also: + return list(top) + have = set(top) + official, atip = [], [] + for idx in top: + c = self.chunks[idx] + if c.get("doc_type") != "commentary": + continue + for j in see_also.get(c.get("section"), ()): + if j in have or (keep is not None and not keep(j)): + continue + have.add(j) + if self.chunks[j].get("source_status") == ATIP_STATUS: + atip.append(j) + else: + official.append(j) + return list(top) + (official + atip)[:READ_NEXT_CAP] + + def _states_named(self, query_lc): + """The slugs of the US-state commentary pages whose state the + lower-cased query names as whole words ('texas', 'new york'); such + a page is exempt from the one-source cap (STATE_ONE_SOURCE).""" + pages = getattr(self, "_state_pages", None) + if pages is None: + pages = [] + for c in self.chunks: + cid = str(c.get("id", "")) + if not cid.startswith("commentary-state-"): + continue + slug = cid[len("commentary-state-"):] + name = (c.get("marginal_note") or "").split("—", 1)[-1].strip().lower() + if not name: + name = slug.replace("-", " ") + pages.append((slug, re.compile(r"(? b - a + for _, (a2, b2) in spans)) + + def _s36_piece(self, query_lc): + """The English IRPA s. 36 chunk the criminality route may inject + when no s. 36 piece was retrieved: the piece carrying s. 36(3) -- + record suspensions and pardons, the youth-record rule, transborder + criminality -- when the query has those words, else the first.""" + pieces = getattr(self, "_s36_pieces", None) + if pieces is None: + pieces = sorted( + (i for i, c in enumerate(self.chunks) + if c.get("act_code") == "I-2.5" and c.get("section") == "36" + and c.get("lang", "en") != "fr"), + key=lambda i: _piece_no(self.chunks[i])) + self._s36_pieces = pieces + if not pieces: + return None + if _S36_3_WORDS.search(query_lc or "") or "transborder" in (query_lc or ""): + for i in reversed(pieces): + if "(3)" in self.chunks[i].get("text", ""): + return i + return pieces[0] + + def _build_appendix_index(self): + """Index directive appendices by (act_code, letter), so a directive + section that cites 'Appendix C' can pull that appendix into the result + set -- a directive's rate-table appendices are bare numbers and rank + poorly on a natural-language query, yet the section citing them is of + little use without them.""" + self._appendix = defaultdict(list) + for idx, c in enumerate(self.chunks): + if c.get("doc_type") != "directive": + continue + m = _APPENDIX_HEAD.match(c["marginal_note"]) + if m: + self._appendix[(c["act_code"], m.group(1).upper())].append(idx) + + def _load_semantic(self): + """Load precomputed embeddings and the query embedder. + + Any failure (missing embeddings, or numpy/model unavailable) degrades the + index to BM25-only rather than breaking retrieval. embedding_missing + is set only once the embedder has loaded: a failure leaves it None + and `semantic` False. Verification 2026-09-22 (A-EMB-2): it was set + before Embedder() and kept when that raised, so an index whose + embedder could not load reported 0 missing and canlex_list_acts said + semantic search was current while the index ran BM25 only. + """ + self.semantic = False + self.embedding_missing = len(self.chunks) + emb_path = PROCESSED_DIR / "embeddings.npz" + if not emb_path.exists(): + print("CanLex index: no embeddings.npz; using BM25 only " + "(run 'canlex.embed' to enable semantic search).", file=sys.stderr) + return + try: + import numpy as np + from .embed import Embedder + # Multi-vector layout: rows share a chunk id when the chunk was + # embedded as several windows (long D-memo sections etc.); a + # chunk's semantic score is the max over its windows. Rows whose + # id matches no loaded chunk are dropped; chunks with no rows + # count as missing (BM25-only for them). + chunk_ordinal = {c["id"]: i for i, c in enumerate(self.chunks)} + with np.load(emb_path) as data: + all_ids = data["ids"].tolist() + # Stored as float16 (see embed.build); compute in float32. + all_vecs = data["vectors"].astype(np.float32) + keep = [(chunk_ordinal[i], r) for r, i in enumerate(all_ids) + if i in chunk_ordinal] + self._np = np + self.vectors = all_vecs[[r for _, r in keep]] + self._row_chunk = np.array([o for o, _ in keep], dtype=np.int64) + covered = set(self._row_chunk.tolist()) + missing = len(self.chunks) - len(covered) + self.embedder = Embedder() + self.embedding_missing = missing + self.semantic = True + if missing: + print(f"CanLex index: {missing}/{len(self.chunks)} sections lack " + f"embeddings; re-run 'canlex.embed' to refresh.", file=sys.stderr) + except Exception as exc: + print(f"CanLex index: semantic search disabled ({type(exc).__name__}: " + f"{exc}); using BM25 only.", file=sys.stderr) + self.semantic = False + self.embedding_missing = None + + def _load_reranker(self): + """Load the cross-encoder reranker; degrade to the fusion order on failure.""" + self.reranker = None + try: + from .rerank import Reranker + self.reranker = Reranker() + except Exception as exc: + print(f"CanLex index: reranker disabled ({type(exc).__name__}: {exc}); " + f"using hybrid fusion order.", file=sys.stderr) + + def _fuzzy_term(self, term): + """The closest vocabulary term for a token the corpus has never seen + -- a typo bridge ('detension' -> 'detention'), scored by character- + trigram Dice similarity with a strict 0.6 floor and a length guard. + Only ever consulted for UNKNOWN tokens, so correctly-spelled queries + are untouched. The trigram index is built lazily on first miss.""" + if len(term) < 4: + return None + tri_index = self._trigram_index() + tris = {term[i:i + 3] for i in range(len(term) - 2)} + counts = Counter() + for t in tris: + counts.update(tri_index.get(t, ())) + best, best_score = None, 0.6 + for cand, n in counts.most_common(50): + # Typos rarely corrupt a word's opening letters; requiring a + # shared 2-char prefix kills drift like 'detens' -> 'pretens' + # while keeping 'seisur' -> 'seizur'. + if abs(len(cand) - len(term)) > 2 or cand[:2] != term[:2]: + continue + score = 2 * n / (len(tris) + len(cand) - 2) + if score > best_score: + best, best_score = cand, score + return best + + def _trigram_index(self): + """The stem vocabulary's character-trigram index (_fuzzy_term), + built once: at load by _prebuild_vocabularies, or on first use by a + synthetic index.""" + if self._tri_index is None: + tri = defaultdict(set) + for vocab_term in self.idf: + for i in range(len(vocab_term) - 2): + tri[vocab_term[i:i + 3]].add(vocab_term) + self._tri_index = tri + return self._tri_index + + def _prebuild_vocabularies(self): + """Build the typo bridge's trigram index and the corrector's two + vocabularies at load rather than inside the first query that needs + them (about 2.7 s on the first typo or proper-noun query after every + restart; audit follow-up 2026-09-21), so no search pays for them and + two concurrent searches never build them twice.""" + self._trigram_index() + for scope in ("en", "all", "fr"): + self._word_index(scope) + + @staticmethod + def _edit_distance(a, b, cap=3): + """Damerau-Levenshtein distance, early-exited past `cap` -- typos + are substitutions, insertions, deletions and transpositions, and a + transposition ('undeclraed') must cost 1, not 2.""" + if abs(len(a) - len(b)) > cap: + return cap + 1 + prev2, prev = None, list(range(len(b) + 1)) + for i, ca in enumerate(a, start=1): + cur = [i] + best_row = i + for j, cb in enumerate(b, start=1): + cost = 0 if ca == cb else 1 + d = min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + cost) + if (prev2 is not None and i > 1 and j > 1 + and ca == b[j - 2] and a[i - 2] == cb): + d = min(d, prev2[j - 2] + 1) + cur.append(d) + best_row = min(best_row, d) + if best_row > cap: + return cap + 1 + prev2, prev = prev, cur + return prev[-1] + + # The corpus frequency a correction needs before it may stand IN for the + # officer's word rather than travel beside it (see _correct_query). + _CORRECTION_FLOOR = 100 + + def _word_index(self, scope): + """The (word Counter, trigram index) the corrector draws on, built + lazily per scope: 'en' from the English chunks, 'all' from both + languages (a French query searches both), 'fr' from the French + chunks with their accents (a French query's accented words). Scoped like _build_bm25's + idf_en, and for the same reason: a corpus-wide vocabulary corrected + English queries into French words -- 'Venmo' -> 'vente', 'Percocet' + -> 'percoit', 'Modelo' -> 'modele' (code audit 2026-09-21).""" + if not self._word_vocab: # None or {}: nothing built yet + self._word_vocab, self._word_tri = {}, {} + if scope not in self._word_vocab: + counted = getattr(self, "_corrector_counts", None) or {} + words = counted.get(scope) + if words is None: # a synthetic index + words = Counter() + for c in self.chunks: + french = c.get("lang", "en") == "fr" + if scope == "fr": + if french: + words.update(_corrector_fr_words(c["text"])) + continue + if scope == "en" and french: + continue + words.update(_corrector_words(c["text"])) + tri = defaultdict(set) + for w, n in words.items(): + if n >= 5: # junk never enters the index + for i in range(len(w) - 2): + tri[w[i:i + 3]].add(w) + self._word_vocab[scope], self._word_tri[scope] = words, tri + return self._word_vocab[scope], self._word_tri[scope] + + def _best_correction(self, word, scope="en"): + """(nearest corpus WORD -- not stem -- for a misspelling, its + Damerau-Levenshtein distance, its corpus frequency), or (None, 0, 0). + + Trigram index generates candidates; Damerau-Levenshtein picks the + winner (distance <= 2, shared 2-char prefix, frequency floor of 5 + so OCR debris in the corpus can never be a 'correction'). The + result is real text: it feeds the semantic embedder and the + cross-encoder, which score words, not stems.""" + words, tri = self._word_index(scope) + tris = {word[i:i + 3] for i in range(len(word) - 2)} + candidates = set() + for t in tris: + candidates |= tri.get(t, set()) + best, best_key = None, (3, 0) # (distance, -frequency) + for cand in candidates: + if abs(len(cand) - len(word)) > 2 or cand[:2] != word[:2]: + continue + dist = self._edit_distance(word, cand, cap=2) + if dist > 2: + continue + key = (dist, -words[cand]) + if key < best_key: + best, best_key = cand, key + if best is None: + return None, 0, 0 + return best, best_key[0], -best_key[1] + + def _correct_word(self, word, scope="en"): + """Nearest corpus word for a misspelling, or None (_best_correction).""" + return self._best_correction(word, scope)[0] + + def _correct_query(self, query, lang="en"): + """The query with misspelled words repaired; the query itself when + nothing needed fixing. A word is only 'misspelled' when its stem is + absent from the whole BM25 vocabulary, so correctly-spelled queries + pass through untouched by construction. `lang` (query_lang) picks + the vocabulary: an English query is corrected into English words. + + A candidate is a correction at all only within one edit, or two for + a word of seven letters or more (two edits rewrite 40% of a short + word). It REPLACES the word only when it is strong -- a common word, + frequency >= _CORRECTION_FLOOR -- and the word is lower-case: a + capitalised token is read as a proper noun, which only ever has a + strong correction APPENDED beside it. A weak correction of a + lower-case unknown word is appended too, and a word the corpus does + hold rarely (df <= 2) gets a strong one appended, so the officer's + own word keeps its recall for BM25, the embedder and the + cross-encoder. Code audit 2026-09-21: the corrector replaced any + unknown word within two edits of a corpus word -- 'Biden' -> 'bids', + 'Tesla' -> 'test', 'Monero' -> 'money', 'Rolex' -> 'role' -- so the + officer's question was silently answered as a different one. + + A French query's ACCENTED word is a French word: it is corrected + only against the French chunks' own words, accents kept ('fr' + scope), never into the unaccented English word -- audit follow-up + 2026-09-21: 'détension du résident' became 'detention du résident'. + A lone accented word in an English query stays English (query_lang's + documented trade-off) and is corrected as before.""" + scope = "all" if lang == "fr" else "en" + out, changed = [], False + for tok in query.split(): + word = _fold_accents(tok.lower()).strip(".,;:!?()'\"") + if len(word) >= 4 and word.isalpha(): + stem = _stem(word) + df = len(self.postings.get(stem, ())) + if df <= 2: + if lang == "fr" and _FR_HINT.search(tok): + word = tok.lower().strip(".,;:!?()'\"«»’") + fix, dist, freq = self._best_correction(word, "fr") + else: + fix, dist, freq = self._best_correction(word, scope) + if fix and fix != word and (dist == 1 or len(word) >= 7): + strong = freq >= self._CORRECTION_FLOOR + proper = tok.lstrip("(\"'")[:1].isupper() + if df == 0 and strong and not proper: + # Unknown everywhere, a common word one slip + # away: stand in for it. + out.append(fix) + changed = True + continue + if strong or (df == 0 and not proper): + # A proper noun, a weak correction, or a word + # the corpus holds rarely (df <= 2 -- often a + # typo that collides with OCR debris, but also + # a chemical in a CDSA schedule, a case name, a + # brand): the correction travels beside the + # word, never in its place. + out.append(f"{tok} {fix}") + changed = True + continue + out.append(tok) + return " ".join(out) if changed else query + + def _bm25_accumulate(self, query, allowed=None, en_only=False): + """(scores, touched): float64 BM25 scores by chunk index, and the + chunks that scored, in first-seen order. + + allowed: optional boolean list by chunk index. Filtered searches + pass it so recall competes only within scope -- filtering after + recall meant a doc_type='caselaw' search drew from whatever case law + happened to survive the open-corpus top-N competition. + + en_only: score with the English-scoped collection statistics (see + _build_bm25). Set for every English query, whose mask already bars + the French twins from the results they would otherwise be weighting. + + Each term's postings are added at once (np.add.at), term by term in + the order the per-posting Python loop used, with the same float64 + arithmetic -- so every score is bit-identical to that loop's and the + tie order is its insertion order (audit follow-up 2026-09-21: a + 12-term question touched 196,652 postings in Python, 108-149 ms).""" + idf_map = self.idf_en if en_only else self.idf + norm = self._bm25_norm_en if en_only else self._bm25_norm + mask = self._mask_array(allowed) + scores = np.zeros(len(self.chunks), dtype=np.float64) + touched = [] + for term in set(tokenize(query)): + idf = idf_map.get(term) + weight = 1.0 + if idf is None: + # Unknown token: try the typo bridge at a discount, so a + # misspelled query still reaches the provision while a real + # match always outranks a fuzzy one. + alt = self._fuzzy_term(term) + if alt is None or alt not in idf_map: + continue + term, idf, weight = alt, idf_map[alt], 0.7 + docs, tfs = self.postings.span(term) + if mask is not None: + keep = mask[docs] + docs, tfs = docs[keep], tfs[keep] + tf = tfs.astype(np.float64) + np.add.at(scores, docs, weight * idf * tf * (K1 + 1) / (tf + norm[docs])) + touched.append(docs) + return scores, bm25.first_seen(touched) + + def _bm25_scores(self, query, allowed=None, en_only=False): + """{chunk index: BM25 score} for the chunks the query's terms reach, + in first-seen order (see _bm25_accumulate).""" + scores, touched = self._bm25_accumulate(query, allowed, en_only) + return dict(zip(touched.tolist(), scores[touched].tolist())) + + def _bm25_ranking(self, query, allowed=None, en_only=False): + """The chunks the query's terms reach, best BM25 score first; ties + in first-seen order -- exactly sorted(scores, key=scores.get, + reverse=True) over _bm25_scores, without the Python sort.""" + scores, touched = self._bm25_accumulate(query, allowed, en_only) + return touched[np.argsort(-scores[touched], kind="stable")].tolist() + + def _mask_array(self, allowed): + """`allowed` (a boolean list by chunk index) as a numpy array, or + None; the last one converted is kept, so the cached English mask is + converted once, not on every search.""" + if allowed is None: + return None + cached = getattr(self, "_mask_cache", None) + if cached is not None and cached[0] is allowed: + return cached[1] + arr = np.asarray(allowed, dtype=bool) + self._mask_cache = (allowed, arr) + return arr + + def _semantic_ranking(self, query, allowed=None): + np = self._np + qv = self.embedder.encode_query(query) + row_sims = self.vectors @ qv + # Aggregate windows to chunks by max: a long section scores as its + # best-matching window, so a fact buried mid-memo is reachable. + sims = np.full(len(self.chunks), -1.0, dtype=np.float32) + np.maximum.at(sims, self._row_chunk, row_sims) + if allowed is not None: + # Mask out-of-scope chunks so a filtered search ranks (and reports + # confidence) within its own scope, not corpus-wide. + sims = np.where(self._mask_array(allowed), sims, -1.0) + # Over-fetch 4x: the recall-stage source cap (_capped_top) drops the + # surplus chunks of any one deep source, and the extra depth is what + # lets other sources backfill the freed CANDIDATES slots. + order = np.argsort(sims)[::-1][:CANDIDATES * 4] + # The top cosine similarity doubles as a corpus-coverage signal: a query + # the corpus cannot answer has no passage close to it. + return [int(i) for i in order], float(sims.max()) + + def _rerank_doc(self, idx): + """What the cross-encoder reads for a chunk: its citation and its + topical title (topical_title: a section's marginal note, a D-memo's + subject in 'part', a decision's proposition in 'heading'), then its + text. The header carried the marginal note, which for a D-memo is a + generic banner and for case law a paragraph range ('paras 11-13'). + Adopted on the eval (audit follow-up 2026-09-21, group A item 8e): + legacy Hit@3 .9583 -> .9688, MRR .8651 -> .8671; holdout (undated) + Hit@3 .9250 -> .9500, MRR .8383 -> .8404; no Hit@1 moved (the + criminality slice's Hit@3 .6333 -> .6000: one question 2 -> 3).""" + c = self.chunks[idx] + return f"{c['citation']} — {topical_title(c)}\n{c['text']}" + + def _source_key(self, idx, exempt=frozenset()): + """The parent document a chunk belongs to, for diversity capping. Returns + None for primary instruments -- legislation, collective agreements and + directives -- whose sections are distinct provisions and are never + capped; case law is keyed by citation, memoranda by memo number. + + The 51 US-state commentary pages are one source, ('commentary', + 'state'), when STATE_ONE_SOURCE is on: uncapped, a general + disposition question filled the officer's top 6 with arbitrary + states (expungement: Iowa, Illinois, Indiana). A page whose state + the query names is passed in `exempt` (slugs, see _states_named) + and stays uncapped, so 'Texas deferred adjudication' keeps Texas. + + Benefit-plan booklets also return None: a booklet's sections are + distinct topics (eligibility, coverage, claims, appeals...) like + statute sections, so capping the whole booklet to SOURCE_CAP chunks + would starve a detailed benefits query. They are not in + PRIMARY_DOC_TYPES, though -- they are guidance, not governing + instruments, so _ensure_primary does not pull them in as if they were + enacted law. + + Exception: the pieces of one long section/article that was split for + embedding (ids ending '-p') share a key, so the diversity cap stops a + single multi-part provision from filling the result set with its own + fragments. This applies to the otherwise-uncapped primary types + (legislation, agreement, directive, delegation): pieces are keyed by + their shared base id, distinct sections keep distinct keys. (Memoranda + and case law are already capped at the parent level, so their pieces are + bounded without a separate key.)""" + c = self.chunks[idx] + doc_type = c.get("doc_type", "legislation") + m = _PIECE_ID.search(c.get("id", "")) + if m and doc_type in PRIMARY_DOC_TYPES: + return ("piece", c["id"][:m.start()]) # pieces of one chunk together + # Commentary entries are like benefits chunks: each is a distinct + # curated topic, so capping the dataset as one source would starve a + # multi-entry question; not PRIMARY, so never force-pulled as law. + # The per-state pages are the exception (see the docstring). + if doc_type == "commentary": + cid = str(c.get("id", "")) + if (STATE_ONE_SOURCE and cid.startswith("commentary-state-") + and cid[len("commentary-state-"):] not in exempt): + return ("commentary", "state") + return None + if doc_type in PRIMARY_DOC_TYPES or doc_type == "benefits": + return None + if doc_type == "memorandum": + # The key must be the parent DOCUMENT, and which field carries that + # identity differs across the memorandum families. ENF chapters, + # PDI pages and the IRB guidelines each have their own act_code and + # use `section` to subdivide it, so act_code is the document. + # _FLAT_MEMO_CODES are the two families where one act_code covers + # many independent items -- ~360 D-memoranda and ~1,000 AMPS + # contraventions -- and `section` (memo number / contravention + # code) is the unit a reader would cite. + # + # Keying every memorandum on `section` (as this did) gave each ENF + # section and each PDI page its own source key, so neither cap ever + # bound at the document level: the ~150-page PDI refugee tree could + # fill the whole candidate pool and push IRPA s. 112 out of a PRRA + # query entirely. + code = c["act_code"] + return ("memorandum", + c["section"] if code in _FLAT_MEMO_CODES else code) + return (doc_type, c["act_code"]) # one decision, keyed by citation + + def _capped_top(self, ordered, exempt=frozenset()): + """The first CANDIDATES indices of `ordered`, keeping at most + RECALL_CAP per capped source. The recall-stage counterpart of + _diversify: where _diversify reorders the fused candidates so a deep + source cannot monopolise the visible results, this stops that source + from monopolising the candidate pool itself -- skipped surplus frees + slots that later, weaker-ranked chunks of OTHER sources backfill.""" + out, counts = [], defaultdict(int) + for idx in ordered: + key = self._source_key(idx, exempt) + if key is not None: + counts[key] += 1 + if counts[key] > RECALL_CAP: + continue + out.append(idx) + if len(out) >= CANDIDATES: + break + return out + + def _fill_pool(self, candidates, exempt=frozenset()): + """The rerank pool: the first RERANK_POOL fused candidates, holding + at most POOL_CAP chunks of any one capped source. Skipped surplus + frees slots that the next-ranked chunks of OTHER sources backfill, + so the cross-encoder sees the field, not one deep decision.""" + pool, counts = [], defaultdict(int) + for idx in candidates: + key = self._source_key(idx, exempt) + # Pieces of ONE long provision (FB art. 17 is eight) are not a + # monopoly to break: the reranker must see every piece to find + # the one that answers, and _diversify shows only two of them + # anyway. Capping them here lost the discipline article outright. + if key is not None and key[0] != "piece": + counts[key] += 1 + if counts[key] > POOL_CAP: + continue + pool.append(idx) + if len(pool) >= RERANK_POOL: + break + return pool + + def _diversify(self, ordered, exempt=frozenset()): + """Reorder so no single case, memorandum, agreement or directive can + monopolise the results: once a source has contributed SOURCE_CAP chunks, + its remaining chunks are deferred below every other candidate. This stops + a heavily paragraph-chunked decision from crowding out the statute it + interprets. Legislation is never capped.""" + kept, deferred, counts = [], [], defaultdict(int) + for idx in ordered: + key = self._source_key(idx, exempt) + if key is None: + kept.append(idx) + continue + counts[key] += 1 + (kept if counts[key] <= SOURCE_CAP else deferred).append(idx) + return kept + deferred + + def _acts_named(self, query_lc): + """The act_codes whose short or full name the lower-cased query + contains as whole words. Names under four characters, and case-law + names, never count.""" + forms = getattr(self, "_act_name_forms", None) + if forms is None: + forms = {} + for c in self.chunks: + if c.get("doc_type") == "caselaw": + continue + code = c.get("act_code") + if not code or code in forms: + continue + names = set() + for key in ("act_short", "act_name"): + n = (c.get(key) or "").strip().lower() + if len(n) >= 4: + names.add(n) + if names: + forms[code] = [re.compile(r"(? 3, twice). So s. 36 takes the + # first slot and the promotions are shown in title order. + # An appeal to the IAD is governed by IRPA ss. 63-64, which the route + # does not promote: on 'can he appeal the removal order to the IAD?' + # after a six-month sentence the route sank IRPA s. 64 from rank 3 to 19 + # (2026-09-15 verifier), so such questions are exempt from the route. + # A bare 'appeal' is not the cue -- 'conviction under appeal (right + # now)' and 'the court of appeal substituted a discharge' are s. 36 + # questions. Nor is a bare 'right of appeal': it is also how an officer + # asks whether a foreign conviction is final ('while the right to appeal + # has not expired', 'appeal rights not exhausted' -- ENF 2 s. 15.2), and + # exempting those put IRPA s. 64 'No appeal' where s. 36 belonged + # (2026-09-15 track B verifier: s. 36 gone from the top 6 on 3 of 4). + # So route_exempt: the IAD or an appeal of the removal order exempts; an + # appeal right exempts only beside a removal/deportation order, or beside + # a permanent resident with no finality word (_FINALITY). + _ROUTE_EXEMPT = re.compile( + r"(? 4 (criminality Hit@3 .5667 -> .6000, MRR .4410 -> + .4493 with this rule; no other eval question moves). + - a section allowlisted only because an equivalency pairing verifies + against it is promoted only if its title shares a distinctive + query word (idf >= _ROUTE_IDF_FLOOR -- not 'of' or 'the'): + counting two IRPA s. 36 pieces once opened a slot that CC s. 173 + 'Indecent acts' (the indecent-exposure pairing) took with no title + overlap on the legacy disorderly-conduct question (2026-09-15 + track B verifier). A governing provision needs no overlap -- IRPA + s. 37 'Organized criminality' answers a patched-gang-member + question that says neither word, and IRPR s. 18 a shoplifting + equivalency question.""" + top, rest = ordered[:top_k], ordered[top_k:] + + def provision(i): + return self._source_key(i) or ("chunk", i) + + def governing(i): + c = self.chunks[i] + return ((c.get("act_code"), c.get("section")) in _CRIMINALITY_STATIC + or c.get("act_code") in _CRIMINALITY_ACTS) + + have = {provision(i) for i in top if is_primary(i)} + need = 2 - len(have) + allow = getattr(self, "_crim_allow", None) or _CRIMINALITY_STATIC + cands = [i for i in rest if is_primary(i) and ( + (self.chunks[i].get("act_code"), self.chunks[i].get("section")) in allow + or self.chunks[i].get("act_code") in _CRIMINALITY_ACTS)] + if q_tokens: + cands.sort(key=lambda i: self._title_score(i, q_tokens), reverse=True) + s36 = self._s36_piece(query_lc) + injected = None + if (s36 is not None and (allowed is None or allowed(s36)) + and not any(self._is_s36(i) for i in top)): + cands = [i for i in cands if not self._is_s36(i)] + injected = s36 + if any(is_primary(i) and governing(i) for i in top): + cands.append(s36) + if q_tokens: + cands.sort(key=lambda i: self._title_score(i, q_tokens), + reverse=True) + else: + cands.insert(0, s36) + if need <= 0: + return ordered + # One promotion per provision -- among the promotions only: a second + # piece of a provision already in the window may still be promoted + # (on 'PR convicted in Canada, six months' IRPA s. 36's second piece + # beats s. 37, whose title merely shares 'criminality'). + promote, seen = [], set() + for i in cands: + if (q_tokens and i != injected and not governing(i) + and not any(t in q_tokens + and self.idf.get(t, 0.0) >= _ROUTE_IDF_FLOOR + for t in self._note_tokens[i])): + continue + if provision(i) not in seen: + seen.add(provision(i)) + promote.append(i) + if len(promote) == need: + break + drop = [i for i in reversed(top) if not is_primary(i)][:len(promote)] + if not drop: + return ordered + promote = promote[:len(drop)] + if q_tokens: + promote.sort(key=lambda i: self._title_score(i, q_tokens), reverse=True) + dropped, promoted = set(drop), set(promote) + kept = [i for i in top if i not in dropped] + at = 1 + while at < len(kept) and is_primary(kept[at]): + at += 1 + return kept[:at] + promote + kept[at:] + drop + [ + i for i in rest if i not in promoted] + + def _cosurface_appendices(self, top, scores=None): + """Append the appendices the directive results cite but that retrieval + missed. A directive's rate tables ('Appendix C') rank poorly on a + natural-language query, yet a section that cites them is of little use + without them -- so the appendix travels with it. When more appendices + are cited than APPENDIX_CAP allows, the ones cited by the most results + win, so a lone off-topic result cannot outvote the relevant ones. + Returns `top` extended by one chunk for each of up to APPENDIX_CAP + distinct appendices: of a split appendix's pieces, the one the fusion + scored best (`scores`, chunk index -> fusion score), else the first. + An appendix with any piece already in `top` is not cited again. Code + audit 2026-09-21: the cap counted pieces, so one long appendix could + fill it and appear twice.""" + have = set(top) + cited = Counter() # (act_code, letter) -> citing results + for idx in top: + c = self.chunks[idx] + if c.get("doc_type") != "directive": + continue + seen = set() # count an appendix once per citing result + for m in _APPENDIX_REF.finditer(c["text"]): + key = (c["act_code"], m.group(1).upper()) + pieces = self._appendix.get(key, ()) + if pieces and key not in seen and not any(p in have for p in pieces): + seen.add(key) + cited[key] += 1 + out = list(top) + for key, _ in cited.most_common(APPENDIX_CAP): + pieces = self._appendix[key] + if scores: + # max() keeps the first of equals: the first piece on a tie. + out.append(max(pieces, key=lambda i: scores.get(i, float("-inf")))) + else: + out.append(pieces[0]) + return out + + def _highlight(self, query, indices): + """For each result chunk, the subsection or paragraph most on point for + the query: {result_position: (citation_suffix, snippet)}. Uses the + cross-encoder; returns {} if it is unavailable or nothing is structured. + Only the first HIGHLIGHT_DEPTH results are scored -- a pinpoint deep + in the list is not worth the cross-encoder cost.""" + if not self.reranker: + return {} + jobs = [] # (result_position, label, scoring_text, snippet) + for pos, idx in enumerate(indices[:HIGHLIGHT_DEPTH]): + c = self.chunks[idx] + if c.get("doc_type", "legislation") != "legislation": + continue + note = c["marginal_note"] + for label, scoring, snippet in _provision_units(c["text"]): + jobs.append((pos, label, f"{note}. {scoring}", snippet)) + if not jobs: + return {} + best = {} # result_position -> (score, label, snippet) + for (pos, label, _, snippet), score in zip( + jobs, self.reranker.score(query, [j[2] for j in jobs])): + if pos not in best or score > best[pos][0]: + best[pos] = (score, label, snippet) + return {pos: (label, " ".join(snippet[:240].split())) + for pos, (score, label, snippet) in best.items()} + + def _filter_ok(self, c, act, doc_type, qlang="en", court=None, + decided_after=None): + """One predicate for both recall masking and the late result filter. + The act filter matches short name, code, or full name -- an agent + passing 'Immigration and Refugee Protection Act' should not get + silence. English queries never see French chunks; French queries see + both languages (guidance/case law exist only in English). court / + decided_after constrain CASE-LAW chunks only (a case-law chunk's + 'part' carries the deciding court and current_to the decision date); + non-caselaw chunks pass those filters untouched so a court-filtered + mixed search still surfaces the governing statute.""" + if qlang == "en" and c.get("lang", "en") == "fr": + return False + if act and act.strip(): + a = act.strip().lower() + if a not in (c["act_short"].lower(), c["act_code"].lower(), + c.get("act_name", "").lower()): + return False + if doc_type and c.get("doc_type", "legislation") != doc_type: + return False + if (court or decided_after) and c.get("doc_type") == "caselaw": + # Match the court by name fragment ('Federal Court of Appeal') + # or by the neutral-citation abbreviation in the citation + # ('FCA', 'ONCA'): callers use both. + if court: + needle = court.strip().lower() + if (needle not in (c.get("part") or "").lower() + and not re.search(r"\b" + re.escape(needle) + r"\b", + (c.get("citation") or "").lower())): + return False + if decided_after: + year = (c.get("current_to") or "")[:4] + if not year.isdigit() or int(year) < decided_after: + return False + return True + + def search(self, query, top_k=6, act=None, doc_type=None, court=None, + decided_after=None): + """Hybrid candidate fusion (BM25 + semantic), then cross-encoder rerank.""" + # A filter value's surrounding whitespace is not part of it: an agent + # sending act='IRPA ' got no results (code audit 2026-09-21). + act = (act.strip() or None) if act else None + court = (court.strip() or None) if court else None + # Repair misspellings first (no-op for correctly-spelled queries; + # the query's language picks the vocabulary), then expand legal + # abbreviations (PRRA, H&C, ...) into statutory wording for the + # recall stages. The corrected text is what feeds the semantic + # embedder and the cross-encoder -- the stem-level fuzzy bridge in + # _bm25_scores only ever helped BM25, which is why typo queries used + # to reach the pool but not the top ranks. + qlang = query_lang(query) + corrected = self._correct_query(query, qlang) + expanded = expand_query(corrected) + # An aerodrome named by IATA code or name reaches its CASR 2012 schedule row, + # which writes only the city and the ICAO indicator (2026-10-15). Read from + # the officer's own words: the corrector turns 'Billy' into 'Billy bill'. + aerodromes = airport_terms(query) + if aerodromes: + expanded = f"{expanded} {' '.join(aerodromes)}" + confidence = None + fused = defaultdict(float) + # A filtered search competes only within its scope: the mask reaches + # both recall stages, so results come from the best of the whole + # filtered corpus rather than whatever survived open recall. An + # unfiltered English query uses the cached English-only mask; a + # French query searches both languages unmasked. + if act or doc_type or court or decided_after: + mask = [self._filter_ok(c, act, doc_type, qlang, court, + decided_after) + for c in self.chunks] + elif qlang == "en": + mask = self._en_mask + else: + mask = None + q_lc = query.lower() + # A US-state commentary page whose state the query names escapes + # the one-source cap on those pages (STATE_ONE_SOURCE, _source_key). + exempt = self._states_named(q_lc) if STATE_ONE_SOURCE else frozenset() + bm25_order = self._bm25_ranking(expanded, allowed=mask, + en_only=(qlang == "en")) + for rank, idx in enumerate(self._capped_top(bm25_order, exempt)): + fused[idx] += 1.0 / (RRF_K + rank) + if self.semantic: + sem_order, confidence = self._semantic_ranking(expanded, allowed=mask) + for rank, idx in enumerate(self._capped_top(sem_order, exempt)): + fused[idx] += W_SEM / (RRF_K + rank) + + # Ensure explicitly-referenced sections are retrieved even if recall + # missed them -- but only for Acts the query actually names. A query + # like "IRPA s. 40 misrepresentation defence" uses the section number + # topically; pulling every Act's s. 40 into the pool would drown out + # the case law that interprets the section the user meant. An Act is + # named by its short name, code or full name as WHOLE WORDS -- + # _acts_named for names of four characters or more, the whole-word + # test below for the three-letter short names (ITA, OLA, FAA) and the + # codes, which _acts_named skips -- never as a substring: 'ita' sits + # inside 'rehabilitation', 'capital' and 'hospital', 'ola' inside + # 'violation', and the substring test pinned the Income Tax Act's + # s. 18 to rank 1 of 'deemed rehabilitation section 18' at score 0.0 + # (code audit 2026-09-21). Whole words rather than token overlap + # because act_codes split into trivial tokens ("A-8.8" -> {a, 8}) + # that spuriously match common query words. + named = self._acts_named(q_lc) + refs = _section_refs(query) | self._act_cited_sections(q_lc) + # A French Act title names its Act for the pin alone (never for the + # ACT_HINT), and is only looked for when a section is cited. + fr_named = self._french_acts_named(q_lc) if refs else set() + whole = {} + + def _whole_word(name): + hit = whole.get(name) + if hit is None: + n = (name or "").lower() + hit = bool(len(n) >= 3 and re.search( + r"(? 0: + fused[idx] += min(MN_WEIGHT * matched * matched / total, MN_CAP) + + # Hierarchy penalties: a topical question should surface the governing + # provision, not the supplementary material around it. An Act creates a + # duty while a regulation only elaborates procedure; a collective + # agreement's numbered articles are its substance while its memoranda + # and letters of understanding are back-matter; a not-in-force + # amendment is pending law, not the law. Each takes a small fusion + # penalty so the governing provision wins a close contest (the NIF + # penalty is larger: pending text should surface only when it is the + # clearly best match for the subject, never beat the live provision). + # Guidance has two sub-tiers below MEMO_PENALTY (unofficial ATIP + # copies of withdrawn chapters; IRCC applicant-facing pages) and + # CanLex's own commentary is a tier of its own (COMMENTARY_PENALTY), + # each 0 until an eval run earns it a value. + # `dt`, not `doc_type`: that name is search()'s doc_type FILTER + # parameter, and shadowing it here left the last-iterated chunk's + # type applied by _filter_ok below as a silent doc_type filter on + # every unfiltered query (legacy Hit@1 0.79 -> 0.47 before the fix). + for idx in list(fused): + dt = self.chunks[idx].get("doc_type") + if self._is_nif[idx]: + fused[idx] -= NIF_PENALTY + elif self._is_repealed[idx]: + fused[idx] -= REPEALED_PENALTY + elif self._inapplicable[idx]: + fused[idx] -= INAPPLICABLE_PENALTY + elif self._is_regulation[idx]: + fused[idx] -= REG_PENALTY + elif dt == "memorandum": + fused[idx] -= MEMO_PENALTY + if self._is_atip[idx]: + fused[idx] -= ATIP_PENALTY + elif self._is_applicant[idx]: + fused[idx] -= APPLICANT_PENALTY + elif dt == "commentary": + fused[idx] -= COMMENTARY_PENALTY + elif self._is_backmatter[idx]: + fused[idx] -= BACKMATTER_PENALTY + + # Court hint: the query names the deciding court, so that court's + # decisions get COURT_HINT (see the constant for the rationale). + hints = court_hints(q_lc) + if hints: + for idx in list(fused): + c = self.chunks[idx] + if (c.get("doc_type") == "caselaw" + and any(pred(c.get("part")) for pred in hints)): + fused[idx] += COURT_HINT + + # Act hint: the query names the instrument (see ACT_HINT). + if named: + for idx in list(fused): + c = self.chunks[idx] + if (c.get("doc_type") != "caselaw" + and c.get("act_code") in named): + fused[idx] += ACT_HINT + + candidates = [i for i in sorted(fused, key=fused.get, reverse=True) + if self._filter_ok(self.chunks[i], act, doc_type, qlang, + court, decided_after)] + if not candidates: + return [] + scores = {i: fused[i] for i in candidates} + + # Precision stage: the cross-encoder rescores the top candidate pool, but + # may only PROMOTE -- each pooled candidate is placed at the better of its + # fusion rank and its rerank rank, never below its fusion rank. The + # reranker reliably surfaces a strong answer the fusion ranked low, yet is + # unreliable on long statutory text (it can score the right section + # negative), so its power to demote a candidate is deliberately removed. + if self.reranker: + pool = self._fill_pool(candidates, exempt) + docs = [self._rerank_doc(i) for i in pool] + ce_orig = self.reranker.score(query, docs) + # A repaired query is scored alongside the original and each + # candidate keeps its better score -- the cross-encoder cannot + # ground 'detension' against text saying 'detention'. Applies + # only when a correction fired, so legacy queries never take + # the second pass. + if corrected != query: + ce_fix = self.reranker.score(corrected, docs) + ce_orig = [max(a, b) for a, b in zip(ce_orig, ce_fix)] + # Optional (CANLEX_RERANK_DUAL=1): likewise for the + # synonym-expanded form. Off by default: the 2026-07 eval could + # not attribute its effect cleanly on the grown corpus. + if expanded != corrected and os.environ.get("CANLEX_RERANK_DUAL"): + ce_exp = self.reranker.score(expanded, docs) + ce_orig = [max(a, b) for a, b in zip(ce_orig, ce_exp)] + ce = dict(zip(pool, ce_orig)) + fusion_rank = {idx: r for r, idx in enumerate(pool)} + rerank_rank = {idx: r for r, idx in enumerate( + sorted(pool, key=ce.get, reverse=True))} + pool.sort(key=lambda i: (min(fusion_rank[i], rerank_rank[i]), + fusion_rank[i])) + in_pool = set(pool) + candidates = pool + [i for i in candidates if i not in in_pool] + + # Explicit section references are pinned to the very top -- using the + # same Act-mentioned constraint as the recall step above, for the same + # reason: a bare "s. 40" without an Act name is usually topical + # (e.g. "the IRPA s. 40 misrepresentation defence"), not a lookup. + if refs: + pinned = [i for i in candidates + if self.chunks[i]["section"] in refs + and _act_in_query(self.chunks[i])] + if pinned: + pinned_set = set(pinned) + candidates = pinned + [i for i in candidates if i not in pinned_set] + + # Cap one-source monopolies, then guarantee a primary instrument on + # the topic is represented. The guarantee operates on a fixed visible + # window of min(top_k, 5), not the full top_k -- with top_k=20 (the + # eval default) the larger window almost always contains incidental + # legislation, so the guarantee never fires even when the governing + # provision is buried at rank 10+. + candidates = self._diversify(candidates, exempt) + # A query that asks what a named court DECIDED is answered by the + # decision; the primary-instrument guarantee would otherwise pull a + # title-matching section into rank 2 regardless of fit (the CBSA Code + # of Conduct's "Use of Personal Electronic Devices at Work" above Pike + # on the device-search question). The court hint is already gated on + # a decision word, so this is the one case where case law is the + # governing source the user asked for. + # A filtered search's results must all pass the caller's filter. + # The steps below that bring in chunks from outside the candidates + # -- the route's IRPA s. 36 piece, directive appendices, read-next + # -- check them against it (2026-09-15 verifier: doc_type + # 'caselaw', 'memorandum' and 'commentary' and act='ENF 2' searches + # each showed IRPA s. 36 at rank 2). + filtered = bool(act or doc_type or court or decided_after) + + def keep(i): + return self._filter_ok(self.chunks[i], act, doc_type, qlang, + court, decided_after) + if not hints: + candidates = self._ensure_primary( + candidates, min(top_k, 5), q_tokens, + allow_policy=employment_query(q_lc), query_lc=q_lc, + criminality=(CRIMINALITY_ROUTE and criminality_query(q_lc) + and not self.route_exempt(q_lc)), + allowed=keep if filtered else None) + + # The ranked window, then the directive appendices it cites (flagged + # 'appendix') and the 'read next' official guidance the curated + # entries in view name (flagged 'read-next'): both appended after + # everything ranked and flagged 'cosurfaced', so a client (and the + # eval) can tell them from a ranked result. Code audit 2026-09-21: + # the appendices were unflagged and counted as ranked, so top_k=6 + # answered with nine results. + ranked = candidates[:top_k] + appended = self._cosurface_appendices(ranked, scores)[len(ranked):] + if filtered: + ranked = [i for i in ranked if keep(i)] + appended = [i for i in appended if keep(i)] + n_ranked, n_appendix = len(ranked), len(appended) + top = self._cosurface_guidance(ranked + appended, + keep=keep if filtered else None) + highlights = self._highlight(query, top) + results = [] + for pos, i in enumerate(top): + result = {**self.chunks[i], "score": round(scores.get(i, 0.0), 4), + "confidence": confidence} + if pos >= n_ranked + n_appendix: + result["cosurfaced"] = "read-next" + elif pos >= n_ranked: + result["cosurfaced"] = "appendix" + if pos in highlights: + result["highlight"] = highlights[pos] + results.append(result) + return results + + def get_section(self, act, section, lang="en"): + """The section chunk in the requested language, falling back to the + other language rather than returning nothing.""" + act = act.strip().lower() + fallback = None + for c in self.chunks: + if c["section"] == section and act in ( + c["act_short"].lower(), c["act_code"].lower(), + c.get("act_name", "").lower()): + if c.get("lang", "en") == lang: + return c + fallback = fallback or c + return fallback + + def get_section_parts(self, act, section, lang="en"): + """Every stored piece of one section, in order -- the chunk + get_section returns plus its '-p2', '-p3', ... siblings (same base + id, same language). Long provisions are split for embedding, and + serving only the first piece silently dropped IRPA s. 36(3) and + IRPR s. 18(2): the subsections the screening tools point at. + Returns [] when the section is not found.""" + first = self.get_section(act, section, lang) + if first is None: + return [] + base = _PIECE_ID.sub("", first["id"]) + first_lang = first.get("lang", "en") + parts = [c for c in self.chunks + if _PIECE_ID.sub("", c.get("id", "")) == base + and c.get("lang", "en") == first_lang] + return sorted(parts, key=_piece_no) + + def _build_xref(self): + """Index legislation by (act, section); find each Act's definitions + section; link every regulation to its enabling Act and every + D-memorandum to the provisions it cites -- all for cross-referencing.""" + self._by_section = {} + self._defs_section = {} + acts, regs = {}, {} # act_code -> (act_short, act_name) + for c in self.chunks: + if c.get("doc_type", "legislation") != "legislation": + continue + if c.get("lang", "en") == "fr": + continue # xref keys are per (act, section); English canonical + self._by_section[(c["act_code"], c["section"])] = c + if c["act_code"] not in self._defs_section and ( + c["marginal_note"].strip().lower() in ( + "definitions", "definition", "interpretation")): + self._defs_section[c["act_code"]] = c + bucket = regs if c["act_code"].startswith(("SOR", "C.R.C")) else acts + bucket.setdefault(c["act_code"], (c["act_short"], c["act_name"])) + + # Link a regulation to the Act it is made under by matching their names + # ("X Regulations" <-> "X Act"). + self._enabling_act = {} # reg code -> (act_short, act_name) + self._regulations = defaultdict(list) # act code -> [(reg_short, reg_name)] + def base(name): + return re.sub(r"\b(?:Act|Regulations)\b", "", name).strip().lower() + act_by_base = {base(n): (code, s, n) for code, (s, n) in acts.items()} + for rcode, (rshort, rname) in regs.items(): + hit = act_by_base.get(base(rname)) + if hit: + self._enabling_act[rcode] = (hit[1], hit[2]) + self._regulations[hit[0]].append((rshort, rname)) + + # Link D-memoranda to the Customs Act / Customs Tariff provisions they + # cite; an unqualified "the Act" in a D-memo means the Customs Act. + by_short = {s.lower(): code for code, (s, n) in acts.items()} + customs, tariff = by_short.get("customs act"), by_short.get("customs tariff") + self._memos_for_section = defaultdict(set) # (act_code, section) -> memos + for c in self.chunks: + if c.get("doc_type") != "memorandum": + continue + for num, which in _MEMO_CITE.findall(c["text"]): + code = tariff if which.lower() == "customs tariff" else customs + if code: + self._memos_for_section[(code, num)].add(c["section"]) + + def related(self, chunk): + """Cross-references for a legislation result, as a dict: 'provisions' + (intra-Act sections it cites, plus the definitions section), + 'regulations' (made under this Act), 'enabling_act' (for a regulation, + the Act it is made under) and 'memoranda' (D-memo numbers citing this + section). Empty dict for case law, memoranda, etc.""" + if chunk.get("doc_type", "legislation") != "legislation": + return {} + act = chunk["act_code"] + provisions, seen = [], {chunk["section"]} + defs = self._defs_section.get(act) + if defs and defs["section"] not in seen: + provisions.append((defs["section"], defs["marginal_note"])) + seen.add(defs["section"]) + for match in _XREF.finditer(chunk["text"]): + sec = match.group(1) + if sec in seen: + continue + target = self._by_section.get((act, sec)) + if target: + provisions.append((sec, target["marginal_note"])) + seen.add(sec) + if len(provisions) >= 8: + break + return { + "provisions": provisions, + "regulations": self._regulations.get(act, []), + "enabling_act": self._enabling_act.get(act), + "memoranda": sorted(self._memos_for_section.get( + (act, chunk["section"]), []))[:6], + } + + +def main(): + if len(sys.argv) < 2: + print('usage: python -m canlex.index "your query"') + return + query = " ".join(sys.argv[1:]) + index = LegislationIndex() + if index.reranker: + mode = "hybrid + cross-encoder rerank" + elif index.semantic: + mode = "hybrid (BM25 + semantic)" + else: + mode = "BM25 only" + print(f"{len(index.chunks)} sections indexed - {mode}. Query: {query!r}\n") + for r in index.search(query): + print(f"[{r['score']:.3f}] {r['citation']} - {r['marginal_note']}") + print(f" {r['text'].replace(chr(10), ' ')[:160]}") + print() + + +if __name__ == "__main__": + main()