diff --git a/scripts/compare_shortlist_ab.py b/scripts/compare_shortlist_ab.py new file mode 100644 index 00000000..d3efee6c --- /dev/null +++ b/scripts/compare_shortlist_ab.py @@ -0,0 +1,112 @@ +#!/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" +# 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", + "ndcg_full@10", + "P@10(eligible)", + "P@10(rel>=1)", + "recall@1000", +) +COST = ("shortlist_size",) +# 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: + 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(" 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.") + 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/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..69d03a79 100644 --- a/src/trialmatchai/main.py +++ b/src/trialmatchai/main.py @@ -15,6 +15,8 @@ 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 from trialmatchai.matching.retrieval.location import ( @@ -107,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", [])) @@ -125,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"] @@ -306,16 +312,23 @@ 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( 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 +336,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", @@ -524,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( @@ -618,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/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/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_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 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 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