Initial release: checkpoints, loader and model card

README.md

@@ -0,0 +1,146 @@
+---
+license: cc-by-nc-4.0
+language: [en]
+tags:
+  - clinical-trials
+  - information-retrieval
+  - reranking
+  - graph-neural-network
+  - biomedical
+library_name: pytorch
+pipeline_tag: text-classification
+---
+
+# CrossGAT Trial Eligibility — 5-fold CV checkpoints
+
+A 3-class eligibility model that **reranks clinical-trial retrieval candidates**.
+It scores a patient description against a trial's inclusion/exclusion criteria as
+`irrelevant / excluded / eligible`, by fusing a graph encoder over parsed
+eligibility criteria with a text branch over whole-document embeddings.
+
+These are the **five cross-validation checkpoints** — one per fold, so each is
+trained on 80% of the data and never saw its own fold's topics. Use them to
+reproduce the reported numbers, or ensemble them. For a single deployment model,
+see [crossgat-trial-eligibility-full](https://huggingface.co/2001jdev/crossgat-trial-eligibility-full).
+
+## ⚠️ This model is not self-contained
+
+It does **not** take raw text. It takes *parsed eligibility graphs* — entity and
+relation structures extracted from trial criteria and patient descriptions by an
+LLM. Producing those needs the generator from the project repo plus an LLM
+endpoint (~38 GB VRAM).
+
+**If you just want to reproduce the results**, the graphs for all 13,229 TREC CT
+trials are published as a dataset — no LLM needed:
+[clinical-trials-eligibility-graphs-rerank](https://huggingface.co/datasets/2001jdev/clinical-trials-eligibility-graphs-rerank).
+
+## What it is for, and what it is not for
+
+**It reranks. It does not retrieve.** Trained on an eligibility target, it reorders
+*within* the candidate set a first-stage retriever already found — promoting
+eligible trials over merely topically-relevant ones.
+
+**It cannot judge disease relevance.** Per-topic AUC for relevant-vs-rest is
+**0.533** — chance. Whether a trial is even on the right disease is entirely the
+first stage's job. This is why it must be *fused* with retrieval, never used to
+replace it.
+
+## Recommended configuration
+
+```
+score = (1 - w) * retrieval_norm + w * p_eligible_norm      with w = 0.5
+```
+
+Both terms min-max normalised **within a topic**. `p_eligible` = P(eligible) alone
+wins over the alternatives in all four runs, at every cutoff, in both metric
+families — even though `expected_gain` models the graded gain directly and is what
+training optimises.
+
+## Results
+
+5-fold cross-validation, TREC Clinical Trials 2021–23, **162 topics**, paired
+Wilcoxon signed-rank against each run's own first stage. Metric is
+**graded NDCG@10** (gain `2^label - 1` = 0/1/3).
+
+First stages: hybrid BM25+`embeddinggemma-300m-medical` (retrieval 0.5973) and
+dense-only `MedEmbed-base-v0.1` (retrieval 0.5187).
+
+| Variant | gemma-med | Δ | medembed | Δ |
+|---|---|---|---|---|
+| w=0.3 `p_eligible` | 0.6340 | +0.0367 *** | 0.5467 | +0.0280 *** |
+| w=0.3 `expected_gain` | 0.6300 | +0.0327 *** | 0.5433 | +0.0246 *** |
+| w=0.3 `one_minus_irrelevant` | 0.6230 | +0.0257 *** | 0.5334 | +0.0147 *** |
+| w=0.5 `p_eligible` | 0.6515 | +0.0542 *** | 0.5644 | +0.0457 *** |
+| w=0.5 `expected_gain` | 0.6485 | +0.0512 *** | 0.5572 | +0.0385 *** |
+| w=0.5 `one_minus_irrelevant` | 0.6302 | +0.0329 *** | 0.5377 | +0.0190 ** |
+| w=0.7 `p_eligible` | 0.6491 | +0.0518 *** | 0.5548 | +0.0361 *** |
+| w=0.7 `expected_gain` | 0.6455 | +0.0482 *** | 0.5536 | +0.0349 *** |
+| w=0.7 `one_minus_irrelevant` | 0.6262 | +0.0289 ** | 0.5405 | +0.0218 ** |
+| w=1 `p_eligible` | 0.6191 | +0.0218 n.s. | 0.5202 | +0.0015 n.s. |
+| w=1 `expected_gain` | 0.6163 | +0.0190 n.s. | 0.5222 | +0.0035 n.s. |
+| w=1 `one_minus_irrelevant` | 0.6049 | +0.0076 n.s. | 0.5148 | -0.0039 n.s. |
+
+`***` p<0.001, `**` p<0.01, `*` p<0.05.
+
+**Read both metric families.** `raw` scores unjudged documents as 0 (punishing a
+system for surfacing trials TREC never pooled); `condensed` deletes them first
+(flattering a system that surfaces many). Neither is honest alone — the gap between
+them, read against `judged@10`, says how far apart they can be. Where the two agree
+the result is robust; at w=1.0 they diverge sharply and `raw` loses significance
+entirely.
+
+## Which checkpoint to load
+
+Each fold ships two: `ndcg.ckpt` (selected on `val_ndcg10`) and `f1.ckpt`
+(selected on `val_macro_f1`).
+
+**Every number above uses `ndcg`.** The `f1` checkpoints classify better
+(macro F1 0.597 ± 0.020 vs 0.561 ± 0.035 on held-out folds) but have **not** been
+benchmarked for reranking — offered for completeness, not recommended blindly.
+
+Note `val_ndcg10` is flat across epochs 1–5, so the per-fold selected epochs
+(1/2/5/2/4) are largely noise around a plateau. Do not read them as convergence.
+
+## Usage
+
+```python
+pip install "crossgat @ git+https://github.com/JDev2001/msc_v2"
+
+from crossgat_hf import load, score, rerank_score
+
+pipe = load("2001jdev/crossgat-trial-eligibility-cv", fold=0, metric="ndcg")
+out  = score(pipe, patient_graph, inc_graph, exc_graph,
+             patient_text="65yo male, type 2 diabetes, HbA1c 8.1% ...",
+             trial_text="Metformin in adults with type 2 diabetes ...")
+
+out["class_probs"]          # {'irrelevant': .., 'excluded': .., 'eligible': ..}
+rerank_score(out)           # p_eligible — fuse this with your retrieval score
+```
+
+`patient_text` and `trial_text` are **required**: the text branch is part of the
+model, and omitting them raises rather than silently scoring an all-zero branch.
+
+## Training
+
+- **Data:** TREC Clinical Trials 2021–23 qrels/topics; eligibility criteria parsed into graphs (CHIA-style entities/relations)
+- **Objective:** cross-entropy + **LambdaRank** (k=10) on topic-grouped batches (4 topics/batch)
+- **Hard negatives:** 40/topic drawn from the retrieval candidate distribution (+5,307 samples), mined from *hybrid gemma-med* run files
+- **Text branch:** `abhinand/MedEmbed-base-v0.1`, whole-document embeddings of patient description and trial title+summary
+- **Splits:** 5-fold CV over patients, seed fixed; each fold's test topics never seen in its training
+
+## Limitations
+
+- **The first stage is the ceiling.** medembed *with* reranking still scores below gemma-med *without* it. A better retriever buys more than more reranker tuning.
+- **Do not use w=1.0.** Pure model ranking loses significance on both retrievers (p=0.65 / 0.85). It ranks better inside the judged pool but drifts out of it — `judged@10` falls from 0.92 to 0.80.
+- **The hard-negative gain is retriever-sensitive.** +0.0175 (p=0.007) on gemma-med, +0.0148 (p=0.050) on medembed — the negatives were mined from gemma-med candidates and transfer only partly.
+- **Within-collection evidence only.** 162 topics, one corpus, one set of assessors. Nothing about generalisation to another trial collection follows.
+- **Not a clinical device.** Research artefact. It does not determine patient eligibility and must not be used to make enrolment decisions.
+
+## Citation & provenance
+
+Trained on TREC Clinical Trials track data (2021–23). Users must comply with the
+TREC data usage agreements for the underlying collection and relevance judgments.
+
+Checkpoint SHA-256 prefixes and the exact training configuration are in
+`checkpoint_manifest.json`. Every benchmark run in the source repository writes a
+`manifest.json` recording checkpoint hashes, git state and sanity checks.

checkpoint_manifest.json

@@ -0,0 +1,61 @@
+{
+  "cv": {
+    "fold0/ndcg": {
+      "original_name": "joint-fold0-ndcg-epoch=01-val_ndcg10=0.6293.ckpt",
+      "sha256_12": "d64902e332cd",
+      "size_mb": 27.3
+    },
+    "fold0/f1": {
+      "original_name": "joint-fold0-f1-epoch=15-val_macro_f1=0.6327.ckpt",
+      "sha256_12": "af1611c7a7a8",
+      "size_mb": 27.3
+    },
+    "fold1/ndcg": {
+      "original_name": "joint-fold1-ndcg-epoch=02-val_ndcg10=0.6497.ckpt",
+      "sha256_12": "8c9e5b222ece",
+      "size_mb": 27.3
+    },
+    "fold1/f1": {
+      "original_name": "joint-fold1-f1-epoch=07-val_macro_f1=0.6015.ckpt",
+      "sha256_12": "69b5668b994c",
+      "size_mb": 27.3
+    },
+    "fold2/ndcg": {
+      "original_name": "joint-fold2-ndcg-epoch=05-val_ndcg10=0.5998.ckpt",
+      "sha256_12": "9ca8c1c3419f",
+      "size_mb": 27.3
+    },
+    "fold2/f1": {
+      "original_name": "joint-fold2-f1-epoch=04-val_macro_f1=0.5922.ckpt",
+      "sha256_12": "c3ee31254421",
+      "size_mb": 27.3
+    },
+    "fold3/ndcg": {
+      "original_name": "joint-fold3-ndcg-epoch=02-val_ndcg10=0.6571.ckpt",
+      "sha256_12": "9cfb1cf3cb33",
+      "size_mb": 27.3
+    },
+    "fold3/f1": {
+      "original_name": "joint-fold3-f1-epoch=05-val_macro_f1=0.5765.ckpt",
+      "sha256_12": "5236ff7beccb",
+      "size_mb": 27.3
+    },
+    "fold4/ndcg": {
+      "original_name": "joint-fold4-ndcg-epoch=04-val_ndcg10=0.7009.ckpt",
+      "sha256_12": "c29677ab75dc",
+      "size_mb": 27.3
+    },
+    "fold4/f1": {
+      "original_name": "joint-fold4-f1-epoch=07-val_macro_f1=0.5839.ckpt",
+      "sha256_12": "355b93e402fd",
+      "size_mb": 27.3
+    }
+  },
+  "full": {
+    "model": {
+      "original_name": "joint-fulldata-epoch=03.ckpt",
+      "sha256_12": "dbd1930dc442",
+      "size_mb": 27.3
+    }
+  }
+}

checkpoints/fold0/f1.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:af1611c7a7a84b0f0998feb076167c99d8b1aea9d0de6771d77249384c07d7ae
+size 28643449

checkpoints/fold0/ndcg.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d64902e332cd9b93757209a56a0557cc3fa11dd12e7b1d351ade9878ca3e6deb
+size 28643449

checkpoints/fold1/f1.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:69b5668b994cb226776bad3d8e48165044d04764e768581c0d320b19f76c8eb2
+size 28643513

checkpoints/fold1/ndcg.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c9e5b222ecec6b3147eba3dc95085f152d45126596862867edb4130ad60e240
+size 28643513

checkpoints/fold2/f1.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c3ee312544215487a8363d46a8e486db9c75c60e2c8a0cbfce57f3ea86f1f22b
+size 28643513

checkpoints/fold2/ndcg.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9ca8c1c3419f5dde902aff1c8a59530e6f259838b8f2866e5357319b890ff4d1
+size 28643513

checkpoints/fold3/f1.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5236ff7beccba7c851f052cbae16eb600211e56e93f47fdecf3b00ff7cc073ae
+size 28643513

checkpoints/fold3/ndcg.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9cfb1cf3cb3347f651f8cf3cc9a0a25c2702a58baaf0f491c177b3009ce90727
+size 28643513

checkpoints/fold4/f1.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:355b93e402fd9617ba806bd439befc66737f3acb2dcb4ae0a2592d7250e5148a
+size 28643513

checkpoints/fold4/ndcg.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c29677ab75dcd3ebd3e0e919d05300af3933c161631ab068bb157488d4d7912e
+size 28643513

crossgat_hf.py

@@ -0,0 +1,93 @@
+"""Load a CrossGAT eligibility checkpoint from this repo and score a pair.
+
+This model is NOT self-contained. It scores a *patient graph* against a trial's
+*inclusion/exclusion graphs* — parsed entity/relation structures, not raw text.
+Producing those graphs needs the eligibility-graph generator from the project
+repository, which in turn needs an LLM endpoint. See the README.
+
+What this file gives you is the loading and scoring half: given graphs (from the
+companion dataset, or from your own generator), it returns the 3-class posterior.
+
+    pip install "crossgat @ git+https://github.com/JDev2001/msc_v2"
+    from crossgat_hf import load, score
+
+    pipe = load("2001jdev/crossgat-trial-eligibility-cv", fold=0, metric="ndcg")
+    out  = score(pipe, patient_graph, inc_graph, exc_graph,
+                 patient_text="65yo male, type 2 diabetes, HbA1c 8.1%",
+                 trial_text="Metformin in adults with type 2 diabetes ...")
+    out["probs"]  # [P(irrelevant), P(excluded), P(eligible)]
+"""
+from __future__ import annotations
+
+from typing import Any
+
+CLASSES = ("irrelevant", "excluded", "eligible")
+
+
+def load(repo_id: str, fold: int | None = None, metric: str = "ndcg",
+         revision: str | None = None):
+    """Download one checkpoint and return a ready InferencePipeline.
+
+    fold=None loads ``checkpoints/model.ckpt`` (the full-data refit repo);
+    otherwise ``checkpoints/fold{fold}/{metric}.ckpt``.
+
+    metric selects which checkpoint the fold's training saved:
+      "ndcg" -- selected on val_ndcg10, used for every benchmark number reported
+      "f1"   -- selected on val_macro_f1; classifies better (macro F1 .597 vs .561)
+                but has NOT been benchmarked for reranking.
+    """
+    from huggingface_hub import hf_hub_download
+    from crossgat.inference.pipeline import InferencePipeline
+
+    if metric not in ("ndcg", "f1"):
+        raise ValueError(f"metric must be 'ndcg' or 'f1', got {metric!r}")
+    name = "checkpoints/model.ckpt" if fold is None else f"checkpoints/fold{fold}/{metric}.ckpt"
+    path = hf_hub_download(repo_id=repo_id, filename=name, revision=revision)
+    return InferencePipeline(path)
+
+
+def score(pipeline, patient_graph, inc_graph, exc_graph,
+          patient_text: str, trial_text: str) -> dict[str, Any]:
+    """Score one patient/trial pair.
+
+    Graphs are objects with ``.entities`` and ``.relations`` (crossgat RawGraph, or
+    anything with those attributes -- e.g. json.load of a file from the companion
+    graph dataset passed through ``crossgat.data.schemas``).
+
+    patient_text / trial_text feed the model's text branch and are REQUIRED: the
+    3-class checkpoints fuse whole-document embeddings with the graph features, and
+    omitting them raises rather than silently scoring against an all-zero branch.
+    Use the TREC topic description and the trial's brief_title + brief_summary.
+    """
+    out = pipeline.predict_chia(
+        patient_entities=patient_graph.entities,
+        patient_relations=patient_graph.relations,
+        inc_trial_entities=inc_graph.entities,
+        inc_trial_relations=inc_graph.relations,
+        exc_trial_entities=exc_graph.entities,
+        exc_trial_relations=exc_graph.relations,
+        patient_text=patient_text,
+        trial_text=trial_text,
+    )
+    out["label"] = CLASSES[out["pred"]]
+    out["class_probs"] = dict(zip(CLASSES, out["probs"]))
+    return out
+
+
+def rerank_score(out: dict, score_fn: str = "p_eligible") -> float:
+    """Reduce the posterior to the scalar a ranking sorts by.
+
+    ``p_eligible`` is the recommended one: it wins on both first stages tested, at
+    every cutoff and in both metric families. Fuse it with the retrieval score as
+    ``(1 - w) * retrieval_norm + w * model_norm`` at **w = 0.5**, min-max
+    normalising both within a topic. Do not replace the first stage outright
+    (w = 1.0): that loses significance on both retrievers.
+    """
+    p_irr, p_exc, p_elig = out["probs"]
+    if score_fn == "p_eligible":
+        return p_elig
+    if score_fn == "one_minus_irrelevant":
+        return 1.0 - p_irr
+    if score_fn == "expected_gain":
+        return 1.0 * p_exc + 3.0 * p_elig
+    raise ValueError(f"unknown score_fn {score_fn!r}")