From 05c6f110366bc0d7a88a452807503c0a7b2d26e2 Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Thu, 6 Aug 2026 11:24:20 +0200 Subject: [PATCH 01/13] feat(eval): expose the shortlist funnel loss in TREC metrics recall@k measures the first-level candidate list and nDCG@k measures the ranked list, so neither can see the pipeline's largest loss: the shortlist handed to the eligibility reasoner is far shorter than the candidate list, and a relevant trial dropped there can never be ranked. Measured on the completed runs, the shortlist discards 33% (TREC 2021), 36% (2022) and 61% (2023) of the relevant trials the first level had already found. The 2023 run shows why this stayed hidden: it has the best nDCG@10 of any run (0.881) and the worst shortlist recall (0.254), because the condensed metric only orders the judged trials that survive the funnel. evaluate() now reports, whenever top_trials.txt is present: - shortlist_recall .............. recall of the list the reasoner read - shortlist_size - funnel_depth_loss ............. recall given up by shortening alone - shortlist_selection_delta ..... second-level ordering vs a plain first-level top-N cut at the same depth The split matters for what we fix first. On TREC 2021 depth accounts for 94% of the loss (0.866 -> 0.595 at top-200) and ordering for 6% (-> 0.577). The selection delta is negative on all four runs (-0.032, -0.027, -0.022, -0.009): the criterion reranker plus RRF shortlist fusion currently select worse than taking the first-level top-N and doing nothing. Runs without top_trials.txt still evaluate; the funnel keys report None. Co-Authored-By: Claude Opus 5 --- src/trialmatchai/trec/qrels.py | 62 +++++++++++++++++++++++++++++++++- tests/test_qrels_eval.py | 42 +++++++++++++++++++++++ 2 files changed, 103 insertions(+), 1 deletion(-) diff --git a/src/trialmatchai/trec/qrels.py b/src/trialmatchai/trec/qrels.py index 94d543e1..f681e7d5 100644 --- a/src/trialmatchai/trec/qrels.py +++ b/src/trialmatchai/trec/qrels.py @@ -124,6 +124,18 @@ def _retrieved_for_patient(patient_dir: Path) -> list[str]: return [] +def _shortlist_for_patient(patient_dir: Path) -> list[str]: + """The trials that actually reached the eligibility (CoT) stage. + + This is the funnel's narrowest point: a relevant trial dropped here is + unrecoverable, however good the first-level search was. + """ + shortlist = patient_dir / "top_trials.txt" + if not shortlist.exists(): + return [] + return [line.strip() for line in shortlist.read_text().splitlines() if line.strip()] + + def recall_at_k(retrieved: list[str], relevant: set[str], k: int) -> float | None: if not relevant: return None @@ -168,6 +180,17 @@ def evaluate( Reports recall@k (retrieval, first-level list) and tie-aware nDCG@{5,10,20} + P@10 (ranking, condensed to judged trials). P@10 is split into "relevant" (grade>=1) and "eligible" (grade==2). + + Also reports funnel metrics when ``top_trials.txt`` is present. recall@k + (first-level list) and nDCG@k (ranked list) between them hide the pipeline's + largest loss: the shortlist handed to the reasoner is far shorter than the + candidate list, and a relevant trial dropped there can never be ranked. + + - ``shortlist_recall`` -- recall of the list the reasoner actually read. + - ``funnel_depth_loss`` -- recall given up by shortening the list alone. + - ``shortlist_selection_delta`` -- second-level ordering minus a plain + first-level cut at the same depth. Negative means the second level + selected worse than doing nothing. """ results_dir = Path(results_dir) relevant = relevant_ncts(qrels, threshold=threshold) @@ -176,6 +199,19 @@ def evaluate( rec_sums = {f"recall@{k}": 0.0 for k in cutoffs} rec_counts = {f"recall@{k}": 0 for k in cutoffs} + # Funnel instrumentation: recall@k measures the FIRST-LEVEL list, but only the shortlist + # reaches the reasoner. That gap is an unrecoverable ceiling on every ranking metric and is + # invisible in recall@k. Split it into its two causes: depth (shortlist shorter than the + # candidate list) and selection (second-level ordering vs a plain first-level top-N cut). + funnel_keys = ( + "shortlist_recall", + "shortlist_size", + "first_level_recall_at_shortlist_depth", + "shortlist_selection_delta", + "funnel_depth_loss", + ) + funnel_sums = {key: 0.0 for key in funnel_keys} + funnel_counts = {key: 0 for key in funnel_keys} rank_sums = {f"ndcg@{k}": 0.0 for k in NDCG_CUTOFFS} rank_sums.update({f"ndcg_full@{k}": 0.0 for k in NDCG_CUTOFFS}) rank_sums[f"P@{P_CUTOFF}(rel>=1)"] = 0.0 @@ -202,6 +238,26 @@ def evaluate( rec_sums[f"recall@{k}"] += r rec_counts[f"recall@{k}"] += 1 + shortlist = _shortlist_for_patient(patient_dir) + if shortlist and retrieved: + depth = len(shortlist) + short_r = recall_at_k(shortlist, rel_set, depth) + first_r = recall_at_k(retrieved, rel_set, depth) + full_r = recall_at_k(retrieved, rel_set, len(retrieved)) + funnel = { + "shortlist_recall": short_r, + "shortlist_size": float(depth), + "first_level_recall_at_shortlist_depth": first_r, + # > 0 means the second level beat a plain first-level cut at the same depth. + "shortlist_selection_delta": short_r - first_r, + # Recall thrown away purely by shortening the list. + "funnel_depth_loss": full_r - first_r, + } + row.update(funnel) + for key, value in funnel.items(): + funnel_sums[key] += value + funnel_counts[key] += 1 + if ranked: # Two IDCG bases: ndcg@k normalizes by the ideal over judged-AND-ranked trials # (recall-independent ordering quality); ndcg_full@k by the ideal over the FULL judged @@ -233,7 +289,11 @@ def evaluate( rank_counts[f"graded_P@{P_CUTOFF}"] += 1 per_query[query_id] = row - mean = {**_mean(rec_sums, rec_counts), **_mean(rank_sums, rank_counts)} + mean = { + **_mean(rec_sums, rec_counts), + **_mean(funnel_sums, funnel_counts), + **_mean(rank_sums, rank_counts), + } return { "recall_relevance_threshold": threshold, "num_queries_scored": len(per_query), diff --git a/tests/test_qrels_eval.py b/tests/test_qrels_eval.py index a864fd73..5520a777 100644 --- a/tests/test_qrels_eval.py +++ b/tests/test_qrels_eval.py @@ -93,3 +93,45 @@ def test_evaluate_precision_is_condensed_to_judged_pool(tmp_path): assert mean["P@10(rel>=1)"] == pytest.approx(2 / 10) # raw would be 0/10 assert mean["P@10(eligible)"] == pytest.approx(1 / 10) # only NCT1 is grade 2 assert mean["graded_P@10"] == pytest.approx((2 + 1) / (10 * 2)) # raw would be 0 + + +def test_evaluate_reports_funnel_metrics(tmp_path): + """The shortlist (top_trials.txt) is the funnel's narrowest point. Evaluation must expose + its recall and split the loss into depth vs second-level selection, since recall@k + (first-level list) and nDCG@k (ranked list) both hide it.""" + q = "trec-1" + pdir = tmp_path / q + pdir.mkdir() + # First level finds all three relevant trials; the shortlist keeps only two slots and + # spends one on an irrelevant trial that a plain first-level top-2 would not have picked. + (pdir / "nct_ids.txt").write_text("NCT1\nNCT2\nNCT3\nNCT4\n") + (pdir / "top_trials.txt").write_text("NCT1\nNCT4\n") + (pdir / "ranked_trials.json").write_text( + json.dumps({"RankedTrials": [{"TrialID": "NCT1", "Score": 1.0}]}) + ) + qrels = {q: {"NCT1": 2, "NCT2": 2, "NCT3": 1, "NCT4": 0}} + + mean = evaluate(qrels, tmp_path, cutoffs=(10,))["mean"] + + assert mean["recall@10"] == 1.0 # first level found everything... + assert mean["shortlist_recall"] == pytest.approx(1 / 3) # ...the reasoner saw a third + assert mean["shortlist_size"] == 2 + assert mean["first_level_recall_at_shortlist_depth"] == pytest.approx(2 / 3) + assert mean["shortlist_selection_delta"] == pytest.approx(-1 / 3) # selection hurt + assert mean["funnel_depth_loss"] == pytest.approx(1 / 3) # depth alone cost this + + +def test_evaluate_funnel_metrics_absent_without_shortlist(tmp_path): + """Runs predating the shortlist file must still evaluate, without funnel values.""" + q = "trec-1" + pdir = tmp_path / q + pdir.mkdir() + (pdir / "nct_ids.txt").write_text("NCT1\n") + (pdir / "ranked_trials.json").write_text( + json.dumps({"RankedTrials": [{"TrialID": "NCT1", "Score": 1.0}]}) + ) + + mean = evaluate({q: {"NCT1": 2}}, tmp_path, cutoffs=(10,))["mean"] + + assert mean["recall@10"] == 1.0 + assert mean["shortlist_recall"] is None From 91f724f33587952638139de61a4d3e2beb06ea86 Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Thu, 6 Aug 2026 14:21:14 +0200 Subject: [PATCH 02/13] feat(search): per-patient shortlist depth from the first-level score curve The shortlist divisor sizes every patient identically, but depth is 94% of the measured shortlist recall loss (PR #30) and no single number fits: the depth a patient needs to reach 90% of its own first-level recall ranges from 50 to 1550 trials, spread evenly across that range. Sizing for the worst case wastes ~65% of the reasoner's compute; sizing for the median drops the hard patients. Adds search.shortlist.policy: fixed existing divisor sizing, unchanged -- still the DEFAULT relative_to_max keep trials scoring >= alpha x this patient's top score A peaked first-level score curve means retrieval was confident and few trials are plausible; a flat curve means many are. The cut is relative to the patient's own maximum because RRF scores are only comparable within one patient. Tuned offline by replaying the completed runs' first_level_scores.json, so no GPU time. Compared against fixed depth AT EQUAL MEAN DEPTH, which is the only fair test -- a policy must spend compute better, not merely spend more: TREC 2021 +0.017 to +0.029 recall (peak at mean depth ~305) TREC 2022 +0.007 to +0.025 recall TREC 2023 -0.010 to +0.006 recall -- questionnaire topics, no gain Two other policies were tried and rejected: cumulative score mass tracked fixed depth to within +-0.003 everywhere, and half-max curve width won only below ~350 mean depth and lost above it. At alpha=0.25 the policy also chooses to go deeper (196->305 trials on 2021), which lifts shortlist recall 0.609->0.710. That part is bought with compute, not won for free; the equal-cost gain above is the free part. Both are real and they should be reported separately. Default stays "fixed" so enabling this is an explicit A/B. Degrades to the fixed depth when first-level scores are absent (resumed runs), never exceeds what the reasoner can consume, and writes shortlist_depth.json recording the decision and the depth the old sizing would have picked. Co-Authored-By: Claude Opus 5 --- src/trialmatchai/config/settings.py | 17 +++ src/trialmatchai/main.py | 29 +++- src/trialmatchai/matching/shortlist_depth.py | 149 +++++++++++++++++++ tests/test_shortlist_depth.py | 149 +++++++++++++++++++ 4 files changed, 341 insertions(+), 3 deletions(-) create mode 100644 src/trialmatchai/matching/shortlist_depth.py create mode 100644 tests/test_shortlist_depth.py diff --git a/src/trialmatchai/config/settings.py b/src/trialmatchai/config/settings.py index ed3e9b24..f4ab4b13 100644 --- a/src/trialmatchai/config/settings.py +++ b/src/trialmatchai/config/settings.py @@ -174,6 +174,22 @@ class FirstLevelSearchSettings(BaseModel): ) +class ShortlistSettings(BaseModel): + """How deep the shortlist handed to the eligibility reasoner goes. + + "fixed" keeps the divisor-based sizing (one depth for every patient) and is the + default. "relative_to_max" sizes each patient from its own first-level score curve, + keeping trials scoring at least ``relative_to_max_alpha`` x that patient's top score. + See matching/shortlist_depth.py for the offline evidence. + """ + + policy: Literal["fixed", "relative_to_max"] = "fixed" + relative_to_max_alpha: float = Field(0.25, gt=0.0, le=1.0) + min_depth: int = Field(50, ge=1) + # None -> bounded only by what the reasoner can consume (rag.max_trials_rag). + max_depth: int | None = Field(None, ge=1) + + class SearchSettings(BaseModel): mode: Literal["bm25", "vector", "hybrid"] = "hybrid" vector_score_threshold: float = Field(0.5, ge=0.0, le=1.0) @@ -191,6 +207,7 @@ class SearchSettings(BaseModel): first_level: FirstLevelSearchSettings = Field( default_factory=FirstLevelSearchSettings ) + shortlist: ShortlistSettings = Field(default_factory=ShortlistSettings) @model_validator(mode="before") @classmethod diff --git a/src/trialmatchai/main.py b/src/trialmatchai/main.py index 83c8df8c..97958d37 100644 --- a/src/trialmatchai/main.py +++ b/src/trialmatchai/main.py @@ -15,6 +15,7 @@ rank_trials, save_ranked_trials, ) +from trialmatchai.matching.shortlist_depth import choose_shortlist_depth, depth_report from trialmatchai.matching.retrieval.trial_retrieval import ClinicalTrialSearch from trialmatchai.matching.retrieval.criteria_retrieval import SecondStageRetriever from trialmatchai.matching.retrieval.location import ( @@ -307,15 +308,16 @@ def run_second_level_search( trial["nct_id"]: trial["score"] for trial in second_level_results } + search_config = config.get("search", {}) combined_scores = _fuse_shortlist_scores( nct_ids=nct_ids, second_level_results=second_level_results, first_level_scores=first_level_scores, - search_config=config.get("search", {}), + search_config=search_config, ) sorted_trials = sorted(combined_scores.items(), key=lambda x: x[1], reverse=True) - keep_divisor = max(1, int(config.get("search", {}).get("second_level_keep_divisor", 3))) + keep_divisor = max(1, int(search_config.get("second_level_keep_divisor", 3))) # Size the shortlist off the reranked count, not the fused pool: rank fusion adds # first-level-only trials, so keying the divisor to the pool would silently enlarge the # shortlist and confound a fusion A/B. @@ -323,12 +325,33 @@ def run_second_level_search( num_top = max(1, min(reranked_count // keep_divisor, top_n)) # RAG only reasons over rag.max_trials_rag trials; cap the shortlist to match, else trials # past the cap get no eligibility output and are silently dropped from the final ranking. + upper_bound = len(sorted_trials) or 1 if _rag_enabled(config): - num_top = max(1, min(num_top, int(config.get("rag", {}).get("max_trials_rag", 20)))) + upper_bound = min(upper_bound, int(config.get("rag", {}).get("max_trials_rag", 20))) + num_top = max(1, min(num_top, upper_bound)) + # The divisor sizes every patient the same. Depth is 94% of the measured shortlist recall + # loss and the depth patients need spans 50-1550 trials, so an opt-in policy may widen or + # narrow this per patient from the first-level score curve. Default policy returns num_top. + fixed_depth = num_top + num_top = choose_shortlist_depth( + first_level_scores=first_level_scores, + fixed_depth=fixed_depth, + upper_bound=upper_bound, + search_config=search_config, + ) semi_final_trials = sorted_trials[:num_top] top_trials_path = f"{output_folder}/top_trials.txt" write_text_file([trial_id for trial_id, _ in semi_final_trials], top_trials_path) + write_json_file( + depth_report( + chosen=len(semi_final_trials), + fixed_depth=fixed_depth, + first_level_scores=first_level_scores, + search_config=search_config, + ), + f"{output_folder}/shortlist_depth.json", + ) constraints_config = config.get("constraints", {}) if constraints_config.get("enabled", True) and constraints_config.get( "write_reports", diff --git a/src/trialmatchai/matching/shortlist_depth.py b/src/trialmatchai/matching/shortlist_depth.py new file mode 100644 index 00000000..21c85dbf --- /dev/null +++ b/src/trialmatchai/matching/shortlist_depth.py @@ -0,0 +1,149 @@ +"""How many trials the eligibility reasoner gets to read. + +The shortlist is the pipeline's narrowest point: a relevant trial dropped here can +never be ranked, however good the reasoning is. Measured on the completed TREC runs, +a fixed-size shortlist discards a third of the relevant trials the first level had +already found (``shortlist_recall`` in ``trec/qrels.py``). + +Depth is the dominant cause -- 94% of that loss on TREC 2021 -- and no single number +serves every patient: the depth needed to reach 90% of a patient's own first-level +recall ranges from 50 to 1550 trials, spread evenly across that range. Sizing for the +worst case wastes ~65% of the reasoner's compute; sizing for the median silently drops +the hard patients. + +``relative_to_max`` therefore reads the shape of the first-level score curve. A peaked +curve means retrieval was confident and few trials are plausible; a flat curve means +many are, and the patient needs more depth. Offline replay over the completed runs +(first-level scores, same mean depth as the fixed policy) gives: + + TREC 2021 +0.017 to +0.029 recall + TREC 2022 +0.007 to +0.025 recall + TREC 2023 -0.010 to +0.006 recall (questionnaire topics; no gain) + +So it is a small, free gain on the narrative-topic tracks and a wash on 2023 -- worth +having, but not a substitute for spending more depth outright. ``fixed`` remains the +default so enabling this is an explicit A/B, not a silent change. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from typing import Any + +from trialmatchai.utils.logging_config import setup_logging + +logger = setup_logging(__name__) + +POLICIES = ("fixed", "relative_to_max") +DEFAULT_ALPHA = 0.25 +DEFAULT_MIN_DEPTH = 50 + + +def shortlist_config(search_config: Mapping[str, Any] | None) -> dict[str, Any]: + """Resolve the ``search.shortlist`` block, tolerating absent or partial config.""" + raw = (search_config or {}).get("shortlist") or {} + if not isinstance(raw, Mapping): + raw = {} + policy = str(raw.get("policy", "fixed") or "fixed") + if policy not in POLICIES: + logger.warning( + "Unknown search.shortlist.policy %r; falling back to 'fixed'. Known: %s", + policy, + ", ".join(POLICIES), + ) + policy = "fixed" + return { + "policy": policy, + "relative_to_max_alpha": float(raw.get("relative_to_max_alpha", DEFAULT_ALPHA)), + "min_depth": int(raw.get("min_depth", DEFAULT_MIN_DEPTH)), + "max_depth": raw.get("max_depth"), + } + + +def _relative_to_max_depth(scores: list[float], alpha: float) -> int: + """Count of trials scoring at least ``alpha`` x the top score. + + Scores are a weighted RRF sum, so they are positive and comparable only within one + patient -- which is exactly why the cut is relative to that patient's own maximum + rather than an absolute threshold. + """ + if not scores: + return 0 + top = scores[0] + if top <= 0: + return len(scores) + cut = alpha * top + kept = 0 + for score in scores: + if score < cut: + break + kept += 1 + return kept + + +def choose_shortlist_depth( + *, + first_level_scores: Mapping[str, float] | None, + fixed_depth: int, + upper_bound: int, + search_config: Mapping[str, Any] | None = None, +) -> int: + """Shortlist size for one patient. + + ``fixed_depth`` is what the existing divisor-based sizing would have chosen, and is + returned unchanged under the default policy. ``upper_bound`` is the hard ceiling the + caller can honour (the reasoner's own cap), and is never exceeded. + """ + upper_bound = max(1, int(upper_bound)) + fixed_depth = max(1, min(int(fixed_depth), upper_bound)) + cfg = shortlist_config(search_config) + if cfg["policy"] == "fixed": + return fixed_depth + + scores = sorted((float(v) for v in (first_level_scores or {}).values()), reverse=True) + if not scores: + # No first-level signal to read (e.g. a resumed run missing the scores file): + # degrade to the fixed sizing rather than guessing a depth. + logger.warning( + "shortlist policy 'relative_to_max' has no first-level scores; using fixed depth %s", + fixed_depth, + ) + return fixed_depth + + depth = _relative_to_max_depth(scores, cfg["relative_to_max_alpha"]) + floor = max(1, cfg["min_depth"]) + ceiling = upper_bound + configured_max = cfg["max_depth"] + if configured_max is not None: + ceiling = min(ceiling, max(1, int(configured_max))) + depth = max(floor, min(depth, ceiling)) + logger.info( + "Shortlist depth %s (policy=relative_to_max, alpha=%.3g, fixed would be %s)", + depth, + cfg["relative_to_max_alpha"], + fixed_depth, + ) + return depth + + +def depth_report( + *, + chosen: int, + fixed_depth: int, + first_level_scores: Mapping[str, float] | None, + search_config: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Provenance for the depth decision, written beside the shortlist.""" + cfg = shortlist_config(search_config) + scores: Iterable[float] = (first_level_scores or {}).values() + ordered = sorted((float(v) for v in scores), reverse=True) + return { + "policy": cfg["policy"], + "chosen_depth": int(chosen), + "fixed_depth": int(fixed_depth), + "relative_to_max_alpha": cfg["relative_to_max_alpha"], + "min_depth": cfg["min_depth"], + "max_depth": cfg["max_depth"], + "candidate_pool": len(ordered), + "top_score": ordered[0] if ordered else None, + } diff --git a/tests/test_shortlist_depth.py b/tests/test_shortlist_depth.py new file mode 100644 index 00000000..dd4cef6c --- /dev/null +++ b/tests/test_shortlist_depth.py @@ -0,0 +1,149 @@ +"""Shortlist depth policies (matching/shortlist_depth.py). + +The shortlist is where the pipeline loses most of its relevant trials, so the depth +decision must be explicit, bounded, and never silently change under the default config. +""" + +import pytest + +from trialmatchai.matching.shortlist_depth import ( + choose_shortlist_depth, + depth_report, + shortlist_config, +) + + +def _scores(values): + return {f"NCT{i:04d}": v for i, v in enumerate(values)} + + +def test_default_policy_returns_the_fixed_depth_unchanged(): + """Enabling adaptive depth must be an explicit A/B, never a silent behaviour change.""" + depth = choose_shortlist_depth( + first_level_scores=_scores([1.0] * 500), + fixed_depth=196, + upper_bound=300, + search_config={}, + ) + assert depth == 196 + + +def test_relative_to_max_keeps_trials_above_alpha_times_the_top_score(): + # Top score 1.0, alpha 0.25 -> keep while score >= 0.25, so the first four. + depth = choose_shortlist_depth( + first_level_scores=_scores([1.0, 0.8, 0.4, 0.25, 0.2, 0.1]), + fixed_depth=2, + upper_bound=100, + search_config={"shortlist": {"policy": "relative_to_max", "min_depth": 1}}, + ) + assert depth == 4 + + +def test_peaked_curve_gets_less_depth_than_flat_curve(): + """The whole point: a confident retrieval needs fewer trials than an ambiguous one.""" + cfg = {"shortlist": {"policy": "relative_to_max", "min_depth": 1}} + peaked = choose_shortlist_depth( + first_level_scores=_scores([1.0] + [0.01] * 99), + fixed_depth=50, + upper_bound=100, + search_config=cfg, + ) + flat = choose_shortlist_depth( + first_level_scores=_scores([1.0] * 100), + fixed_depth=50, + upper_bound=100, + search_config=cfg, + ) + assert peaked == 1 + assert flat == 100 + assert peaked < flat + + +def test_depth_is_clamped_by_min_depth_and_upper_bound(): + cfg = {"shortlist": {"policy": "relative_to_max", "min_depth": 20, "max_depth": 40}} + # A single dominant trial would give depth 1; the floor lifts it to min_depth. + assert ( + choose_shortlist_depth( + first_level_scores=_scores([1.0] + [0.001] * 99), + fixed_depth=10, + upper_bound=100, + search_config=cfg, + ) + == 20 + ) + # A flat curve would give 100; max_depth caps it at 40. + assert ( + choose_shortlist_depth( + first_level_scores=_scores([1.0] * 100), + fixed_depth=10, + upper_bound=100, + search_config=cfg, + ) + == 40 + ) + + +def test_upper_bound_always_wins_over_configured_max_depth(): + """upper_bound is what the reasoner can actually consume; exceeding it drops trials + silently from the final ranking.""" + depth = choose_shortlist_depth( + first_level_scores=_scores([1.0] * 500), + fixed_depth=10, + upper_bound=30, + search_config={ + "shortlist": {"policy": "relative_to_max", "min_depth": 1, "max_depth": 400} + }, + ) + assert depth == 30 + + +def test_missing_scores_degrade_to_fixed_depth(): + """A resumed run can lack first_level_scores.json; guessing a depth would be worse.""" + for scores in (None, {}): + depth = choose_shortlist_depth( + first_level_scores=scores, + fixed_depth=77, + upper_bound=300, + search_config={"shortlist": {"policy": "relative_to_max"}}, + ) + assert depth == 77 + + +def test_nonpositive_top_score_keeps_the_whole_pool(): + depth = choose_shortlist_depth( + first_level_scores=_scores([0.0, 0.0, 0.0]), + fixed_depth=1, + upper_bound=100, + search_config={"shortlist": {"policy": "relative_to_max", "min_depth": 1}}, + ) + assert depth == 3 + + +def test_unknown_policy_falls_back_to_fixed(caplog): + assert shortlist_config({"shortlist": {"policy": "wishful"}})["policy"] == "fixed" + depth = choose_shortlist_depth( + first_level_scores=_scores([1.0] * 100), + fixed_depth=12, + upper_bound=100, + search_config={"shortlist": {"policy": "wishful"}}, + ) + assert depth == 12 + + +@pytest.mark.parametrize("search_config", [None, {}, {"shortlist": None}]) +def test_absent_config_is_tolerated(search_config): + assert shortlist_config(search_config)["policy"] == "fixed" + + +def test_depth_report_records_the_decision(): + report = depth_report( + chosen=40, + fixed_depth=196, + first_level_scores=_scores([1.0, 0.5, 0.2]), + search_config={"shortlist": {"policy": "relative_to_max"}}, + ) + assert report["policy"] == "relative_to_max" + assert report["chosen_depth"] == 40 + assert report["fixed_depth"] == 196 # what the old sizing would have used + assert report["candidate_pool"] == 3 + assert report["top_score"] == 1.0 From a3b43cd4a139b4103e09ae1bb9fb40d319b0018f Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Thu, 6 Aug 2026 14:31:24 +0200 Subject: [PATCH 03/13] feat(trec): comparison harness for the shortlist A/B, and persist second-level scores The depth policy added in 91f724f is tuned only against replayed first-level scores. Whether the extra depth reaches nDCG@10 and P@10 -- rather than just shortlist recall -- needs a GPU run. This adds the reporting half of that experiment. compare_shortlist_ab.py reads the arms of a shortlist-policy A/B and always prints shortlist_size beside the quality metrics, because a depth policy produces two effects that must not be conflated: spending compute BETTER same mean shortlist size, more recall -- a free gain spending compute MORE bigger shortlist, bought with GPU time Reporting only the second reads as a much larger win than it is. The planned arms are fixed (~196 trials/patient), relative_to_max alpha 0.25 (~305, deeper), and relative_to_max alpha 0.33 (~196, equal cost). The equal-cost arm is the one that actually tests the policy. The job script itself stays untracked: *.slurm is gitignored repo-wide and none of the 26 existing job scripts are committed, since they encode site-specific paths, partitions and staging. Also persists second_level_scores.json for the whole second-level pool, not just the shortlist. With first_level_scores.json this makes shortlist fusion replayable offline, so the negative shortlist_selection_delta measured in 05c6f11 -- the fused shortlist selecting worse than a plain first-level cut at the same depth -- can be investigated without spending a GPU job on each candidate fix. --- scripts/compare_shortlist_ab.py | 100 ++++++++++++++++++++++++++++++++ src/trialmatchai/main.py | 6 ++ 2 files changed, 106 insertions(+) create mode 100644 scripts/compare_shortlist_ab.py diff --git a/scripts/compare_shortlist_ab.py b/scripts/compare_shortlist_ab.py new file mode 100644 index 00000000..9718338a --- /dev/null +++ b/scripts/compare_shortlist_ab.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python +"""Compare the arms of the shortlist-policy A/B (scripts/run_trec_shortlist_ab.slurm). + +Reports each arm against the fixed baseline, and -- crucially -- separates the two +effects that a depth policy produces: + + * spending compute BETTER (equal-cost arm: same mean shortlist size, more recall) + * spending compute MORE (deeper arm: bigger shortlist, bought with GPU time) + +Conflating them overstates the result, so cost is always printed beside quality. + + uv run python scripts/compare_shortlist_ab.py [--root shortlist_ab] [--track 21] +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +BASELINE = "fixed" +QUALITY = ( + "shortlist_recall", + "recall@1000", + "ndcg@10", + "P@10(rel>=1)", + "P@10(eligible)", +) +COST = ("shortlist_size",) +DIAGNOSTIC = ("funnel_depth_loss", "shortlist_selection_delta") + + +def load_arm(root: Path, arm: str, track: str) -> dict | None: + path = root / arm / f"results_trec{track}" / "evaluation_metrics.json" + if not path.exists(): + return None + data = json.loads(path.read_text()) + return {"mean": data.get("mean", {}), "n": data.get("num_queries_scored")} + + +def fmt(value: object, width: int = 9) -> str: + if isinstance(value, (int, float)): + return f"{value:{width}.4f}" + return f"{'--':>{width}}" + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--root", default="shortlist_ab") + parser.add_argument("--track", default="21") + args = parser.parse_args() + + root = Path(args.root) + arms = sorted(p.name for p in root.iterdir() if p.is_dir()) if root.is_dir() else [] + if BASELINE not in arms: + print(f"No '{BASELINE}' arm under {root}/ -- nothing to compare against.") + return 1 + + loaded = {arm: load_arm(root, arm, args.track) for arm in arms} + missing = [arm for arm, data in loaded.items() if data is None] + for arm in missing: + print(f"NOTE: arm '{arm}' has no evaluation_metrics.json yet (still running?)") + loaded.pop(arm) + if BASELINE not in loaded: + return 1 + + base = loaded[BASELINE]["mean"] + order = [BASELINE] + [a for a in loaded if a != BASELINE] + + print(f"\nTREC {args.track} shortlist-policy A/B (n={loaded[BASELINE]['n']} topics)\n") + label_w = max(len(a) for a in order) + 2 + for group, keys in (("COST", COST), ("QUALITY", QUALITY), ("DIAGNOSTIC", DIAGNOSTIC)): + print(f" {group}") + print(f" {'arm':<{label_w}}" + "".join(f"{k:>26s}" for k in keys)) + for arm in order: + mean = loaded[arm]["mean"] + cells = "" + for key in keys: + value = mean.get(key) + cell = fmt(value) + if arm != BASELINE and isinstance(value, (int, float)): + ref = base.get(key) + if isinstance(ref, (int, float)): + cell += f" ({value - ref:+.4f})" + cells += f"{cell:>26s}" + print(f" {arm:<{label_w}}" + cells) + print() + + print(" Reading this:") + print(" An arm at the SAME shortlist_size as 'fixed' with higher shortlist_recall") + print(" spends its compute better -- that gain is free.") + print(" An arm with a LARGER shortlist_size bought its gain with GPU time; compare") + print(" it against 'fixed' only after noting the extra cost.") + print(" shortlist_selection_delta < 0 means the second level is still selecting worse") + print(" than a plain first-level cut at that depth, independent of the policy.\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/trialmatchai/main.py b/src/trialmatchai/main.py index 97958d37..336b42a2 100644 --- a/src/trialmatchai/main.py +++ b/src/trialmatchai/main.py @@ -307,6 +307,12 @@ def run_second_level_search( second_level_scores = { trial["nct_id"]: trial["score"] for trial in second_level_results } + # Persist the whole second-level pool, not just the shortlist. Together with + # first_level_scores.json this makes shortlist fusion replayable offline, so fusion + # weights can be retuned against completed runs instead of costing a GPU job each time. + # Measured motivation: shortlist_selection_delta is negative on every run so far, i.e. + # the fused shortlist selects worse than a plain first-level cut at the same depth. + write_json_file(second_level_scores, f"{output_folder}/second_level_scores.json") search_config = config.get("search", {}) combined_scores = _fuse_shortlist_scores( From 3ce7d606b4008ddfa1235d327389c3c923c5f67f Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Thu, 6 Aug 2026 14:54:36 +0200 Subject: [PATCH 04/13] feat(search): implement the first-level LLM query expansion backend The llm_expansion search channel was dead. LLMQueryExpansionBackend (the Protocol), parse_llm_query_expansion (the parser), LLMQueryExpansion (the schema), the channel's 0.5 weight and search.first_level.llm_expansion_enabled all existed, but nothing ever built a backend and main.py never passed one to ClinicalTrialSearch. Setting the flag logged "no expander is enabled but no expander is configured" and contributed no terms. Adds FirstLevelQueryExpander, which subclasses QueryExpander to reuse its engine, chat template and structured-output machinery -- so it shares the one cached vLLM engine rather than loading a second copy -- and swaps in a retrieval-query prompt and schema. This is a genuinely different task from the existing expander, not a rename. QueryExpander enriches the patient SUMMARY (conditions plus narrative sentences) for the expand stage. This one writes RETRIEVAL QUERIES: short noun phrases that should match a trial's title, condition list or eligibility text, bucketed into the six fields the planner turns into weighted query channels. To support both, QueryExpander's prompt and schema became overridable class attributes; the base class's values are unchanged. The per-field maxItems caps are deliberately uneven. llm_max_terms is a SHARED budget across the five query fields, spent in field order, so a model that pads primary_queries starves biomarker_queries and treatment_queries entirely -- a test pins that behaviour. primary_queries is therefore capped at 3 (a patient has one main disease), leaving the budget for the later fields. Off by default: the builder returns None unless llm_expansion_enabled is set, and returns None rather than raising if the expander cannot be constructed. A failed expansion degrades to empty and is logged, because this is one channel of eight -- losing it should cost recall, not the run. Wired once per run in main_pipeline, not per patient, so the engine is resolved once. --- src/trialmatchai/main.py | 12 ++ src/trialmatchai/matching/query_expansion.py | 169 ++++++++++++++++- tests/test_first_level_expansion.py | 182 +++++++++++++++++++ 3 files changed, 361 insertions(+), 2 deletions(-) create mode 100644 tests/test_first_level_expansion.py diff --git a/src/trialmatchai/main.py b/src/trialmatchai/main.py index 336b42a2..69d03a79 100644 --- a/src/trialmatchai/main.py +++ b/src/trialmatchai/main.py @@ -15,6 +15,7 @@ rank_trials, save_ranked_trials, ) +from trialmatchai.matching.query_expansion import build_first_level_expander from trialmatchai.matching.shortlist_depth import choose_shortlist_depth, depth_report from trialmatchai.matching.retrieval.trial_retrieval import ClinicalTrialSearch from trialmatchai.matching.retrieval.criteria_retrieval import SecondStageRetriever @@ -108,6 +109,7 @@ def run_first_level_search( config: Dict, search_backend, patient_profile: PatientProfile | None = None, + llm_query_expander=None, ) -> Optional[Tuple]: main_conditions = list(keywords.get("main_conditions", [])) other_conditions = list(keywords.get("other_conditions", [])) @@ -126,6 +128,9 @@ def run_first_level_search( search_backend=search_backend, embedder=embedder, entity_annotator=entity_annotator, + # Without this the llm_expansion channel is dead: the planner logs "no expander is + # configured" and returns [], however the config flag is set. + llm_query_expander=llm_query_expander, ) search_cfg = config["search"] @@ -553,6 +558,12 @@ def main_pipeline( embedder = build_embedder(config) entity_annotator = build_entity_annotator(config, embedder=embedder) + # Built once for the whole run, not per patient: it shares the cached CoT engine, and + # rebuilding per patient would re-resolve that engine 75 times. None when + # search.first_level.llm_expansion_enabled is off, which is the default. + llm_query_expander = build_first_level_expander(config) + if llm_query_expander is not None: + logger.info("First-level LLM query expansion is ON (llm_expansion channel active).") with warnings.catch_warnings(): warnings.filterwarnings( @@ -647,6 +658,7 @@ def main_pipeline( config, search_backend, patient_profile=profile, + llm_query_expander=llm_query_expander, ) if not result: logger.error("First-level search failed for %s", patient_id) diff --git a/src/trialmatchai/matching/query_expansion.py b/src/trialmatchai/matching/query_expansion.py index 5ffbcd00..3d36c1b6 100644 --- a/src/trialmatchai/matching/query_expansion.py +++ b/src/trialmatchai/matching/query_expansion.py @@ -105,6 +105,11 @@ def _resolve_settings(config: Dict[str, Any]) -> Dict[str, Any]: class QueryExpander: """CoT expander; loads its model lazily so import stays base-deps safe.""" + # Overridable by subclasses that reuse this engine/template machinery for a different + # extraction task (see FirstLevelQueryExpander). + system_prompt: str = SYSTEM_PROMPT + json_schema: Dict[str, Any] = _KEYWORDS_JSON_SCHEMA + def __init__(self, settings: Dict[str, Any], config: Dict[str, Any]): self.settings = settings self.config = config @@ -160,7 +165,7 @@ def _generate(self, narrative: str) -> str: no_think = bool(self.settings.get("no_think")) user_content = ("/no_think\n" + narrative) if no_think else narrative messages = [ - {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "system", "content": self.system_prompt}, {"role": "user", "content": user_content}, ] # Qwen3.x-style templates take enable_thinking; harmless-and-ignored elsewhere (guarded). @@ -199,7 +204,7 @@ def _apply_template(tokenize): if self.settings.get("guided_json"): from vllm.sampling_params import StructuredOutputsParams # type: ignore - structured = StructuredOutputsParams(json=_KEYWORDS_JSON_SCHEMA, disable_any_whitespace=True) + structured = StructuredOutputsParams(json=self.json_schema, disable_any_whitespace=True) params = SamplingParams( temperature=0.0, max_tokens=self.settings["max_new_tokens"], @@ -257,3 +262,163 @@ def enrich_summary( if sentences: out["patient_narrative"] = sentences return out + + +# --- first-level retrieval query expansion (the llm_expansion search channel) ------------ # + +# Distinct from SYSTEM_PROMPT above. That one enriches the patient SUMMARY (conditions and +# narrative sentences). This one writes RETRIEVAL QUERIES: short noun phrases that should +# match trial titles, conditions and eligibility text. The two are not interchangeable -- +# the first-level planner buckets these six fields into weighted query channels. +FIRST_LEVEL_SYSTEM_PROMPT = """ +You expand a patient description into search queries for a clinical trial index. + +Write SHORT NOUN PHRASES that would appear in a trial's title, condition list or +eligibility criteria. Do not write sentences, questions or explanations. + +Fill these six fields: + +1. "primary_queries": the patient's main disease as a trial would name it. Include the + staging or subtype only when the patient description states it. +2. "disease_aliases": other names for that same disease -- synonyms, abbreviations, older + or regional terminology, and the expanded form of any abbreviation. +3. "broader_queries": the parent disease categories a trial might recruit under, from + narrower to wider. These deliberately trade precision for coverage. +4. "biomarker_queries": genes, mutations, fusions, receptor and expression status, and + other molecular markers stated for this patient. +5. "treatment_queries": drugs, drug classes, procedures and prior therapies stated for + this patient. +6. "discarded_or_uncertain": terms you considered but rejected, and anything you are not + confident the patient description supports. + +Rules: +- Use ONLY what the patient description states. Never infer a diagnosis, stage, biomarker + or therapy that is not written there. Put anything doubtful in "discarded_or_uncertain". +- Leave a field as an empty list when the description supports nothing for it. An empty + list is correct; an invented term is not. +- No duplicates within a field. + +Return a JSON object with exactly those six keys and no other commentary. +""".strip() + +_FIRST_LEVEL_FIELDS = ( + "primary_queries", + "disease_aliases", + "broader_queries", + "biomarker_queries", + "treatment_queries", + "discarded_or_uncertain", +) + +# maxItems bounds the array COUNT for the same reason as _KEYWORDS_JSON_SCHEMA: it forces a +# verbose model to close each array instead of emitting terms until max_tokens runs out. +# These are noun phrases, so a short maxLength is safe here (unlike expanded_sentences). +# +# The per-field caps are deliberately uneven. search.first_level.llm_max_terms is a SHARED +# budget across the five query fields, spent in field order (first_level_planner +# parse_llm_query_expansion), so a model that pads primary_queries starves the biomarker and +# treatment channels entirely. Capping primary_queries tightly -- a patient has one main +# disease, not twelve -- keeps the budget available for the later fields. +_FIRST_LEVEL_MAX_ITEMS = { + "primary_queries": 3, + "disease_aliases": 8, + "broader_queries": 5, + "biomarker_queries": 8, + "treatment_queries": 8, + "discarded_or_uncertain": 12, +} +_FIRST_LEVEL_JSON_SCHEMA = { + "type": "object", + "properties": { + field: { + "type": "array", + "maxItems": _FIRST_LEVEL_MAX_ITEMS[field], + "items": {"type": "string", "maxLength": 120}, + } + for field in _FIRST_LEVEL_FIELDS + }, + "required": list(_FIRST_LEVEL_FIELDS), +} + +_FIRST_LEVEL_EMPTY: Dict[str, List[str]] = {field: [] for field in _FIRST_LEVEL_FIELDS} + + +def _first_level_patient_text(profile: Any, matching_summary: Dict[str, Any]) -> str: + """Compact patient description for the expander prompt. + + Built from the matching summary rather than the raw profile so it stays in step with + what first-level retrieval actually searches on. + """ + summary = matching_summary or {} + parts: List[str] = [] + main = [c for c in _as_list(summary.get("main_conditions")) if c][:12] + other = [c for c in _as_list(summary.get("other_conditions")) if c][:30] + narrative = [s for s in _as_list(summary.get("patient_narrative")) if s][:12] + if main: + parts.append("Main conditions: " + "; ".join(main)) + if other: + parts.append("Other conditions and factors: " + "; ".join(other)) + age, gender = summary.get("age"), summary.get("gender") + demographics = [ + f"Age: {age}" for _ in (1,) if age not in (None, "", "all") + ] + [f"Sex: {gender}" for _ in (1,) if gender not in (None, "", "all")] + if demographics: + parts.append(", ".join(demographics)) + if narrative: + parts.append("Description: " + " ".join(narrative)) + return "\n".join(parts).strip() + + +class FirstLevelQueryExpander(QueryExpander): + """Implements ``LLMQueryExpansionBackend`` for the first-level ``llm_expansion`` channel. + + Reuses QueryExpander's engine, chat-template and structured-output machinery -- so it + shares the one cached vLLM engine rather than loading a second copy -- but swaps in the + retrieval-query prompt and schema. + """ + + system_prompt = FIRST_LEVEL_SYSTEM_PROMPT + json_schema = _FIRST_LEVEL_JSON_SCHEMA + + def expand_first_level_queries( + self, + *, + profile: Any, + matching_summary: Dict[str, Any], + ) -> Dict[str, Any]: + patient_text = _first_level_patient_text(profile, matching_summary) + if not patient_text: + return dict(_FIRST_LEVEL_EMPTY) + try: + raw = self._generate(patient_text) + parsed = extract_json_object(BaseTrialProcessor._strip_thinking_tags(raw)) + if not isinstance(parsed, dict): + raise ValueError("first-level expansion output was not a JSON object") + return {field: _as_list(parsed.get(field)) for field in _FIRST_LEVEL_FIELDS} + except Exception as exc: + # Retrieval must not fail because expansion did: the channel is one of eight and + # carries weight 0.5, so an empty expansion degrades recall rather than the run. + logger.error( + "First-level query expansion failed; continuing without that channel: %s", exc + ) + return dict(_FIRST_LEVEL_EMPTY) + + +def build_first_level_expander(config: Dict[str, Any]) -> "FirstLevelQueryExpander | None": + """Construct the expander when ``search.first_level.llm_expansion_enabled`` is true. + + Independent of ``query_expansion.enabled``: that flag governs the separate summary + enrichment stage. Both may run, and they share one engine. + """ + first_level = (config.get("search") or {}).get("first_level") or {} + if not first_level.get("llm_expansion_enabled"): + return None + try: + return FirstLevelQueryExpander(_resolve_settings(config), config) + except Exception as exc: + logger.error( + "search.first_level.llm_expansion_enabled is set but the expander could not be " + "built; first-level search continues without that channel: %s", + exc, + ) + return None diff --git a/tests/test_first_level_expansion.py b/tests/test_first_level_expansion.py new file mode 100644 index 00000000..e7a900da --- /dev/null +++ b/tests/test_first_level_expansion.py @@ -0,0 +1,182 @@ +"""First-level LLM query expansion (the llm_expansion search channel). + +This channel was dead before: the protocol, parser, schema and config flag all existed, +but nothing ever constructed a backend, so enabling the flag logged "no expander is +configured" and returned no terms. These tests cover the backend and the config gate. +""" + +import json + +import pytest + +from trialmatchai.matching.query_expansion import ( + _FIRST_LEVEL_FIELDS, + FirstLevelQueryExpander, + _first_level_patient_text, + build_first_level_expander, +) +from trialmatchai.matching.retrieval.first_level_planner import parse_llm_query_expansion + +SUMMARY = { + "main_conditions": ["metastatic breast cancer"], + "other_conditions": ["hypertension", "HER2 positive"], + "patient_narrative": ["A 54 year old woman with metastatic breast cancer."], + "age": 54, + "gender": "female", +} + + +class _FakeExpander(FirstLevelQueryExpander): + """Bypasses __init__ so no model or GPU is touched; _generate returns a canned reply.""" + + def __init__(self, reply): + self._reply = reply + self.settings = {"guided_json": True, "max_new_tokens": 512} + self.config = {} + self.backend = "vllm" + + def _generate(self, narrative): + if isinstance(self._reply, Exception): + raise self._reply + return self._reply + + +def test_expands_into_the_six_planner_fields(): + payload = { + "primary_queries": ["metastatic breast cancer"], + "disease_aliases": ["breast carcinoma", "mammary carcinoma"], + "broader_queries": ["solid tumor"], + "biomarker_queries": ["HER2 positive"], + "treatment_queries": ["trastuzumab"], + "discarded_or_uncertain": ["hypertension"], + } + result = _FakeExpander(json.dumps(payload)).expand_first_level_queries( + profile=None, matching_summary=SUMMARY + ) + assert result == payload + assert set(result) == set(_FIRST_LEVEL_FIELDS) + + +def test_output_is_consumable_by_the_planner_parser(): + """The backend's contract is the planner's parser, not just valid JSON.""" + payload = { + "primary_queries": ["metastatic breast cancer"], + "disease_aliases": ["breast carcinoma"], + "broader_queries": ["solid tumor"], + "biomarker_queries": ["HER2 positive"], + "treatment_queries": ["trastuzumab"], + "discarded_or_uncertain": [], + } + raw = _FakeExpander(json.dumps(payload)).expand_first_level_queries( + profile=None, matching_summary=SUMMARY + ) + parsed = parse_llm_query_expansion(raw, max_terms=12) + assert parsed.primary_queries == ["metastatic breast cancer"] + assert parsed.biomarker_queries == ["HER2 positive"] + + +def test_max_terms_is_a_shared_budget_spent_primary_first(): + """llm_max_terms caps the TOTAL across the five query fields, not each one, and is spent + in field order. A model that fills primary_queries can starve the later channels, so the + prompt must keep primary_queries to the actual disease rather than padding it.""" + payload = {field: [] for field in _FIRST_LEVEL_FIELDS} + payload["primary_queries"] = [f"q{i}" for i in range(5)] + payload["disease_aliases"] = ["alias1", "alias2"] + payload["biomarker_queries"] = ["EGFR"] + + parsed = parse_llm_query_expansion(payload, max_terms=6) + + assert parsed.primary_queries == [f"q{i}" for i in range(5)] + assert parsed.disease_aliases == ["alias1"] # only one slot left + assert parsed.biomarker_queries == [] # budget exhausted before this field + + +def test_reasoning_tags_are_stripped_before_json_extraction(): + """Reasoning models emit containing an echo of the schema; extracting from that + would return the schema instead of the answer.""" + payload = {field: [] for field in _FIRST_LEVEL_FIELDS} + payload["primary_queries"] = ["glioblastoma"] + reply = ( + "The schema wants primary_queries, disease_aliases, ..." + + json.dumps(payload) + ) + result = _FakeExpander(reply).expand_first_level_queries( + profile=None, matching_summary=SUMMARY + ) + assert result["primary_queries"] == ["glioblastoma"] + + +@pytest.mark.parametrize( + "reply", + ["not json at all", json.dumps(["a", "list"]), RuntimeError("engine died")], +) +def test_failures_degrade_to_empty_not_raise(reply): + """Retrieval must survive a failed expansion: this is 1 of 8 channels, weight 0.5.""" + result = _FakeExpander(reply).expand_first_level_queries( + profile=None, matching_summary=SUMMARY + ) + assert result == {field: [] for field in _FIRST_LEVEL_FIELDS} + + +def test_a_bare_string_field_is_not_shredded_into_characters(): + payload = {field: [] for field in _FIRST_LEVEL_FIELDS} + payload["primary_queries"] = "glioblastoma" + result = _FakeExpander(json.dumps(payload)).expand_first_level_queries( + profile=None, matching_summary=SUMMARY + ) + assert result["primary_queries"] == ["glioblastoma"] + + +def test_empty_summary_skips_the_model_entirely(): + expander = _FakeExpander(RuntimeError("must not be called")) + assert expander.expand_first_level_queries(profile=None, matching_summary={}) == { + field: [] for field in _FIRST_LEVEL_FIELDS + } + + +def test_patient_text_includes_conditions_and_demographics(): + text = _first_level_patient_text(None, SUMMARY) + assert "metastatic breast cancer" in text + assert "HER2 positive" in text + assert "54" in text and "female" in text + + +def test_patient_text_omits_placeholder_demographics(): + text = _first_level_patient_text( + None, {"main_conditions": ["asthma"], "age": "all", "gender": "all"} + ) + assert "asthma" in text + assert "Age:" not in text and "Sex:" not in text + + +def test_builder_returns_none_unless_the_flag_is_set(): + assert build_first_level_expander({}) is None + assert build_first_level_expander({"search": {"first_level": {}}}) is None + assert ( + build_first_level_expander( + {"search": {"first_level": {"llm_expansion_enabled": False}}} + ) + is None + ) + + +def test_builder_degrades_to_none_when_construction_fails(): + """A misconfigured expander must not abort the run; the channel just stays empty.""" + config = { + "search": {"first_level": {"llm_expansion_enabled": True}}, + "model": {}, # no base_model -> QueryExpander raises + "query_expansion": {}, + } + assert build_first_level_expander(config) is None + + +def test_schema_caps_primary_queries_tightly_to_protect_the_shared_budget(): + """Guards the interaction pinned above: primary_queries is spent first out of + llm_max_terms, so its schema cap must leave room for the later channels.""" + from trialmatchai.matching.query_expansion import _FIRST_LEVEL_JSON_SCHEMA + + props = _FIRST_LEVEL_JSON_SCHEMA["properties"] + primary = props["primary_queries"]["maxItems"] + assert primary <= 3 + for field in ("biomarker_queries", "treatment_queries", "disease_aliases"): + assert props[field]["maxItems"] > primary From 57b5ab59d16ab878ad8d3fe69d67229090dd3856 Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Thu, 6 Aug 2026 17:13:56 +0200 Subject: [PATCH 05/13] fix(trec): judge shortlist depth on recall-aware metrics, not ndcg@10 ndcg@10 cannot answer a question about depth, and reading the A/B on it would have given exactly the wrong answer. It normalizes by the ideal over judged-AND-RANKED trials, which makes it recall-independent by construction: ranking fewer trials shrinks the ideal along with the DCG. Replaying the finished runs at reduced shortlist depth (re-ranking from cached CoT, no inference) shows it moving the WRONG way as depth grows, on all four runs. TREC 2021 / qwen36_medcpt: depth shortRec ndcg@10 ndcg_full@10 P@10(elig) 10 0.0787 0.8948 0.6028 0.3853 50 0.2565 0.8134 0.7916 0.6800 100 0.3924 0.8210 0.8113 0.7360 196 0.5774 0.8215 0.8197 0.7680 ndcg@10 reads BEST at depth 10, where the shortlist holds 8% of the relevant trials. Meanwhile ndcg_full@10 (+0.2169) and P@10(eligible) (+0.3827) both rise strongly with depth, because they are normalized against the full judged pool and so count a trial that never entered the shortlist as a miss. Same pattern on h100_baichuan/21, qwen36/22 and l40/23. So the funnel premise holds: depth does reach the top ten. It is simply invisible in the metric the comparison led with. ndcg_full@10 and P@10(eligible) move to QUALITY; ndcg@10 moves to DIAGNOSTIC with the reason recorded beside it. The replay also shows clear diminishing returns -- 25->50 buys +0.051 ndcg_full@10 while 100->196 buys +0.008 -- so the deeper arm should be expected to gain modestly, not doubled. TREC 2023 is the exception and is still climbing steeply at its full depth (shortlist recall only 0.254), which makes it the strongest case for more depth rather than 2021. --- scripts/compare_shortlist_ab.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/scripts/compare_shortlist_ab.py b/scripts/compare_shortlist_ab.py index 9718338a..d3efee6c 100644 --- a/scripts/compare_shortlist_ab.py +++ b/scripts/compare_shortlist_ab.py @@ -19,15 +19,24 @@ from pathlib import Path BASELINE = "fixed" +# Ordered deliberately. ndcg_full@10 normalizes by the ideal over the FULL judged pool, so a +# relevant trial that never entered the shortlist counts as a miss -- it is recall-aware, and +# it is the metric a depth change should be judged on. QUALITY = ( "shortlist_recall", - "recall@1000", - "ndcg@10", - "P@10(rel>=1)", + "ndcg_full@10", "P@10(eligible)", + "P@10(rel>=1)", + "recall@1000", ) COST = ("shortlist_size",) -DIAGNOSTIC = ("funnel_depth_loss", "shortlist_selection_delta") +# ndcg@10 sits here, NOT in QUALITY. It normalizes by the ideal over judged-AND-RANKED trials, +# which makes it recall-INDEPENDENT by construction: ranking fewer trials shrinks the ideal +# too. Replaying finished runs at reduced depth shows it moving the wrong way -- on TREC 2021 +# it reads 0.8948 at depth 10 and 0.8215 at depth 196, while ndcg_full@10 goes 0.6028 -> 0.8197 +# and P@10(eligible) 0.3853 -> 0.7680 over the same range. Judging a depth policy on ndcg@10 +# would conclude that more depth hurts, which is exactly backwards. +DIAGNOSTIC = ("ndcg@10", "funnel_depth_loss", "shortlist_selection_delta") def load_arm(root: Path, arm: str, track: str) -> dict | None: @@ -87,7 +96,10 @@ def main() -> int: print() print(" Reading this:") - print(" An arm at the SAME shortlist_size as 'fixed' with higher shortlist_recall") + print(" Judge a depth change on ndcg_full@10 and P@10(eligible). Both are recall-aware.") + print(" Do NOT judge it on ndcg@10 -- that is normalized over judged-AND-ranked trials,") + print(" so it is recall-independent and moves the WRONG way as depth grows.") + print(" An arm at the SAME shortlist_size as 'fixed' with higher recall-aware quality") print(" spends its compute better -- that gain is free.") print(" An arm with a LARGER shortlist_size bought its gain with GPU time; compare") print(" it against 'fixed' only after noting the extra cost.") From 51428d4135ae28220a9ed872f786804a069ad06a Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Fri, 7 Aug 2026 10:32:23 +0200 Subject: [PATCH 06/13] feat(agent): deterministic verification of the reasoner's eligibility claims First tool-grounded step of the agent: the reasoner's per-criterion verdicts are now checkable against the deterministic constraint engine instead of being taken on trust. The eligibility stage decides Met / Not Met / Violated by reading criterion text, which is where language models are weakest -- comparing an age or a lab value against a numeric threshold. GPT-4 scores ~51% on MedCalc-Bench with chain-of-thought and 81-85% when given the computation instead; AgentMD reports 87.7% against 40.9% for CoT on the same shape of task. TrialMatchAI already had a deterministic constraint engine, but it was consulted only during second-level retrieval scoring, blended in at weight 0.25 -- the reasoner never saw it. verify_trial_output() re-evaluates each criterion's constraints against the patient and compares the verdict with what the reasoner claimed, reporting every disagreement. The split is the whole design. Deterministic verdicts are authoritative ONLY for age, sex, lab and performance_status -- decidable comparisons against a structured patient value. Condition, phenotype, medication, procedure and biomarker identity stay with the reasoner, which beats a regex at them. That is the neuro-symbolic division alphaNeSy-CTM measured on TREC CT 2021-2023: adding a symbolic verifier moved specificity -- the ability to correctly REJECT a trial -- from 24.7% to 75.7%, while the agentic loop around it added 1.4 points. The verifier is where the value is, so it lands before any loop. Deliberately conservative. It abstains unless a constraint is decisive (status matched or violated, confidence >= 0.75); an unknown patient age decides nothing; a criterion whose text does not pair with a stored criterion is left alone rather than guessed at; malformed reasoner output cannot raise. Corrections record ReasonerClassification and VerifiedBy so a corrected run stays auditable, and apply_corrections=false gives a measure-only mode to size the effect before letting it act. Off by default (verification.enabled). Calls no model, so it is testable and tunable without a GPU -- being replayed over the completed runs now. --- src/trialmatchai/matching/verification.py | 208 ++++++++++++++++++++++ tests/test_verification.py | 192 ++++++++++++++++++++ 2 files changed, 400 insertions(+) create mode 100644 src/trialmatchai/matching/verification.py create mode 100644 tests/test_verification.py diff --git a/src/trialmatchai/matching/verification.py b/src/trialmatchai/matching/verification.py new file mode 100644 index 00000000..1281deea --- /dev/null +++ b/src/trialmatchai/matching/verification.py @@ -0,0 +1,208 @@ +"""Deterministic verification of the reasoner's per-criterion eligibility claims. + +The eligibility stage reads criterion text and emits Met / Not Met / Unclear / Irrelevant +(inclusion) or Violated / Not Violated / ... (exclusion). It does that in free text, which +is exactly where language models are weakest: comparing a patient's age or a lab value +against a numeric threshold. On MedCalc-Bench, GPT-4 scores ~51% on medical calculation +with chain-of-thought, and giving it the computation instead of asking it to reason raises +that to 81-85%. AgentMD reports 87.7% against 40.9% for CoT on the same shape of task. + +TrialMatchAI already has a deterministic constraint engine (``constraints/``), but it is +only consulted during second-level retrieval scoring, blended in at weight 0.25. The +reasoner never sees it. This module closes that gap: it re-derives the constraints for each +criterion, evaluates them against the patient, and compares the verdict with what the +reasoner claimed. + +**The split matters.** Deterministic verdicts are authoritative ONLY for constraint kinds +that are genuinely decidable -- age, sex, labs, performance status. Semantic kinds +(condition, phenotype, medication, procedure, biomarker identity) stay with the reasoner, +which is better at them than a regex. This is the neuro-symbolic division that alphaNeSy-CTM +measured on TREC CT 2021-2023: adding a symbolic verifier moved specificity -- the ability +to correctly REJECT a trial -- from 24.7% to 75.7%, while the agentic loop around it added +only 1.4 accuracy points. The verifier is where the value is. + +Disabled by default (``verification.enabled``). Nothing here calls a model. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +from trialmatchai.constraints import ( + CriterionConstraintEvaluation, + PatientConstraintContext, + evaluate_constraint_set, + extract_constraint_set, +) +from trialmatchai.matching.trial_ranker import _normalize_classification +from trialmatchai.utils.logging_config import setup_logging + +logger = setup_logging(__name__) + +# Constraint kinds whose deterministic verdict outranks the reasoner. These are decidable +# comparisons against a structured patient value: "age >= 18", "ANC > 1500", "ECOG 0-2". +# Everything else -- condition, phenotype, medication, procedure, biomarker, temporal -- is +# semantic or needs clinical judgement, and the reasoner keeps the final word there. +AUTHORITATIVE_KINDS = frozenset({"age", "sex", "lab", "performance_status"}) + +# Below this the extractor is guessing; a regex that half-matched must not overturn the model. +MIN_CONFIDENCE = 0.75 + +_INCLUSION_MET = "met" +_INCLUSION_NOT_MET = "not met" +_EXCLUSION_VIOLATED = "violated" +_EXCLUSION_NOT_VIOLATED = "not violated" + + +def verification_config(config: Mapping[str, Any] | None) -> dict[str, Any]: + raw = (config or {}).get("verification") or {} + if not isinstance(raw, Mapping): + raw = {} + return { + "enabled": bool(raw.get("enabled", False)), + "min_confidence": float(raw.get("min_confidence", MIN_CONFIDENCE)), + "authoritative_kinds": frozenset( + raw.get("authoritative_kinds") or AUTHORITATIVE_KINDS + ), + # When false, disagreements are recorded but the classification is left alone -- + # useful for measuring how often the verifier fires before letting it act. + "apply_corrections": bool(raw.get("apply_corrections", True)), + } + + +def _decisive_evaluations( + evaluation: CriterionConstraintEvaluation, + *, + authoritative_kinds: frozenset[str], + min_confidence: float, +) -> list[Any]: + """Constraint evaluations that are allowed to overrule the reasoner.""" + decisive = [] + for item in evaluation.evaluations: + constraint = item.constraint + if constraint.kind not in authoritative_kinds: + continue + if constraint.confidence < min_confidence: + continue + if item.status not in ("matched", "violated"): + continue # unknown / not_applicable decide nothing + decisive.append(item) + return decisive + + +def _verdict_for(polarity: str, decisive: Sequence[Any]) -> str | None: + """The label the deterministic engine implies, or None if it implies nothing. + + A violated constraint dominates: for an inclusion criterion the patient fails to meet + it, and for an exclusion criterion the patient is excluded by it. Requiring every + decisive constraint to match before claiming the positive label keeps the verifier + conservative -- it is far more willing to say "this does not hold" than "this holds". + """ + if not decisive: + return None + any_violated = any(item.status == "violated" for item in decisive) + all_matched = all(item.status == "matched" for item in decisive) + if polarity == "exclusion": + if any_violated: + return _EXCLUSION_VIOLATED + return _EXCLUSION_NOT_VIOLATED if all_matched else None + if any_violated: + return _INCLUSION_NOT_MET + return _INCLUSION_MET if all_matched else None + + +def verify_trial_output( + *, + trial_output: Mapping[str, Any], + criteria: Sequence[Mapping[str, Any]], + patient_context: PatientConstraintContext, + nct_id: str, + config: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Check one trial's reasoner output against the deterministic constraint engine. + + ``criteria`` supplies the criterion text and inclusion/exclusion type, keyed so each + reasoner claim can be paired with the criterion it judged. Returns a report and, when + corrections are enabled, a corrected copy of the trial output. + """ + cfg = verification_config(config) + by_text = { + str(c.get("criterion") or c.get("text") or "").strip(): c + for c in criteria + if (c.get("criterion") or c.get("text")) + } + + disagreements: list[dict[str, Any]] = [] + corrected = {k: (list(v) if isinstance(v, list) else v) for k, v in trial_output.items()} + + for section, polarity in ( + ("Inclusion_Criteria_Evaluation", "inclusion"), + ("Exclusion_Criteria_Evaluation", "exclusion"), + ): + rows = corrected.get(section) + if not isinstance(rows, list): + continue + new_rows = [] + for row in rows: + if not isinstance(row, dict): + new_rows.append(row) + continue + row = dict(row) + text = str(row.get("Criterion") or "").strip() + claimed = _normalize_classification(row.get("Classification")) + source = by_text.get(text) + if not text or not source: + new_rows.append(row) + continue + try: + constraint_set = extract_constraint_set( + nct_id=nct_id, + criteria_id=str(source.get("criteria_id") or text[:64]), + criterion=text, + eligibility_type=str(source.get("eligibility_type") or polarity), + entities=source.get("entities"), + ) + evaluation = evaluate_constraint_set(constraint_set, patient_context) + except Exception as exc: # never let verification break a finished run + logger.warning("Verification skipped for %s criterion %r: %s", nct_id, text[:60], exc) + new_rows.append(row) + continue + + decisive = _decisive_evaluations( + evaluation, + authoritative_kinds=cfg["authoritative_kinds"], + min_confidence=cfg["min_confidence"], + ) + verdict = _verdict_for(polarity, decisive) + if verdict is None or verdict == claimed: + new_rows.append(row) + continue + + disagreements.append( + { + "nct_id": nct_id, + "polarity": polarity, + "criterion": text, + "reasoner_said": claimed, + "verifier_said": verdict, + "kinds": sorted({d.constraint.kind for d in decisive}), + "reasons": [d.reason for d in decisive if d.reason][:3], + } + ) + if cfg["apply_corrections"]: + # Record what the reasoner said BEFORE overwriting it, so a corrected run + # stays auditable: score_trial only reads Classification. + row["ReasonerClassification"] = row.get("Classification") + row["VerifiedBy"] = "deterministic-constraints" + row["Classification"] = verdict.title() + new_rows.append(row) + corrected[section] = new_rows + + return { + "nct_id": nct_id, + "disagreements": disagreements, + "n_disagreements": len(disagreements), + "corrected_output": corrected if cfg["apply_corrections"] else trial_output, + "applied": bool(cfg["apply_corrections"] and disagreements), + } diff --git a/tests/test_verification.py b/tests/test_verification.py new file mode 100644 index 00000000..b8c6374d --- /dev/null +++ b/tests/test_verification.py @@ -0,0 +1,192 @@ +"""Deterministic verification of reasoner eligibility claims (matching/verification.py). + +The contract that matters: the verifier is allowed to overrule the reasoner ONLY on +decidable numeric/structured comparisons, and must stay silent everywhere else. +""" + +import pytest + +from trialmatchai.constraints import PatientConstraintContext, PatientConstraintFact +from trialmatchai.matching.verification import ( + AUTHORITATIVE_KINDS, + verification_config, + verify_trial_output, +) + + +def _ctx(age=None, sex=None, facts=None): + return PatientConstraintContext( + patient_id="p1", age_years=age, sex=sex, gender=sex, facts=facts or [] + ) + + +def _output(section, criterion, classification): + other = ( + "Exclusion_Criteria_Evaluation" + if section == "Inclusion_Criteria_Evaluation" + else "Inclusion_Criteria_Evaluation" + ) + return { + section: [{"Criterion": criterion, "Classification": classification}], + other: [], + "Final Decision": "Eligible", + } + + +def _criteria(text, eligibility_type): + return [{"criterion": text, "criteria_id": "c1", "eligibility_type": eligibility_type}] + + +ENABLED = {"verification": {"enabled": True, "apply_corrections": True}} + + +def test_disabled_by_default_makes_no_corrections(): + assert verification_config({})["apply_corrections"] is True + assert verification_config({})["enabled"] is False + + +def test_age_floor_overrules_a_wrong_met_claim(): + """A 12-year-old cannot meet 'age >= 18'. The reasoner said Met; the verifier must not + accept that, because it is a decidable comparison.""" + text = "Patients must be at least 18 years of age" + report = verify_trial_output( + trial_output=_output("Inclusion_Criteria_Evaluation", text, "Met"), + criteria=_criteria(text, "inclusion"), + patient_context=_ctx(age=12), + nct_id="NCT1", + config=ENABLED, + ) + assert report["n_disagreements"] == 1 + d = report["disagreements"][0] + assert d["reasoner_said"] == "met" + assert d["verifier_said"] == "not met" + assert "age" in d["kinds"] + row = report["corrected_output"]["Inclusion_Criteria_Evaluation"][0] + assert row["Classification"] == "Not Met" + assert row["ReasonerClassification"] == "Met" # provenance kept for audit + assert row["VerifiedBy"] == "deterministic-constraints" + + +def test_no_disagreement_when_the_reasoner_is_right(): + text = "Patients must be at least 18 years of age" + report = verify_trial_output( + trial_output=_output("Inclusion_Criteria_Evaluation", text, "Met"), + criteria=_criteria(text, "inclusion"), + patient_context=_ctx(age=40), + nct_id="NCT1", + config=ENABLED, + ) + assert report["n_disagreements"] == 0 + assert report["applied"] is False + + +def test_verifier_stays_silent_on_semantic_criteria(): + """Condition/phenotype matching is the reasoner's job -- a regex must not overrule it.""" + text = "Histologically confirmed anaplastic astrocytoma" + report = verify_trial_output( + trial_output=_output("Inclusion_Criteria_Evaluation", text, "Met"), + criteria=_criteria(text, "inclusion"), + patient_context=_ctx(age=40), + nct_id="NCT1", + config=ENABLED, + ) + assert report["n_disagreements"] == 0 + + +def test_authoritative_kinds_are_only_the_decidable_ones(): + assert AUTHORITATIVE_KINDS == {"age", "sex", "lab", "performance_status"} + for semantic in ("condition", "phenotype", "medication", "procedure", "biomarker"): + assert semantic not in AUTHORITATIVE_KINDS + + +def test_low_confidence_extraction_never_overrules(): + """A half-matched regex must not flip a label; confidence gates the override.""" + text = "Patients must be at least 18 years of age" + strict = {"verification": {"enabled": True, "min_confidence": 1.01}} + report = verify_trial_output( + trial_output=_output("Inclusion_Criteria_Evaluation", text, "Met"), + criteria=_criteria(text, "inclusion"), + patient_context=_ctx(age=12), + nct_id="NCT1", + config=strict, + ) + assert report["n_disagreements"] == 0 + + +def test_apply_corrections_false_reports_without_changing_labels(): + """Measure-only mode: see how often the verifier fires before letting it act.""" + text = "Patients must be at least 18 years of age" + report = verify_trial_output( + trial_output=_output("Inclusion_Criteria_Evaluation", text, "Met"), + criteria=_criteria(text, "inclusion"), + patient_context=_ctx(age=12), + nct_id="NCT1", + config={"verification": {"enabled": True, "apply_corrections": False}}, + ) + assert report["n_disagreements"] == 1 + assert report["applied"] is False + row = report["corrected_output"]["Inclusion_Criteria_Evaluation"][0] + assert row["Classification"] == "Met" # untouched + + +def test_unmatched_criterion_text_is_left_alone(): + """The reasoner sometimes paraphrases; no pairing means no verdict, not a guess.""" + report = verify_trial_output( + trial_output=_output("Inclusion_Criteria_Evaluation", "some paraphrase", "Met"), + criteria=_criteria("Patients must be at least 18 years of age", "inclusion"), + patient_context=_ctx(age=12), + nct_id="NCT1", + config=ENABLED, + ) + assert report["n_disagreements"] == 0 + + +def test_missing_age_decides_nothing(): + """Unknown patient age must produce 'unknown', never a violation.""" + text = "Patients must be at least 18 years of age" + report = verify_trial_output( + trial_output=_output("Inclusion_Criteria_Evaluation", text, "Met"), + criteria=_criteria(text, "inclusion"), + patient_context=_ctx(age=None), + nct_id="NCT1", + config=ENABLED, + ) + assert report["n_disagreements"] == 0 + + +@pytest.mark.parametrize("bad", [{}, {"Inclusion_Criteria_Evaluation": "not a list"}]) +def test_malformed_reasoner_output_does_not_raise(bad): + """Verification must never break a finished run.""" + report = verify_trial_output( + trial_output=bad, + criteria=_criteria("Patients must be at least 18 years of age", "inclusion"), + patient_context=_ctx(age=12), + nct_id="NCT1", + config=ENABLED, + ) + assert report["n_disagreements"] == 0 + + +def test_lab_threshold_flags_an_exclusion_the_reasoner_missed(): + """The case the literature says models fail: a numeric lab comparison.""" + text = "Platelet count less than 100,000/mm3" + facts = [ + PatientConstraintFact( + kind="lab", + label="platelet count", + value="50000", + unit="/mm3", + evidence_text="platelets 50000", + ) + ] + report = verify_trial_output( + trial_output=_output("Exclusion_Criteria_Evaluation", text, "Not Violated"), + criteria=_criteria(text, "exclusion"), + patient_context=_ctx(age=40, facts=facts), + nct_id="NCT1", + config=ENABLED, + ) + # Either the verifier catches it, or it abstains -- it must never confirm the wrong label. + for d in report["disagreements"]: + assert d["verifier_said"] in ("violated", "not violated") + assert d["reasoner_said"] == "not violated" From b1aeef6ac05a3e36753519db58cc405064e9621d Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Fri, 7 Aug 2026 12:21:19 +0200 Subject: [PATCH 07/13] fix(agent): make verification sound, and record that it does not pay on TREC Two findings from replaying the verifier over the completed runs, no inference. 1. The first version was UNSOUND. It asserted the positive label whenever every authoritative constraint matched. But a criterion is a conjunction and the verifier sees only its authoritative part: refuting the age clause of "age >= 18 with confirmed glioma" refutes the whole criterion, while matching it proves nothing about the glioma clause. That produced 484 "unclear -> met" and 322 "not met -> met" flips on TREC 2021 and cost -0.0098 ndcg_full@10 / -0.0133 P@10(eligible). It now abstains unless something is actually refuted; two regression tests pin the asymmetry. 2. Corrected, it is flat, and the reason is structural: TREC 2021 671 disagreements (4.6/100) ndcg_full@10 +0.0002 P@10(elig) -0.0040 TREC 2022 535 disagreements (6.9/100) ndcg_full@10 -0.0007 P@10(elig) -0.0020 TREC 2023 3 disagreements (0.0/100) no change Nearly every firing constraint is sex (473) or age (226) -- and hard_filters already enforces both at retrieval, so a trial whose age or sex excludes the patient never reaches the reasoner. The verifier is re-litigating a filter that already ran, against criterion prose rather than structured trial fields. Meanwhile lab and performance_status almost never fire, because TREC profiles are narrative summaries with no structured lab or ECOG values: 3 disagreements across 8,987 trials on 2023. The ceiling is set by the data, not the logic. Kept off by default. Still worth having for deployments with real EHR records, where labs and performance status exist and the filters may be looser -- but it is not a benchmark win and should not be sold as one. Also corrects an overreach in my earlier framing: alphaNeSy-CTM's specificity numbers come from a balanced binary eligible/ineligible task (100/100 per year), not corpus ranking, so they never predicted an ndcg/P@10 gain here. Its missing half -- LLM abduction of structured patient attributes from sparse notes -- is what would give a verifier material to work with on narrative data. --- src/trialmatchai/matching/verification.py | 67 ++++++++++++++++++----- tests/test_verification.py | 32 +++++++++++ 2 files changed, 84 insertions(+), 15 deletions(-) diff --git a/src/trialmatchai/matching/verification.py b/src/trialmatchai/matching/verification.py index 1281deea..329c67f2 100644 --- a/src/trialmatchai/matching/verification.py +++ b/src/trialmatchai/matching/verification.py @@ -21,6 +21,36 @@ to correctly REJECT a trial -- from 24.7% to 75.7%, while the agentic loop around it added only 1.4 accuracy points. The verifier is where the value is. +**Measured verdict on TREC: this does not pay, and should stay off.** Replayed over the +completed runs (corrected logic, no inference): + + TREC 2021 671 disagreements (4.6/100 trials) ndcg_full@10 +0.0002 P@10(elig) -0.0040 + TREC 2022 535 disagreements (6.9/100 trials) ndcg_full@10 -0.0007 P@10(elig) -0.0020 + TREC 2023 3 disagreements (0.0/100 trials) no change + +Two reasons, both structural rather than fixable here: + +1. **Redundant with the hard filters.** Nearly every firing constraint is sex (473 on 2021) + or age (226), and ``search.first_level.hard_filters`` already enforces exactly those at + retrieval (``lancedb_backend._trial_passes_filters``). A trial whose age bounds or sex + exclude the patient never reaches the reasoner, so the verifier is re-litigating a filter + that already ran -- against criterion prose instead of the structured trial fields, which + is strictly worse and occasionally wrong. +2. **Nothing else to check.** lab and performance_status almost never fire, because TREC + patient profiles are narrative summaries with no structured lab or ECOG values. TREC 2023 + produced 3 disagreements across 8,987 trials. + +So the ceiling here is set by the DATA, not the logic. The component is sound and cheap, and +is worth keeping for deployments with real EHR records -- where labs and performance status +exist and the hard filters may be looser -- but it is not a benchmark win, and enabling it on +TREC costs a little precision for nothing. + +The missing half is abduction: alphaNeSy-CTM pairs its symbolic verifier with an LLM step +that INFERS structured patient attributes from sparse notes. That is what would give a +verifier something to check on narrative data. Note also that its headline numbers come from +a balanced binary eligible/ineligible task (100/100 per year), not corpus ranking, so they do +not transfer to ndcg/P@10 directly. + Disabled by default (``verification.enabled``). Nothing here calls a model. """ @@ -49,10 +79,8 @@ # Below this the extractor is guessing; a regex that half-matched must not overturn the model. MIN_CONFIDENCE = 0.75 -_INCLUSION_MET = "met" _INCLUSION_NOT_MET = "not met" _EXCLUSION_VIOLATED = "violated" -_EXCLUSION_NOT_VIOLATED = "not violated" def verification_config(config: Mapping[str, Any] | None) -> dict[str, Any]: @@ -94,22 +122,31 @@ def _decisive_evaluations( def _verdict_for(polarity: str, decisive: Sequence[Any]) -> str | None: """The label the deterministic engine implies, or None if it implies nothing. - A violated constraint dominates: for an inclusion criterion the patient fails to meet - it, and for an exclusion criterion the patient is excluded by it. Requiring every - decisive constraint to match before claiming the positive label keeps the verifier - conservative -- it is far more willing to say "this does not hold" than "this holds". + ONLY the negative direction is sound, and this asymmetry is the whole correctness + argument. A criterion is a conjunction, and this verifier sees only the part of it that + falls in AUTHORITATIVE_KINDS. + + - A violated authoritative constraint refutes the whole conjunction. "Age >= 18 with + confirmed glioma" cannot be Met by a 12-year-old, whatever the glioma status. Sound. + - A matched authoritative constraint proves nothing about the rest. That same criterion + is NOT Met merely because the patient is 40; the glioma clause is unexamined, and it + belongs to the reasoner. + + An earlier version returned the positive labels when every decisive constraint matched. + Replayed over the completed runs that produced 484 "unclear -> met" and 322 "not met -> + met" flips on TREC 2021 alone, and cost -0.0098 ndcg_full@10 / -0.0133 P@10(eligible): + it was asserting whole criteria held on the strength of an age or sex match. Hence + abstention unless something is actually refuted. + + Exclusion polarity is already normalized by the engine: status "violated" means the + patient HAS the excluded item, so the criterion is Violated (see + constraints/evaluation.py _status_and_signal). """ if not decisive: return None - any_violated = any(item.status == "violated" for item in decisive) - all_matched = all(item.status == "matched" for item in decisive) - if polarity == "exclusion": - if any_violated: - return _EXCLUSION_VIOLATED - return _EXCLUSION_NOT_VIOLATED if all_matched else None - if any_violated: - return _INCLUSION_NOT_MET - return _INCLUSION_MET if all_matched else None + if not any(item.status == "violated" for item in decisive): + return None + return _EXCLUSION_VIOLATED if polarity == "exclusion" else _INCLUSION_NOT_MET def verify_trial_output( diff --git a/tests/test_verification.py b/tests/test_verification.py index b8c6374d..81a300fa 100644 --- a/tests/test_verification.py +++ b/tests/test_verification.py @@ -190,3 +190,35 @@ def test_lab_threshold_flags_an_exclusion_the_reasoner_missed(): for d in report["disagreements"]: assert d["verifier_said"] in ("violated", "not violated") assert d["reasoner_said"] == "not violated" + + +def test_verifier_never_asserts_the_positive_label(): + """Only the negative direction is sound. A criterion is a conjunction and the verifier + sees only its authoritative part, so a matching age proves nothing about the rest. + + Regression: an earlier version returned Met when every decisive constraint matched. On + TREC 2021 that produced 484 'unclear -> met' and 322 'not met -> met' flips and cost + -0.0098 ndcg_full@10, because it asserted whole criteria on the strength of an age match. + """ + text = "Patients aged 18 years or older with histologically confirmed glioma" + for claimed in ("Unclear", "Not Met", "Irrelevant"): + report = verify_trial_output( + trial_output=_output("Inclusion_Criteria_Evaluation", text, claimed), + criteria=_criteria(text, "inclusion"), + patient_context=_ctx(age=40), # age clause matches; glioma clause unexamined + nct_id="NCT1", + config=ENABLED, + ) + assert report["n_disagreements"] == 0, f"must not upgrade {claimed!r} to Met" + + +def test_exclusion_is_not_cleared_by_a_matching_authoritative_constraint(): + text = "Age over 75 years, or any active infection" + report = verify_trial_output( + trial_output=_output("Exclusion_Criteria_Evaluation", text, "Unclear"), + criteria=_criteria(text, "exclusion"), + patient_context=_ctx(age=40), # not over 75, but infection status is unexamined + nct_id="NCT1", + config=ENABLED, + ) + assert report["n_disagreements"] == 0 From fe960e46e133fea9573598e8f8959e09ad23c381 Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Fri, 7 Aug 2026 15:25:20 +0200 Subject: [PATCH 08/13] docs(ranker): record that softening the disqualification rule was measured and is worse The hard disqualification looks indefensible on inspection -- one Violated label out of a mean 9 (max 63) exclusion criteria irreversibly sinks a trial, and it fires on 19% of trials the TREC judges rated ELIGIBLE, 25% on 2023. That invites a future "fix". It was measured by replaying the completed runs from cached CoT, and every relaxation loses: rule 2021 ndcg_full@10 / P@10(elig) 2022 2023 >=1 violated (kept) 0.8197 / 0.7680 0.7481/0.6940 0.8806/0.8270 >=2 violated 0.7907 / 0.7280 0.7317/0.6760 0.8592/0.7676 >=3 violated 0.7839 / 0.7147 0.7293/0.6760 0.8515/0.7568 never disqualify 0.7837 / 0.7133 0.7295/0.6780 0.8502/0.7541 The rule pays because disqualification is ~2.3x more likely on an irrelevant trial than an eligible one (44% vs 19% on 2021): removing true violations buys more precision than the false ones cost. Reading the disagreements also shows many are not model errors but genuine clinical ambiguity -- the reasoner excluding a patient for prior bevacizumab under "no prior biologic therapy" is defensible even though the assessor graded that trial eligible. Comment only; no behaviour change. --- src/trialmatchai/matching/trial_ranker.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/src/trialmatchai/matching/trial_ranker.py b/src/trialmatchai/matching/trial_ranker.py index 61eb060a..77072b8e 100644 --- a/src/trialmatchai/matching/trial_ranker.py +++ b/src/trialmatchai/matching/trial_ranker.py @@ -36,6 +36,24 @@ def load_trial_data( # Eligibility scoring contract (REFACTOR_PLAN.md PR1, audit C1): a single Violated exclusion # HARD-DISQUALIFIES (not averaged in, which let a violated trial outrank an eligible one); # eligible trials score [0, 1] by the fraction of decided inclusion criteria (Met/Not Met) Met. +# +# DO NOT SOFTEN THIS. It looks far too brittle -- one Violated label out of a mean 9 (max 63) +# exclusion criteria irreversibly sinks a trial, and it fires on 19% of trials the TREC judges +# rated ELIGIBLE (25% on 2023). It was measured anyway, by replaying the completed runs from +# cached CoT, and every relaxation is WORSE: +# +# rule TREC 2021 ndcg_full@10 / P@10(elig) 2022 2023 +# >=1 violated (this) 0.8197 / 0.7680 0.7481/0.6940 0.8806/0.8270 +# >=2 violated 0.7907 / 0.7280 0.7317/0.6760 0.8592/0.7676 +# >=3 violated 0.7839 / 0.7147 0.7293/0.6760 0.8515/0.7568 +# never disqualify 0.7837 / 0.7133 0.7295/0.6780 0.8502/0.7541 +# +# The false-positive rate is real but the rule still pays, because disqualification is ~2.3x +# more likely on an irrelevant trial than an eligible one (44% vs 19% on 2021). Removing true +# violations buys more precision than the false ones cost. Inspecting the disagreements also +# shows many are not model errors at all but genuine clinical ambiguity -- e.g. the reasoner +# excluding a patient for prior bevacizumab under "no prior biologic therapy", which is +# defensible even though the assessor graded the trial eligible. DISQUALIFIED_SCORE = -1.0 # "Unclear" (info insufficient to decide) is the dominant classification. Partial credit — vs From cb335f7a5b242f89c732a138d235a347689a348b Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Thu, 13 Aug 2026 11:16:49 +0200 Subject: [PATCH 09/13] feat(reranker): add graded relevance scoring alongside the binary Yes/No mode Two defects in the reranker prompt, both cheap to fix, both pointing at the same measured symptom: shortlist_selection_delta is negative on every completed run, i.e. the second level selects WORSE than a plain first-level cut at the same depth. 1. It asks the wrong question. The prompt asks whether the patient text contains "sufficient information to evaluate" the criterion -- answerability, not relevance and not eligibility. It correlates with relevance (a matching trial's criteria are more discussable) but is a different quantity, and the downstream 0.5 cut in aggregate_to_trials therefore discards trials whose criteria the patient text merely fails to mention, however eligible the patient is. 2. It answers in a binary that does not aggregate. Criterion scores are summed and averaged into a trial score, so magnitude matters as much as order, and binary Yes/No logprob is the formulation known to lack that discrimination. Three independent sources converge: Setwise (SIGIR 2024) measures pointwise Yes/No at 0.386 BEIR nDCG@10, BELOW BM25's 0.436; "Don't Overthink Passage Reranking" (2025) shows it collapsing the partial-relevance band to ~0% of scores in 0.1-0.9 against 11.4% for a well-behaved reranker; ERank and TFRank (2025) were both built specifically because it "lacks the necessary scoring discrimination". "Beyond Yes and No" (NAACL 2024) measured the fix at +2.2 BEIR average and +7.1 on SciFact, for one forward pass and zero generated tokens. Adds LLM_reranker.scoring: binary (default) P(Yes) over Yes/No -- historical behaviour, byte-identical prompt graded Expected Relevance Value sum(p_k * y_k) over Not/Somewhat/Highly, normalized to [0, 1], with a prompt that asks for relevance Deliberately kept as one switch: the graded label space is what makes the relevance question aggregable, so splitting them would produce a scale the 0.5 threshold cannot read. Same cost -- one forward pass, one token, no generation. Default is unchanged and the binary prompt is byte-identical, because the LoRA adapter was tuned against that exact wording; tests assert both. _yes_probability falls back to binary when label_token_ids is absent, so the pre-existing test_model_helpers contract test passes unmodified, which is the evidence the old path is untouched. Not yet validated on GPU. It needs only the 2B reranker, not the 35B CoT model, so a second-level-only ablation is cheap: measure how many judged-relevant trials survive to the shortlist under each mode. --- src/trialmatchai/config/settings.py | 4 + src/trialmatchai/main.py | 1 + src/trialmatchai/models/llm/llm_reranker.py | 168 +++++++++++++++++--- tests/test_llm_reranker_scoring.py | 135 ++++++++++++++++ 4 files changed, 285 insertions(+), 23 deletions(-) create mode 100644 tests/test_llm_reranker_scoring.py diff --git a/src/trialmatchai/config/settings.py b/src/trialmatchai/config/settings.py index f4ab4b13..25053296 100644 --- a/src/trialmatchai/config/settings.py +++ b/src/trialmatchai/config/settings.py @@ -289,6 +289,10 @@ class LLMRerankerSettings(BaseModel): # to fit both engines on a smaller card (e.g. 48GB A40/L40). gpu_memory_utilization: float = Field(0.4, gt=0.0, le=1.0) tensor_parallel_size: int = Field(1, ge=1) + # "binary" = historical P(Yes) over Yes/No. "graded" = Expected Relevance Value over a + # 3-level scale, which yields an aggregable magnitude instead of a saturated binary. + # See models/llm/llm_reranker.py for the evidence; default keeps existing behaviour. + scoring: Literal["binary", "graded"] = "binary" model_config = ConfigDict(extra="allow") diff --git a/src/trialmatchai/main.py b/src/trialmatchai/main.py index 69d03a79..54378c90 100644 --- a/src/trialmatchai/main.py +++ b/src/trialmatchai/main.py @@ -603,6 +603,7 @@ def main_pipeline( # to the shared vllm section, so a large reranker fits the same way the CoT does. quantization=reranker_cfg.get("quantization", vllm_cfg.get("quantization", "")), kv_cache_dtype=reranker_cfg.get("kv_cache_dtype", vllm_cfg.get("kv_cache_dtype")), + scoring=str(reranker_cfg.get("scoring", "binary")), ) else: raise ValueError( diff --git a/src/trialmatchai/models/llm/llm_reranker.py b/src/trialmatchai/models/llm/llm_reranker.py index 7f1fb5a6..f1b54beb 100644 --- a/src/trialmatchai/models/llm/llm_reranker.py +++ b/src/trialmatchai/models/llm/llm_reranker.py @@ -9,9 +9,41 @@ logger = setup_logging(__name__) +# Graded label space for "graded" scoring mode. The score is the Expected Relevance Value +# sum(p_k * y_k) over these label tokens, normalized to [0, 1]. +# +# Why this exists. The default "binary" mode scores P(Yes) over a Yes/No pair, and three +# independent lines of evidence say that specific formulation is the weak link when scores are +# AGGREGATED rather than merely ordered -- which is exactly what aggregate_to_trials does: +# * Setwise (SIGIR 2024): pointwise Yes/No reaches 0.386 BEIR nDCG@10, BELOW BM25 at 0.436. +# * "Don't Overthink Passage Reranking" (2025): binary Yes/No collapses the partial-relevance +# band -- ~0% of scores land in 0.1-0.9 versus 11.4% for a well-behaved reranker. +# * ERank and TFRank (2025) were both built specifically because binary yes/no logprob +# "lacks the necessary scoring discrimination", and both emit graded scores instead. +# "Beyond Yes and No" (NAACL 2024) measured the fix: replacing Yes/No with a 3-level scale and +# taking sum(p_k * y_k) gained +2.2 BEIR nDCG@10 on average and +7.1 on SciFact, at identical +# cost -- one forward pass, zero generated tokens. Gains flatten past ~7 levels. +# +# This matches an independently measured defect in this pipeline: shortlist_selection_delta is +# negative on every completed run, i.e. the second level currently selects WORSE than a plain +# first-level cut at the same depth. +GRADED_LABELS: tuple[str, ...] = ("Not", "Somewhat", "Highly") + + class LLMReranker: - """vLLM-backed pointwise reranker: scores each (patient, criterion) pair by the constrained - Yes/No next-token probability, optionally via a LoRA adapter.""" + """vLLM-backed pointwise reranker over (patient, criterion) pairs, optionally LoRA-adapted. + + Two scoring modes, selected by ``LLM_reranker.scoring``: + + ``binary`` (default) + P(Yes) over a constrained Yes/No next token. Preserves the historical contract exactly. + + ``graded`` + Expected Relevance Value over ``GRADED_LABELS``: ``sum(p_k * y_k)`` normalized to + [0, 1]. Same single forward pass and still zero generated tokens, but it yields a + continuous magnitude instead of a saturated binary -- which is what the downstream + aggregation and its 0.5 cut actually need. + """ def __init__( self, @@ -29,6 +61,7 @@ def __init__( tensor_parallel_size: int = 1, quantization: str = "", kv_cache_dtype: str | None = None, + scoring: str = "binary", ): from vllm import SamplingParams # type: ignore @@ -62,26 +95,79 @@ def __init__( model_config=model_config, vllm_cfg=vllm_cfg ) + self.scoring = str(scoring or "binary").lower() + if self.scoring not in ("binary", "graded"): + logger.warning( + "Unknown LLM_reranker.scoring %r; falling back to 'binary'.", scoring + ) + self.scoring = "binary" + self.applicable_token_id, self.not_applicable_token_id = self._yes_no_token_ids() + if self.scoring == "graded": + self.label_token_ids = [self._first_token_id(w) for w in GRADED_LABELS] + allowed = list(dict.fromkeys(self.label_token_ids)) + if len(allowed) != len(GRADED_LABELS): + # Distinct labels must map to distinct first tokens or the softmax is degenerate. + logger.error( + "Graded labels %s collide under this tokenizer; using binary scoring.", + GRADED_LABELS, + ) + self.scoring = "binary" + if self.scoring == "binary": + self.label_token_ids = [ + self.not_applicable_token_id, + self.applicable_token_id, + ] + allowed = list(self.label_token_ids) + self.sampling_params = SamplingParams( temperature=0.0, max_tokens=1, logprobs=20, - allowed_token_ids=[self.applicable_token_id, self.not_applicable_token_id], + allowed_token_ids=allowed, ) + logger.info("Reranker scoring mode: %s", self.scoring) + + def _first_token_id(self, word: str) -> int: + return self.tokenizer(word, add_special_tokens=False)["input_ids"][0] def _yes_no_token_ids(self) -> tuple[int, int]: - yes = self.tokenizer("Yes", add_special_tokens=False)["input_ids"] - no = self.tokenizer("No", add_special_tokens=False)["input_ids"] - return yes[0], no[0] + return self._first_token_id("Yes"), self._first_token_id("No") - @staticmethod - def create_messages(patient_text: str, trial_text: str) -> List[Dict]: + # Kept verbatim: this is the historical prompt and the LoRA adapter was tuned against it. + # Note what it actually asks -- whether the patient text contains ENOUGH INFORMATION TO + # EVALUATE the criterion. That is answerability, not relevance and not eligibility. It + # correlates with relevance (a matching trial's criteria are more discussable) but is a + # different quantity, and the downstream 0.5 cut therefore discards trials whose criteria + # the patient text simply does not mention, however eligible the patient may be. + BINARY_SYSTEM_PROMPT = ( + "You are a clinical assistant tasked with determining whether the patient information (Statement A) " + "provides enough details to evaluate whether the patient satisfies or violates the clinical " + "trial eligibility criterion (Statement B). Respond with 'Yes' if Statement A contains sufficient " + "information to make this evaluation, or 'No' if it does not." + ) + + # Asks for RELEVANCE on a graded scale. Two changes from the above, deliberately bundled + # because the graded label space is what makes the relevance question aggregable. + GRADED_SYSTEM_PROMPT = ( + "You are a clinical assistant. Judge how relevant the clinical trial eligibility " + "criterion (Statement B) is to the patient described in Statement A -- that is, how " + "much this criterion bears on whether this particular patient could join the trial.\n" + "Answer with exactly one word:\n" + "'Highly' - the criterion concerns this patient's condition, treatment or " + "characteristics, and clearly bears on their eligibility.\n" + "'Somewhat' - the criterion is related to the patient's clinical picture but is " + "peripheral or only partly applicable.\n" + "'Not' - the criterion concerns a different disease, population or context and does " + "not bear on this patient." + ) + + @classmethod + def create_messages( + cls, patient_text: str, trial_text: str, *, scoring: str = "binary" + ) -> List[Dict]: system_prompt = ( - "You are a clinical assistant tasked with determining whether the patient information (Statement A) " - "provides enough details to evaluate whether the patient satisfies or violates the clinical " - "trial eligibility criterion (Statement B). Respond with 'Yes' if Statement A contains sufficient " - "information to make this evaluation, or 'No' if it does not." + cls.GRADED_SYSTEM_PROMPT if scoring == "graded" else cls.BINARY_SYSTEM_PROMPT ) return [ {"role": "user", "content": system_prompt}, @@ -99,27 +185,63 @@ def preprocess_text(self, text: str) -> str: def _build_prompt(self, patient_text: str, trial_text: str) -> str: messages = self.create_messages( - self.preprocess_text(patient_text), self.preprocess_text(trial_text) + self.preprocess_text(patient_text), + self.preprocess_text(trial_text), + scoring=self.scoring, ) return self.tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) - def _yes_probability(self, output: Any) -> float: + def _label_ids(self) -> list[int]: + """Label tokens in ascending-relevance order. + + Falls back to [No, Yes] when ``label_token_ids`` is absent so an instance built + without the full __init__ (older callers, hand-made test doubles) still scores as + plain binary rather than raising. + """ + ids = getattr(self, "label_token_ids", None) + if ids: + return list(ids) + return [self.not_applicable_token_id, self.applicable_token_id] + + def _label_probabilities(self, output: Any) -> list[float] | None: + """Softmax over the label tokens only, in ascending-relevance order.""" try: token_logprobs = output.outputs[0].logprobs[0] except (AttributeError, IndexError, TypeError): - return 0.0 - yes = token_logprobs.get(self.applicable_token_id) - no = token_logprobs.get(self.not_applicable_token_id) - yes_lp = yes.logprob if yes is not None else float("-inf") - no_lp = no.logprob if no is not None else float("-inf") - highest = max(yes_lp, no_lp) + return None + logprobs = [] + for token_id in self._label_ids(): + entry = token_logprobs.get(token_id) + logprobs.append(entry.logprob if entry is not None else float("-inf")) + highest = max(logprobs) if highest == float("-inf"): + return None + exps = [math.exp(lp - highest) for lp in logprobs] + total = sum(exps) + if total <= 0: + return None + return [e / total for e in exps] + + def _yes_probability(self, output: Any) -> float: + """Score in [0, 1]. + + binary: P(Yes). graded: Expected Relevance Value sum(p_k * y_k) with y_k = k, divided + by (K-1) so both modes share the [0, 1] range that aggregate_to_trials and its 0.5 + threshold assume. Renormalizing over the label tokens (rather than the full vocabulary) + keeps the score comparable across prompts of different lengths. + """ + probabilities = self._label_probabilities(output) + if probabilities is None: + return 0.0 + if getattr(self, "scoring", "binary") == "binary": + # label_token_ids is [No, Yes]; P(Yes) reproduces the historical value exactly. + return probabilities[1] + k = len(probabilities) - 1 + if k <= 0: return 0.0 - ey = math.exp(yes_lp - highest) - en = math.exp(no_lp - highest) - return ey / (ey + en) + return sum(p * i for i, p in enumerate(probabilities)) / k def rank_pairs(self, patient_trial_pairs: List[tuple]) -> List[Dict]: results: List[Dict] = [] diff --git a/tests/test_llm_reranker_scoring.py b/tests/test_llm_reranker_scoring.py new file mode 100644 index 00000000..a0b1f192 --- /dev/null +++ b/tests/test_llm_reranker_scoring.py @@ -0,0 +1,135 @@ +"""Reranker scoring modes (models/llm/llm_reranker.py). + +The scores this produces are SUMMED and AVERAGED into a trial score and then cut at 0.5 +(criteria_retrieval.aggregate_to_trials), so what matters is not just the ordering but the +magnitude. These tests pin the [0, 1] range, the binary contract, and the graded spread. +""" + +import math +import types + +import pytest + +from trialmatchai.models.llm.llm_reranker import GRADED_LABELS, LLMReranker + + +def _reranker(scoring, label_token_ids): + """An LLMReranker without __init__ -- no vLLM engine, no GPU.""" + r = LLMReranker.__new__(LLMReranker) + r.scoring = scoring + r.label_token_ids = label_token_ids + r.applicable_token_id, r.not_applicable_token_id = 1, 0 + return r + + +def _output(logprob_by_token): + entries = {t: types.SimpleNamespace(logprob=lp) for t, lp in logprob_by_token.items()} + return types.SimpleNamespace(outputs=[types.SimpleNamespace(logprobs=[entries])]) + + +# --------------------------------------------------------------- binary contract +def test_binary_returns_probability_of_yes(): + r = _reranker("binary", [0, 1]) # [No, Yes] + # Equal logprobs -> 0.5 exactly. + assert r._yes_probability(_output({0: math.log(0.5), 1: math.log(0.5)})) == pytest.approx(0.5) + # Yes dominant. + assert r._yes_probability(_output({0: math.log(0.1), 1: math.log(0.9)})) == pytest.approx(0.9) + # No dominant. + assert r._yes_probability(_output({0: math.log(0.8), 1: math.log(0.2)})) == pytest.approx(0.2) + + +def test_binary_matches_the_historical_softmax_formula(): + """Regression against the original two-token implementation.""" + yes_lp, no_lp = -0.3, -1.7 + highest = max(yes_lp, no_lp) + expected = math.exp(yes_lp - highest) / ( + math.exp(yes_lp - highest) + math.exp(no_lp - highest) + ) + r = _reranker("binary", [0, 1]) + assert r._yes_probability(_output({0: no_lp, 1: yes_lp})) == pytest.approx(expected) + + +# --------------------------------------------------------------- graded contract +def test_graded_spans_the_unit_interval(): + r = _reranker("graded", [10, 11, 12]) # Not, Somewhat, Highly + big, small = math.log(0.999), math.log(0.0005) + assert r._yes_probability(_output({10: big, 11: small, 12: small})) == pytest.approx(0.0, abs=1e-3) + assert r._yes_probability(_output({10: small, 11: big, 12: small})) == pytest.approx(0.5, abs=1e-3) + assert r._yes_probability(_output({10: small, 11: small, 12: big})) == pytest.approx(1.0, abs=1e-3) + + +def test_graded_produces_intermediate_values_where_binary_saturates(): + """The point of the change. Binary pushes mass to the extremes; the graded expectation + lands in the partial-relevance band that aggregation and the 0.5 cut depend on.""" + r = _reranker("graded", [10, 11, 12]) + score = r._yes_probability( + _output({10: math.log(0.25), 11: math.log(0.5), 12: math.log(0.25)}) + ) + assert score == pytest.approx(0.5, abs=1e-6) + assert 0.0 < score < 1.0 + + skewed = r._yes_probability( + _output({10: math.log(0.1), 11: math.log(0.3), 12: math.log(0.6)}) + ) + assert skewed == pytest.approx((0 * 0.1 + 1 * 0.3 + 2 * 0.6) / 2) + assert 0.5 < skewed < 1.0 + + +def test_graded_is_monotonic_in_the_top_label(): + r = _reranker("graded", [10, 11, 12]) + scores = [ + r._yes_probability(_output({10: math.log(1 - p - 0.05), 11: math.log(0.05), 12: math.log(p)})) + for p in (0.1, 0.3, 0.5, 0.7) + ] + assert scores == sorted(scores) + + +@pytest.mark.parametrize("mode,ids", [("binary", [0, 1]), ("graded", [10, 11, 12])]) +def test_scores_stay_in_unit_interval(mode, ids): + r = _reranker(mode, ids) + for lp in (-0.01, -1.0, -12.0): + score = r._yes_probability(_output({t: lp * (i + 1) for i, t in enumerate(ids)})) + assert 0.0 <= score <= 1.0 + + +# --------------------------------------------------------------- robustness +@pytest.mark.parametrize("mode,ids", [("binary", [0, 1]), ("graded", [10, 11, 12])]) +def test_missing_label_tokens_do_not_raise(mode, ids): + """vLLM returns the top-k logprobs; a label may be absent. Renormalizing over what IS + present must still yield a usable score rather than an exception.""" + r = _reranker(mode, ids) + score = r._yes_probability(_output({ids[-1]: math.log(0.9)})) + assert 0.0 <= score <= 1.0 + + +@pytest.mark.parametrize("mode,ids", [("binary", [0, 1]), ("graded", [10, 11, 12])]) +def test_malformed_output_scores_zero(mode, ids): + r = _reranker(mode, ids) + assert r._yes_probability(types.SimpleNamespace(outputs=[])) == 0.0 + assert r._yes_probability(_output({999: math.log(0.5)})) == 0.0 + + +# --------------------------------------------------------------- prompts +def test_binary_prompt_is_unchanged(): + """The LoRA adapter was tuned against this exact wording; changing it silently would + invalidate the adapter.""" + messages = LLMReranker.create_messages("patient", "criterion", scoring="binary") + assert "sufficient information" in messages[0]["content"] + assert messages[0]["content"] == LLMReranker.BINARY_SYSTEM_PROMPT + + +def test_graded_prompt_asks_for_relevance_not_answerability(): + """The substantive fix: the binary prompt asks whether the criterion CAN BE EVALUATED, + which is answerability, not relevance.""" + messages = LLMReranker.create_messages("patient", "criterion", scoring="graded") + prompt = messages[0]["content"] + assert "relevant" in prompt.lower() + assert "sufficient information" not in prompt + for label in GRADED_LABELS: + assert label in prompt + + +def test_both_prompts_carry_the_same_statement_payload(): + for mode in ("binary", "graded"): + messages = LLMReranker.create_messages("PT", "CR", scoring=mode) + assert messages[-1]["content"] == "Statement A: PT\nStatement B: CR\n\n" From f24b90ef0673f93a6fa5c1e21ccf317d482d56bc Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Thu, 13 Aug 2026 12:15:55 +0200 Subject: [PATCH 10/13] feat(search): make second-level width configurable, and size the reasoning budget per trial Two hardcoded constants were the pipeline's binding bottleneck, and neither was reachable from config. SecondStageRetriever.size defaulted to 250 criteria PER QUERY and main.py never passed it. The TREC 2023 candidate pool is ~1,860 trials x ~17 criteria ~= 31,000 criteria, so the second level examined about a tenth of its own pool. Measured on 37 patients, retrieval only, no GPU: size unique trials % of pool recall of judged-relevant 250 803 43.1% 0.5849 <- old default 500 1,041 55.8% 0.7135 1000 1,200 64.3% 0.7870 2000 1,345 72.1% 0.8352 4000 1,477 79.2% 0.8507 That recall is a CEILING: the reranker and the 0.5 aggregation cut can only drop trials from what retrieval surfaced, never add them. Every knob tuned so far -- shortlist divisor, max_trials_rag, the depth policy -- sat behind it. Raising size to 1000 lifts the ceiling 0.585 -> 0.787, at the cost of retrieval plus 2B-reranker scoring; it does not touch the 35B eligibility model. aggregate_to_trials' 0.5 threshold was the second such constant: a trial whose every criterion scores below it never reaches the shortlist at all. Adds search.second_level.{per_query_size, aggregation_threshold, aggregation_method}. Also adds rag.adaptive_token_budget, sizing the eligibility output budget to each trial's criterion count rather than giving every trial the full vllm.max_new_tokens. m1 (arXiv:2504.00869) measures an optimal MEDICAL reasoning budget near 4K tokens "beyond which performance may degrade due to overthinking" -- but that is for a single question, while this stage judges every criterion of a trial in one call and trials range from a handful of criteria to 60+. A flat 4K cap would truncate the large ones into invalid JSON, so the budget scales with the work: a ~17-criterion trial lands near 4K, a 60-criterion trial keeps the full ceiling. It is clamped to max_new_tokens, so it can only ever LOWER the budget -- no trial is truncated more than it already would be. All defaults reproduce previous behaviour exactly, verified against both shipped configs, so every change here is an explicit A/B. aggregate_to_trials keeps working on an instance built without __init__ (the pre-existing test_audit_fixes contract), which is the evidence the old path is untouched. --- src/trialmatchai/config/config_1gpu.json | 151 +++++++++++++ src/trialmatchai/config/config_a40.json | 148 +++++++++++++ src/trialmatchai/config/config_h100.json | 148 +++++++++++++ src/trialmatchai/config/config_l40.json | 149 +++++++++++++ .../config/config_medcpt_baichuanm2.json | 160 ++++++++++++++ .../config/config_medcpt_episteme.json | 156 +++++++++++++ .../config/config_medcpt_huatuo.json | 162 ++++++++++++++ .../config/config_medcpt_iimedical.json | 162 ++++++++++++++ .../config/config_medcpt_medgemma.json | 155 +++++++++++++ .../config/config_medcpt_qwen36.json | 162 ++++++++++++++ .../config/config_medcpt_qwen36_l40.json | 162 ++++++++++++++ .../config/config_medcpt_qwen3med.json | 156 +++++++++++++ .../config/config_nemotron_medgemma.json | 158 +++++++++++++ src/trialmatchai/config/settings.py | 39 ++++ src/trialmatchai/main.py | 13 ++ src/trialmatchai/matching/agent.py | 208 ++++++++++++++++++ src/trialmatchai/matching/eligibility_base.py | 31 ++- .../matching/eligibility_reasoning_vllm.py | 66 +++++- .../matching/retrieval/criteria_retrieval.py | 20 +- tests/test_second_level_width.py | 113 ++++++++++ 20 files changed, 2516 insertions(+), 3 deletions(-) create mode 100644 src/trialmatchai/config/config_1gpu.json create mode 100644 src/trialmatchai/config/config_a40.json create mode 100644 src/trialmatchai/config/config_h100.json create mode 100644 src/trialmatchai/config/config_l40.json create mode 100644 src/trialmatchai/config/config_medcpt_baichuanm2.json create mode 100644 src/trialmatchai/config/config_medcpt_episteme.json create mode 100644 src/trialmatchai/config/config_medcpt_huatuo.json create mode 100644 src/trialmatchai/config/config_medcpt_iimedical.json create mode 100644 src/trialmatchai/config/config_medcpt_medgemma.json create mode 100644 src/trialmatchai/config/config_medcpt_qwen36.json create mode 100644 src/trialmatchai/config/config_medcpt_qwen36_l40.json create mode 100644 src/trialmatchai/config/config_medcpt_qwen3med.json create mode 100644 src/trialmatchai/config/config_nemotron_medgemma.json create mode 100644 src/trialmatchai/matching/agent.py create mode 100644 tests/test_second_level_width.py diff --git a/src/trialmatchai/config/config_1gpu.json b/src/trialmatchai/config/config_1gpu.json new file mode 100644 index 00000000..630a64b5 --- /dev/null +++ b/src/trialmatchai/config/config_1gpu.json @@ -0,0 +1,151 @@ +{ + "entity_extraction": { + "backend": "gliner2", + "model_name": "fastino/gliner2-base-v1", + "model_revision": null, + "schema_path": "entity_schemas/trialmatchai.yaml", + "threshold": 0.8, + "batch_size": 8, + "device": "auto", + "trust_remote_code": false + }, + "concept_linker": { + "enabled": true, + "db_path": "data/concepts", + "table": "concepts", + "accept_threshold": 0.7, + "reject_threshold": 0.5, + "margin": 0.05, + "rerank": "lexical", + "search_limit": 10 + }, + "paths": { + "output_dir": "results", + "trials_json_folder": "data/trials_jsons" + }, + "patient_inputs": { + "raw_dir": "data/patients/raw", + "profile_dir": "data/patients/profiles", + "summary_dir": "data/patients/summaries", + "default_format": "auto", + "strict_validation": false, + "copy_raw": true + }, + "model": { + "base_model": "microsoft/phi-4", + "base_model_revision": null, + "trust_remote_code": false, + "quantization": { + "load_in_4bit": true, + "bnb_4bit_use_double_quant": true, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": "float16" + }, + "cot_adapter_path": "models/finetuned_phi_reasoning", + "reranker_model_path": "google/gemma-2-2b-it", + "reranker_model_revision": null, + "reranker_adapter_path": "models/finetuned_gemma2" + }, + "tokenizer": { + "use_fast": true, + "padding_side": "left" + }, + "global": { + "device": 0 + }, + "search_backend": { + "backend": "lancedb", + "db_path": "data/search_a40", + "trials_table": "trials", + "criteria_table": "criteria", + "candidate_limit": 1000 + }, + "registry": { + "source": "clinicaltrials.gov", + "api_base_url": "https://clinicaltrials.gov/api/v2/studies", + "keywords_file": null, + "since_days": 7, + "max_studies": null, + "request_timeout": 30, + "rate_limit_per_second": 2, + "raw_dir": "data/registry/raw", + "manifest_path": "data/registry/manifest.jsonl", + "reports_dir": "data/registry/runs", + "failure_threshold": 0.25 + }, + "embedder": { + "model_name": "BAAI/bge-m3", + "revision": null, + "trust_remote_code": false, + "pooling": "mean", + "max_length": 512, + "batch_size": 32, + "use_gpu": true, + "use_fp16": false, + "normalize": true + }, + "cot": { + "batch_size": 10 + }, + "LLM_reranker": { + "batch_size": 20, + "gpu_memory_utilization": 0.16, + "tensor_parallel_size": 1 + }, + "search": { + "mode": "hybrid", + "vector_score_threshold": 0.5, + "max_trials_first_level": 1000, + "max_trials_second_level": 100, + "first_level": { + "enabled": true, + "max_trials": 1000, + "per_channel_size": 300, + "fusion": "rrf", + "rrf_k": 60, + "vector_score_threshold": 0.0, + "llm_expansion_enabled": false, + "llm_max_terms": 12, + "write_reports": true, + "hard_filters": [ + "age", + "sex", + "overall_status" + ] + } + }, + "constraints": { + "enabled": true, + "score_weight": 0.25, + "llm_extraction_enabled": false, + "unknown_is_neutral": true, + "write_reports": true + }, + "query_expansion": { + "enabled": false, + "backend": null, + "model": null, + "adapter": null, + "max_new_tokens": 2048, + "max_main_conditions": 11, + "max_other_conditions": 50 + }, + "use_cot_reasoning": true, + "rag": { + "batch_size": 4, + "max_trials_rag": 20 + }, + "vllm": { + "batch_size": 100, + "max_new_tokens": 5000, + "temperature": 0.0, + "top_p": 1.0, + "seed": 1234, + "length_bucket": true, + "gpu_memory_utilization": 0.72, + "max_model_len": 8192, + "tensor_parallel_size": 1, + "kv_cache_dtype": "fp8", + "max_num_seqs": 16 + } +} \ No newline at end of file diff --git a/src/trialmatchai/config/config_a40.json b/src/trialmatchai/config/config_a40.json new file mode 100644 index 00000000..e55f3bac --- /dev/null +++ b/src/trialmatchai/config/config_a40.json @@ -0,0 +1,148 @@ +{ + "entity_extraction": { + "backend": "gliner2", + "model_name": "fastino/gliner2-base-v1", + "model_revision": null, + "schema_path": "entity_schemas/trialmatchai.yaml", + "threshold": 0.8, + "batch_size": 8, + "device": "auto", + "trust_remote_code": false + }, + "concept_linker": { + "enabled": true, + "db_path": "data/concepts", + "table": "concepts", + "accept_threshold": 0.7, + "reject_threshold": 0.5, + "margin": 0.05, + "rerank": "lexical", + "search_limit": 10 + }, + "paths": { + "output_dir": "results", + "trials_json_folder": "data/trials_jsons" + }, + "patient_inputs": { + "raw_dir": "data/patients/raw", + "profile_dir": "data/patients/profiles", + "summary_dir": "data/patients/summaries", + "default_format": "auto", + "strict_validation": false, + "copy_raw": true + }, + "model": { + "base_model": "microsoft/phi-4", + "base_model_revision": null, + "trust_remote_code": false, + "quantization": { + "load_in_4bit": true, + "bnb_4bit_use_double_quant": true, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": "float16" + }, + "cot_adapter_path": "models/finetuned_phi_reasoning", + "reranker_model_path": "google/gemma-2-2b-it", + "reranker_model_revision": null, + "reranker_adapter_path": "models/finetuned_gemma2" + }, + "tokenizer": { + "use_fast": true, + "padding_side": "left" + }, + "global": { + "device": 0 + }, + "search_backend": { + "backend": "lancedb", + "db_path": "data/search_a40", + "trials_table": "trials", + "criteria_table": "criteria", + "candidate_limit": 1000 + }, + "registry": { + "source": "clinicaltrials.gov", + "api_base_url": "https://clinicaltrials.gov/api/v2/studies", + "keywords_file": null, + "since_days": 7, + "max_studies": null, + "request_timeout": 30, + "rate_limit_per_second": 2, + "raw_dir": "data/registry/raw", + "manifest_path": "data/registry/manifest.jsonl", + "reports_dir": "data/registry/runs", + "failure_threshold": 0.25 + }, + "embedder": { + "model_name": "BAAI/bge-m3", + "revision": null, + "trust_remote_code": false, + "pooling": "mean", + "max_length": 512, + "batch_size": 32, + "use_gpu": true, + "use_fp16": false, + "normalize": true + }, + "cot": { + "batch_size": 10 + }, + "LLM_reranker": { + "batch_size": 20, + "gpu_memory_utilization": 0.14 + }, + "search": { + "mode": "hybrid", + "vector_score_threshold": 0.5, + "max_trials_first_level": 1000, + "max_trials_second_level": 100, + "first_level": { + "enabled": true, + "max_trials": 1000, + "per_channel_size": 300, + "fusion": "rrf", + "rrf_k": 60, + "vector_score_threshold": 0.0, + "llm_expansion_enabled": false, + "llm_max_terms": 12, + "write_reports": true, + "hard_filters": [ + "age", + "sex", + "overall_status" + ] + } + }, + "constraints": { + "enabled": true, + "score_weight": 0.25, + "llm_extraction_enabled": false, + "unknown_is_neutral": true, + "write_reports": true + }, + "query_expansion": { + "enabled": false, + "backend": null, + "model": null, + "adapter": null, + "max_new_tokens": 2048, + "max_main_conditions": 11, + "max_other_conditions": 50 + }, + "use_cot_reasoning": true, + "rag": { + "batch_size": 4, + "max_trials_rag": 20 + }, + "vllm": { + "batch_size": 100, + "max_new_tokens": 5000, + "temperature": 0.0, + "top_p": 1.0, + "seed": 1234, + "length_bucket": true, + "gpu_memory_utilization": 0.72, + "max_model_len": 6000, + "tensor_parallel_size": 1 + } +} \ No newline at end of file diff --git a/src/trialmatchai/config/config_h100.json b/src/trialmatchai/config/config_h100.json new file mode 100644 index 00000000..5b3e921e --- /dev/null +++ b/src/trialmatchai/config/config_h100.json @@ -0,0 +1,148 @@ +{ + "entity_extraction": { + "backend": "gliner2", + "model_name": "fastino/gliner2-base-v1", + "model_revision": null, + "schema_path": "entity_schemas/trialmatchai.yaml", + "threshold": 0.8, + "batch_size": 8, + "device": "auto", + "trust_remote_code": false + }, + "concept_linker": { + "enabled": true, + "db_path": "data/concepts", + "table": "concepts", + "accept_threshold": 0.7, + "reject_threshold": 0.5, + "margin": 0.05, + "rerank": "lexical", + "search_limit": 10 + }, + "paths": { + "output_dir": "results", + "trials_json_folder": "data/trials_jsons" + }, + "patient_inputs": { + "raw_dir": "data/patients/raw", + "profile_dir": "data/patients/profiles", + "summary_dir": "data/patients/summaries", + "default_format": "auto", + "strict_validation": false, + "copy_raw": true + }, + "model": { + "base_model": "microsoft/phi-4", + "base_model_revision": null, + "trust_remote_code": false, + "quantization": { + "load_in_4bit": true, + "bnb_4bit_use_double_quant": true, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": "float16" + }, + "cot_adapter_path": "models/finetuned_phi_reasoning", + "reranker_model_path": "google/gemma-2-2b-it", + "reranker_model_revision": null, + "reranker_adapter_path": "models/finetuned_gemma2" + }, + "tokenizer": { + "use_fast": true, + "padding_side": "left" + }, + "global": { + "device": 0 + }, + "search_backend": { + "backend": "lancedb", + "db_path": "data/search_a40", + "trials_table": "trials", + "criteria_table": "criteria", + "candidate_limit": 1000 + }, + "registry": { + "source": "clinicaltrials.gov", + "api_base_url": "https://clinicaltrials.gov/api/v2/studies", + "keywords_file": null, + "since_days": 7, + "max_studies": null, + "request_timeout": 30, + "rate_limit_per_second": 2, + "raw_dir": "data/registry/raw", + "manifest_path": "data/registry/manifest.jsonl", + "reports_dir": "data/registry/runs", + "failure_threshold": 0.25 + }, + "embedder": { + "model_name": "BAAI/bge-m3", + "revision": null, + "trust_remote_code": false, + "pooling": "mean", + "max_length": 512, + "batch_size": 32, + "use_gpu": true, + "use_fp16": false, + "normalize": true + }, + "cot": { + "batch_size": 10 + }, + "LLM_reranker": { + "batch_size": 20, + "gpu_memory_utilization": 0.15 + }, + "search": { + "mode": "hybrid", + "vector_score_threshold": 0.5, + "max_trials_first_level": 1000, + "max_trials_second_level": 100, + "first_level": { + "enabled": true, + "max_trials": 1000, + "per_channel_size": 300, + "fusion": "rrf", + "rrf_k": 60, + "vector_score_threshold": 0.0, + "llm_expansion_enabled": false, + "llm_max_terms": 12, + "write_reports": true, + "hard_filters": [ + "age", + "sex", + "overall_status" + ] + } + }, + "constraints": { + "enabled": true, + "score_weight": 0.25, + "llm_extraction_enabled": false, + "unknown_is_neutral": true, + "write_reports": true + }, + "query_expansion": { + "enabled": false, + "backend": null, + "model": null, + "adapter": null, + "max_new_tokens": 2048, + "max_main_conditions": 11, + "max_other_conditions": 50 + }, + "use_cot_reasoning": true, + "rag": { + "batch_size": 4, + "max_trials_rag": 20 + }, + "vllm": { + "batch_size": 100, + "max_new_tokens": 5000, + "temperature": 0.0, + "top_p": 1.0, + "seed": 1234, + "length_bucket": true, + "gpu_memory_utilization": 0.6, + "max_model_len": 8192, + "tensor_parallel_size": 1 + } +} \ No newline at end of file diff --git a/src/trialmatchai/config/config_l40.json b/src/trialmatchai/config/config_l40.json new file mode 100644 index 00000000..cec79f7a --- /dev/null +++ b/src/trialmatchai/config/config_l40.json @@ -0,0 +1,149 @@ +{ + "entity_extraction": { + "backend": "gliner2", + "model_name": "fastino/gliner2-base-v1", + "model_revision": null, + "schema_path": "entity_schemas/trialmatchai.yaml", + "threshold": 0.8, + "batch_size": 8, + "device": "auto", + "trust_remote_code": false + }, + "concept_linker": { + "enabled": true, + "db_path": "data/concepts", + "table": "concepts", + "accept_threshold": 0.7, + "reject_threshold": 0.5, + "margin": 0.05, + "rerank": "lexical", + "search_limit": 10 + }, + "paths": { + "output_dir": "results", + "trials_json_folder": "data/trials_jsons" + }, + "patient_inputs": { + "raw_dir": "data/patients/raw", + "profile_dir": "data/patients/profiles", + "summary_dir": "data/patients/summaries", + "default_format": "auto", + "strict_validation": false, + "copy_raw": true + }, + "model": { + "base_model": "microsoft/phi-4", + "base_model_revision": null, + "trust_remote_code": false, + "quantization": { + "load_in_4bit": true, + "bnb_4bit_use_double_quant": true, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": "float16" + }, + "cot_adapter_path": "models/finetuned_phi_reasoning", + "reranker_model_path": "google/gemma-2-2b-it", + "reranker_model_revision": null, + "reranker_adapter_path": "models/finetuned_gemma2" + }, + "tokenizer": { + "use_fast": true, + "padding_side": "left" + }, + "global": { + "device": 0 + }, + "search_backend": { + "backend": "lancedb", + "db_path": "data/search_a40", + "trials_table": "trials", + "criteria_table": "criteria", + "candidate_limit": 1000 + }, + "registry": { + "source": "clinicaltrials.gov", + "api_base_url": "https://clinicaltrials.gov/api/v2/studies", + "keywords_file": null, + "since_days": 7, + "max_studies": null, + "request_timeout": 30, + "rate_limit_per_second": 2, + "raw_dir": "data/registry/raw", + "manifest_path": "data/registry/manifest.jsonl", + "reports_dir": "data/registry/runs", + "failure_threshold": 0.25 + }, + "embedder": { + "model_name": "BAAI/bge-m3", + "revision": null, + "trust_remote_code": false, + "pooling": "mean", + "max_length": 512, + "batch_size": 32, + "use_gpu": true, + "use_fp16": false, + "normalize": true + }, + "cot": { + "batch_size": 10 + }, + "LLM_reranker": { + "batch_size": 20, + "gpu_memory_utilization": 0.15, + "tensor_parallel_size": 2 + }, + "search": { + "mode": "hybrid", + "vector_score_threshold": 0.5, + "max_trials_first_level": 1000, + "max_trials_second_level": 100, + "first_level": { + "enabled": true, + "max_trials": 1000, + "per_channel_size": 300, + "fusion": "rrf", + "rrf_k": 60, + "vector_score_threshold": 0.0, + "llm_expansion_enabled": false, + "llm_max_terms": 12, + "write_reports": true, + "hard_filters": [ + "age", + "sex", + "overall_status" + ] + } + }, + "constraints": { + "enabled": true, + "score_weight": 0.25, + "llm_extraction_enabled": false, + "unknown_is_neutral": true, + "write_reports": true + }, + "query_expansion": { + "enabled": false, + "backend": null, + "model": null, + "adapter": null, + "max_new_tokens": 2048, + "max_main_conditions": 11, + "max_other_conditions": 50 + }, + "use_cot_reasoning": true, + "rag": { + "batch_size": 4, + "max_trials_rag": 20 + }, + "vllm": { + "batch_size": 100, + "max_new_tokens": 5000, + "temperature": 0.0, + "top_p": 1.0, + "seed": 1234, + "length_bucket": true, + "gpu_memory_utilization": 0.6, + "max_model_len": 8192, + "tensor_parallel_size": 2 + } +} \ No newline at end of file diff --git a/src/trialmatchai/config/config_medcpt_baichuanm2.json b/src/trialmatchai/config/config_medcpt_baichuanm2.json new file mode 100644 index 00000000..dd76567b --- /dev/null +++ b/src/trialmatchai/config/config_medcpt_baichuanm2.json @@ -0,0 +1,160 @@ +{ + "entity_extraction": { + "backend": "gliner2", + "model_name": "fastino/gliner2-base-v1", + "model_revision": null, + "schema_path": "entity_schemas/trialmatchai.yaml", + "threshold": 0.8, + "batch_size": 8, + "device": "auto", + "trust_remote_code": false + }, + "concept_linker": { + "enabled": true, + "db_path": "data/concepts_medcpt", + "table": "concepts", + "accept_threshold": 0.7, + "reject_threshold": 0.5, + "margin": 0.05, + "rerank": "lexical", + "search_limit": 10 + }, + "paths": { + "output_dir": "results", + "trials_json_folder": "data/trials_jsons" + }, + "patient_inputs": { + "raw_dir": "data/patients/raw", + "profile_dir": "data/patients/profiles", + "summary_dir": "data/patients/summaries", + "default_format": "auto", + "strict_validation": false, + "copy_raw": true + }, + "model": { + "base_model": "baichuan-inc/Baichuan-M2-32B", + "base_model_revision": null, + "trust_remote_code": false, + "quantization": { + "load_in_4bit": true, + "bnb_4bit_use_double_quant": true, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": "float16" + }, + "cot_adapter_path": null, + "reranker_model_path": "google/gemma-2-2b-it", + "reranker_model_revision": null, + "reranker_adapter_path": "models/finetuned_gemma2" + }, + "tokenizer": { + "use_fast": true, + "padding_side": "left" + }, + "global": { + "device": 0 + }, + "search_backend": { + "backend": "lancedb", + "db_path": "data/search_medcpt", + "trials_table": "trials", + "criteria_table": "criteria", + "candidate_limit": 1000, + "vector_metric": "dot", + "vector_weight": 0.6, + "reembed_index": true + }, + "registry": { + "source": "clinicaltrials.gov", + "api_base_url": "https://clinicaltrials.gov/api/v2/studies", + "keywords_file": null, + "since_days": 7, + "max_studies": null, + "request_timeout": 30, + "rate_limit_per_second": 2, + "raw_dir": "data/registry/raw", + "manifest_path": "data/registry/manifest.jsonl", + "reports_dir": "data/registry/runs", + "failure_threshold": 0.25 + }, + "embedder": { + "backend": "hf", + "model_name": "ncbi/MedCPT-Article-Encoder", + "query_model_name": "ncbi/MedCPT-Query-Encoder", + "pooling": "cls", + "max_length": 512, + "query_max_length": 64, + "normalize": false, + "use_gpu": true, + "batch_size": 32, + "use_fp16": false + }, + "cot": { + "batch_size": 10 + }, + "LLM_reranker": { + "batch_size": 20, + "gpu_memory_utilization": 0.15, + "tensor_parallel_size": 1 + }, + "search": { + "mode": "hybrid", + "vector_score_threshold": 0.5, + "max_trials_first_level": 2000, + "max_trials_second_level": 1000, + "first_level": { + "enabled": true, + "max_trials": 2000, + "per_channel_size": 600, + "fusion": "rrf", + "rrf_k": 60, + "vector_score_threshold": 0.0, + "llm_expansion_enabled": false, + "llm_max_terms": 12, + "write_reports": true, + "hard_filters": [ + "age", + "sex", + "overall_status" + ] + } + }, + "constraints": { + "enabled": true, + "score_weight": 0.25, + "llm_extraction_enabled": false, + "unknown_is_neutral": true, + "write_reports": true + }, + "query_expansion": { + "enabled": true, + "backend": null, + "model": null, + "adapter": null, + "max_new_tokens": 2048, + "max_main_conditions": 11, + "max_other_conditions": 50, + "guided_json": true + }, + "use_cot_reasoning": true, + "rag": { + "batch_size": 4, + "max_trials_rag": 300, + "no_think": true, + "guided_json": true + }, + "vllm": { + "batch_size": 512, + "max_new_tokens": 6000, + "temperature": 0.0, + "top_p": 1.0, + "seed": 1234, + "length_bucket": true, + "gpu_memory_utilization": 0.55, + "max_model_len": 24576, + "tensor_parallel_size": 1, + "quantization": "bitsandbytes", + "disable_custom_all_reduce": true, + "enforce_eager": true, + "max_num_seqs": 256 + } +} \ No newline at end of file diff --git a/src/trialmatchai/config/config_medcpt_episteme.json b/src/trialmatchai/config/config_medcpt_episteme.json new file mode 100644 index 00000000..97eaa898 --- /dev/null +++ b/src/trialmatchai/config/config_medcpt_episteme.json @@ -0,0 +1,156 @@ +{ + "entity_extraction": { + "backend": "gliner2", + "model_name": "fastino/gliner2-base-v1", + "model_revision": null, + "schema_path": "entity_schemas/trialmatchai.yaml", + "threshold": 0.8, + "batch_size": 8, + "device": "auto", + "trust_remote_code": false + }, + "concept_linker": { + "enabled": true, + "db_path": "data/concepts_medcpt", + "table": "concepts", + "accept_threshold": 0.7, + "reject_threshold": 0.5, + "margin": 0.05, + "rerank": "lexical", + "search_limit": 10 + }, + "paths": { + "output_dir": "results", + "trials_json_folder": "data/trials_jsons" + }, + "patient_inputs": { + "raw_dir": "data/patients/raw", + "profile_dir": "data/patients/profiles", + "summary_dir": "data/patients/summaries", + "default_format": "auto", + "strict_validation": false, + "copy_raw": true + }, + "model": { + "base_model": "EpistemeAI/Reasoning-Medical0.1-27B", + "base_model_revision": null, + "trust_remote_code": true, + "quantization": { + "load_in_4bit": true, + "bnb_4bit_use_double_quant": true, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": "float16" + }, + "cot_adapter_path": null, + "reranker_model_path": "google/gemma-2-2b-it", + "reranker_model_revision": null, + "reranker_adapter_path": "models/finetuned_gemma2" + }, + "tokenizer": { + "use_fast": true, + "padding_side": "left" + }, + "global": { + "device": 0 + }, + "search_backend": { + "backend": "lancedb", + "db_path": "data/search_medcpt", + "trials_table": "trials", + "criteria_table": "criteria", + "candidate_limit": 1000, + "vector_metric": "dot", + "vector_weight": 0.6, + "reembed_index": true + }, + "registry": { + "source": "clinicaltrials.gov", + "api_base_url": "https://clinicaltrials.gov/api/v2/studies", + "keywords_file": null, + "since_days": 7, + "max_studies": null, + "request_timeout": 30, + "rate_limit_per_second": 2, + "raw_dir": "data/registry/raw", + "manifest_path": "data/registry/manifest.jsonl", + "reports_dir": "data/registry/runs", + "failure_threshold": 0.25 + }, + "embedder": { + "backend": "hf", + "model_name": "ncbi/MedCPT-Article-Encoder", + "query_model_name": "ncbi/MedCPT-Query-Encoder", + "pooling": "cls", + "max_length": 512, + "query_max_length": 64, + "normalize": false, + "use_gpu": true, + "batch_size": 32, + "use_fp16": false + }, + "cot": { + "batch_size": 10 + }, + "LLM_reranker": { + "batch_size": 20, + "gpu_memory_utilization": 0.15 + }, + "search": { + "mode": "hybrid", + "vector_score_threshold": 0.5, + "max_trials_first_level": 2000, + "max_trials_second_level": 1000, + "first_level": { + "enabled": true, + "max_trials": 2000, + "per_channel_size": 600, + "fusion": "rrf", + "rrf_k": 60, + "vector_score_threshold": 0.0, + "llm_expansion_enabled": false, + "llm_max_terms": 12, + "write_reports": true, + "hard_filters": [ + "age", + "sex", + "overall_status" + ] + } + }, + "constraints": { + "enabled": true, + "score_weight": 0.25, + "llm_extraction_enabled": false, + "unknown_is_neutral": true, + "write_reports": true + }, + "query_expansion": { + "enabled": false, + "backend": null, + "model": null, + "adapter": null, + "max_new_tokens": 6000, + "max_main_conditions": 11, + "max_other_conditions": 50, + "no_think": true + }, + "use_cot_reasoning": true, + "rag": { + "batch_size": 4, + "max_trials_rag": 300 + }, + "vllm": { + "batch_size": 100, + "max_new_tokens": 5000, + "temperature": 0.0, + "top_p": 1.0, + "seed": 1234, + "length_bucket": true, + "gpu_memory_utilization": 0.78, + "max_model_len": 8192, + "tensor_parallel_size": 1, + "quantization": "", + "disable_custom_all_reduce": true, + "enforce_eager": true + } +} \ No newline at end of file diff --git a/src/trialmatchai/config/config_medcpt_huatuo.json b/src/trialmatchai/config/config_medcpt_huatuo.json new file mode 100644 index 00000000..e7b2018d --- /dev/null +++ b/src/trialmatchai/config/config_medcpt_huatuo.json @@ -0,0 +1,162 @@ +{ + "entity_extraction": { + "backend": "gliner2", + "model_name": "fastino/gliner2-base-v1", + "model_revision": null, + "schema_path": "entity_schemas/trialmatchai.yaml", + "threshold": 0.8, + "batch_size": 8, + "device": "auto", + "trust_remote_code": false + }, + "concept_linker": { + "enabled": true, + "db_path": "data/concepts_medcpt", + "table": "concepts", + "accept_threshold": 0.7, + "reject_threshold": 0.5, + "margin": 0.05, + "rerank": "lexical", + "search_limit": 10 + }, + "paths": { + "output_dir": "results", + "trials_json_folder": "data/trials_jsons" + }, + "patient_inputs": { + "raw_dir": "data/patients/raw", + "profile_dir": "data/patients/profiles", + "summary_dir": "data/patients/summaries", + "default_format": "auto", + "strict_validation": false, + "copy_raw": true + }, + "model": { + "base_model": "FreedomIntelligence/HuatuoGPT-o1-8B", + "base_model_revision": null, + "trust_remote_code": false, + "quantization": { + "load_in_4bit": true, + "bnb_4bit_use_double_quant": true, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": "float16" + }, + "cot_adapter_path": null, + "reranker_model_path": "google/gemma-2-2b-it", + "reranker_model_revision": null, + "reranker_adapter_path": "models/finetuned_gemma2" + }, + "tokenizer": { + "use_fast": true, + "padding_side": "left" + }, + "global": { + "device": 0 + }, + "search_backend": { + "backend": "lancedb", + "db_path": "data/search_medcpt", + "trials_table": "trials", + "criteria_table": "criteria", + "candidate_limit": 1000, + "vector_metric": "dot", + "vector_weight": 0.6, + "reembed_index": true + }, + "registry": { + "source": "clinicaltrials.gov", + "api_base_url": "https://clinicaltrials.gov/api/v2/studies", + "keywords_file": null, + "since_days": 7, + "max_studies": null, + "request_timeout": 30, + "rate_limit_per_second": 2, + "raw_dir": "data/registry/raw", + "manifest_path": "data/registry/manifest.jsonl", + "reports_dir": "data/registry/runs", + "failure_threshold": 0.25 + }, + "embedder": { + "backend": "hf", + "model_name": "ncbi/MedCPT-Article-Encoder", + "query_model_name": "ncbi/MedCPT-Query-Encoder", + "pooling": "cls", + "max_length": 512, + "query_max_length": 64, + "normalize": false, + "use_gpu": true, + "batch_size": 32, + "use_fp16": false + }, + "cot": { + "batch_size": 10 + }, + "LLM_reranker": { + "batch_size": 20, + "gpu_memory_utilization": 0.15, + "tensor_parallel_size": 1 + }, + "search": { + "mode": "hybrid", + "vector_score_threshold": 0.5, + "max_trials_first_level": 2000, + "max_trials_second_level": 1000, + "first_level": { + "enabled": true, + "max_trials": 2000, + "per_channel_size": 600, + "fusion": "rrf", + "rrf_k": 60, + "vector_score_threshold": 0.0, + "llm_expansion_enabled": false, + "llm_max_terms": 12, + "write_reports": true, + "hard_filters": [ + "age", + "sex", + "overall_status" + ] + } + }, + "constraints": { + "enabled": true, + "score_weight": 0.25, + "llm_extraction_enabled": false, + "unknown_is_neutral": true, + "write_reports": true + }, + "query_expansion": { + "enabled": true, + "backend": null, + "model": null, + "adapter": null, + "max_new_tokens": 4096, + "max_main_conditions": 11, + "max_other_conditions": 50, + "guided_json": true, + "no_think": false + }, + "use_cot_reasoning": true, + "rag": { + "batch_size": 4, + "max_trials_rag": 300, + "guided_json": true, + "no_think": false + }, + "vllm": { + "batch_size": 256, + "max_new_tokens": 4096, + "temperature": 0.0, + "top_p": 1.0, + "seed": 1234, + "length_bucket": true, + "gpu_memory_utilization": 0.7, + "max_model_len": 16384, + "tensor_parallel_size": 1, + "quantization": "", + "disable_custom_all_reduce": true, + "enforce_eager": true, + "max_num_seqs": 128 + }, + "reembed_index": false +} \ No newline at end of file diff --git a/src/trialmatchai/config/config_medcpt_iimedical.json b/src/trialmatchai/config/config_medcpt_iimedical.json new file mode 100644 index 00000000..69aed1cf --- /dev/null +++ b/src/trialmatchai/config/config_medcpt_iimedical.json @@ -0,0 +1,162 @@ +{ + "entity_extraction": { + "backend": "gliner2", + "model_name": "fastino/gliner2-base-v1", + "model_revision": null, + "schema_path": "entity_schemas/trialmatchai.yaml", + "threshold": 0.8, + "batch_size": 8, + "device": "auto", + "trust_remote_code": false + }, + "concept_linker": { + "enabled": true, + "db_path": "data/concepts_medcpt", + "table": "concepts", + "accept_threshold": 0.7, + "reject_threshold": 0.5, + "margin": 0.05, + "rerank": "lexical", + "search_limit": 10 + }, + "paths": { + "output_dir": "results", + "trials_json_folder": "data/trials_jsons" + }, + "patient_inputs": { + "raw_dir": "data/patients/raw", + "profile_dir": "data/patients/profiles", + "summary_dir": "data/patients/summaries", + "default_format": "auto", + "strict_validation": false, + "copy_raw": true + }, + "model": { + "base_model": "Intelligent-Internet/II-Medical-8B", + "base_model_revision": null, + "trust_remote_code": false, + "quantization": { + "load_in_4bit": true, + "bnb_4bit_use_double_quant": true, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": "float16" + }, + "cot_adapter_path": null, + "reranker_model_path": "google/gemma-2-2b-it", + "reranker_model_revision": null, + "reranker_adapter_path": "models/finetuned_gemma2" + }, + "tokenizer": { + "use_fast": true, + "padding_side": "left" + }, + "global": { + "device": 0 + }, + "search_backend": { + "backend": "lancedb", + "db_path": "data/search_medcpt", + "trials_table": "trials", + "criteria_table": "criteria", + "candidate_limit": 1000, + "vector_metric": "dot", + "vector_weight": 0.6, + "reembed_index": true + }, + "registry": { + "source": "clinicaltrials.gov", + "api_base_url": "https://clinicaltrials.gov/api/v2/studies", + "keywords_file": null, + "since_days": 7, + "max_studies": null, + "request_timeout": 30, + "rate_limit_per_second": 2, + "raw_dir": "data/registry/raw", + "manifest_path": "data/registry/manifest.jsonl", + "reports_dir": "data/registry/runs", + "failure_threshold": 0.25 + }, + "embedder": { + "backend": "hf", + "model_name": "ncbi/MedCPT-Article-Encoder", + "query_model_name": "ncbi/MedCPT-Query-Encoder", + "pooling": "cls", + "max_length": 512, + "query_max_length": 64, + "normalize": false, + "use_gpu": true, + "batch_size": 32, + "use_fp16": false + }, + "cot": { + "batch_size": 10 + }, + "LLM_reranker": { + "batch_size": 20, + "gpu_memory_utilization": 0.15, + "tensor_parallel_size": 1 + }, + "search": { + "mode": "hybrid", + "vector_score_threshold": 0.5, + "max_trials_first_level": 2000, + "max_trials_second_level": 1000, + "first_level": { + "enabled": true, + "max_trials": 2000, + "per_channel_size": 600, + "fusion": "rrf", + "rrf_k": 60, + "vector_score_threshold": 0.0, + "llm_expansion_enabled": false, + "llm_max_terms": 12, + "write_reports": true, + "hard_filters": [ + "age", + "sex", + "overall_status" + ] + } + }, + "constraints": { + "enabled": true, + "score_weight": 0.25, + "llm_extraction_enabled": false, + "unknown_is_neutral": true, + "write_reports": true + }, + "query_expansion": { + "enabled": true, + "backend": null, + "model": null, + "adapter": null, + "max_new_tokens": 4096, + "max_main_conditions": 11, + "max_other_conditions": 50, + "guided_json": true, + "no_think": true + }, + "use_cot_reasoning": true, + "rag": { + "batch_size": 4, + "max_trials_rag": 300, + "guided_json": true, + "no_think": true + }, + "vllm": { + "batch_size": 256, + "max_new_tokens": 8192, + "temperature": 0.0, + "top_p": 1.0, + "seed": 1234, + "length_bucket": true, + "gpu_memory_utilization": 0.7, + "max_model_len": 32768, + "tensor_parallel_size": 1, + "quantization": "", + "disable_custom_all_reduce": true, + "enforce_eager": true, + "max_num_seqs": 256 + }, + "reembed_index": false +} \ No newline at end of file diff --git a/src/trialmatchai/config/config_medcpt_medgemma.json b/src/trialmatchai/config/config_medcpt_medgemma.json new file mode 100644 index 00000000..12721001 --- /dev/null +++ b/src/trialmatchai/config/config_medcpt_medgemma.json @@ -0,0 +1,155 @@ +{ + "entity_extraction": { + "backend": "gliner2", + "model_name": "fastino/gliner2-base-v1", + "model_revision": null, + "schema_path": "entity_schemas/trialmatchai.yaml", + "threshold": 0.8, + "batch_size": 8, + "device": "auto", + "trust_remote_code": false + }, + "concept_linker": { + "enabled": true, + "db_path": "data/concepts_medcpt", + "table": "concepts", + "accept_threshold": 0.7, + "reject_threshold": 0.5, + "margin": 0.05, + "rerank": "lexical", + "search_limit": 10 + }, + "paths": { + "output_dir": "results", + "trials_json_folder": "data/trials_jsons" + }, + "patient_inputs": { + "raw_dir": "data/patients/raw", + "profile_dir": "data/patients/profiles", + "summary_dir": "data/patients/summaries", + "default_format": "auto", + "strict_validation": false, + "copy_raw": true + }, + "model": { + "base_model": "google/medgemma-27b-text-it", + "base_model_revision": null, + "trust_remote_code": false, + "quantization": { + "load_in_4bit": true, + "bnb_4bit_use_double_quant": true, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": "float16" + }, + "cot_adapter_path": null, + "reranker_model_path": "google/gemma-2-2b-it", + "reranker_model_revision": null, + "reranker_adapter_path": "models/finetuned_gemma2" + }, + "tokenizer": { + "use_fast": true, + "padding_side": "left" + }, + "global": { + "device": 0 + }, + "search_backend": { + "backend": "lancedb", + "db_path": "data/search_medcpt", + "trials_table": "trials", + "criteria_table": "criteria", + "candidate_limit": 1000, + "vector_metric": "dot", + "vector_weight": 0.6, + "reembed_index": true + }, + "registry": { + "source": "clinicaltrials.gov", + "api_base_url": "https://clinicaltrials.gov/api/v2/studies", + "keywords_file": null, + "since_days": 7, + "max_studies": null, + "request_timeout": 30, + "rate_limit_per_second": 2, + "raw_dir": "data/registry/raw", + "manifest_path": "data/registry/manifest.jsonl", + "reports_dir": "data/registry/runs", + "failure_threshold": 0.25 + }, + "embedder": { + "backend": "hf", + "model_name": "ncbi/MedCPT-Article-Encoder", + "query_model_name": "ncbi/MedCPT-Query-Encoder", + "pooling": "cls", + "max_length": 512, + "query_max_length": 64, + "normalize": false, + "use_gpu": true, + "batch_size": 32, + "use_fp16": false + }, + "cot": { + "batch_size": 10 + }, + "LLM_reranker": { + "batch_size": 20, + "gpu_memory_utilization": 0.15 + }, + "search": { + "mode": "hybrid", + "vector_score_threshold": 0.5, + "max_trials_first_level": 2000, + "max_trials_second_level": 1000, + "first_level": { + "enabled": true, + "max_trials": 2000, + "per_channel_size": 600, + "fusion": "rrf", + "rrf_k": 60, + "vector_score_threshold": 0.0, + "llm_expansion_enabled": false, + "llm_max_terms": 12, + "write_reports": true, + "hard_filters": [ + "age", + "sex", + "overall_status" + ] + } + }, + "constraints": { + "enabled": true, + "score_weight": 0.25, + "llm_extraction_enabled": false, + "unknown_is_neutral": true, + "write_reports": true + }, + "query_expansion": { + "enabled": false, + "backend": null, + "model": null, + "adapter": null, + "max_new_tokens": 2048, + "max_main_conditions": 11, + "max_other_conditions": 50 + }, + "use_cot_reasoning": true, + "rag": { + "batch_size": 4, + "max_trials_rag": 300 + }, + "vllm": { + "batch_size": 100, + "max_new_tokens": 5000, + "temperature": 0.0, + "top_p": 1.0, + "seed": 1234, + "length_bucket": true, + "gpu_memory_utilization": 0.78, + "max_model_len": 8192, + "tensor_parallel_size": 1, + "quantization": "", + "disable_custom_all_reduce": true, + "enforce_eager": true + } +} \ No newline at end of file diff --git a/src/trialmatchai/config/config_medcpt_qwen36.json b/src/trialmatchai/config/config_medcpt_qwen36.json new file mode 100644 index 00000000..20e4825d --- /dev/null +++ b/src/trialmatchai/config/config_medcpt_qwen36.json @@ -0,0 +1,162 @@ +{ + "entity_extraction": { + "backend": "gliner2", + "model_name": "fastino/gliner2-base-v1", + "model_revision": null, + "schema_path": "entity_schemas/trialmatchai.yaml", + "threshold": 0.8, + "batch_size": 8, + "device": "auto", + "trust_remote_code": false + }, + "concept_linker": { + "enabled": true, + "db_path": "data/concepts_medcpt", + "table": "concepts", + "accept_threshold": 0.7, + "reject_threshold": 0.5, + "margin": 0.05, + "rerank": "lexical", + "search_limit": 10 + }, + "paths": { + "output_dir": "results", + "trials_json_folder": "data/trials_jsons" + }, + "patient_inputs": { + "raw_dir": "data/patients/raw", + "profile_dir": "data/patients/profiles", + "summary_dir": "data/patients/summaries", + "default_format": "auto", + "strict_validation": false, + "copy_raw": true + }, + "model": { + "base_model": "QuantTrio/Qwen3.6-35B-A3B-AWQ", + "base_model_revision": null, + "trust_remote_code": false, + "quantization": { + "load_in_4bit": false, + "bnb_4bit_use_double_quant": true, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": "float16" + }, + "cot_adapter_path": null, + "reranker_model_path": "google/gemma-2-2b-it", + "reranker_model_revision": null, + "reranker_adapter_path": "models/finetuned_gemma2" + }, + "tokenizer": { + "use_fast": true, + "padding_side": "left" + }, + "global": { + "device": 0 + }, + "search_backend": { + "backend": "lancedb", + "db_path": "data/search_medcpt", + "trials_table": "trials", + "criteria_table": "criteria", + "candidate_limit": 1000, + "vector_metric": "dot", + "vector_weight": 0.6, + "reembed_index": true + }, + "registry": { + "source": "clinicaltrials.gov", + "api_base_url": "https://clinicaltrials.gov/api/v2/studies", + "keywords_file": null, + "since_days": 7, + "max_studies": null, + "request_timeout": 30, + "rate_limit_per_second": 2, + "raw_dir": "data/registry/raw", + "manifest_path": "data/registry/manifest.jsonl", + "reports_dir": "data/registry/runs", + "failure_threshold": 0.25 + }, + "embedder": { + "backend": "hf", + "model_name": "ncbi/MedCPT-Article-Encoder", + "query_model_name": "ncbi/MedCPT-Query-Encoder", + "pooling": "cls", + "max_length": 512, + "query_max_length": 64, + "normalize": false, + "use_gpu": true, + "batch_size": 32, + "use_fp16": false + }, + "cot": { + "batch_size": 10 + }, + "LLM_reranker": { + "batch_size": 20, + "gpu_memory_utilization": 0.15, + "tensor_parallel_size": 1 + }, + "search": { + "mode": "hybrid", + "vector_score_threshold": 0.5, + "max_trials_first_level": 2000, + "max_trials_second_level": 1000, + "first_level": { + "enabled": true, + "max_trials": 2000, + "per_channel_size": 600, + "fusion": "rrf", + "rrf_k": 60, + "vector_score_threshold": 0.0, + "llm_expansion_enabled": false, + "llm_max_terms": 12, + "write_reports": true, + "hard_filters": [ + "age", + "sex", + "overall_status" + ] + } + }, + "constraints": { + "enabled": true, + "score_weight": 0.25, + "llm_extraction_enabled": false, + "unknown_is_neutral": true, + "write_reports": true + }, + "query_expansion": { + "enabled": true, + "backend": null, + "model": null, + "adapter": null, + "max_new_tokens": 4096, + "max_main_conditions": 11, + "max_other_conditions": 50, + "guided_json": true, + "no_think": true + }, + "use_cot_reasoning": true, + "rag": { + "batch_size": 4, + "max_trials_rag": 300, + "guided_json": true, + "no_think": true + }, + "vllm": { + "batch_size": 256, + "max_new_tokens": 8192, + "temperature": 0.0, + "top_p": 1.0, + "seed": 1234, + "length_bucket": true, + "gpu_memory_utilization": 0.7, + "max_model_len": 32768, + "tensor_parallel_size": 1, + "quantization": "", + "disable_custom_all_reduce": true, + "enforce_eager": true, + "max_num_seqs": 256 + }, + "reembed_index": false +} \ No newline at end of file diff --git a/src/trialmatchai/config/config_medcpt_qwen36_l40.json b/src/trialmatchai/config/config_medcpt_qwen36_l40.json new file mode 100644 index 00000000..f20e7ea9 --- /dev/null +++ b/src/trialmatchai/config/config_medcpt_qwen36_l40.json @@ -0,0 +1,162 @@ +{ + "entity_extraction": { + "backend": "gliner2", + "model_name": "fastino/gliner2-base-v1", + "model_revision": null, + "schema_path": "entity_schemas/trialmatchai.yaml", + "threshold": 0.8, + "batch_size": 8, + "device": "auto", + "trust_remote_code": false + }, + "concept_linker": { + "enabled": true, + "db_path": "data/concepts_medcpt", + "table": "concepts", + "accept_threshold": 0.7, + "reject_threshold": 0.5, + "margin": 0.05, + "rerank": "lexical", + "search_limit": 10 + }, + "paths": { + "output_dir": "results", + "trials_json_folder": "data/trials_jsons" + }, + "patient_inputs": { + "raw_dir": "data/patients/raw", + "profile_dir": "data/patients/profiles", + "summary_dir": "data/patients/summaries", + "default_format": "auto", + "strict_validation": false, + "copy_raw": true + }, + "model": { + "base_model": "QuantTrio/Qwen3.6-35B-A3B-AWQ", + "base_model_revision": null, + "trust_remote_code": false, + "quantization": { + "load_in_4bit": false, + "bnb_4bit_use_double_quant": true, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": "float16" + }, + "cot_adapter_path": null, + "reranker_model_path": "google/gemma-2-2b-it", + "reranker_model_revision": null, + "reranker_adapter_path": "models/finetuned_gemma2" + }, + "tokenizer": { + "use_fast": true, + "padding_side": "left" + }, + "global": { + "device": 0 + }, + "search_backend": { + "backend": "lancedb", + "db_path": "data/search_medcpt", + "trials_table": "trials", + "criteria_table": "criteria", + "candidate_limit": 1000, + "vector_metric": "dot", + "vector_weight": 0.6, + "reembed_index": true + }, + "registry": { + "source": "clinicaltrials.gov", + "api_base_url": "https://clinicaltrials.gov/api/v2/studies", + "keywords_file": null, + "since_days": 7, + "max_studies": null, + "request_timeout": 30, + "rate_limit_per_second": 2, + "raw_dir": "data/registry/raw", + "manifest_path": "data/registry/manifest.jsonl", + "reports_dir": "data/registry/runs", + "failure_threshold": 0.25 + }, + "embedder": { + "backend": "hf", + "model_name": "ncbi/MedCPT-Article-Encoder", + "query_model_name": "ncbi/MedCPT-Query-Encoder", + "pooling": "cls", + "max_length": 512, + "query_max_length": 64, + "normalize": false, + "use_gpu": false, + "batch_size": 32, + "use_fp16": false + }, + "cot": { + "batch_size": 10 + }, + "LLM_reranker": { + "batch_size": 20, + "gpu_memory_utilization": 0.15, + "tensor_parallel_size": 1 + }, + "search": { + "mode": "hybrid", + "vector_score_threshold": 0.5, + "max_trials_first_level": 2000, + "max_trials_second_level": 1000, + "first_level": { + "enabled": true, + "max_trials": 2000, + "per_channel_size": 600, + "fusion": "rrf", + "rrf_k": 60, + "vector_score_threshold": 0.0, + "llm_expansion_enabled": false, + "llm_max_terms": 12, + "write_reports": true, + "hard_filters": [ + "age", + "sex", + "overall_status" + ] + } + }, + "constraints": { + "enabled": true, + "score_weight": 0.25, + "llm_extraction_enabled": false, + "unknown_is_neutral": true, + "write_reports": true + }, + "query_expansion": { + "enabled": true, + "backend": null, + "model": null, + "adapter": null, + "max_new_tokens": 4096, + "max_main_conditions": 11, + "max_other_conditions": 50, + "guided_json": true, + "no_think": true + }, + "use_cot_reasoning": true, + "rag": { + "batch_size": 4, + "max_trials_rag": 300, + "guided_json": true, + "no_think": true + }, + "vllm": { + "batch_size": 256, + "max_new_tokens": 8192, + "temperature": 0.0, + "top_p": 1.0, + "seed": 1234, + "length_bucket": true, + "gpu_memory_utilization": 0.78, + "max_model_len": 16384, + "tensor_parallel_size": 1, + "quantization": "", + "disable_custom_all_reduce": true, + "enforce_eager": true, + "max_num_seqs": 256 + }, + "reembed_index": false +} \ No newline at end of file diff --git a/src/trialmatchai/config/config_medcpt_qwen3med.json b/src/trialmatchai/config/config_medcpt_qwen3med.json new file mode 100644 index 00000000..e8726a96 --- /dev/null +++ b/src/trialmatchai/config/config_medcpt_qwen3med.json @@ -0,0 +1,156 @@ +{ + "entity_extraction": { + "backend": "gliner2", + "model_name": "fastino/gliner2-base-v1", + "model_revision": null, + "schema_path": "entity_schemas/trialmatchai.yaml", + "threshold": 0.8, + "batch_size": 8, + "device": "auto", + "trust_remote_code": false + }, + "concept_linker": { + "enabled": true, + "db_path": "data/concepts_medcpt", + "table": "concepts", + "accept_threshold": 0.7, + "reject_threshold": 0.5, + "margin": 0.05, + "rerank": "lexical", + "search_limit": 10 + }, + "paths": { + "output_dir": "results", + "trials_json_folder": "data/trials_jsons" + }, + "patient_inputs": { + "raw_dir": "data/patients/raw", + "profile_dir": "data/patients/profiles", + "summary_dir": "data/patients/summaries", + "default_format": "auto", + "strict_validation": false, + "copy_raw": true + }, + "model": { + "base_model": "nicoboss/Qwen-3-32B-Medical-Reasoning", + "base_model_revision": null, + "trust_remote_code": false, + "quantization": { + "load_in_4bit": true, + "bnb_4bit_use_double_quant": true, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": "float16" + }, + "cot_adapter_path": null, + "reranker_model_path": "google/gemma-2-2b-it", + "reranker_model_revision": null, + "reranker_adapter_path": "models/finetuned_gemma2" + }, + "tokenizer": { + "use_fast": true, + "padding_side": "left" + }, + "global": { + "device": 0 + }, + "search_backend": { + "backend": "lancedb", + "db_path": "data/search_medcpt", + "trials_table": "trials", + "criteria_table": "criteria", + "candidate_limit": 1000, + "vector_metric": "dot", + "vector_weight": 0.6, + "reembed_index": true + }, + "registry": { + "source": "clinicaltrials.gov", + "api_base_url": "https://clinicaltrials.gov/api/v2/studies", + "keywords_file": null, + "since_days": 7, + "max_studies": null, + "request_timeout": 30, + "rate_limit_per_second": 2, + "raw_dir": "data/registry/raw", + "manifest_path": "data/registry/manifest.jsonl", + "reports_dir": "data/registry/runs", + "failure_threshold": 0.25 + }, + "embedder": { + "backend": "hf", + "model_name": "ncbi/MedCPT-Article-Encoder", + "query_model_name": "ncbi/MedCPT-Query-Encoder", + "pooling": "cls", + "max_length": 512, + "query_max_length": 64, + "normalize": false, + "use_gpu": true, + "batch_size": 32, + "use_fp16": false + }, + "cot": { + "batch_size": 10 + }, + "LLM_reranker": { + "batch_size": 20, + "gpu_memory_utilization": 0.15 + }, + "search": { + "mode": "hybrid", + "vector_score_threshold": 0.5, + "max_trials_first_level": 2000, + "max_trials_second_level": 1000, + "first_level": { + "enabled": true, + "max_trials": 2000, + "per_channel_size": 600, + "fusion": "rrf", + "rrf_k": 60, + "vector_score_threshold": 0.0, + "llm_expansion_enabled": false, + "llm_max_terms": 12, + "write_reports": true, + "hard_filters": [ + "age", + "sex", + "overall_status" + ] + } + }, + "constraints": { + "enabled": true, + "score_weight": 0.25, + "llm_extraction_enabled": false, + "unknown_is_neutral": true, + "write_reports": true + }, + "query_expansion": { + "enabled": false, + "backend": null, + "model": null, + "adapter": null, + "max_new_tokens": 2048, + "max_main_conditions": 11, + "max_other_conditions": 50 + }, + "use_cot_reasoning": true, + "rag": { + "batch_size": 4, + "max_trials_rag": 300, + "no_think": true + }, + "vllm": { + "batch_size": 100, + "max_new_tokens": 5000, + "temperature": 0.0, + "top_p": 1.0, + "seed": 1234, + "length_bucket": true, + "gpu_memory_utilization": 0.55, + "max_model_len": 16384, + "tensor_parallel_size": 1, + "quantization": "bitsandbytes", + "disable_custom_all_reduce": true, + "enforce_eager": true + } +} diff --git a/src/trialmatchai/config/config_nemotron_medgemma.json b/src/trialmatchai/config/config_nemotron_medgemma.json new file mode 100644 index 00000000..59640984 --- /dev/null +++ b/src/trialmatchai/config/config_nemotron_medgemma.json @@ -0,0 +1,158 @@ +{ + "entity_extraction": { + "backend": "gliner2", + "model_name": "fastino/gliner2-base-v1", + "model_revision": null, + "schema_path": "entity_schemas/trialmatchai.yaml", + "threshold": 0.8, + "batch_size": 8, + "device": "auto", + "trust_remote_code": false + }, + "concept_linker": { + "enabled": true, + "db_path": "data/concepts_medcpt", + "table": "concepts", + "accept_threshold": 0.7, + "reject_threshold": 0.5, + "margin": 0.05, + "rerank": "lexical", + "search_limit": 10 + }, + "paths": { + "output_dir": "results", + "trials_json_folder": "data/trials_jsons" + }, + "patient_inputs": { + "raw_dir": "data/patients/raw", + "profile_dir": "data/patients/profiles", + "summary_dir": "data/patients/summaries", + "default_format": "auto", + "strict_validation": false, + "copy_raw": true + }, + "model": { + "base_model": "google/medgemma-27b-text-it", + "base_model_revision": null, + "trust_remote_code": false, + "quantization": { + "load_in_4bit": true, + "bnb_4bit_use_double_quant": true, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": "float16" + }, + "cot_adapter_path": null, + "reranker_model_path": "google/gemma-2-2b-it", + "reranker_model_revision": null, + "reranker_adapter_path": "models/finetuned_gemma2" + }, + "tokenizer": { + "use_fast": true, + "padding_side": "left" + }, + "global": { + "device": 0 + }, + "search_backend": { + "backend": "lancedb", + "db_path": "data/search_medcpt", + "trials_table": "trials", + "criteria_table": "criteria", + "candidate_limit": 1000, + "vector_metric": "cosine", + "vector_weight": 0.6, + "reembed_index": true + }, + "registry": { + "source": "clinicaltrials.gov", + "api_base_url": "https://clinicaltrials.gov/api/v2/studies", + "keywords_file": null, + "since_days": 7, + "max_studies": null, + "request_timeout": 30, + "rate_limit_per_second": 2, + "raw_dir": "data/registry/raw", + "manifest_path": "data/registry/manifest.jsonl", + "reports_dir": "data/registry/runs", + "failure_threshold": 0.25 + }, + "embedder": { + "backend": "hf", + "model_name": "nvidia/Nemotron-3-Embed-1B-BF16", + "pooling": "mean", + "max_length": 512, + "normalize": true, + "native_metric": "cosine", + "query_prefix": "query: ", + "document_prefix": "passage: ", + "trust_remote_code": false, + "use_gpu": true, + "batch_size": 32, + "use_fp16": false + }, + "cot": { + "batch_size": 10 + }, + "LLM_reranker": { + "batch_size": 20, + "gpu_memory_utilization": 0.15 + }, + "search": { + "mode": "hybrid", + "vector_score_threshold": 0.5, + "max_trials_first_level": 2000, + "max_trials_second_level": 1000, + "first_level": { + "enabled": true, + "max_trials": 2000, + "per_channel_size": 600, + "fusion": "rrf", + "rrf_k": 60, + "vector_score_threshold": 0.0, + "llm_expansion_enabled": false, + "llm_max_terms": 12, + "write_reports": true, + "hard_filters": [ + "age", + "sex", + "overall_status" + ] + } + }, + "constraints": { + "enabled": true, + "score_weight": 0.25, + "llm_extraction_enabled": false, + "unknown_is_neutral": true, + "write_reports": true + }, + "query_expansion": { + "enabled": false, + "backend": null, + "model": null, + "adapter": null, + "max_new_tokens": 2048, + "max_main_conditions": 11, + "max_other_conditions": 50 + }, + "use_cot_reasoning": true, + "rag": { + "batch_size": 4, + "max_trials_rag": 300 + }, + "vllm": { + "batch_size": 100, + "max_new_tokens": 5000, + "temperature": 0.0, + "top_p": 1.0, + "seed": 1234, + "length_bucket": true, + "gpu_memory_utilization": 0.78, + "max_model_len": 8192, + "tensor_parallel_size": 1, + "quantization": "", + "disable_custom_all_reduce": true, + "enforce_eager": true + }, + "reembed_index": true +} \ No newline at end of file diff --git a/src/trialmatchai/config/settings.py b/src/trialmatchai/config/settings.py index 25053296..ae82d564 100644 --- a/src/trialmatchai/config/settings.py +++ b/src/trialmatchai/config/settings.py @@ -190,6 +190,37 @@ class ShortlistSettings(BaseModel): max_depth: int | None = Field(None, ge=1) +class SecondLevelSearchSettings(BaseModel): + """Width of the second level — the pipeline's true bottleneck. + + Both values were hardcoded and unreachable from config. Measured on TREC 2023 (37 + patients, retrieval only), ``per_query_size`` alone sets a hard ceiling on everything + downstream, because the reranker and the aggregation threshold can only DROP trials + from what retrieval surfaced: + + size unique trials % of candidate pool recall of judged-relevant + 250 803 43.1% 0.5849 <- old default + 500 1,041 55.8% 0.7135 + 1000 1,200 64.3% 0.7870 + 2000 1,345 72.1% 0.8352 + 4000 1,477 79.2% 0.8507 + + The candidate pool is ~1,860 trials x ~17 criteria ~= 31,000 criteria, so at 250 the + second level examined about a tenth of its own pool. Raising it costs retrieval plus + 2B-reranker scoring; it does not touch the 35B eligibility model. + + Defaults reproduce the previous hardcoded behaviour exactly, so changing width is an + explicit A/B. + """ + + # Criteria retrieved PER QUERY (~10-13 queries per patient), not per patient. + per_query_size: int = Field(250, ge=1) + # Criteria scoring below this are dropped in aggregate_to_trials, so a trial whose every + # criterion falls short never reaches the shortlist at all. + aggregation_threshold: float = Field(0.5, ge=0.0, le=1.0) + aggregation_method: Literal["weighted", "avg", "sqrt", "log"] = "weighted" + + class SearchSettings(BaseModel): mode: Literal["bm25", "vector", "hybrid"] = "hybrid" vector_score_threshold: float = Field(0.5, ge=0.0, le=1.0) @@ -208,6 +239,9 @@ class SearchSettings(BaseModel): default_factory=FirstLevelSearchSettings ) shortlist: ShortlistSettings = Field(default_factory=ShortlistSettings) + second_level: SecondLevelSearchSettings = Field( + default_factory=SecondLevelSearchSettings + ) @model_validator(mode="before") @classmethod @@ -241,6 +275,11 @@ class RagSettings(BaseModel): backend: Literal["vllm", "transformers"] = "vllm" batch_size: int = Field(4, ge=1) max_trials_rag: int = Field(20, ge=1) + # Size the eligibility output budget to each trial's criterion count instead of giving + # every trial the full vllm.max_new_tokens. Clamped to that ceiling, so it can only lower + # the budget for small trials -- large trials keep today's headroom and cannot be + # truncated further. See matching/eligibility_reasoning_vllm.adaptive_max_tokens. + adaptive_token_budget: bool = False # Suppress chain-of-thought in the eligibility stage for reasoning models (Qwen3): # sends enable_thinking=False / a /no_think prefix and strips residual think tags. no_think: bool = False diff --git a/src/trialmatchai/main.py b/src/trialmatchai/main.py index 54378c90..0a499fd5 100644 --- a/src/trialmatchai/main.py +++ b/src/trialmatchai/main.py @@ -457,6 +457,7 @@ def run_rag_processing( seed=vllm_cfg.get("seed", 1234), length_bucket=vllm_cfg.get("length_bucket", True), max_model_len=vllm_cfg.get("max_model_len"), + adaptive_token_budget=bool(rag_cfg.get("adaptive_token_budget", False)), lora_request=lora_request, chat_template_kwargs=rag_cfg.get("chat_template_kwargs"), guided_json=rag_cfg.get("guided_json", False), @@ -612,12 +613,24 @@ def main_pipeline( else: llm_reranker = None + # search.second_level controls the width of the pipeline's binding bottleneck. Previously + # both values were hardcoded in SecondStageRetriever and unreachable from config, which + # capped the second level at ~43% of its own candidate pool. + second_level_cfg = config["search"].get("second_level") or {} gemma_retriever = SecondStageRetriever( search_backend=search_backend, llm_reranker=llm_reranker, embedder=embedder, entity_annotator=entity_annotator, search_mode=config["search"].get("mode", "hybrid"), + size=int(second_level_cfg.get("per_query_size", 250)), + aggregation_threshold=float(second_level_cfg.get("aggregation_threshold", 0.5)), + aggregation_method=str(second_level_cfg.get("aggregation_method", "weighted")), + ) + logger.info( + "Second level width: per_query_size=%s, aggregation_threshold=%s", + gemma_retriever.size, + gemma_retriever.aggregation_threshold, ) completed_patients = 0 diff --git a/src/trialmatchai/matching/agent.py b/src/trialmatchai/matching/agent.py new file mode 100644 index 00000000..fb4ddd61 --- /dev/null +++ b/src/trialmatchai/matching/agent.py @@ -0,0 +1,208 @@ +"""A bounded per-patient controller over retrieval width and shortlist depth. + +The pipeline is otherwise a single pass with hardcoded widths: the second level retrieves +``size=250`` criteria per query and the shortlist takes a fixed slice of whatever that +produced. Measurement says both numbers are wrong, and wrong in a patient-dependent way: + + * The depth a patient needs to reach 90% of its own first-level recall ranges from 50 to + 1550 trials, spread evenly across that range (17 of 75 TREC-2021 patients need <=200, + 21 need >700). One number cannot serve both ends; sizing for the worst case wastes ~65% + of the reasoner's compute and sizing for the median drops the hard patients. + * The second level surfaces ~446 unique trials from a ~1860-trial candidate pool, and the + shortlist can only choose among those. Every downstream knob sits behind that. + +So the agent's job is ALLOCATION, not extra reasoning. Three attempts to add a reasoning +component to this pipeline measured neutral or worse (constraint verification, uncertainty +gating, softened disqualification), while changing how compute is distributed is the one +intervention that has held up. This module does only that. + +Design constraints, taken from the failure literature rather than invented: + + * **Closed action set.** Four actions, no free-form tool choice. The MAST study of 1,600+ + traces of popular multi-agent frameworks found 41-87% task failure, dominated by + specification (41.8%) and coordination (36.9%) problems -- i.e. system design, not model + quality. A small fixed action set is what avoids that class of failure. + * **Hard budget.** Rounds and total criteria examined are both capped. The agent cannot + spend without bound however promising things look. + * **Monotone widening.** Each round only ever ADDS candidates. No round removes something + an earlier round accepted, so errors cannot compound across rounds. + * **No LLM in the control loop.** Every decision is a pure function of counts and scores + the pipeline already computes. The agent is auditable, reproducible, and free. + +Disabled by default (``agent.enabled``). +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, Literal + +from trialmatchai.utils.logging_config import setup_logging + +logger = setup_logging(__name__) + +Action = Literal["widen_retrieval", "deepen_shortlist", "stop"] + +# Defaults chosen from the measured curves, not tuned. per_query_size starts at the current +# hardcoded 250 so round 1 reproduces today's behaviour exactly; widening is what the agent +# adds on top. +DEFAULT_MAX_ROUNDS = 3 +DEFAULT_START_SIZE = 250 +DEFAULT_SIZE_MULTIPLIER = 2.0 +DEFAULT_MAX_SIZE = 2000 +# Stop widening when a round's extra retrieval buys less than this fraction of new trials +# relative to what was already surfaced. At 0.10 a round must grow the pool by >=10% to +# justify the next one. +DEFAULT_MIN_YIELD = 0.10 +# Never let one patient consume more than this many criteria across all rounds. +DEFAULT_MAX_CRITERIA = 40_000 + + +@dataclass(frozen=True) +class AgentBudget: + max_rounds: int = DEFAULT_MAX_ROUNDS + max_criteria: int = DEFAULT_MAX_CRITERIA + max_size: int = DEFAULT_MAX_SIZE + + def exhausted(self, *, rounds_done: int, criteria_examined: int) -> bool: + return rounds_done >= self.max_rounds or criteria_examined >= self.max_criteria + + +@dataclass +class RoundObservation: + """What the agent can see after one second-level round. Counts only -- no model calls.""" + + round_index: int + per_query_size: int + criteria_examined: int + trials_surfaced: int + new_trials: int + candidate_pool: int + + @property + def yield_rate(self) -> float: + """New trials this round as a fraction of what was already surfaced.""" + previously = self.trials_surfaced - self.new_trials + if previously <= 0: + return 1.0 if self.new_trials else 0.0 + return self.new_trials / previously + + @property + def pool_coverage(self) -> float: + if self.candidate_pool <= 0: + return 1.0 + return self.trials_surfaced / self.candidate_pool + + +@dataclass +class AgentTrace: + """Provenance for the whole episode, written beside the shortlist.""" + + rounds: list[dict[str, Any]] = field(default_factory=list) + actions: list[str] = field(default_factory=list) + stopped_because: str = "" + + def record(self, observation: RoundObservation, action: Action, reason: str) -> None: + self.rounds.append( + { + "round": observation.round_index, + "per_query_size": observation.per_query_size, + "criteria_examined": observation.criteria_examined, + "trials_surfaced": observation.trials_surfaced, + "new_trials": observation.new_trials, + "yield_rate": round(observation.yield_rate, 4), + "pool_coverage": round(observation.pool_coverage, 4), + "action": action, + "reason": reason, + } + ) + self.actions.append(action) + if action == "stop": + self.stopped_because = reason + + def as_dict(self) -> dict[str, Any]: + return { + "rounds": self.rounds, + "actions": self.actions, + "stopped_because": self.stopped_because, + "total_criteria_examined": sum(r["criteria_examined"] for r in self.rounds), + "final_trials_surfaced": self.rounds[-1]["trials_surfaced"] if self.rounds else 0, + } + + +def agent_config(config: Mapping[str, Any] | None) -> dict[str, Any]: + raw = (config or {}).get("agent") or {} + if not isinstance(raw, Mapping): + raw = {} + return { + "enabled": bool(raw.get("enabled", False)), + "max_rounds": int(raw.get("max_rounds", DEFAULT_MAX_ROUNDS)), + "start_size": int(raw.get("start_size", DEFAULT_START_SIZE)), + "size_multiplier": float(raw.get("size_multiplier", DEFAULT_SIZE_MULTIPLIER)), + "max_size": int(raw.get("max_size", DEFAULT_MAX_SIZE)), + "min_yield": float(raw.get("min_yield", DEFAULT_MIN_YIELD)), + "max_criteria": int(raw.get("max_criteria", DEFAULT_MAX_CRITERIA)), + } + + +def budget_from_config(config: Mapping[str, Any] | None) -> AgentBudget: + cfg = agent_config(config) + return AgentBudget( + max_rounds=cfg["max_rounds"], + max_criteria=cfg["max_criteria"], + max_size=cfg["max_size"], + ) + + +def decide( + observation: RoundObservation, + budget: AgentBudget, + *, + min_yield: float = DEFAULT_MIN_YIELD, + criteria_examined: int = 0, +) -> tuple[Action, str]: + """The whole policy. A pure function of counts -- no model, no hidden state. + + Widening is justified only while it is still finding trials the previous rounds missed. + ``yield_rate`` is measured against what was already surfaced rather than against the + candidate pool, because the pool is a poor denominator: a patient with 2000 candidates + and 400 genuinely plausible trials should stop at 400, not chase 20% coverage. + """ + rounds_done = observation.round_index + 1 + if budget.exhausted(rounds_done=rounds_done, criteria_examined=criteria_examined): + return "stop", ( + f"budget exhausted (rounds {rounds_done}/{budget.max_rounds}, " + f"criteria {criteria_examined}/{budget.max_criteria})" + ) + if observation.trials_surfaced >= observation.candidate_pool > 0: + return "stop", "whole candidate pool already surfaced" + if observation.per_query_size >= budget.max_size: + return "deepen_shortlist", f"per-query size at the cap ({budget.max_size})" + if observation.yield_rate < min_yield: + return "stop", ( + f"marginal yield {observation.yield_rate:.3f} below {min_yield:.3f}; " + "widening is no longer finding new trials" + ) + return "widen_retrieval", ( + f"yield {observation.yield_rate:.3f} >= {min_yield:.3f} and budget remains" + ) + + +def next_size(current: int, *, multiplier: float, cap: int) -> int: + return max(current + 1, min(cap, int(round(current * multiplier)))) + + +def plan_sizes(config: Mapping[str, Any] | None) -> Sequence[int]: + """The size ladder the agent would climb if every round justified widening. + + Exposed so a run can be costed before it is launched. + """ + cfg = agent_config(config) + sizes, size = [], cfg["start_size"] + for _ in range(max(1, cfg["max_rounds"])): + sizes.append(size) + if size >= cfg["max_size"]: + break + size = next_size(size, multiplier=cfg["size_multiplier"], cap=cfg["max_size"]) + return sizes diff --git a/src/trialmatchai/matching/eligibility_base.py b/src/trialmatchai/matching/eligibility_base.py index 2966b874..7040828c 100644 --- a/src/trialmatchai/matching/eligibility_base.py +++ b/src/trialmatchai/matching/eligibility_base.py @@ -9,7 +9,7 @@ import json import os import re -from typing import Dict, List +from typing import Any, Dict, List from trialmatchai.utils.file_utils import read_json_file, write_json_file, write_text_file from trialmatchai.utils.json_utils import extract_json_object @@ -19,6 +19,32 @@ logger = setup_logging(__name__) +def count_criteria(criteria_text: Any) -> int: + """Rough count of eligibility criteria in a trial's criteria block. + + Criteria arrive as one free-text block, one criterion per line (occasionally bulleted). + Counting non-empty lines and ignoring section headers ("Inclusion Criteria:") is close + enough to size an output budget -- it does not need to be exact, only proportional. + """ + if isinstance(criteria_text, (list, tuple)): + return sum(1 for item in criteria_text if str(item).strip()) + text = str(criteria_text or "") + count = 0 + for raw in text.splitlines(): + line = raw.strip().lstrip("-*• \t") + if not line: + continue + # Section headers introduce criteria rather than being one. + if line.rstrip(":").strip().lower() in { + "inclusion criteria", + "exclusion criteria", + "eligibility criteria", + }: + continue + count += 1 + return count + + def _is_error_output(path: str) -> bool: """True if a per-trial output is a recorded failure or unparseable, so the resume worklist retries it instead of locking a transient failure into the ranking.""" @@ -260,6 +286,9 @@ def process_trials( "nct_id": nct_id, "prompt": prompt, "tok_len": self._token_length(prompt, nct_id), + # Trials range from a handful of criteria to 60+, and the model must emit a + # verdict per criterion, so this drives the adaptive output budget. + "n_criteria": count_criteria(criteria_text), } ) diff --git a/src/trialmatchai/matching/eligibility_reasoning_vllm.py b/src/trialmatchai/matching/eligibility_reasoning_vllm.py index 156ff2ff..cb02c917 100644 --- a/src/trialmatchai/matching/eligibility_reasoning_vllm.py +++ b/src/trialmatchai/matching/eligibility_reasoning_vllm.py @@ -31,6 +31,40 @@ def _criteria_array(labels: list[str]) -> dict: } +# Adaptive output budget. m1 (arXiv:2504.00869) measures an optimal MEDICAL reasoning budget +# near 4K tokens, "beyond which performance may degrade due to overthinking" -- but that is for +# a single question. This stage judges every criterion of a trial in ONE call, and trials range +# from a handful of criteria to 60+, so a flat 4K cap would truncate the large ones into invalid +# JSON. Scaling with the actual work keeps the per-criterion budget near the measured optimum at +# both ends. +# +# The result is CLAMPED to the configured max_new_tokens, so this can only ever LOWER the budget +# for small trials -- large trials keep exactly the headroom they have today, and no trial can be +# truncated more than it already would be. +_BUDGET_BASE_TOKENS = 1024 # recap, final decision, JSON scaffolding +_BUDGET_PER_CRITERION_TOKENS = 192 # one verdict + justification +_BUDGET_FLOOR_TOKENS = 1024 + + +def adaptive_max_tokens( + n_criteria: int, + *, + ceiling: int, + base: int = _BUDGET_BASE_TOKENS, + per_criterion: int = _BUDGET_PER_CRITERION_TOKENS, + floor: int = _BUDGET_FLOOR_TOKENS, +) -> int: + """Output-token budget for a trial with ``n_criteria`` criteria. + + A ~17-criterion trial (the TREC 2023 mean) lands near 4K; a 60-criterion trial asks for + more and is capped by ``ceiling``. Never exceeds ``ceiling``, never drops below ``floor``. + """ + if n_criteria <= 0: + return min(ceiling, max(floor, base)) + want = base + per_criterion * n_criteria + return max(min(floor, ceiling), min(ceiling, want)) + + ELIGIBILITY_JSON_SCHEMA = { "type": "object", "properties": { @@ -62,6 +96,7 @@ def __init__( lora_request: Optional[Any] = None, chat_template_kwargs: Optional[dict] = None, guided_json: bool = False, + adaptive_token_budget: bool = False, ): """vLLM-backed CoT eligibility processor with optional LoRA adapter.""" self.llm = llm @@ -71,6 +106,7 @@ def __init__( self.no_think = no_think self.chat_template_kwargs = chat_template_kwargs or {} self.max_new_tokens = max_new_tokens + self.adaptive_token_budget = adaptive_token_budget self.temperature = temperature self.top_p = top_p self.seed = seed @@ -179,6 +215,29 @@ def _init_validate_lora_request(self, lora_request): # ---------------------- Core batch path (vLLM) ---------------------- + def _sampling_params_for(self, batch: List[Dict]) -> Any: + """Shared SamplingParams, or one per prompt sized to that trial's criterion count.""" + if not self.adaptive_token_budget: + return self.sampling_params + from copy import copy + + params = [] + for item in batch: + budget = adaptive_max_tokens( + int(item.get("n_criteria") or 0), ceiling=self.max_new_tokens + ) + per_request = copy(self.sampling_params) + per_request.max_tokens = budget + params.append(per_request) + if params: + logger.debug( + "Adaptive output budget over batch: min=%s max=%s (ceiling %s)", + min(p.max_tokens for p in params), + max(p.max_tokens for p in params), + self.max_new_tokens, + ) + return params + def _process_batch(self, batch: List[Dict], output_folder: str): try: prompts = [item["prompt"] for item in batch] @@ -188,10 +247,15 @@ def _process_batch(self, batch: List[Dict], output_folder: str): safe_lora_request = self._validate_lora_request() + # One SamplingParams per prompt when the budget is adaptive, so a 5-criterion trial + # does not get the same 8K allowance as a 60-criterion one. Falls back to the shared + # object when disabled, which is the historical behaviour. + params = self._sampling_params_for(batch) + try: results = self.llm.generate( prompts, - self.sampling_params, + params, lora_request=safe_lora_request, ) except TypeError as e: diff --git a/src/trialmatchai/matching/retrieval/criteria_retrieval.py b/src/trialmatchai/matching/retrieval/criteria_retrieval.py index e6625950..8a96e229 100644 --- a/src/trialmatchai/matching/retrieval/criteria_retrieval.py +++ b/src/trialmatchai/matching/retrieval/criteria_retrieval.py @@ -28,6 +28,8 @@ def __init__( llm_reranker: Optional[LLMReranker], embedder: Optional[TextEmbedder], size: int = 250, + aggregation_threshold: float = 0.5, + aggregation_method: str = "weighted", inclusion_weight: float = 1.0, exclusion_weight: float = 0.25, entity_annotator=None, @@ -36,7 +38,13 @@ def __init__( self.search_backend = search_backend self.llm_reranker = llm_reranker self.embedder = embedder + # Criteria retrieved PER QUERY. This is the pipeline's binding width constraint: + # measured on TREC 2023, 250 surfaces 43% of the candidate pool at 0.585 recall of + # judged-relevant trials, while 1000 reaches 64%/0.787. Everything downstream can + # only drop trials from what this retrieves. self.size = size + self.aggregation_threshold = aggregation_threshold + self.aggregation_method = aggregation_method self.inclusion_weight = inclusion_weight self.exclusion_weight = exclusion_weight self.entity_annotator = entity_annotator @@ -201,8 +209,18 @@ def apply_constraint_adjustments( return criteria def aggregate_to_trials( - self, criteria: List[Dict], threshold: float = 0.5, method: str = "weighted" + self, + criteria: List[Dict], + threshold: float | None = None, + method: str | None = None, ) -> List[Dict]: + # None -> the configured values; explicit args still win so callers and tests can pin + # them. getattr keeps the historical contract that this method uses no instance state, + # so an instance built without __init__ still aggregates at the old 0.5/"weighted". + if threshold is None: + threshold = getattr(self, "aggregation_threshold", 0.5) + if method is None: + method = getattr(self, "aggregation_method", "weighted") # A criterion matched by several paraphrases appears once per query; keep only the best # score per UNIQUE criterion so a trial isn't inflated by query overlap (skews sqrt/weighted). best_by_criterion: dict[tuple[str, str], float] = {} diff --git a/tests/test_second_level_width.py b/tests/test_second_level_width.py new file mode 100644 index 00000000..b3dd1dec --- /dev/null +++ b/tests/test_second_level_width.py @@ -0,0 +1,113 @@ +"""Second-level width and adaptive output budget. + +Two hardcoded constants were the pipeline's binding bottleneck: the per-query criteria +retrieval size (250) and the aggregation threshold (0.5). Measured on TREC 2023, size alone +capped the second level at 43% of its candidate pool. These tests pin the new config surface +and, importantly, that the defaults reproduce the old behaviour. +""" + +import pytest + +from trialmatchai.matching.eligibility_base import count_criteria +from trialmatchai.matching.eligibility_reasoning_vllm import adaptive_max_tokens +from trialmatchai.matching.retrieval.criteria_retrieval import SecondStageRetriever + + +def _retriever(**kwargs): + return SecondStageRetriever( + search_backend=object(), llm_reranker=None, embedder=None, **kwargs + ) + + +# ----------------------------------------------------------------- width defaults +def test_defaults_reproduce_the_previous_hardcoded_behaviour(): + r = _retriever() + assert r.size == 250 + assert r.aggregation_threshold == 0.5 + assert r.aggregation_method == "weighted" + + +def test_width_is_configurable(): + r = _retriever(size=1000, aggregation_threshold=0.2, aggregation_method="sqrt") + assert r.size == 1000 + assert r.aggregation_threshold == 0.2 + assert r.aggregation_method == "sqrt" + + +def _criterion(nct, score, cid): + return {"_source": {"nct_id": nct, "criteria_id": cid}, "llm_score": score} + + +def test_aggregation_threshold_gates_which_trials_survive(): + """The 0.5 cut drops a trial whose every criterion scores below it -- that trial then + never reaches the shortlist, whatever the reranker thought.""" + criteria = [_criterion("NCT1", 0.9, "a"), _criterion("NCT2", 0.3, "b")] + + strict = _retriever().aggregate_to_trials(criteria) + assert {t["nct_id"] for t in strict} == {"NCT1"} # NCT2 dropped at 0.5 + + lenient = _retriever(aggregation_threshold=0.2).aggregate_to_trials(criteria) + assert {t["nct_id"] for t in lenient} == {"NCT1", "NCT2"} + + +def test_explicit_argument_still_overrides_the_configured_threshold(): + criteria = [_criterion("NCT2", 0.3, "b")] + r = _retriever(aggregation_threshold=0.5) + assert r.aggregate_to_trials(criteria) == [] + assert len(r.aggregate_to_trials(criteria, threshold=0.1)) == 1 + + +# ----------------------------------------------------------------- criterion counting +def test_count_criteria_ignores_headers_and_bullets(): + text = ( + "Inclusion Criteria:\n" + "- Age 18 or older\n" + "- Histologically confirmed glioma\n" + "\n" + "Exclusion Criteria:\n" + "* Prior systemic therapy\n" + ) + assert count_criteria(text) == 3 + + +@pytest.mark.parametrize("value", ["", None, " \n \n"]) +def test_count_criteria_handles_empty(value): + assert count_criteria(value) == 0 + + +def test_count_criteria_accepts_a_list(): + assert count_criteria(["a", "b", ""]) == 2 + + +# ----------------------------------------------------------------- adaptive budget +def test_budget_scales_with_criterion_count(): + small = adaptive_max_tokens(5, ceiling=8192) + typical = adaptive_max_tokens(17, ceiling=8192) + large = adaptive_max_tokens(60, ceiling=8192) + assert small < typical < large + + +def test_typical_trial_lands_near_the_measured_medical_optimum(): + """m1 puts the medical reasoning optimum near 4K tokens. The TREC 2023 mean trial carries + ~17 criteria, so that case should land in the same neighbourhood -- not at an 8K default.""" + assert 3000 <= adaptive_max_tokens(17, ceiling=8192) <= 5000 + + +def test_budget_never_exceeds_the_ceiling(): + """The whole safety argument: this can only LOWER the budget, so no trial is truncated + more than it already would be under the fixed setting.""" + for n in (0, 1, 17, 60, 500): + assert adaptive_max_tokens(n, ceiling=8192) <= 8192 + assert adaptive_max_tokens(n, ceiling=2048) <= 2048 + + +def test_large_trials_keep_full_headroom(): + """A 60-criterion trial must still get the whole configured budget rather than a 4K cap + that would truncate it into invalid JSON.""" + assert adaptive_max_tokens(60, ceiling=8192) == 8192 + + +def test_budget_has_a_floor_for_degenerate_input(): + assert adaptive_max_tokens(0, ceiling=8192) >= 1024 + # A tiny ceiling still wins -- the floor never pushes past what the caller allows. + assert adaptive_max_tokens(0, ceiling=256) == 256 From 9457c30790c09c93436c3e35bb1d38978446f06d Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Thu, 13 Aug 2026 14:11:11 +0200 Subject: [PATCH 11/13] fix(trec): stop the funnel preset silently overwriting an explicit config _track_config assigned the TREC funnel unconditionally: search["max_trials_first_level"] = 1000 search["max_trials_second_level"] = 500 search["second_level_keep_divisor"] = 1 rag["max_trials_rag"] = 250 Whatever the config file said was discarded, with no warning. That made funnel experiments impossible to run and produced results that looked like findings: * A run configured for a 600-trial shortlist got 250. The "depth 600" TREC 2023 run therefore never changed depth at all -- its shortlist came out at 243, which I attributed to the second-level pool binding. It was this line. * A config asking for max_trials_second_level=1000 got 500. The width A/B was measuring a truncation constant: at per_query_size=1000 the second level aggregated 1,149 trials and at 250 it aggregated 573, and BOTH were cut to exactly 500 before the shortlist saw them. Across 120 completed patients, 56% land on exactly 500 and nothing ever exceeds it. So this cap sat upstream of every knob tuned so far -- per_query_size, the aggregation threshold, the shortlist divisor, max_trials_rag -- which is a large part of why each of them moved so little. The preset is now a DEFAULT rather than an override: each value is applied only where the base config left it at the schema default, i.e. the user expressed no opinion. setdefault alone is insufficient because a shipped config always carries these keys, so the preset would never apply at all. The resolved funnel is logged with which values came from the preset and which from the config, so this can never again be silent. Note this changes the effective funnel for the shipped configs, which set max_trials_second_level=1000: TREC runs will now honour that instead of 500. That is the intended fix, but it means new runs are not comparable to earlier ones without re-baselining. --- src/trialmatchai/trec/runner.py | 55 +++++++++++++++++++++--- tests/test_trec_funnel_config.py | 74 ++++++++++++++++++++++++++++++++ 2 files changed, 124 insertions(+), 5 deletions(-) create mode 100644 tests/test_trec_funnel_config.py diff --git a/src/trialmatchai/trec/runner.py b/src/trialmatchai/trec/runner.py index a521beb6..c14a6aba 100644 --- a/src/trialmatchai/trec/runner.py +++ b/src/trialmatchai/trec/runner.py @@ -41,13 +41,58 @@ def _track_config(base_config: Dict[str, Any], spec: TrackSpec) -> Dict[str, Any cfg["patient_inputs"]["summary_dir"] = str(spec.summary_dir) cfg.setdefault("paths", {})["output_dir"] = str(spec.output_dir) cfg.setdefault("query_expansion", {})["enabled"] = True - # TREC funnel 1000 -> 500 -> 250 (CoT), no second->CoT thinning (keep_divisor=1); + # TREC funnel preset: 1000 -> 500 -> 250 (CoT), no second->CoT thinning (keep_divisor=1); # deeper than the interactive defaults because TREC scores the whole ranked list. + # + # These are DEFAULTS, not overrides. They used to be assigned unconditionally, which + # silently discarded whatever the config file said and made funnel experiments + # unrunnable: a config asking for max_trials_second_level=1000 still got 500, and a run + # configured for a 600-trial shortlist still got 250. Several experiments measured the + # preset rather than the variable they were manipulating before this was caught. + # + # setdefault is not enough here because a shipped config always carries these keys, so + # the preset would never apply. Instead each value is applied only when the base config + # left it at the schema default, i.e. the user did not express an opinion. search = cfg.setdefault("search", {}) - search["max_trials_first_level"] = 1000 - search["max_trials_second_level"] = 500 - search["second_level_keep_divisor"] = 1 - cfg.setdefault("rag", {})["max_trials_rag"] = 250 + rag = cfg.setdefault("rag", {}) + preset = { + ("search", "max_trials_first_level"): 1000, + ("search", "max_trials_second_level"): 500, + ("search", "second_level_keep_divisor"): 3, + ("rag", "max_trials_rag"): 250, + } + schema_defaults = { + ("search", "max_trials_first_level"): 1000, + ("search", "max_trials_second_level"): 100, + ("search", "second_level_keep_divisor"): 3, + ("rag", "max_trials_rag"): 20, + } + sections = {"search": search, "rag": rag} + applied, respected = {}, {} + for key, preset_value in preset.items(): + section, name = key + current = sections[section].get(name) + if current is None or current == schema_defaults[key]: + sections[section][name] = preset_value + applied[f"{section}.{name}"] = preset_value + else: + respected[f"{section}.{name}"] = current + # keep_divisor=1 is part of the preset's intent (no second->CoT thinning), and the schema + # default happens to equal the preset value, so set it explicitly when untouched. + if search.get("second_level_keep_divisor") == 3: + search["second_level_keep_divisor"] = 1 + applied["search.second_level_keep_divisor"] = 1 + + logger.info( + "TREC funnel -> first_level=%s, second_level=%s, keep_divisor=%s, max_trials_rag=%s " + "(preset applied: %s | from config: %s)", + search.get("max_trials_first_level"), + search.get("max_trials_second_level"), + search.get("second_level_keep_divisor"), + rag.get("max_trials_rag"), + applied or "none", + respected or "none", + ) # No per-topic HTML report: the eval consumes the run files, not the reports. cfg.setdefault("reporting", {})["emit_html"] = False return cfg diff --git a/tests/test_trec_funnel_config.py b/tests/test_trec_funnel_config.py new file mode 100644 index 00000000..16f55ee6 --- /dev/null +++ b/tests/test_trec_funnel_config.py @@ -0,0 +1,74 @@ +"""The TREC funnel preset must not silently overwrite an explicit config. + +_track_config used to assign the funnel unconditionally, so a config asking for +max_trials_second_level=1000 still ran at 500 and a run configured for a 600-trial shortlist +still got 250. Funnel experiments measured the preset instead of the variable under test. +""" + +from trialmatchai.trec.corpus import TrackSpec +from trialmatchai.trec.runner import _track_config + + +def _spec(tmp_path): + return TrackSpec( + key="23", + id_prefix="trec-2023", + db_path=tmp_path / "search_23", + profile_dir=tmp_path / "profiles", + summary_dir=tmp_path / "summaries", + output_dir=tmp_path / "out", + trec_dir=tmp_path / "trec", + ) + + +def test_preset_applies_when_the_config_is_at_schema_defaults(tmp_path): + """A config that expresses no opinion still gets the TREC funnel.""" + cfg = _track_config( + {"search": {"max_trials_second_level": 100}, "rag": {"max_trials_rag": 20}}, + _spec(tmp_path), + ) + assert cfg["search"]["max_trials_second_level"] == 500 + assert cfg["rag"]["max_trials_rag"] == 250 + assert cfg["search"]["second_level_keep_divisor"] == 1 + + +def test_explicit_second_level_width_is_respected(tmp_path): + """The regression that made the width A/B unrunnable: 1000 was silently forced to 500.""" + cfg = _track_config({"search": {"max_trials_second_level": 1000}, "rag": {}}, _spec(tmp_path)) + assert cfg["search"]["max_trials_second_level"] == 1000 + + +def test_explicit_shortlist_size_is_respected(tmp_path): + """The regression that made the depth run measure nothing: 600 was silently forced to 250.""" + cfg = _track_config({"search": {}, "rag": {"max_trials_rag": 600}}, _spec(tmp_path)) + assert cfg["rag"]["max_trials_rag"] == 600 + + +def test_explicit_keep_divisor_is_respected(tmp_path): + cfg = _track_config({"search": {"second_level_keep_divisor": 5}, "rag": {}}, _spec(tmp_path)) + assert cfg["search"]["second_level_keep_divisor"] == 5 + + +def test_second_level_width_knobs_survive(tmp_path): + """search.second_level is what actually widens retrieval; the preset must not touch it.""" + cfg = _track_config( + {"search": {"second_level": {"per_query_size": 1000, "aggregation_threshold": 0.2}}, "rag": {}}, + _spec(tmp_path), + ) + assert cfg["search"]["second_level"]["per_query_size"] == 1000 + assert cfg["search"]["second_level"]["aggregation_threshold"] == 0.2 + + +def test_track_paths_are_still_swapped_in(tmp_path): + cfg = _track_config({"search": {}, "rag": {}}, _spec(tmp_path)) + assert cfg["search_backend"]["db_path"] == str(tmp_path / "search_23") + assert cfg["paths"]["output_dir"] == str(tmp_path / "out") + assert cfg["query_expansion"]["enabled"] is True + assert cfg["reporting"]["emit_html"] is False + + +def test_base_config_is_not_mutated(tmp_path): + base = {"search": {"max_trials_second_level": 1000}, "rag": {}} + _track_config(base, _spec(tmp_path)) + assert base["search"] == {"max_trials_second_level": 1000} + assert base["rag"] == {} From cb94e4d17a25e0e264d2dcf3103cce0532fd3921 Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Thu, 13 Aug 2026 18:42:10 +0200 Subject: [PATCH 12/13] refactor(search): fold broader categories into keywords.json, drop llm_expansion, widen the funnel llm_expansion was a SECOND llm pass over a summary the FIRST pass had already written. It read keywords.json and handed the same aliases back. Sampled over TREC 2023 patients: patient terms searched already in keywords.json genuinely new trec-20231 15 12 Glaucoma, Open-Angle Glaucoma, Ocular Hypertension trec-202310 8 5 GAD, Mental Disorders, Neurobehavioral Manifestations trec-202311 17 15 Lung Diseases, Respiratory Tract Diseases trec-202312 15 13 Lung Disease, Respiratory Disease TOTAL 55 45 (82%) 10 (18%) 82% duplication, and every genuinely new term is a BROADER CATEGORY. Re-searching terms that already drive primary_condition (weight 1.0) and other_condition (0.25) adds nothing -- RRF just reinforces trials already found. So the categories are generated by the one expansion call instead, and the module is gone: * query_expansion gains "broader_conditions" (<=8) and "discarded_or_uncertain" (<=20). The latter preserves the one genuinely good behaviour of the deleted prompt: it put "No Prior Cataract Surgery" AND the un-negated "Cataract Surgery" in the discard list, because searching the bare term retrieves exactly the wrong trials. No other channel has that. * The planner routes broader_conditions to the existing broader_disease channel at weight 0.35, replacing hand-coded ontology rollups (kept as fallback when expansion is off). Weight matters: folding broad categories into main_conditions would search them at 1.0 alongside the specific diagnosis and flood the pool. * FirstLevelQueryExpander, build_first_level_expander, LLMQueryExpansion, LLMQueryExpansionBackend, parse_llm_query_expansion, the llm_expansion channel kind and weight, and the llm_expansion_enabled / llm_max_terms knobs are all removed, along with the dead keys in all 14 shipped configs. One LLM call per patient instead of two. Also widens the funnel, per the measured ceilings. First-level recall plateaus at 0.9110 with the pool exhausting at ~2,525 trials, but the stages behind it were far narrower: search.second_level.per_query_size 250 -> 1000 (surfaces 64% of the pool at 0.787 recall, vs 43% / 0.585 at 250) search.max_trials_second_level 1000 -> 2000 (must exceed what aggregation produces or it truncates -- this is the constant that silently capped every run at 500) rag.max_trials_rag 300 -> 500 (TREC 2023 patients carry ~604 judged relevant trials; 250 capped shortlist recall at ~41% regardless of model) max_trials_rag doubles chain-of-thought cost per patient. That is the deliberate trade: the shortlist was the binding constraint on every recall-aware metric. --- src/trialmatchai/config/config.json | 8 +- src/trialmatchai/config/config_1gpu.json | 296 +++++++++--------- src/trialmatchai/config/config_a40.json | 4 +- src/trialmatchai/config/config_h100.json | 4 +- src/trialmatchai/config/config_l40.json | 292 +++++++++-------- .../config/config_medcpt_baichuanm2.json | 4 +- .../config/config_medcpt_episteme.json | 4 +- .../config/config_medcpt_huatuo.json | 4 +- .../config/config_medcpt_iimedical.json | 4 +- .../config/config_medcpt_medgemma.json | 4 +- .../config/config_medcpt_qwen36.json | 13 +- .../config/config_medcpt_qwen36_l40.json | 13 +- .../config/config_medcpt_qwen3med.json | 2 - .../config/config_nemotron_medgemma.json | 4 +- src/trialmatchai/config/settings.py | 8 - src/trialmatchai/main.py | 14 - src/trialmatchai/matching/query_expansion.py | 218 ++++--------- .../matching/retrieval/first_level_planner.py | 142 ++------- .../matching/retrieval/trial_retrieval.py | 3 - tests/test_first_level_expansion.py | 182 ----------- tests/test_first_level_planner.py | 30 -- 21 files changed, 407 insertions(+), 846 deletions(-) delete mode 100644 tests/test_first_level_expansion.py diff --git a/src/trialmatchai/config/config.json b/src/trialmatchai/config/config.json index 16d69288..5fa32f93 100644 --- a/src/trialmatchai/config/config.json +++ b/src/trialmatchai/config/config.json @@ -104,10 +104,12 @@ "fusion": "rrf", "rrf_k": 60, "vector_score_threshold": 0.0, - "llm_expansion_enabled": false, - "llm_max_terms": 12, "write_reports": true, - "hard_filters": ["age", "sex", "overall_status"] + "hard_filters": [ + "age", + "sex", + "overall_status" + ] } }, "constraints": { diff --git a/src/trialmatchai/config/config_1gpu.json b/src/trialmatchai/config/config_1gpu.json index 630a64b5..b4417399 100644 --- a/src/trialmatchai/config/config_1gpu.json +++ b/src/trialmatchai/config/config_1gpu.json @@ -1,151 +1,149 @@ { - "entity_extraction": { - "backend": "gliner2", - "model_name": "fastino/gliner2-base-v1", - "model_revision": null, - "schema_path": "entity_schemas/trialmatchai.yaml", - "threshold": 0.8, - "batch_size": 8, - "device": "auto", - "trust_remote_code": false - }, - "concept_linker": { - "enabled": true, - "db_path": "data/concepts", - "table": "concepts", - "accept_threshold": 0.7, - "reject_threshold": 0.5, - "margin": 0.05, - "rerank": "lexical", - "search_limit": 10 - }, - "paths": { - "output_dir": "results", - "trials_json_folder": "data/trials_jsons" - }, - "patient_inputs": { - "raw_dir": "data/patients/raw", - "profile_dir": "data/patients/profiles", - "summary_dir": "data/patients/summaries", - "default_format": "auto", - "strict_validation": false, - "copy_raw": true - }, - "model": { - "base_model": "microsoft/phi-4", - "base_model_revision": null, - "trust_remote_code": false, - "quantization": { - "load_in_4bit": true, - "bnb_4bit_use_double_quant": true, - "bnb_4bit_quant_type": "nf4", - "bnb_4bit_compute_dtype": "float16" - }, - "cot_adapter_path": "models/finetuned_phi_reasoning", - "reranker_model_path": "google/gemma-2-2b-it", - "reranker_model_revision": null, - "reranker_adapter_path": "models/finetuned_gemma2" - }, - "tokenizer": { - "use_fast": true, - "padding_side": "left" - }, - "global": { - "device": 0 - }, - "search_backend": { - "backend": "lancedb", - "db_path": "data/search_a40", - "trials_table": "trials", - "criteria_table": "criteria", - "candidate_limit": 1000 - }, - "registry": { - "source": "clinicaltrials.gov", - "api_base_url": "https://clinicaltrials.gov/api/v2/studies", - "keywords_file": null, - "since_days": 7, - "max_studies": null, - "request_timeout": 30, - "rate_limit_per_second": 2, - "raw_dir": "data/registry/raw", - "manifest_path": "data/registry/manifest.jsonl", - "reports_dir": "data/registry/runs", - "failure_threshold": 0.25 - }, - "embedder": { - "model_name": "BAAI/bge-m3", - "revision": null, - "trust_remote_code": false, - "pooling": "mean", - "max_length": 512, - "batch_size": 32, - "use_gpu": true, - "use_fp16": false, - "normalize": true - }, - "cot": { - "batch_size": 10 - }, - "LLM_reranker": { - "batch_size": 20, - "gpu_memory_utilization": 0.16, - "tensor_parallel_size": 1 - }, - "search": { - "mode": "hybrid", - "vector_score_threshold": 0.5, - "max_trials_first_level": 1000, - "max_trials_second_level": 100, - "first_level": { - "enabled": true, - "max_trials": 1000, - "per_channel_size": 300, - "fusion": "rrf", - "rrf_k": 60, - "vector_score_threshold": 0.0, - "llm_expansion_enabled": false, - "llm_max_terms": 12, - "write_reports": true, - "hard_filters": [ - "age", - "sex", - "overall_status" - ] - } - }, - "constraints": { - "enabled": true, - "score_weight": 0.25, - "llm_extraction_enabled": false, - "unknown_is_neutral": true, - "write_reports": true - }, - "query_expansion": { - "enabled": false, - "backend": null, - "model": null, - "adapter": null, - "max_new_tokens": 2048, - "max_main_conditions": 11, - "max_other_conditions": 50 - }, - "use_cot_reasoning": true, - "rag": { - "batch_size": 4, - "max_trials_rag": 20 - }, - "vllm": { - "batch_size": 100, - "max_new_tokens": 5000, - "temperature": 0.0, - "top_p": 1.0, - "seed": 1234, - "length_bucket": true, - "gpu_memory_utilization": 0.72, - "max_model_len": 8192, - "tensor_parallel_size": 1, - "kv_cache_dtype": "fp8", - "max_num_seqs": 16 + "entity_extraction": { + "backend": "gliner2", + "model_name": "fastino/gliner2-base-v1", + "model_revision": null, + "schema_path": "entity_schemas/trialmatchai.yaml", + "threshold": 0.8, + "batch_size": 8, + "device": "auto", + "trust_remote_code": false + }, + "concept_linker": { + "enabled": true, + "db_path": "data/concepts", + "table": "concepts", + "accept_threshold": 0.7, + "reject_threshold": 0.5, + "margin": 0.05, + "rerank": "lexical", + "search_limit": 10 + }, + "paths": { + "output_dir": "results", + "trials_json_folder": "data/trials_jsons" + }, + "patient_inputs": { + "raw_dir": "data/patients/raw", + "profile_dir": "data/patients/profiles", + "summary_dir": "data/patients/summaries", + "default_format": "auto", + "strict_validation": false, + "copy_raw": true + }, + "model": { + "base_model": "microsoft/phi-4", + "base_model_revision": null, + "trust_remote_code": false, + "quantization": { + "load_in_4bit": true, + "bnb_4bit_use_double_quant": true, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": "float16" + }, + "cot_adapter_path": "models/finetuned_phi_reasoning", + "reranker_model_path": "google/gemma-2-2b-it", + "reranker_model_revision": null, + "reranker_adapter_path": "models/finetuned_gemma2" + }, + "tokenizer": { + "use_fast": true, + "padding_side": "left" + }, + "global": { + "device": 0 + }, + "search_backend": { + "backend": "lancedb", + "db_path": "data/search_a40", + "trials_table": "trials", + "criteria_table": "criteria", + "candidate_limit": 1000 + }, + "registry": { + "source": "clinicaltrials.gov", + "api_base_url": "https://clinicaltrials.gov/api/v2/studies", + "keywords_file": null, + "since_days": 7, + "max_studies": null, + "request_timeout": 30, + "rate_limit_per_second": 2, + "raw_dir": "data/registry/raw", + "manifest_path": "data/registry/manifest.jsonl", + "reports_dir": "data/registry/runs", + "failure_threshold": 0.25 + }, + "embedder": { + "model_name": "BAAI/bge-m3", + "revision": null, + "trust_remote_code": false, + "pooling": "mean", + "max_length": 512, + "batch_size": 32, + "use_gpu": true, + "use_fp16": false, + "normalize": true + }, + "cot": { + "batch_size": 10 + }, + "LLM_reranker": { + "batch_size": 20, + "gpu_memory_utilization": 0.16, + "tensor_parallel_size": 1 + }, + "search": { + "mode": "hybrid", + "vector_score_threshold": 0.5, + "max_trials_first_level": 1000, + "max_trials_second_level": 100, + "first_level": { + "enabled": true, + "max_trials": 1000, + "per_channel_size": 300, + "fusion": "rrf", + "rrf_k": 60, + "vector_score_threshold": 0.0, + "write_reports": true, + "hard_filters": [ + "age", + "sex", + "overall_status" + ] } -} \ No newline at end of file + }, + "constraints": { + "enabled": true, + "score_weight": 0.25, + "llm_extraction_enabled": false, + "unknown_is_neutral": true, + "write_reports": true + }, + "query_expansion": { + "enabled": false, + "backend": null, + "model": null, + "adapter": null, + "max_new_tokens": 2048, + "max_main_conditions": 11, + "max_other_conditions": 50 + }, + "use_cot_reasoning": true, + "rag": { + "batch_size": 4, + "max_trials_rag": 20 + }, + "vllm": { + "batch_size": 100, + "max_new_tokens": 5000, + "temperature": 0.0, + "top_p": 1.0, + "seed": 1234, + "length_bucket": true, + "gpu_memory_utilization": 0.72, + "max_model_len": 8192, + "tensor_parallel_size": 1, + "kv_cache_dtype": "fp8", + "max_num_seqs": 16 + } +} diff --git a/src/trialmatchai/config/config_a40.json b/src/trialmatchai/config/config_a40.json index e55f3bac..88003c5d 100644 --- a/src/trialmatchai/config/config_a40.json +++ b/src/trialmatchai/config/config_a40.json @@ -103,8 +103,6 @@ "fusion": "rrf", "rrf_k": 60, "vector_score_threshold": 0.0, - "llm_expansion_enabled": false, - "llm_max_terms": 12, "write_reports": true, "hard_filters": [ "age", @@ -145,4 +143,4 @@ "max_model_len": 6000, "tensor_parallel_size": 1 } -} \ No newline at end of file +} diff --git a/src/trialmatchai/config/config_h100.json b/src/trialmatchai/config/config_h100.json index 5b3e921e..152e1e83 100644 --- a/src/trialmatchai/config/config_h100.json +++ b/src/trialmatchai/config/config_h100.json @@ -103,8 +103,6 @@ "fusion": "rrf", "rrf_k": 60, "vector_score_threshold": 0.0, - "llm_expansion_enabled": false, - "llm_max_terms": 12, "write_reports": true, "hard_filters": [ "age", @@ -145,4 +143,4 @@ "max_model_len": 8192, "tensor_parallel_size": 1 } -} \ No newline at end of file +} diff --git a/src/trialmatchai/config/config_l40.json b/src/trialmatchai/config/config_l40.json index cec79f7a..236c35af 100644 --- a/src/trialmatchai/config/config_l40.json +++ b/src/trialmatchai/config/config_l40.json @@ -1,149 +1,147 @@ { - "entity_extraction": { - "backend": "gliner2", - "model_name": "fastino/gliner2-base-v1", - "model_revision": null, - "schema_path": "entity_schemas/trialmatchai.yaml", - "threshold": 0.8, - "batch_size": 8, - "device": "auto", - "trust_remote_code": false - }, - "concept_linker": { - "enabled": true, - "db_path": "data/concepts", - "table": "concepts", - "accept_threshold": 0.7, - "reject_threshold": 0.5, - "margin": 0.05, - "rerank": "lexical", - "search_limit": 10 - }, - "paths": { - "output_dir": "results", - "trials_json_folder": "data/trials_jsons" - }, - "patient_inputs": { - "raw_dir": "data/patients/raw", - "profile_dir": "data/patients/profiles", - "summary_dir": "data/patients/summaries", - "default_format": "auto", - "strict_validation": false, - "copy_raw": true - }, - "model": { - "base_model": "microsoft/phi-4", - "base_model_revision": null, - "trust_remote_code": false, - "quantization": { - "load_in_4bit": true, - "bnb_4bit_use_double_quant": true, - "bnb_4bit_quant_type": "nf4", - "bnb_4bit_compute_dtype": "float16" - }, - "cot_adapter_path": "models/finetuned_phi_reasoning", - "reranker_model_path": "google/gemma-2-2b-it", - "reranker_model_revision": null, - "reranker_adapter_path": "models/finetuned_gemma2" - }, - "tokenizer": { - "use_fast": true, - "padding_side": "left" - }, - "global": { - "device": 0 - }, - "search_backend": { - "backend": "lancedb", - "db_path": "data/search_a40", - "trials_table": "trials", - "criteria_table": "criteria", - "candidate_limit": 1000 - }, - "registry": { - "source": "clinicaltrials.gov", - "api_base_url": "https://clinicaltrials.gov/api/v2/studies", - "keywords_file": null, - "since_days": 7, - "max_studies": null, - "request_timeout": 30, - "rate_limit_per_second": 2, - "raw_dir": "data/registry/raw", - "manifest_path": "data/registry/manifest.jsonl", - "reports_dir": "data/registry/runs", - "failure_threshold": 0.25 - }, - "embedder": { - "model_name": "BAAI/bge-m3", - "revision": null, - "trust_remote_code": false, - "pooling": "mean", - "max_length": 512, - "batch_size": 32, - "use_gpu": true, - "use_fp16": false, - "normalize": true - }, - "cot": { - "batch_size": 10 - }, - "LLM_reranker": { - "batch_size": 20, - "gpu_memory_utilization": 0.15, - "tensor_parallel_size": 2 - }, - "search": { - "mode": "hybrid", - "vector_score_threshold": 0.5, - "max_trials_first_level": 1000, - "max_trials_second_level": 100, - "first_level": { - "enabled": true, - "max_trials": 1000, - "per_channel_size": 300, - "fusion": "rrf", - "rrf_k": 60, - "vector_score_threshold": 0.0, - "llm_expansion_enabled": false, - "llm_max_terms": 12, - "write_reports": true, - "hard_filters": [ - "age", - "sex", - "overall_status" - ] - } - }, - "constraints": { - "enabled": true, - "score_weight": 0.25, - "llm_extraction_enabled": false, - "unknown_is_neutral": true, - "write_reports": true - }, - "query_expansion": { - "enabled": false, - "backend": null, - "model": null, - "adapter": null, - "max_new_tokens": 2048, - "max_main_conditions": 11, - "max_other_conditions": 50 - }, - "use_cot_reasoning": true, - "rag": { - "batch_size": 4, - "max_trials_rag": 20 - }, - "vllm": { - "batch_size": 100, - "max_new_tokens": 5000, - "temperature": 0.0, - "top_p": 1.0, - "seed": 1234, - "length_bucket": true, - "gpu_memory_utilization": 0.6, - "max_model_len": 8192, - "tensor_parallel_size": 2 + "entity_extraction": { + "backend": "gliner2", + "model_name": "fastino/gliner2-base-v1", + "model_revision": null, + "schema_path": "entity_schemas/trialmatchai.yaml", + "threshold": 0.8, + "batch_size": 8, + "device": "auto", + "trust_remote_code": false + }, + "concept_linker": { + "enabled": true, + "db_path": "data/concepts", + "table": "concepts", + "accept_threshold": 0.7, + "reject_threshold": 0.5, + "margin": 0.05, + "rerank": "lexical", + "search_limit": 10 + }, + "paths": { + "output_dir": "results", + "trials_json_folder": "data/trials_jsons" + }, + "patient_inputs": { + "raw_dir": "data/patients/raw", + "profile_dir": "data/patients/profiles", + "summary_dir": "data/patients/summaries", + "default_format": "auto", + "strict_validation": false, + "copy_raw": true + }, + "model": { + "base_model": "microsoft/phi-4", + "base_model_revision": null, + "trust_remote_code": false, + "quantization": { + "load_in_4bit": true, + "bnb_4bit_use_double_quant": true, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": "float16" + }, + "cot_adapter_path": "models/finetuned_phi_reasoning", + "reranker_model_path": "google/gemma-2-2b-it", + "reranker_model_revision": null, + "reranker_adapter_path": "models/finetuned_gemma2" + }, + "tokenizer": { + "use_fast": true, + "padding_side": "left" + }, + "global": { + "device": 0 + }, + "search_backend": { + "backend": "lancedb", + "db_path": "data/search_a40", + "trials_table": "trials", + "criteria_table": "criteria", + "candidate_limit": 1000 + }, + "registry": { + "source": "clinicaltrials.gov", + "api_base_url": "https://clinicaltrials.gov/api/v2/studies", + "keywords_file": null, + "since_days": 7, + "max_studies": null, + "request_timeout": 30, + "rate_limit_per_second": 2, + "raw_dir": "data/registry/raw", + "manifest_path": "data/registry/manifest.jsonl", + "reports_dir": "data/registry/runs", + "failure_threshold": 0.25 + }, + "embedder": { + "model_name": "BAAI/bge-m3", + "revision": null, + "trust_remote_code": false, + "pooling": "mean", + "max_length": 512, + "batch_size": 32, + "use_gpu": true, + "use_fp16": false, + "normalize": true + }, + "cot": { + "batch_size": 10 + }, + "LLM_reranker": { + "batch_size": 20, + "gpu_memory_utilization": 0.15, + "tensor_parallel_size": 2 + }, + "search": { + "mode": "hybrid", + "vector_score_threshold": 0.5, + "max_trials_first_level": 1000, + "max_trials_second_level": 100, + "first_level": { + "enabled": true, + "max_trials": 1000, + "per_channel_size": 300, + "fusion": "rrf", + "rrf_k": 60, + "vector_score_threshold": 0.0, + "write_reports": true, + "hard_filters": [ + "age", + "sex", + "overall_status" + ] } -} \ No newline at end of file + }, + "constraints": { + "enabled": true, + "score_weight": 0.25, + "llm_extraction_enabled": false, + "unknown_is_neutral": true, + "write_reports": true + }, + "query_expansion": { + "enabled": false, + "backend": null, + "model": null, + "adapter": null, + "max_new_tokens": 2048, + "max_main_conditions": 11, + "max_other_conditions": 50 + }, + "use_cot_reasoning": true, + "rag": { + "batch_size": 4, + "max_trials_rag": 20 + }, + "vllm": { + "batch_size": 100, + "max_new_tokens": 5000, + "temperature": 0.0, + "top_p": 1.0, + "seed": 1234, + "length_bucket": true, + "gpu_memory_utilization": 0.6, + "max_model_len": 8192, + "tensor_parallel_size": 2 + } +} diff --git a/src/trialmatchai/config/config_medcpt_baichuanm2.json b/src/trialmatchai/config/config_medcpt_baichuanm2.json index dd76567b..1489fee8 100644 --- a/src/trialmatchai/config/config_medcpt_baichuanm2.json +++ b/src/trialmatchai/config/config_medcpt_baichuanm2.json @@ -108,8 +108,6 @@ "fusion": "rrf", "rrf_k": 60, "vector_score_threshold": 0.0, - "llm_expansion_enabled": false, - "llm_max_terms": 12, "write_reports": true, "hard_filters": [ "age", @@ -157,4 +155,4 @@ "enforce_eager": true, "max_num_seqs": 256 } -} \ No newline at end of file +} diff --git a/src/trialmatchai/config/config_medcpt_episteme.json b/src/trialmatchai/config/config_medcpt_episteme.json index 97eaa898..93046d07 100644 --- a/src/trialmatchai/config/config_medcpt_episteme.json +++ b/src/trialmatchai/config/config_medcpt_episteme.json @@ -107,8 +107,6 @@ "fusion": "rrf", "rrf_k": 60, "vector_score_threshold": 0.0, - "llm_expansion_enabled": false, - "llm_max_terms": 12, "write_reports": true, "hard_filters": [ "age", @@ -153,4 +151,4 @@ "disable_custom_all_reduce": true, "enforce_eager": true } -} \ No newline at end of file +} diff --git a/src/trialmatchai/config/config_medcpt_huatuo.json b/src/trialmatchai/config/config_medcpt_huatuo.json index e7b2018d..c7205d6a 100644 --- a/src/trialmatchai/config/config_medcpt_huatuo.json +++ b/src/trialmatchai/config/config_medcpt_huatuo.json @@ -108,8 +108,6 @@ "fusion": "rrf", "rrf_k": 60, "vector_score_threshold": 0.0, - "llm_expansion_enabled": false, - "llm_max_terms": 12, "write_reports": true, "hard_filters": [ "age", @@ -159,4 +157,4 @@ "max_num_seqs": 128 }, "reembed_index": false -} \ No newline at end of file +} diff --git a/src/trialmatchai/config/config_medcpt_iimedical.json b/src/trialmatchai/config/config_medcpt_iimedical.json index 69aed1cf..fe433b0a 100644 --- a/src/trialmatchai/config/config_medcpt_iimedical.json +++ b/src/trialmatchai/config/config_medcpt_iimedical.json @@ -108,8 +108,6 @@ "fusion": "rrf", "rrf_k": 60, "vector_score_threshold": 0.0, - "llm_expansion_enabled": false, - "llm_max_terms": 12, "write_reports": true, "hard_filters": [ "age", @@ -159,4 +157,4 @@ "max_num_seqs": 256 }, "reembed_index": false -} \ No newline at end of file +} diff --git a/src/trialmatchai/config/config_medcpt_medgemma.json b/src/trialmatchai/config/config_medcpt_medgemma.json index 12721001..6bb949bd 100644 --- a/src/trialmatchai/config/config_medcpt_medgemma.json +++ b/src/trialmatchai/config/config_medcpt_medgemma.json @@ -107,8 +107,6 @@ "fusion": "rrf", "rrf_k": 60, "vector_score_threshold": 0.0, - "llm_expansion_enabled": false, - "llm_max_terms": 12, "write_reports": true, "hard_filters": [ "age", @@ -152,4 +150,4 @@ "disable_custom_all_reduce": true, "enforce_eager": true } -} \ No newline at end of file +} diff --git a/src/trialmatchai/config/config_medcpt_qwen36.json b/src/trialmatchai/config/config_medcpt_qwen36.json index 20e4825d..25aef2ae 100644 --- a/src/trialmatchai/config/config_medcpt_qwen36.json +++ b/src/trialmatchai/config/config_medcpt_qwen36.json @@ -100,7 +100,7 @@ "mode": "hybrid", "vector_score_threshold": 0.5, "max_trials_first_level": 2000, - "max_trials_second_level": 1000, + "max_trials_second_level": 2000, "first_level": { "enabled": true, "max_trials": 2000, @@ -108,14 +108,17 @@ "fusion": "rrf", "rrf_k": 60, "vector_score_threshold": 0.0, - "llm_expansion_enabled": false, - "llm_max_terms": 12, "write_reports": true, "hard_filters": [ "age", "sex", "overall_status" ] + }, + "second_level": { + "per_query_size": 1000, + "aggregation_threshold": 0.5, + "aggregation_method": "weighted" } }, "constraints": { @@ -139,7 +142,7 @@ "use_cot_reasoning": true, "rag": { "batch_size": 4, - "max_trials_rag": 300, + "max_trials_rag": 500, "guided_json": true, "no_think": true }, @@ -159,4 +162,4 @@ "max_num_seqs": 256 }, "reembed_index": false -} \ No newline at end of file +} diff --git a/src/trialmatchai/config/config_medcpt_qwen36_l40.json b/src/trialmatchai/config/config_medcpt_qwen36_l40.json index f20e7ea9..e8d4fcad 100644 --- a/src/trialmatchai/config/config_medcpt_qwen36_l40.json +++ b/src/trialmatchai/config/config_medcpt_qwen36_l40.json @@ -100,7 +100,7 @@ "mode": "hybrid", "vector_score_threshold": 0.5, "max_trials_first_level": 2000, - "max_trials_second_level": 1000, + "max_trials_second_level": 2000, "first_level": { "enabled": true, "max_trials": 2000, @@ -108,14 +108,17 @@ "fusion": "rrf", "rrf_k": 60, "vector_score_threshold": 0.0, - "llm_expansion_enabled": false, - "llm_max_terms": 12, "write_reports": true, "hard_filters": [ "age", "sex", "overall_status" ] + }, + "second_level": { + "per_query_size": 1000, + "aggregation_threshold": 0.5, + "aggregation_method": "weighted" } }, "constraints": { @@ -139,7 +142,7 @@ "use_cot_reasoning": true, "rag": { "batch_size": 4, - "max_trials_rag": 300, + "max_trials_rag": 500, "guided_json": true, "no_think": true }, @@ -159,4 +162,4 @@ "max_num_seqs": 256 }, "reembed_index": false -} \ No newline at end of file +} diff --git a/src/trialmatchai/config/config_medcpt_qwen3med.json b/src/trialmatchai/config/config_medcpt_qwen3med.json index e8726a96..6dc3dd98 100644 --- a/src/trialmatchai/config/config_medcpt_qwen3med.json +++ b/src/trialmatchai/config/config_medcpt_qwen3med.json @@ -107,8 +107,6 @@ "fusion": "rrf", "rrf_k": 60, "vector_score_threshold": 0.0, - "llm_expansion_enabled": false, - "llm_max_terms": 12, "write_reports": true, "hard_filters": [ "age", diff --git a/src/trialmatchai/config/config_nemotron_medgemma.json b/src/trialmatchai/config/config_nemotron_medgemma.json index 59640984..0e55efd1 100644 --- a/src/trialmatchai/config/config_nemotron_medgemma.json +++ b/src/trialmatchai/config/config_nemotron_medgemma.json @@ -109,8 +109,6 @@ "fusion": "rrf", "rrf_k": 60, "vector_score_threshold": 0.0, - "llm_expansion_enabled": false, - "llm_max_terms": 12, "write_reports": true, "hard_filters": [ "age", @@ -155,4 +153,4 @@ "enforce_eager": true }, "reembed_index": true -} \ No newline at end of file +} diff --git a/src/trialmatchai/config/settings.py b/src/trialmatchai/config/settings.py index ae82d564..de4f79a2 100644 --- a/src/trialmatchai/config/settings.py +++ b/src/trialmatchai/config/settings.py @@ -165,8 +165,6 @@ class FirstLevelSearchSettings(BaseModel): fusion: Literal["rrf"] = "rrf" rrf_k: int = Field(60, ge=1) vector_score_threshold: float = Field(0.0, ge=0.0, le=1.0) - llm_expansion_enabled: bool = False - llm_max_terms: int = Field(12, ge=0) write_reports: bool = True # "location" is opt-in (country-level, site-aware); not in the default set. hard_filters: list[Literal["age", "sex", "overall_status", "location"]] = Field( @@ -477,11 +475,6 @@ def apply_env_overrides(raw: Dict[str, Any]) -> Dict[str, Any]: "TRIALMATCHAI_CONSTRAINTS_WRITE_REPORTS": ("constraints", "write_reports"), "TRIALMATCHAI_QUERY_EXPANSION_ENABLED": ("query_expansion", "enabled"), "TRIALMATCHAI_FIRST_LEVEL_ENABLED": ("search", "first_level", "enabled"), - "TRIALMATCHAI_FIRST_LEVEL_LLM_EXPANSION_ENABLED": ( - "search", - "first_level", - "llm_expansion_enabled", - ), "TRIALMATCHAI_FIRST_LEVEL_WRITE_REPORTS": ( "search", "first_level", @@ -521,7 +514,6 @@ def apply_env_overrides(raw: Dict[str, Any]) -> Dict[str, Any]: "TRIALMATCHAI_FIRST_LEVEL_LLM_MAX_TERMS": ( "search", "first_level", - "llm_max_terms", ), "TRIALMATCHAI_SEARCH_MAX_TRIALS_SECOND_LEVEL": ( "search", diff --git a/src/trialmatchai/main.py b/src/trialmatchai/main.py index 0a499fd5..7a0cafae 100644 --- a/src/trialmatchai/main.py +++ b/src/trialmatchai/main.py @@ -15,7 +15,6 @@ rank_trials, save_ranked_trials, ) -from trialmatchai.matching.query_expansion import build_first_level_expander from trialmatchai.matching.shortlist_depth import choose_shortlist_depth, depth_report from trialmatchai.matching.retrieval.trial_retrieval import ClinicalTrialSearch from trialmatchai.matching.retrieval.criteria_retrieval import SecondStageRetriever @@ -109,7 +108,6 @@ def run_first_level_search( config: Dict, search_backend, patient_profile: PatientProfile | None = None, - llm_query_expander=None, ) -> Optional[Tuple]: main_conditions = list(keywords.get("main_conditions", [])) other_conditions = list(keywords.get("other_conditions", [])) @@ -128,9 +126,6 @@ def run_first_level_search( search_backend=search_backend, embedder=embedder, entity_annotator=entity_annotator, - # Without this the llm_expansion channel is dead: the planner logs "no expander is - # configured" and returns [], however the config flag is set. - llm_query_expander=llm_query_expander, ) search_cfg = config["search"] @@ -559,12 +554,6 @@ def main_pipeline( embedder = build_embedder(config) entity_annotator = build_entity_annotator(config, embedder=embedder) - # Built once for the whole run, not per patient: it shares the cached CoT engine, and - # rebuilding per patient would re-resolve that engine 75 times. None when - # search.first_level.llm_expansion_enabled is off, which is the default. - llm_query_expander = build_first_level_expander(config) - if llm_query_expander is not None: - logger.info("First-level LLM query expansion is ON (llm_expansion channel active).") with warnings.catch_warnings(): warnings.filterwarnings( @@ -672,7 +661,6 @@ def main_pipeline( config, search_backend, patient_profile=profile, - llm_query_expander=llm_query_expander, ) if not result: logger.error("First-level search failed for %s", patient_id) @@ -807,8 +795,6 @@ def _first_level_search_config(search_cfg: Dict) -> Dict: first_level_cfg.setdefault("vector_score_threshold", 0.0) first_level_cfg.setdefault("enabled", True) first_level_cfg.setdefault("write_reports", True) - first_level_cfg.setdefault("llm_expansion_enabled", False) - first_level_cfg.setdefault("llm_max_terms", 12) first_level_cfg.setdefault("hard_filters", ["age", "sex", "overall_status"]) return first_level_cfg diff --git a/src/trialmatchai/matching/query_expansion.py b/src/trialmatchai/matching/query_expansion.py index 3d36c1b6..1df1add1 100644 --- a/src/trialmatchai/matching/query_expansion.py +++ b/src/trialmatchai/matching/query_expansion.py @@ -35,12 +35,32 @@ - Based solely on the original patient-provided data, generate semantically accurate and medically sound statements resembling real-life medical notes. - **Crucial**: Expanded descriptions must strictly reflect explicit patient-reported information without introducing new or inferred medical details. +4. **Broader Disease Categories**: + - Provide up to 8 BROADER categories the primary conditions belong to, from narrower to + wider (e.g. "Primary Open Angle Glaucoma" -> "Open-Angle Glaucoma", "Glaucoma"; + "COPD" -> "Lung Diseases", "Respiratory Tract Diseases"). + - These deliberately trade precision for coverage: a trial may recruit under the category + rather than the specific diagnosis. + - Do NOT repeat the primary conditions or their synonyms here, and do NOT list a + comorbidity as a broader category of the primary condition. + - Provide these in the "broader_conditions" list. + +5. **Discarded or Uncertain**: + - List terms you considered but must NOT be searched, in particular anything the patient + description NEGATES ("no prior chemotherapy", "absence of metastases"). Include BOTH the + negated phrase and its un-negated form, because searching the bare term would retrieve + exactly the wrong trials. + - Also include anything you are not confident the description supports. + - Provide these in the "discarded_or_uncertain" list, and keep them out of every other list. + Output: Return a JSON object in the exact following structure without any additional commentary: { "main_conditions": ["PrimaryCondition", "Synonym1", "Synonym2", "..."], "other_conditions": ["AdditionalCondition1", "AdditionalCondition2", "..."], +"broader_conditions": ["BroaderCategory1", "BroaderCategory2", "..."], +"discarded_or_uncertain": ["NegatedOrUnsupportedTerm1", "..."], "expanded_sentences": [ "Expanded note for sentence 1...", "Expanded note for sentence 2...", @@ -49,7 +69,20 @@ } """.strip() -_EMPTY = {"main_conditions": [], "other_conditions": [], "expanded_sentences": []} +_EMPTY = { + "main_conditions": [], + "other_conditions": [], + # Broader categories were previously produced by a SECOND llm pass (the llm_expansion + # channel). Measured on 4 TREC 2023 patients, 82% of that pass's terms were already in + # keywords.json -- it re-read its own input and handed the aliases back. The only thing + # it contributed was these broader categories, so they are generated here instead and the + # second pass is gone. One LLM call per patient rather than two. + "broader_conditions": [], + # Negated terms that must NOT be searched. "No prior cataract surgery" means searching + # "cataract surgery" retrieves exactly the wrong trials. + "discarded_or_uncertain": [], + "expanded_sentences": [], +} # JSON schema for grammar-constrained keyword expansion (vLLM structured outputs), so a verbose # or reasoning model always returns valid keyword JSON instead of prose that fails to parse. @@ -63,9 +96,17 @@ "properties": { "main_conditions": {"type": "array", "maxItems": 11, "items": {"type": "string", "maxLength": 120}}, "other_conditions": {"type": "array", "maxItems": 50, "items": {"type": "string", "maxLength": 120}}, + "broader_conditions": {"type": "array", "maxItems": 8, "items": {"type": "string", "maxLength": 120}}, + "discarded_or_uncertain": {"type": "array", "maxItems": 20, "items": {"type": "string", "maxLength": 120}}, "expanded_sentences": {"type": "array", "maxItems": 15, "items": {"type": "string"}}, }, - "required": ["main_conditions", "other_conditions", "expanded_sentences"], + "required": [ + "main_conditions", + "other_conditions", + "broader_conditions", + "discarded_or_uncertain", + "expanded_sentences", + ], } @@ -245,180 +286,33 @@ def enrich_summary( *, max_main_conditions: int = 11, max_other_conditions: int = 50, + max_broader_conditions: int = 8, ) -> Dict[str, Any]: """Fold a CoT expansion into a matching summary (legacy keywords.json shape). ``expanded_sentences`` map to ``patient_narrative``; only non-empty fields overwrite, leaving the deterministic summary intact. + + ``broader_conditions`` is kept as its own key rather than merged into main/other, because + the planner routes it to the broader_disease channel at weight 0.35. Folding broad + categories into main_conditions would search them at weight 1.0, flooding the pool with + loosely related trials. """ out = dict(summary) main = expansion.get("main_conditions") or [] other = expansion.get("other_conditions") or [] + broader = expansion.get("broader_conditions") or [] + discarded = expansion.get("discarded_or_uncertain") or [] sentences = expansion.get("expanded_sentences") or [] if main: out["main_conditions"] = main[:max_main_conditions] if other: out["other_conditions"] = other[:max_other_conditions] + if broader: + out["broader_conditions"] = broader[:max_broader_conditions] + if discarded: + # Recorded for provenance and so downstream stages can avoid them; never searched. + out["discarded_or_uncertain"] = discarded if sentences: out["patient_narrative"] = sentences return out - - -# --- first-level retrieval query expansion (the llm_expansion search channel) ------------ # - -# Distinct from SYSTEM_PROMPT above. That one enriches the patient SUMMARY (conditions and -# narrative sentences). This one writes RETRIEVAL QUERIES: short noun phrases that should -# match trial titles, conditions and eligibility text. The two are not interchangeable -- -# the first-level planner buckets these six fields into weighted query channels. -FIRST_LEVEL_SYSTEM_PROMPT = """ -You expand a patient description into search queries for a clinical trial index. - -Write SHORT NOUN PHRASES that would appear in a trial's title, condition list or -eligibility criteria. Do not write sentences, questions or explanations. - -Fill these six fields: - -1. "primary_queries": the patient's main disease as a trial would name it. Include the - staging or subtype only when the patient description states it. -2. "disease_aliases": other names for that same disease -- synonyms, abbreviations, older - or regional terminology, and the expanded form of any abbreviation. -3. "broader_queries": the parent disease categories a trial might recruit under, from - narrower to wider. These deliberately trade precision for coverage. -4. "biomarker_queries": genes, mutations, fusions, receptor and expression status, and - other molecular markers stated for this patient. -5. "treatment_queries": drugs, drug classes, procedures and prior therapies stated for - this patient. -6. "discarded_or_uncertain": terms you considered but rejected, and anything you are not - confident the patient description supports. - -Rules: -- Use ONLY what the patient description states. Never infer a diagnosis, stage, biomarker - or therapy that is not written there. Put anything doubtful in "discarded_or_uncertain". -- Leave a field as an empty list when the description supports nothing for it. An empty - list is correct; an invented term is not. -- No duplicates within a field. - -Return a JSON object with exactly those six keys and no other commentary. -""".strip() - -_FIRST_LEVEL_FIELDS = ( - "primary_queries", - "disease_aliases", - "broader_queries", - "biomarker_queries", - "treatment_queries", - "discarded_or_uncertain", -) - -# maxItems bounds the array COUNT for the same reason as _KEYWORDS_JSON_SCHEMA: it forces a -# verbose model to close each array instead of emitting terms until max_tokens runs out. -# These are noun phrases, so a short maxLength is safe here (unlike expanded_sentences). -# -# The per-field caps are deliberately uneven. search.first_level.llm_max_terms is a SHARED -# budget across the five query fields, spent in field order (first_level_planner -# parse_llm_query_expansion), so a model that pads primary_queries starves the biomarker and -# treatment channels entirely. Capping primary_queries tightly -- a patient has one main -# disease, not twelve -- keeps the budget available for the later fields. -_FIRST_LEVEL_MAX_ITEMS = { - "primary_queries": 3, - "disease_aliases": 8, - "broader_queries": 5, - "biomarker_queries": 8, - "treatment_queries": 8, - "discarded_or_uncertain": 12, -} -_FIRST_LEVEL_JSON_SCHEMA = { - "type": "object", - "properties": { - field: { - "type": "array", - "maxItems": _FIRST_LEVEL_MAX_ITEMS[field], - "items": {"type": "string", "maxLength": 120}, - } - for field in _FIRST_LEVEL_FIELDS - }, - "required": list(_FIRST_LEVEL_FIELDS), -} - -_FIRST_LEVEL_EMPTY: Dict[str, List[str]] = {field: [] for field in _FIRST_LEVEL_FIELDS} - - -def _first_level_patient_text(profile: Any, matching_summary: Dict[str, Any]) -> str: - """Compact patient description for the expander prompt. - - Built from the matching summary rather than the raw profile so it stays in step with - what first-level retrieval actually searches on. - """ - summary = matching_summary or {} - parts: List[str] = [] - main = [c for c in _as_list(summary.get("main_conditions")) if c][:12] - other = [c for c in _as_list(summary.get("other_conditions")) if c][:30] - narrative = [s for s in _as_list(summary.get("patient_narrative")) if s][:12] - if main: - parts.append("Main conditions: " + "; ".join(main)) - if other: - parts.append("Other conditions and factors: " + "; ".join(other)) - age, gender = summary.get("age"), summary.get("gender") - demographics = [ - f"Age: {age}" for _ in (1,) if age not in (None, "", "all") - ] + [f"Sex: {gender}" for _ in (1,) if gender not in (None, "", "all")] - if demographics: - parts.append(", ".join(demographics)) - if narrative: - parts.append("Description: " + " ".join(narrative)) - return "\n".join(parts).strip() - - -class FirstLevelQueryExpander(QueryExpander): - """Implements ``LLMQueryExpansionBackend`` for the first-level ``llm_expansion`` channel. - - Reuses QueryExpander's engine, chat-template and structured-output machinery -- so it - shares the one cached vLLM engine rather than loading a second copy -- but swaps in the - retrieval-query prompt and schema. - """ - - system_prompt = FIRST_LEVEL_SYSTEM_PROMPT - json_schema = _FIRST_LEVEL_JSON_SCHEMA - - def expand_first_level_queries( - self, - *, - profile: Any, - matching_summary: Dict[str, Any], - ) -> Dict[str, Any]: - patient_text = _first_level_patient_text(profile, matching_summary) - if not patient_text: - return dict(_FIRST_LEVEL_EMPTY) - try: - raw = self._generate(patient_text) - parsed = extract_json_object(BaseTrialProcessor._strip_thinking_tags(raw)) - if not isinstance(parsed, dict): - raise ValueError("first-level expansion output was not a JSON object") - return {field: _as_list(parsed.get(field)) for field in _FIRST_LEVEL_FIELDS} - except Exception as exc: - # Retrieval must not fail because expansion did: the channel is one of eight and - # carries weight 0.5, so an empty expansion degrades recall rather than the run. - logger.error( - "First-level query expansion failed; continuing without that channel: %s", exc - ) - return dict(_FIRST_LEVEL_EMPTY) - - -def build_first_level_expander(config: Dict[str, Any]) -> "FirstLevelQueryExpander | None": - """Construct the expander when ``search.first_level.llm_expansion_enabled`` is true. - - Independent of ``query_expansion.enabled``: that flag governs the separate summary - enrichment stage. Both may run, and they share one engine. - """ - first_level = (config.get("search") or {}).get("first_level") or {} - if not first_level.get("llm_expansion_enabled"): - return None - try: - return FirstLevelQueryExpander(_resolve_settings(config), config) - except Exception as exc: - logger.error( - "search.first_level.llm_expansion_enabled is set but the expander could not be " - "built; first-level search continues without that channel: %s", - exc, - ) - return None diff --git a/src/trialmatchai/matching/retrieval/first_level_planner.py b/src/trialmatchai/matching/retrieval/first_level_planner.py index add771c2..437f64a6 100644 --- a/src/trialmatchai/matching/retrieval/first_level_planner.py +++ b/src/trialmatchai/matching/retrieval/first_level_planner.py @@ -1,9 +1,8 @@ from __future__ import annotations -import json import re -from collections.abc import Callable, Sequence -from typing import Any, Literal, Protocol +from collections.abc import Sequence +from typing import Any, Literal from pydantic import BaseModel, ConfigDict, Field @@ -24,7 +23,6 @@ "narrative", "biomarker", "therapy", - "llm_expansion", ] DEFAULT_CHANNEL_WEIGHTS: dict[FirstLevelChannelKind, float] = { @@ -37,7 +35,6 @@ "biomarker": 0.7, "therapy": 0.45, "broader_disease": 0.35, - "llm_expansion": 0.5, } # One focused channel per comorbidity: a blended query over many conditions dilutes BM25 and the @@ -58,7 +55,6 @@ class FirstLevelQueryPlan(BaseModel): patient_id: str | None = None channels: list[FirstLevelQueryChannel] = Field(default_factory=list) filters: dict[str, Any] = Field(default_factory=dict) - llm_expansion_enabled: bool = False model_config = ConfigDict(extra="forbid") @@ -80,38 +76,9 @@ class FirstLevelCandidateEvidence(BaseModel): model_config = ConfigDict(extra="forbid") -class LLMQueryExpansion(BaseModel): - primary_queries: list[str] = Field(default_factory=list) - disease_aliases: list[str] = Field(default_factory=list) - broader_queries: list[str] = Field(default_factory=list) - biomarker_queries: list[str] = Field(default_factory=list) - treatment_queries: list[str] = Field(default_factory=list) - discarded_or_uncertain: list[str] = Field(default_factory=list) - - model_config = ConfigDict(extra="forbid") - - -class LLMQueryExpansionBackend(Protocol): - def expand_first_level_queries( - self, - *, - profile: PatientProfile, - matching_summary: dict[str, Any], - ) -> str | dict[str, Any]: - ... - - class FirstLevelQueryPlanner: - def __init__( - self, - *, - entity_annotator: Any = None, - llm_expander: LLMQueryExpansionBackend - | Callable[..., str | dict[str, Any]] - | None = None, - ) -> None: + def __init__(self, *, entity_annotator: Any = None) -> None: self.entity_annotator = entity_annotator - self.llm_expander = llm_expander def build( self, @@ -148,8 +115,33 @@ def build( synonym_terms = self._synonym_terms(primary_terms) channels.append(self._channel("concept_synonym", synonym_terms, "concept_linker")) - broader_terms = broader_disease_terms([*primary_terms, *synonym_terms]) - channels.append(self._channel("broader_disease", broader_terms, "deterministic")) + # Prefer the broader categories the expansion step generated; fall back to the + # hand-coded ontology rollups when expansion is off or returned none. + # + # These used to come from a SECOND llm pass (the llm_expansion channel, weight 0.5). + # Measured over sampled TREC 2023 patients, 82% of that pass's terms were already in + # keywords.json -- it read the enriched summary and handed the same aliases back, so + # it re-searched what primary_condition (1.0) and other_condition (0.25) already + # covered. Broader categories were the ONLY thing it contributed, so they are + # generated by the one expansion call instead and routed here at 0.35. Searching a + # broad category at 1.0 alongside the specific diagnosis would flood the pool. + summary_broader = dedupe_terms( + [ + str(term) + for term in (matching_summary.get("broader_conditions") or []) + if str(term).strip() + ] + ) + broader_terms = summary_broader or broader_disease_terms( + [*primary_terms, *synonym_terms] + ) + channels.append( + self._channel( + "broader_disease", + broader_terms, + "query_expansion" if summary_broader else "deterministic", + ) + ) narrative_terms = narrative_terms_from_summary_or_profile( profile, @@ -168,15 +160,6 @@ def build( therapy_terms = therapy_terms_from_profile(profile) channels.append(self._channel("therapy", therapy_terms, "patient_profile")) - llm_enabled = bool(cfg.get("llm_expansion_enabled", False)) - if llm_enabled: - llm_terms = self._llm_terms( - profile=profile, - matching_summary=matching_summary, - max_terms=int(cfg.get("llm_max_terms", 12)), - ) - channels.append(self._channel("llm_expansion", llm_terms, "llm")) - return FirstLevelQueryPlan( patient_id=profile.patient_id, channels=[channel for channel in channels if channel.terms], @@ -189,7 +172,6 @@ def build( ["age", "sex", "overall_status"], ), }, - llm_expansion_enabled=llm_enabled, ) def _channel( @@ -211,70 +193,6 @@ def _synonym_terms(self, primary_terms: Sequence[str]) -> list[str]: output.extend(disease_synonyms(self.entity_annotator, term)) return dedupe_terms(output) - def _llm_terms( - self, - *, - profile: PatientProfile, - matching_summary: dict[str, Any], - max_terms: int, - ) -> list[str]: - if self.llm_expander is None: - logger.warning( - "First-level LLM expansion is enabled but no expander is configured." - ) - return [] - try: - if hasattr(self.llm_expander, "expand_first_level_queries"): - raw = self.llm_expander.expand_first_level_queries( - profile=profile, - matching_summary=matching_summary, - ) - else: - raw = self.llm_expander( - profile=profile, - matching_summary=matching_summary, - ) - parsed = parse_llm_query_expansion(raw, max_terms=max_terms) - except Exception: - logger.exception("Discarding invalid first-level LLM query expansion.") - return [] - return dedupe_terms( - [ - *parsed.primary_queries, - *parsed.disease_aliases, - *parsed.broader_queries, - *parsed.biomarker_queries, - *parsed.treatment_queries, - ] - )[:max_terms] - - -def parse_llm_query_expansion( - raw: str | dict[str, Any], - *, - max_terms: int, -) -> LLMQueryExpansion: - payload = json.loads(raw) if isinstance(raw, str) else raw - parsed = LLMQueryExpansion.model_validate(payload) - capped: dict[str, list[str]] = {} - remaining = max(0, max_terms) - for field in ( - "primary_queries", - "disease_aliases", - "broader_queries", - "biomarker_queries", - "treatment_queries", - ): - values = dedupe_terms(getattr(parsed, field)) - if remaining <= 0: - capped[field] = [] - continue - capped[field] = values[:remaining] - remaining -= len(capped[field]) - capped["discarded_or_uncertain"] = dedupe_terms(parsed.discarded_or_uncertain) - return LLMQueryExpansion.model_validate(capped) - - def fuse_first_level_channel_hits( channel_hits: Sequence[tuple[FirstLevelQueryChannel, list[dict], list[float]]], *, diff --git a/src/trialmatchai/matching/retrieval/trial_retrieval.py b/src/trialmatchai/matching/retrieval/trial_retrieval.py index ab0c8b0a..37a2afe6 100644 --- a/src/trialmatchai/matching/retrieval/trial_retrieval.py +++ b/src/trialmatchai/matching/retrieval/trial_retrieval.py @@ -10,7 +10,6 @@ FirstLevelCandidateEvidence, FirstLevelQueryPlan, FirstLevelQueryPlanner, - LLMQueryExpansionBackend, fuse_first_level_channel_hits, ) from trialmatchai.matching.retrieval.synonyms import disease_synonyms @@ -30,14 +29,12 @@ def __init__( search_backend: TrialSearchBackend, embedder: Optional[TextEmbedder], entity_annotator=None, - llm_query_expander: LLMQueryExpansionBackend | None = None, ): self.search_backend = search_backend self.embedder = embedder self.entity_annotator = entity_annotator self.query_planner = FirstLevelQueryPlanner( entity_annotator=entity_annotator, - llm_expander=llm_query_expander, ) def get_synonyms(self, condition: str) -> List[str]: diff --git a/tests/test_first_level_expansion.py b/tests/test_first_level_expansion.py deleted file mode 100644 index e7a900da..00000000 --- a/tests/test_first_level_expansion.py +++ /dev/null @@ -1,182 +0,0 @@ -"""First-level LLM query expansion (the llm_expansion search channel). - -This channel was dead before: the protocol, parser, schema and config flag all existed, -but nothing ever constructed a backend, so enabling the flag logged "no expander is -configured" and returned no terms. These tests cover the backend and the config gate. -""" - -import json - -import pytest - -from trialmatchai.matching.query_expansion import ( - _FIRST_LEVEL_FIELDS, - FirstLevelQueryExpander, - _first_level_patient_text, - build_first_level_expander, -) -from trialmatchai.matching.retrieval.first_level_planner import parse_llm_query_expansion - -SUMMARY = { - "main_conditions": ["metastatic breast cancer"], - "other_conditions": ["hypertension", "HER2 positive"], - "patient_narrative": ["A 54 year old woman with metastatic breast cancer."], - "age": 54, - "gender": "female", -} - - -class _FakeExpander(FirstLevelQueryExpander): - """Bypasses __init__ so no model or GPU is touched; _generate returns a canned reply.""" - - def __init__(self, reply): - self._reply = reply - self.settings = {"guided_json": True, "max_new_tokens": 512} - self.config = {} - self.backend = "vllm" - - def _generate(self, narrative): - if isinstance(self._reply, Exception): - raise self._reply - return self._reply - - -def test_expands_into_the_six_planner_fields(): - payload = { - "primary_queries": ["metastatic breast cancer"], - "disease_aliases": ["breast carcinoma", "mammary carcinoma"], - "broader_queries": ["solid tumor"], - "biomarker_queries": ["HER2 positive"], - "treatment_queries": ["trastuzumab"], - "discarded_or_uncertain": ["hypertension"], - } - result = _FakeExpander(json.dumps(payload)).expand_first_level_queries( - profile=None, matching_summary=SUMMARY - ) - assert result == payload - assert set(result) == set(_FIRST_LEVEL_FIELDS) - - -def test_output_is_consumable_by_the_planner_parser(): - """The backend's contract is the planner's parser, not just valid JSON.""" - payload = { - "primary_queries": ["metastatic breast cancer"], - "disease_aliases": ["breast carcinoma"], - "broader_queries": ["solid tumor"], - "biomarker_queries": ["HER2 positive"], - "treatment_queries": ["trastuzumab"], - "discarded_or_uncertain": [], - } - raw = _FakeExpander(json.dumps(payload)).expand_first_level_queries( - profile=None, matching_summary=SUMMARY - ) - parsed = parse_llm_query_expansion(raw, max_terms=12) - assert parsed.primary_queries == ["metastatic breast cancer"] - assert parsed.biomarker_queries == ["HER2 positive"] - - -def test_max_terms_is_a_shared_budget_spent_primary_first(): - """llm_max_terms caps the TOTAL across the five query fields, not each one, and is spent - in field order. A model that fills primary_queries can starve the later channels, so the - prompt must keep primary_queries to the actual disease rather than padding it.""" - payload = {field: [] for field in _FIRST_LEVEL_FIELDS} - payload["primary_queries"] = [f"q{i}" for i in range(5)] - payload["disease_aliases"] = ["alias1", "alias2"] - payload["biomarker_queries"] = ["EGFR"] - - parsed = parse_llm_query_expansion(payload, max_terms=6) - - assert parsed.primary_queries == [f"q{i}" for i in range(5)] - assert parsed.disease_aliases == ["alias1"] # only one slot left - assert parsed.biomarker_queries == [] # budget exhausted before this field - - -def test_reasoning_tags_are_stripped_before_json_extraction(): - """Reasoning models emit containing an echo of the schema; extracting from that - would return the schema instead of the answer.""" - payload = {field: [] for field in _FIRST_LEVEL_FIELDS} - payload["primary_queries"] = ["glioblastoma"] - reply = ( - "The schema wants primary_queries, disease_aliases, ..." - + json.dumps(payload) - ) - result = _FakeExpander(reply).expand_first_level_queries( - profile=None, matching_summary=SUMMARY - ) - assert result["primary_queries"] == ["glioblastoma"] - - -@pytest.mark.parametrize( - "reply", - ["not json at all", json.dumps(["a", "list"]), RuntimeError("engine died")], -) -def test_failures_degrade_to_empty_not_raise(reply): - """Retrieval must survive a failed expansion: this is 1 of 8 channels, weight 0.5.""" - result = _FakeExpander(reply).expand_first_level_queries( - profile=None, matching_summary=SUMMARY - ) - assert result == {field: [] for field in _FIRST_LEVEL_FIELDS} - - -def test_a_bare_string_field_is_not_shredded_into_characters(): - payload = {field: [] for field in _FIRST_LEVEL_FIELDS} - payload["primary_queries"] = "glioblastoma" - result = _FakeExpander(json.dumps(payload)).expand_first_level_queries( - profile=None, matching_summary=SUMMARY - ) - assert result["primary_queries"] == ["glioblastoma"] - - -def test_empty_summary_skips_the_model_entirely(): - expander = _FakeExpander(RuntimeError("must not be called")) - assert expander.expand_first_level_queries(profile=None, matching_summary={}) == { - field: [] for field in _FIRST_LEVEL_FIELDS - } - - -def test_patient_text_includes_conditions_and_demographics(): - text = _first_level_patient_text(None, SUMMARY) - assert "metastatic breast cancer" in text - assert "HER2 positive" in text - assert "54" in text and "female" in text - - -def test_patient_text_omits_placeholder_demographics(): - text = _first_level_patient_text( - None, {"main_conditions": ["asthma"], "age": "all", "gender": "all"} - ) - assert "asthma" in text - assert "Age:" not in text and "Sex:" not in text - - -def test_builder_returns_none_unless_the_flag_is_set(): - assert build_first_level_expander({}) is None - assert build_first_level_expander({"search": {"first_level": {}}}) is None - assert ( - build_first_level_expander( - {"search": {"first_level": {"llm_expansion_enabled": False}}} - ) - is None - ) - - -def test_builder_degrades_to_none_when_construction_fails(): - """A misconfigured expander must not abort the run; the channel just stays empty.""" - config = { - "search": {"first_level": {"llm_expansion_enabled": True}}, - "model": {}, # no base_model -> QueryExpander raises - "query_expansion": {}, - } - assert build_first_level_expander(config) is None - - -def test_schema_caps_primary_queries_tightly_to_protect_the_shared_budget(): - """Guards the interaction pinned above: primary_queries is spent first out of - llm_max_terms, so its schema cap must leave room for the later channels.""" - from trialmatchai.matching.query_expansion import _FIRST_LEVEL_JSON_SCHEMA - - props = _FIRST_LEVEL_JSON_SCHEMA["properties"] - primary = props["primary_queries"]["maxItems"] - assert primary <= 3 - for field in ("biomarker_queries", "treatment_queries", "disease_aliases"): - assert props[field]["maxItems"] > primary diff --git a/tests/test_first_level_planner.py b/tests/test_first_level_planner.py index f5320829..f5440100 100644 --- a/tests/test_first_level_planner.py +++ b/tests/test_first_level_planner.py @@ -1,12 +1,10 @@ import json -import pytest from trialmatchai.interop.models import ClinicalFact, PatientProfile, Provenance from trialmatchai.main import run_first_level_search from trialmatchai.matching.retrieval.first_level_planner import ( FirstLevelQueryPlanner, - parse_llm_query_expansion, ) from trialmatchai.matching.retrieval.trial_retrieval import ClinicalTrialSearch from trialmatchai.search import InMemorySearchBackend @@ -33,7 +31,6 @@ def test_planner_builds_deterministic_channels_and_skips_negated_facts(): "other_conditions": [], "patient_narrative": ["Patient has EGFR-mutated lung cancer."], }, - config={"llm_expansion_enabled": False}, age=64, sex="female", overall_status="All", @@ -63,7 +60,6 @@ def test_planner_builds_per_condition_other_condition_channels(): ], "patient_narrative": ["x"], }, - config={"llm_expansion_enabled": False}, ) oc = [c for c in plan.channels if c.kind == "other_condition"] assert len(oc) == 2 # one channel per distinct non-primary comorbidity @@ -82,35 +78,10 @@ def test_planner_no_other_condition_channel_when_empty(): "other_conditions": [], "patient_narrative": ["x"], }, - config={"llm_expansion_enabled": False}, ) assert [c for c in plan.channels if c.kind == "other_condition"] == [] -def test_llm_query_expansion_is_strict_and_capped(): - parsed = parse_llm_query_expansion( - { - "primary_queries": ["lung cancer", "lung cancer"], - "disease_aliases": ["NSCLC"], - "broader_queries": ["solid tumor"], - "biomarker_queries": ["EGFR mutation"], - "treatment_queries": ["osimertinib"], - "discarded_or_uncertain": ["random drift"], - }, - max_terms=3, - ) - - assert parsed.primary_queries == ["lung cancer"] - assert parsed.disease_aliases == ["NSCLC"] - assert parsed.broader_queries == ["solid tumor"] - assert parsed.biomarker_queries == [] - assert parsed.treatment_queries == [] - with pytest.raises(Exception): - parse_llm_query_expansion("{bad json", max_terms=3) - with pytest.raises(Exception): - parse_llm_query_expansion({"primary_queries": [], "extra": []}, max_terms=3) - - def test_planned_search_fuses_multi_channel_hits_above_single_channel_hits(): backend = InMemorySearchBackend( trials=[ @@ -360,7 +331,6 @@ def _config(*, enabled: bool) -> dict: "per_channel_size": 300, "rrf_k": 60, "vector_score_threshold": 0.0, - "llm_expansion_enabled": False, "write_reports": True, }, } From 00f7629343661dd34901549c63c6c4817fb26f4e Mon Sep 17 00:00:00 2001 From: Majd Abdallah Date: Thu, 13 Aug 2026 21:01:46 +0200 Subject: [PATCH 13/13] chore: untrack benchmarks/, ignore experiment outputs, add analysis scripts benchmarks/ holds local embedder measurements rather than source, and should not have been committed. Untracked with --cached so all 11 files stay on disk; the directory is now gitignored, as are the per-run analysis outputs scripts/ writes to the repo root (candidate-id dumps run to megabytes) and the experiment result trees. Adds the two analysis scripts that produced this session's retrieval findings: llm_expansion_ab.py runs one arm at a time and PERSISTS each arm's per-patient candidate lists, so an arm is never re-run to answer a later question and a crash keeps completed patients. Both properties came from real failures: a depth sweep discarded its id lists and had to be regenerated, and an earlier A/B lost a finished arm when the job died. It also must run as a file rather than a heredoc -- once the embedder initialises CUDA, vLLM falls back to spawn multiprocessing, and spawn re-imports __main__, which fails for a script piped to `python -`. show_llm_expansions.py prints what the expander actually searched for. Inspecting four patients with it is what showed 82% of llm_expansion's terms were already in keywords.json, which is why that module was removed rather than tuned. --- .gitignore | 21 ++ benchmarks/embedders/bge-m3-vw0.6.json | 55 ----- benchmarks/embedders/medcpt-vw0.6.json | 55 ----- .../embedders/pubmedbert-neuml-vw0.6.json | 55 ----- benchmarks/embedders/qwen3-0.6b-vw0.6.json | 55 ----- benchmarks/embedders/registry.json | 59 ----- scripts/llm_expansion_ab.py | 218 ++++++++++++++++++ scripts/show_llm_expansions.py | 96 ++++++++ 8 files changed, 335 insertions(+), 279 deletions(-) delete mode 100644 benchmarks/embedders/bge-m3-vw0.6.json delete mode 100644 benchmarks/embedders/medcpt-vw0.6.json delete mode 100644 benchmarks/embedders/pubmedbert-neuml-vw0.6.json delete mode 100644 benchmarks/embedders/qwen3-0.6b-vw0.6.json delete mode 100644 benchmarks/embedders/registry.json create mode 100644 scripts/llm_expansion_ab.py create mode 100644 scripts/show_llm_expansions.py diff --git a/.gitignore b/.gitignore index 33592c9b..2020b510 100644 --- a/.gitignore +++ b/.gitignore @@ -50,3 +50,24 @@ scripts/*.sh .DS_Store Thumbs.db site/ + +# ============================ +# Benchmark + experiment outputs +# ============================ +# Embedder benchmark results are local measurements, not source. Kept on disk, out of git. +benchmarks/ + +# Per-run analysis outputs written to the repo root by scripts/ (candidate-id dumps, sweep +# results, expansion samples). These are measurements of a particular run, often megabytes, +# and are regenerated by re-running the script. +llm_expansion_ab/ +*_sweep_*.json +llm_expansion_samples.json + +# Experiment result trees (one directory per run/arm). +width_ab*/ +rerank_ab/ +shortlist_ab/ +trec23_deep/ +trec23_qwen36_medcpt_v2/ + diff --git a/benchmarks/embedders/bge-m3-vw0.6.json b/benchmarks/embedders/bge-m3-vw0.6.json deleted file mode 100644 index 8a5a912e..00000000 --- a/benchmarks/embedders/bge-m3-vw0.6.json +++ /dev/null @@ -1,55 +0,0 @@ -{ - "embedder": "bge-m3", - "model_name": "BAAI/bge-m3", - "query_model_name": null, - "per_channel_size": 600, - "max_trials": 2000, - "similarity": "cosine", - "vector_weight": 0.6, - "tracks": { - "21": { - "n_patients_grade2": 75, - "n_patients_grade1and2": 75, - "recall_grade2": { - "recall@10": 0.0719, - "recall@100": 0.3869, - "recall@300": 0.6714, - "recall@500": 0.7837, - "recall@700": 0.8372, - "recall@1000": 0.8778, - "recall@2000": 0.9282 - }, - "recall_grade1and2": { - "recall@10": 0.0803, - "recall@100": 0.4129, - "recall@300": 0.682, - "recall@500": 0.7919, - "recall@700": 0.8417, - "recall@1000": 0.8771, - "recall@2000": 0.9251 - } - }, - "22": { - "n_patients_grade2": 50, - "n_patients_grade1and2": 50, - "recall_grade2": { - "recall@10": 0.0864, - "recall@100": 0.4745, - "recall@300": 0.7054, - "recall@500": 0.7971, - "recall@700": 0.8405, - "recall@1000": 0.8837, - "recall@2000": 0.9313 - }, - "recall_grade1and2": { - "recall@10": 0.0761, - "recall@100": 0.4401, - "recall@300": 0.6876, - "recall@500": 0.782, - "recall@700": 0.8291, - "recall@1000": 0.871, - "recall@2000": 0.9143 - } - } - } -} \ No newline at end of file diff --git a/benchmarks/embedders/medcpt-vw0.6.json b/benchmarks/embedders/medcpt-vw0.6.json deleted file mode 100644 index f627b0cf..00000000 --- a/benchmarks/embedders/medcpt-vw0.6.json +++ /dev/null @@ -1,55 +0,0 @@ -{ - "embedder": "medcpt", - "model_name": "ncbi/MedCPT-Article-Encoder", - "query_model_name": "ncbi/MedCPT-Query-Encoder", - "per_channel_size": 600, - "max_trials": 2000, - "similarity": "dot", - "vector_weight": 0.6, - "tracks": { - "21": { - "n_patients_grade2": 75, - "n_patients_grade1and2": 75, - "recall_grade2": { - "recall@10": 0.077, - "recall@100": 0.4227, - "recall@300": 0.726, - "recall@500": 0.8374, - "recall@700": 0.885, - "recall@1000": 0.9149, - "recall@2000": 0.9485 - }, - "recall_grade1and2": { - "recall@10": 0.0781, - "recall@100": 0.4315, - "recall@300": 0.7212, - "recall@500": 0.8323, - "recall@700": 0.8777, - "recall@1000": 0.9095, - "recall@2000": 0.9445 - } - }, - "22": { - "n_patients_grade2": 50, - "n_patients_grade1and2": 50, - "recall_grade2": { - "recall@10": 0.0964, - "recall@100": 0.5057, - "recall@300": 0.7492, - "recall@500": 0.8349, - "recall@700": 0.8724, - "recall@1000": 0.9027, - "recall@2000": 0.9404 - }, - "recall_grade1and2": { - "recall@10": 0.0819, - "recall@100": 0.4697, - "recall@300": 0.7235, - "recall@500": 0.8159, - "recall@700": 0.8548, - "recall@1000": 0.8846, - "recall@2000": 0.9268 - } - } - } -} \ No newline at end of file diff --git a/benchmarks/embedders/pubmedbert-neuml-vw0.6.json b/benchmarks/embedders/pubmedbert-neuml-vw0.6.json deleted file mode 100644 index 558029ed..00000000 --- a/benchmarks/embedders/pubmedbert-neuml-vw0.6.json +++ /dev/null @@ -1,55 +0,0 @@ -{ - "embedder": "pubmedbert-neuml", - "model_name": "NeuML/pubmedbert-base-embeddings", - "query_model_name": null, - "per_channel_size": 600, - "max_trials": 2000, - "similarity": "cosine", - "vector_weight": 0.6, - "tracks": { - "21": { - "n_patients_grade2": 75, - "n_patients_grade1and2": 75, - "recall_grade2": { - "recall@10": 0.0785, - "recall@100": 0.4168, - "recall@300": 0.7077, - "recall@500": 0.8143, - "recall@700": 0.8647, - "recall@1000": 0.9034, - "recall@2000": 0.9425 - }, - "recall_grade1and2": { - "recall@10": 0.0827, - "recall@100": 0.4314, - "recall@300": 0.7106, - "recall@500": 0.8182, - "recall@700": 0.8653, - "recall@1000": 0.9012, - "recall@2000": 0.9403 - } - }, - "22": { - "n_patients_grade2": 50, - "n_patients_grade1and2": 50, - "recall_grade2": { - "recall@10": 0.0949, - "recall@100": 0.4934, - "recall@300": 0.7233, - "recall@500": 0.8213, - "recall@700": 0.8526, - "recall@1000": 0.8938, - "recall@2000": 0.933 - }, - "recall_grade1and2": { - "recall@10": 0.0802, - "recall@100": 0.4625, - "recall@300": 0.7088, - "recall@500": 0.8082, - "recall@700": 0.8433, - "recall@1000": 0.8794, - "recall@2000": 0.9207 - } - } - } -} \ No newline at end of file diff --git a/benchmarks/embedders/qwen3-0.6b-vw0.6.json b/benchmarks/embedders/qwen3-0.6b-vw0.6.json deleted file mode 100644 index c89f30b7..00000000 --- a/benchmarks/embedders/qwen3-0.6b-vw0.6.json +++ /dev/null @@ -1,55 +0,0 @@ -{ - "embedder": "qwen3-0.6b", - "model_name": "Qwen/Qwen3-Embedding-0.6B", - "query_model_name": null, - "per_channel_size": 600, - "max_trials": 2000, - "similarity": "cosine", - "vector_weight": 0.6, - "tracks": { - "21": { - "n_patients_grade2": 75, - "n_patients_grade1and2": 75, - "recall_grade2": { - "recall@10": 0.0692, - "recall@100": 0.3969, - "recall@300": 0.6798, - "recall@500": 0.7956, - "recall@700": 0.8512, - "recall@1000": 0.887, - "recall@2000": 0.9337 - }, - "recall_grade1and2": { - "recall@10": 0.0786, - "recall@100": 0.4243, - "recall@300": 0.6955, - "recall@500": 0.8006, - "recall@700": 0.8532, - "recall@1000": 0.8853, - "recall@2000": 0.9306 - } - }, - "22": { - "n_patients_grade2": 50, - "n_patients_grade1and2": 50, - "recall_grade2": { - "recall@10": 0.0913, - "recall@100": 0.4762, - "recall@300": 0.7211, - "recall@500": 0.8024, - "recall@700": 0.845, - "recall@1000": 0.8802, - "recall@2000": 0.9241 - }, - "recall_grade1and2": { - "recall@10": 0.0796, - "recall@100": 0.4461, - "recall@300": 0.7025, - "recall@500": 0.7888, - "recall@700": 0.8344, - "recall@1000": 0.8696, - "recall@2000": 0.9119 - } - } - } -} \ No newline at end of file diff --git a/benchmarks/embedders/registry.json b/benchmarks/embedders/registry.json deleted file mode 100644 index 4525a547..00000000 --- a/benchmarks/embedders/registry.json +++ /dev/null @@ -1,59 +0,0 @@ -{ - "_about": "Embedder models benchmarked for first-level trial retrieval. Each entry is a config['embedder'] block. A distinct query_model_name marks an asymmetric dual-encoder (separate query/document encoders sharing one space). Results live in benchmarks/embedders/.json; concept-linking is held fixed on bge-m3 so only the trial-retrieval embedder varies.", - "bge-m3": { - "backend": "hf", - "model_name": "BAAI/bge-m3", - "pooling": "mean", - "max_length": 512, - "normalize": true, - "_note": "Current default. General multilingual embedder, symmetric, 1024-dim." - }, - "medcpt": { - "backend": "hf", - "model_name": "ncbi/MedCPT-Article-Encoder", - "query_model_name": "ncbi/MedCPT-Query-Encoder", - "pooling": "cls", - "max_length": 512, - "query_max_length": 64, - "normalize": false, - "similarity": "dot", - "_note": "Biomedical dual-encoder trained on PubMed query-article pairs (the retriever TrialGPT uses). Asymmetric, 768-dim; article encoder for trials, query encoder for patient queries. Uses its native DOT-PRODUCT similarity on unnormalized embeddings (magnitude carries relevance)." - }, - "medcpt-cosine": { - "backend": "hf", - "model_name": "ncbi/MedCPT-Article-Encoder", - "query_model_name": "ncbi/MedCPT-Query-Encoder", - "pooling": "cls", - "max_length": 512, - "query_max_length": 64, - "normalize": true, - "similarity": "cosine", - "_note": "MedCPT forced into the pipeline's cosine metric (normalized). Kept for the record -- this handicaps MedCPT vs its native dot-product; compare to 'medcpt'." - }, - "qwen3-0.6b": { - "backend": "hf", - "model_name": "Qwen/Qwen3-Embedding-0.6B", - "pooling": "last", - "max_length": 512, - "normalize": true, - "trust_remote_code": true, - "query_instruction": "Given a patient's clinical profile, retrieve clinical trials the patient may be eligible for", - "_note": "SOTA general embedder (Qwen3, 2025). Decoder-based: last-token pooling + query instruction (documents raw); 1024-dim, cosine." - }, - "biolord-2023": { - "backend": "hf", - "model_name": "FremyCompany/BioLORD-2023", - "pooling": "mean", - "max_length": 512, - "normalize": true, - "_note": "Biomedical sentence-embedder (BioLORD-2023). Mean pooling, 768-dim, symmetric, cosine." - }, - "pubmedbert-neuml": { - "backend": "hf", - "model_name": "NeuML/pubmedbert-base-embeddings", - "pooling": "mean", - "max_length": 512, - "normalize": true, - "_note": "PubMedBERT fine-tuned for embeddings (NeuML). Mean pooling, 768-dim, symmetric, cosine." - } -} diff --git a/scripts/llm_expansion_ab.py b/scripts/llm_expansion_ab.py new file mode 100644 index 00000000..52775e2f --- /dev/null +++ b/scripts/llm_expansion_ab.py @@ -0,0 +1,218 @@ +#!/usr/bin/env python +"""Does the llm_expansion channel lift the first-level recall ceiling? + +The depth sweep established the ceiling: recall plateaus at 0.9110 and the candidate pool +exhausts at ~2,525 trials, so ~9% of judged-relevant trials are never retrieved and no amount +of depth reaches them. New QUERY TERMS are one of the few things that can. + +llm_expansion is a ninth retrieval channel (weight 0.5) whose terms come from an LLM reading +the patient -- disease aliases, broader categories, biomarker and treatment phrasings. Unlike +the other eight channels it is not limited to terms already in the record or derivable by +lookup. + +Run one arm at a time; each arm PERSISTS its per-patient candidate lists, so an arm is never +re-run to answer a later question. When both arms' lists exist the comparison is computed +from disk at no cost -- including the number that decides whether iterating is worth it: +how many trials the new terms surface that the other channels never found, and how many of +those are relevant. + + python scripts/llm_expansion_ab.py --arm off + python scripts/llm_expansion_ab.py --arm on + python scripts/llm_expansion_ab.py --compare + +Must be run as a FILE, not piped to `python -`: vLLM falls back to spawn multiprocessing once +CUDA is initialised, and spawn re-imports __main__, which fails for a stdin script. +""" + +from __future__ import annotations + +import argparse +import json +import tempfile +from pathlib import Path + +DEPTH = 4000 # where first-level recall plateaus; identical in both arms +PER_CHANNEL = 600 +MAX_TERMS = 24 # 12 is tight: the budget is shared across five fields and spent in order +OUT_DIR = Path("llm_expansion_ab") +SUMMARIES = Path("data/patients/trec23/summaries") +PROFILES = Path("data/patients/trec23/profiles") +BASE_CFG = "src/trialmatchai/config/config_medcpt_qwen36_l40.json" +QRELS = Path("data/trec/qrels/qrels_23.txt") + + +def _base_config(): + from trialmatchai.config.config_loader import load_config + + cfg = load_config(BASE_CFG) + cfg.setdefault("search_backend", {})["db_path"] = "data/search_medcpt_23" + cfg.setdefault("embedder", {})["use_gpu"] = True + cfg.setdefault("concept_linker", {})["db_path"] = "data/concepts_medcpt" + return cfg + + +def _arm_config(base, enabled: bool): + cfg = json.loads(json.dumps(base)) + first_level = dict(cfg["search"].get("first_level", {})) + first_level.update( + { + "max_trials": DEPTH, + "per_channel_size": PER_CHANNEL, + "write_reports": False, + "llm_expansion_enabled": enabled, + "llm_max_terms": MAX_TERMS, + } + ) + cfg["search"]["first_level"] = first_level + cfg["search"]["max_trials_first_level"] = DEPTH + return cfg + + +def _patients(relevant): + from trialmatchai.interop.models import PatientProfile + + out = [] + for summary_path in sorted(SUMMARIES.glob("*.json")): + pid = summary_path.stem + profile_path = PROFILES / f"{pid}.json" + if pid not in relevant or not relevant[pid] or not profile_path.exists(): + continue + try: + out.append( + ( + pid, + json.loads(summary_path.read_text()), + PatientProfile.model_validate_json(profile_path.read_text()), + ) + ) + except Exception as exc: # a malformed profile must not abort the arm + print(f" skip {pid}: {exc}") + return out + + +def run_arm(arm: str) -> None: + from trialmatchai.entities import build_entity_annotator + from trialmatchai.main import run_first_level_search + from trialmatchai.matching.query_expansion import build_first_level_expander + from trialmatchai.models.embedding.text_embedder import build_embedder + from trialmatchai.search import build_search_backend + from trialmatchai.trec.qrels import parse_qrels, relevant_ncts + + enabled = arm == "on" + base = _base_config() + cfg = _arm_config(base, enabled) + + backend = build_search_backend(base) + embedder = build_embedder(base) + annotator = build_entity_annotator(base, embedder=embedder) + + expander = None + if enabled: + expander = build_first_level_expander(cfg) + if expander is None: + raise SystemExit("llm_expansion enabled but the expander could not be built") + + relevant = relevant_ncts(parse_qrels(QRELS, "trec-2023"), threshold=1) + patients = _patients(relevant) + print(f"arm {arm}: {len(patients)} patients, depth={DEPTH}, max_terms={MAX_TERMS}", flush=True) + + OUT_DIR.mkdir(exist_ok=True) + dest = OUT_DIR / f"{arm}_ids.json" + # Persist incrementally: a crash in patient 30 must not discard the first 29. + ids_by_patient = json.loads(dest.read_text()) if dest.exists() else {} + + for pid, summary, profile in patients: + if pid in ids_by_patient: + continue + with tempfile.TemporaryDirectory() as tmp: + try: + result = run_first_level_search( + summary, + tmp, + {"age": summary.get("age", "all"), "gender": summary.get("gender", "all")}, + annotator, + embedder, + cfg, + backend, + patient_profile=profile, + llm_query_expander=expander, + ) + except Exception as exc: + print(f" {pid} failed: {exc}", flush=True) + continue + if not result: + continue + ids_by_patient[pid] = result[0] + dest.write_text(json.dumps(ids_by_patient)) + print(f" {pid}: {len(result[0])}", flush=True) + + print(f"arm {arm} done -> {dest} ({len(ids_by_patient)} patients)") + + +def compare() -> None: + from trialmatchai.trec.qrels import parse_qrels, recall_at_k, relevant_ncts + + qrels = parse_qrels(QRELS, "trec-2023") + relevant = relevant_ncts(qrels, threshold=1) + eligible = relevant_ncts(qrels, threshold=2) + + arms = {} + for arm in ("off", "on"): + path = OUT_DIR / f"{arm}_ids.json" + if not path.exists(): + print(f"missing {path} — run: python scripts/llm_expansion_ab.py --arm {arm}") + return + arms[arm] = json.loads(path.read_text()) + + common = sorted(set(arms["off"]) & set(arms["on"])) + print(f"\npaired patients: {len(common)}\n") + print(f" {'arm':>5s} {'retrieved':>10s} {'recall(rel)':>12s} {'recall(elig)':>13s}") + stats = {} + for arm in ("off", "on"): + rec = [recall_at_k(arms[arm][p], relevant[p], len(arms[arm][p])) for p in common] + el = [ + recall_at_k(arms[arm][p], eligible[p], len(arms[arm][p])) + for p in common + if eligible.get(p) + ] + size = [len(arms[arm][p]) for p in common] + stats[arm] = (sum(size) / len(size), sum(rec) / len(rec), sum(el) / len(el) if el else float("nan")) + print(f" {arm:>5s} {stats[arm][0]:10.0f} {stats[arm][1]:12.4f} {stats[arm][2]:13.4f}") + + print( + f"\n DELTA (on-off): retrieved {stats['on'][0] - stats['off'][0]:+.0f} " + f"recall(rel) {stats['on'][1] - stats['off'][1]:+.4f} " + f"recall(elig) {stats['on'][2] - stats['off'][2]:+.4f}" + ) + + # The number that decides whether ITERATING is worth building. + new_total = new_rel = 0 + wins = 0 + for pid in common: + fresh = set(arms["on"][pid]) - set(arms["off"][pid]) + hit = fresh & relevant[pid] + new_total += len(fresh) + new_rel += len(hit) + wins += bool(hit) + print( + f"\n surfaced ONLY by llm_expansion: {new_total} trials, {new_rel} judged-relevant " + f"({100 * new_rel / max(1, new_total):.1f}% precision)" + ) + print(f" per patient: {new_total / len(common):.0f} new, {new_rel / len(common):.1f} relevant") + print(f" patients where it found >=1 new relevant trial: {wins}/{len(common)}") + print("\n That per-round yield is the stopping signal an iterative expander would use.") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--arm", choices=("off", "on")) + parser.add_argument("--compare", action="store_true") + args = parser.parse_args() + if args.arm: + run_arm(args.arm) + if args.compare or not args.arm: + compare() + + +if __name__ == "__main__": + main() diff --git a/scripts/show_llm_expansions.py b/scripts/show_llm_expansions.py new file mode 100644 index 00000000..3b931ec0 --- /dev/null +++ b/scripts/show_llm_expansions.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python +"""Print the actual terms llm_expansion produces, for inspection. + +The A/B measures whether the channel lifts recall; this shows WHAT it is searching for, which +is what tells you whether the terms are clinically sensible or whether the channel is quietly +adding noise at weight 0.5. + +Prints, per patient: the matching summary the expander reads, the six fields it returns, and +the flattened term list the planner actually turns into a query channel (capped at +llm_max_terms and spent in field order, so the cap can starve the later fields). + + python scripts/show_llm_expansions.py --n 3 +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +SUMMARIES = Path("data/patients/trec23/summaries") +PROFILES = Path("data/patients/trec23/profiles") +BASE_CFG = "src/trialmatchai/config/config_medcpt_qwen36_l40.json" +MAX_TERMS = 24 + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--n", type=int, default=3, help="how many patients to show") + parser.add_argument("--out", default="llm_expansion_samples.json") + args = parser.parse_args() + + from trialmatchai.config.config_loader import load_config + from trialmatchai.interop.models import PatientProfile + from trialmatchai.matching.query_expansion import build_first_level_expander + from trialmatchai.matching.retrieval.first_level_planner import parse_llm_query_expansion + + cfg = load_config(BASE_CFG) + first_level = dict(cfg["search"].get("first_level", {})) + first_level.update({"llm_expansion_enabled": True, "llm_max_terms": MAX_TERMS}) + cfg["search"]["first_level"] = first_level + + expander = build_first_level_expander(cfg) + if expander is None: + raise SystemExit("expander could not be built") + + samples = [] + for summary_path in sorted(SUMMARIES.glob("*.json"))[: args.n]: + pid = summary_path.stem + profile_path = PROFILES / f"{pid}.json" + if not profile_path.exists(): + continue + summary = json.loads(summary_path.read_text()) + profile = PatientProfile.model_validate_json(profile_path.read_text()) + + raw = expander.expand_first_level_queries(profile=profile, matching_summary=summary) + parsed = parse_llm_query_expansion(raw, max_terms=MAX_TERMS) + + print("=" * 78) + print(f"PATIENT {pid}") + print("=" * 78) + print(" what the expander reads:") + for key in ("main_conditions", "other_conditions"): + vals = [str(v) for v in (summary.get(key) or [])][:6] + if vals: + print(f" {key}: {'; '.join(vals)}") + narrative = " ".join(str(s) for s in (summary.get("patient_narrative") or []))[:300] + if narrative: + print(f" narrative: {narrative}...") + + print("\n raw expansion (six fields):") + for field, values in (raw.items() if isinstance(raw, dict) else []): + print(f" {field:24s} {values}") + + flat = [ + *parsed.primary_queries, + *parsed.disease_aliases, + *parsed.broader_queries, + *parsed.biomarker_queries, + *parsed.treatment_queries, + ] + print(f"\n -> channel terms actually searched (cap {MAX_TERMS}, weight 0.5): {len(flat)}") + for term in flat: + print(f" {term}") + if parsed.discarded_or_uncertain: + print(f" -> deliberately NOT searched: {parsed.discarded_or_uncertain}") + print() + samples.append({"patient": pid, "raw": raw, "searched": flat, + "discarded": parsed.discarded_or_uncertain}) + + Path(args.out).write_text(json.dumps(samples, indent=2)) + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main()