feat: config-driven model swapping (registries, self-describing embedders, index provenance)

Make embedder / search-backend / reranker / CoT swaps config-only, with no code edits.

- Non-lossy config load: overlay the validated dump onto the raw config so undeclared
  knobs reach their `.get()` consumers instead of being dropped by model_dump; a new
  runtime knob no longer needs a pydantic field. Revives rag/query_expansion.no_think.
- Plugin registry (trialmatchai.plugins) with build_search_backend / build_embedder
  factories; retires by-name construction across 6 call sites; selection Literals widen
  to str with a fail-fast "registered: [...]" error.
- Self-describing embedders (dim / native_metric / is_asymmetric / pooling / fingerprint);
  the vector metric now follows the embedder (vector_metric defaults to None); fixes the
  asymmetric embed_queries path dropping query_instruction.
- Index embedder-provenance sidecar (_embedder.json) + auto-reembed on an embedder swap;
  pre-provenance indexes are trusted (no rebuild).
- Reranker builds through the shared vLLM loader, inheriting quantization / kv_cache_dtype
  / max_model_len / LoRA-fallback and the shared engine cache; scoring is unchanged.
- Model catalog: `embedder: {model: "medcpt"}` one-line swap resolves the full spec.
- Tighten verbose comments across the package.

Adds 27 tests; full suite green; ruff clean; existing configs and persisted indexes
resolve identically (no behavior drift). Bumps version 0.5.0 -> 0.6.0.

Author: majdabd

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "trialmatchai"
-version = "0.5.0"
+version = "0.6.0"
 description = "AI-driven patient-to-clinical-trial matching: hybrid retrieval + LLM eligibility reasoning."
 readme = "README.md"
 requires-python = ">=3.11,<3.12"

scripts/benchmark_embedder.py

@@ -116,8 +116,13 @@ def main() -> None:
     ap.add_argument("--out", required=True)
     ap.add_argument("--reuse-index", action="store_true", help="Reuse an existing per-track index instead of rebuilding.")
     ap.add_argument("--vector-weight", type=float, default=0.5, help="Hybrid blend: score = (1-w)*text + w*vector (default 0.5).")
+    ap.add_argument("--cutoffs", default=None, help="Comma-separated recall@k cutoffs (default 100,200,300,500,700,1000,2000).")
     args = ap.parse_args()
 
+    if args.cutoffs:
+        global KS
+        KS = [int(x) for x in args.cutoffs.split(",")]
+
     registry = json.loads(Path(args.registry).read_text())
     if args.embedder not in registry:
         raise SystemExit(f"Unknown embedder '{args.embedder}'. Known: {[k for k in registry if not k.startswith('_')]}")

src/trialmatchai/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.5.0"
+__version__ = "0.6.0"

src/trialmatchai/cli/build.py

@@ -76,8 +76,7 @@ def main() -> int:
         return 1
 
     # Build the concept store BEFORE prepare/index so build_system can link entities
-    # between prepare and index, giving the search tables concept IDs instead of
-    # leaving every entity at concept_store_unavailable.
+    # in between, giving search tables concept IDs instead of concept_store_unavailable.
     link_concepts = bool(args.concepts or args.concepts_csv)
     if link_concepts:
         logger.info("=== build: concepts stage ===")

src/trialmatchai/cli/build_concepts.py

@@ -117,9 +117,8 @@ def _concepts_fingerprint(
     embedder_revision: str | None = None,
     skip_embeddings: bool = False,
 ) -> str:
-    # Embedder identity must be in the fingerprint: stored vectors must share the
-    # query-time embedding space, so a model swap (or toggling embeddings) invalidates
-    # the store. Mirrors _prepare_signature in orchestration.py.
+    # Embedder identity must be in the fingerprint: stored vectors must match the
+    # query-time embedding space. Mirrors _prepare_signature in orchestration.py.
     return digest(
         _CONCEPTS_STATE_VERSION,
         sources,
@@ -179,9 +178,8 @@ def run_build_concepts(
     if not force and not concept_csv:
         ready, rows_present = _concept_table_ready(db_path, table_name)
         recorded = _read_concepts_fingerprint(db_path)
-        # Backward-compatible: a pre-marker store (recorded is None) still skips on
-        # presence; only an explicit fingerprint mismatch rebuilds, so a large existing
-        # store is never needlessly re-embedded.
+        # Backward-compatible: a pre-marker store (recorded is None) skips on presence;
+        # only an explicit fingerprint mismatch rebuilds, sparing a needless re-embed.
         if ready and (recorded is None or recorded == concepts_fp):
             logger.info(
                 "Concept store already present at %s/%s (%s concepts); skipping. "

src/trialmatchai/cli/healthcheck.py

@@ -4,7 +4,7 @@
 import sys
 
 from trialmatchai.config.config_loader import load_config
-from trialmatchai.search import LanceDBSearchBackend
+from trialmatchai.search import build_search_backend
 from trialmatchai.services.preflight import run_preflight_checks
 from trialmatchai.utils.logging_config import setup_logging
 
@@ -37,7 +37,7 @@ def main() -> int:
 
     config = load_config(args.config)
     issues = 0
-    search_backend = LanceDBSearchBackend.from_config(config)
+    search_backend = build_search_backend(config)
 
     preflight_issues = run_preflight_checks(
         config,

src/trialmatchai/cli/update_registry.py

@@ -17,7 +17,7 @@
     RegistryUpdater,
     normalize_keywords,
 )
-from trialmatchai.search import InMemorySearchBackend, LanceDBSearchBackend
+from trialmatchai.search import InMemorySearchBackend, build_search_backend
 from trialmatchai.utils.logging_config import setup_logging
 
 logger = setup_logging(__name__)
@@ -118,7 +118,7 @@ def main() -> int:
         timeout=float(registry_cfg.get("request_timeout", 30.0)),
         rate_limit_per_second=float(registry_cfg.get("rate_limit_per_second", 2.0)),
     )
-    backend = InMemorySearchBackend() if args.dry_run else LanceDBSearchBackend.from_config(config)
+    backend = InMemorySearchBackend() if args.dry_run else build_search_backend(config)
     embedder = _NullEmbedder() if args.dry_run else _build_embedder(config)
     entity_annotator = None if args.dry_run else _build_entity_annotator(config, embedder)
 

src/trialmatchai/config/catalog/embedders.json

@@ -0,0 +1,20 @@
+{
+  "bge-m3": {
+    "backend": "hf",
+    "model_name": "BAAI/bge-m3",
+    "pooling": "cls",
+    "normalize": true,
+    "native_metric": "cosine",
+    "max_length": 512
+  },
+  "medcpt": {
+    "backend": "hf",
+    "model_name": "ncbi/MedCPT-Article-Encoder",
+    "query_model_name": "ncbi/MedCPT-Query-Encoder",
+    "pooling": "cls",
+    "normalize": false,
+    "native_metric": "dot",
+    "max_length": 512,
+    "query_max_length": 64
+  }
+}

src/trialmatchai/config/config_loader.py

@@ -35,12 +35,106 @@ def load_config(config_path: str | os.PathLike[str] | None = None) -> Dict[str,
     with resolved_config.open("r", encoding="utf-8") as f:
         raw: Dict[str, Any] = json.load(f)
     raw = apply_env_overrides(deepcopy(raw))
+    raw = _expand_catalog(raw)
     settings = TrialMatchSettings.model_validate(raw)
-    cfg = settings.to_dict()
+    # Non-lossy load: overlay the validated dump ONTO raw rather than replacing it, so declared
+    # fields take validated values (coercions + defaults win) while undeclared knobs survive to
+    # `.get()` consumers instead of being dropped by model_dump (a pure-config knob needs no field).
+    cfg = _deep_merge(raw, settings.to_dict())
+    _warn_unknown_config_keys(raw, settings)
     cfg = normalize_config_paths(cfg, resolved_config)
     return cfg
 
 
+_CATALOGS = {"embedder": "embedders.json"}
+
+
+def _load_catalog(filename: str) -> Dict[str, Any]:
+    path = Path(__file__).resolve().parent / "catalog" / filename
+    if not path.exists():
+        return {}
+    try:
+        return json.loads(path.read_text(encoding="utf-8"))
+    except Exception:
+        logger.warning("Could not read model catalog %s (ignoring).", path)
+        return {}
+
+
+def _expand_catalog(raw: Dict[str, Any]) -> Dict[str, Any]:
+    """Expand a ``model: "<name>"`` shorthand into the full catalog spec (e.g.
+    ``embedder: {model: "medcpt"}`` resolves dim/metric/pooling/templates).
+
+    Resolution order: inline > catalog > code defaults. An unknown name falls back to
+    ``model_name`` so the raw-path form (``{model_name: "org/model"}``) keeps working. Runs before
+    validation and the non-lossy merge, so downstream sees only the expanded section.
+    """
+    for section, catalog_file in _CATALOGS.items():
+        block = raw.get(section)
+        if not isinstance(block, dict):
+            continue
+        name = block.get("model")
+        if not name:
+            continue
+        entry = _load_catalog(catalog_file).get(name)
+        if entry is None:
+            logger.warning(
+                "config[%s].model=%r not in catalog; using it as a raw model_name.", section, name
+            )
+            block.setdefault("model_name", name)
+            block.pop("model", None)
+            continue
+        merged = dict(entry)
+        for key, value in block.items():
+            if key != "model":
+                merged[key] = value  # inline wins over catalog
+        raw[section] = merged
+    return raw
+
+
+def _deep_merge(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]:
+    """Recursively merge ``override`` onto ``base``; ``override`` wins on leaf conflicts.
+
+    Overlays the validated dump onto raw config so validated values win for declared keys while
+    undeclared keys survive. Only dict/dict pairs recurse; other types (lists included) are
+    replaced wholesale by ``override``'s value.
+    """
+    merged = dict(base)
+    for key, over_val in override.items():
+        base_val = merged.get(key)
+        if isinstance(base_val, dict) and isinstance(over_val, dict):
+            merged[key] = _deep_merge(base_val, over_val)
+        else:
+            merged[key] = over_val
+    return merged
+
+
+def _warn_unknown_config_keys(raw: Dict[str, Any], settings: TrialMatchSettings) -> None:
+    """Log config keys that survive passthrough but aren't in the schema, so a typo is visible.
+
+    Non-lossy loading keeps mistyped keys rather than dropping them; this warning replaces the
+    previously silent drop and never removes anything.
+    """
+    fields = type(settings).model_fields
+    # A section is "known" under its field name OR its serialization alias (e.g. ``global_`` is
+    # written ``global`` in JSON), so an aliased section isn't mis-flagged as a typo.
+    known_sections = set(fields) | {f.alias for f in fields.values() if f.alias}
+    for section in raw:
+        if section not in known_sections:
+            logger.warning(
+                "config: unknown top-level section %r (kept but not schema-validated — check for a typo)",
+                section,
+            )
+    for name in fields:
+        sub = getattr(settings, name, None)
+        extra = getattr(sub, "model_extra", None)
+        if extra:
+            logger.warning(
+                "config[%s]: passthrough keys not in schema (kept): %s",
+                name,
+                ", ".join(sorted(extra)),
+            )
+
+
 def resolve_config_path(
     config_path: str | os.PathLike[str] | None = None,
 ) -> Path:

src/trialmatchai/config/settings.py

@@ -66,7 +66,8 @@ class ModelQuantizationSettings(BaseModel):
 class ModelSettings(BaseModel):
     base_model: str
     quantization: ModelQuantizationSettings
-    cot_adapter_path: str
+    # None runs the base CoT model with no LoRA adapter (e.g. MedGemma, already instruction-tuned).
+    cot_adapter_path: str | None = None
     reranker_model_path: str
     reranker_adapter_path: str
     trust_remote_code: bool = False
@@ -84,16 +85,25 @@ class GlobalSettings(BaseModel):
 
 
 class SearchBackendSettings(BaseModel):
-    backend: Literal["lancedb"] = "lancedb"
+    # Resolved via the plugin registry at build_search_backend (raises on a typo), so a plain str
+    # not a Literal that needs widening per backend. ``type`` preferred; ``backend`` is legacy.
+    type: str | None = None
+    backend: str = "lancedb"
     db_path: str = "data/search"
     trials_table: str = "trials"
     criteria_table: str = "criteria"
     candidate_limit: int = Field(1000, ge=1)
-    # Vector similarity for the ANN index and re-ranker. "cosine" (default) suits normalized
-    # embedders (bge-m3); "dot" suits inner-product dual-encoders (e.g. MedCPT).
-    vector_metric: Literal["cosine", "dot"] = "cosine"
+    # ANN index/re-ranker similarity. None follows the embedder (cosine if normalized, dot if not,
+    # e.g. MedCPT); set explicitly to pin it regardless of the embedder.
+    vector_metric: Literal["cosine", "dot"] | None = None
     # Hybrid blend of the first-level score: (1 - vector_weight) * text + vector_weight * vector.
     vector_weight: float = Field(0.5, ge=0.0, le=1.0)
+    # Re-embed the corpus with the config embedder at index-build time (vs reusing prepare-time
+    # vectors). Required to switch embedders: otherwise the new-model query dim-mismatches the old
+    # index vectors and silently falls back to BM25.
+    reembed_index: bool = False
+    # Keep undeclared backend knobs; they reach the tolerant .get() consumer via the non-lossy loader.
+    model_config = ConfigDict(extra="allow")
 
 
 class RegistrySettings(BaseModel):
@@ -111,7 +121,10 @@ class RegistrySettings(BaseModel):
 
 
 class EmbedderSettings(BaseModel):
-    backend: Literal["hf", "hashing"] = "hf"
+    # Embedder family resolved via the plugin registry at build_embedder; plain str (not a Literal)
+    # so a new family needs no schema edit. ``type`` preferred; ``backend`` is legacy (hf | hashing).
+    type: str | None = None
+    backend: str = "hf"
     model_name: str = "BAAI/bge-m3"
     # A distinct query_model_name selects an asymmetric dual-encoder (separate document/query
     # encoders sharing one space, e.g. MedCPT's article/query encoders); None = symmetric.
@@ -128,6 +141,11 @@ class EmbedderSettings(BaseModel):
     hashing_dimensions: int = Field(64, ge=1)
     # Instruction prepended to queries only, for instruction-tuned embedders (e.g. Qwen3-Embedding).
     query_instruction: str | None = None
+    # Vector metric this embedder's space is trained for; None derives from ``normalize`` (cosine
+    # if normalized, dot if not). The search backend reads it so the metric follows the embedder.
+    native_metric: Literal["cosine", "dot"] | None = None
+    # Keep undeclared embedder knobs (they reach the tolerant .get() consumer via the non-lossy loader).
+    model_config = ConfigDict(extra="allow")
 
     @field_validator("pooling")
     @classmethod
@@ -163,11 +181,9 @@ class SearchSettings(BaseModel):
     max_trials_second_level: int = Field(100, ge=1)
     # Keep the top 1/N of reranked second-level trials before CoT (N=1 keeps all).
     second_level_keep_divisor: int = Field(3, ge=1)
-    # How the second-level shortlist combines the first-level (retrieval) and second-level
-    # (reranker) rankings. "rrf" fuses the two by rank, so a strong retrieval hit keeps a
-    # floor and is not evicted when the reranker fails to score its criteria; "score_sum"
-    # is the earlier behaviour of adding the two raw scores (whose scales differ, so one
-    # signal dominates and can drop retrieval-ranked trials).
+    # How the shortlist fuses first-level (retrieval) and second-level (reranker) rankings. "rrf"
+    # fuses by rank, so a strong retrieval hit isn't evicted when the reranker fails to score it;
+    # "score_sum" adds raw scores (mismatched scales let one signal dominate and drop retrieval hits).
     shortlist_fusion: Literal["rrf", "score_sum"] = "rrf"
     shortlist_rrf_k: int = Field(60, ge=1)
     shortlist_first_level_weight: float = Field(1.0, ge=0.0)
@@ -208,6 +224,10 @@ class RagSettings(BaseModel):
     backend: Literal["vllm", "transformers"] = "vllm"
     batch_size: int = Field(4, ge=1)
     max_trials_rag: int = Field(20, ge=1)
+    # Suppress chain-of-thought <think> 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
+    model_config = ConfigDict(extra="allow")
 
 
 class VllmSettings(BaseModel):
@@ -220,10 +240,24 @@ class VllmSettings(BaseModel):
     gpu_memory_utilization: float = Field(0.5, gt=0.0, le=1.0)
     max_model_len: int = Field(8192, ge=256)
     tensor_parallel_size: int = Field(1, ge=1)
+    # In-flight vLLM weight quantization (e.g. "bitsandbytes" NF4 4-bit) so a large base model (32B
+    # CoT) fits one card without a pre-quantized checkpoint. "" = none (bf16). Must be a field or
+    # it is dropped before the loader.
+    quantization: str = ""
+    # Disable vLLM's custom all-reduce kernel. Required for tensor_parallel_size>1 on multi-GPU
+    # nodes WITHOUT NVLink (A40/L40 over PCIe), where it dies with CUDA 'invalid argument'; falls
+    # back to NCCL all-reduce (pair with NCCL_P2P_DISABLE=1).
+    disable_custom_all_reduce: bool = False
+    # Skip CUDA graph capture, freeing ~6-7GB per GPU for KV cache. Useful when a large model
+    # (e.g. a 27B tensor-parallel across 48GB cards) would otherwise starve the KV cache.
+    enforce_eager: bool = False
     # fp8 KV cache halves KV memory so a large window (e.g. 8192) fits on one 48GB card;
     # max_num_seqs caps concurrency so that tight KV budget does not thrash.
     kv_cache_dtype: Literal["auto", "fp8", "fp8_e4m3", "fp8_e5m2"] | None = None
     max_num_seqs: int | None = Field(None, ge=1)
+    # Keep undeclared vLLM runtime knobs (e.g. swap_space, enable_prefix_caching) — they reach
+    # load_vllm_engine's tolerant .get() via the non-lossy loader, no new field required.
+    model_config = ConfigDict(extra="allow")
 
 
 class CotSettings(BaseModel):
@@ -238,6 +272,7 @@ 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)
+    model_config = ConfigDict(extra="allow")
 
 
 class QueryExpansionSettings(BaseModel):
@@ -251,6 +286,9 @@ class QueryExpansionSettings(BaseModel):
     max_main_conditions: int = Field(11, ge=1)
     max_other_conditions: int = Field(50, ge=1)
     trust_remote_code: bool = False
+    # Suppress <think> during expansion for reasoning models (read at query_expansion.py:_resolve_settings).
+    no_think: bool = False
+    model_config = ConfigDict(extra="allow")
 
 
 class ReportingSettings(BaseModel):