Ragzilla: per-passage keep/drop for retrieved context

Ragzilla sits after your retriever. It reads the question and each of the top-10 retrieved passages and returns keep or drop for every passage: several can be kept, or none, and passage 7 can be kept while passage 2 is dropped. Your generator then reads fewer tokens and still finds the answer.

It is answerdotai/ModernBERT-large fine-tuned end to end with the public openjev SDK (openjev-sdk 0.5.0, openjev tod train --method encoder) on a per-passage 0/1 target from eleven public support-labelled sets. Everything here was made with the SDK's command line: training, calibration, export, evaluation and this folder. You can load it with load_policy, fine-tune it on your own retrieval logs with openjev tod train --model, and re-choose its keep thresholds with openjev tod calibrate.

What it does, in numbers

Two numbers are reported throughout. Answer survival: among questions whose top-10 holds an answer-bearing passage, the share that still hold one after the filter. Tokens kept: kept passage tokens as a share of top-10 passage tokens, over every question. All at the default point rho_99.

set	questions with an answer-bearing passage	all questions	answer survival	tokens kept
In-distribution test, all 11 sources	64,948	84,874	99.1%	36.8%
Rebuilt retrieval lists: FiQA, SciFact, NQ, FEVER (bge-large top-10)	7,310	7,714	98.6%	22.7%
TREC DL 2019, never trained on	22	43	100%	59.1%
TREC DL 2020, never trained on	38	54	100%	61.9%
TREC DL 2023, never trained on	39	82	94.9%	56.7%
TREC-COVID, never trained on	47	50	100%	81.1%

With the generator reading only the kept passages, answers stay as good: on RGB (1,500 lists, gemini-3.8-flash, RAGAS answer_correctness) 0.736 against 0.732 with every document, at 80% of the context tokens; on CRAG (500 lists) 0.397 against 0.404 at 60% of the tokens. Details under "Answer quality".

How it works

• Input: one question and 2 to 10 passages per call (trained on 5 to 10). Passages over 4,000 characters are cut. Longer lists are scored in windows of 10.

• One sequence per question (the SDK's packed layout): [CLS] + Which passages are relevant to answering the question? + the question as JSON + [SEP], then each passage as its own span ending in [SEP]. The prefix attends to itself; each passage's tokens attend to the prefix and to that passage only, never to another passage, and positions restart at the end of the prefix for every passage. A passage's probability therefore does not depend on which other passages come with it: reordering passages reorders the scores exactly, and windows give the same scores as one long call.

• Score: mean of the final hidden states over the passage's span, then Linear-GELU-LayerNorm-Linear to one logit, divided by a temperature of 0.407 fitted on the calibration split, through a sigmoid.

• Keep rule (rho, relative): keep passage i when p_i >= rho x the question's highest p; the top passage is always kept. An absolute rule (tau: keep p_i >= tau) is stored too.

• Shipped point: rho_99 = 0.1410, the largest rho that reaches 99% answer survival on the calibration split (81,796 questions never trained on), so the fewest tokens for that target.

• Other stored points (operating_points.json, all chosen on calibration):

  point	value	answer survival on calibration	tokens kept	passages kept of 8.9
  tau_95	0.0874	95.0%	29.3%	2.67
  tau_98	0.0376	98.0%	38.2%	3.53
  tau_99	0.0195	99.0%	45.4%	4.20
  rho_95	0.4897	95.0%	24.2%	2.21
  rho_98	0.2484	98.0%	32.4%	3.03
  rho_99 (default)	0.1410	99.0%	38.4%	3.64

  At every target the rho rule keeps fewer tokens and fewer passages than the tau rule. A rho rule always keeps at least the top passage; a tau rule can keep nothing, which is the right answer when no passage is relevant.

How to use

pip install openjev-sdk==0.5.0
hf download parsecai/ragzilla --local-dir ragzilla

load_policy + decide. The state is the question plus its passages as {'id', 'text'} items; the kind is relevant.

from openjev.tod.rl import load_policy

policy = load_policy('ragzilla', device='cpu')       # load once; 'cuda' serves the trunk in bf16
state = {'question': 'Who wrote the novel Dune?',
         'passages': [{'id': 'p0', 'text': 'Paris is the capital and most populous city of France.'},
                      {'id': 'p1', 'text': 'Dune is a 1965 epic science fiction novel by American author Frank Herbert.'},
                      {'id': 'p2', 'text': 'The Sahara is a desert spanning North Africa.'}]}
pred = policy.decide(state, 'relevant')               # the default point, rho_99
pred.keep                                             # ids kept, e.g. ('p1',)
pred.probabilities                                    # {'p0': ..., 'p1': ..., 'p2': ...}
policy.decide(state, 'relevant', point='rho_98')      # another stored point
policy.decide_many([(state, 'relevant'), (state, 'relevant')])   # several questions per forward (GPU)

Without the SDK. The same folder loads with transformers alone (modeling_openjev.py is the SDK's packing code, copied in by openjev tod export --hf). The query must be rendered exactly as the SDK renders it: the kind's question, a blank line, then {"question": ...} as JSON with indent 1.

import json
from transformers import AutoModel, AutoTokenizer
m = AutoModel.from_pretrained('ragzilla', trust_remote_code=True).eval()
tok = AutoTokenizer.from_pretrained('ragzilla')
q = 'Who wrote the novel Dune?'
passages = ['Paris is the capital and most populous city of France.',
            'Dune is a 1965 epic science fiction novel by American author Frank Herbert.']
query = 'Which passages are relevant to answering the question?\n\n' + json.dumps({'question': q}, ensure_ascii=False, indent=1)
m.score(tok, query, passages, 'relevant')             # calibrated probability per passage, e.g. [0.00.., 0.9..]
m.keep(tok, query, passages, 'relevant')              # indices kept at rho_99, e.g. [1]

On 100 held-out test questions (896 passages) at the default point, the Hugging Face folder and the SDK export give identical keep-sets on 99 of 100 questions on an A100 (largest probability difference 0.0026) and on 99 of 100 on a CPU (0.021); loaded in fp32 the folder agrees with the SDK to 7.3e-8. The SDK's own export --hf check reports 0.0298 and exits 1 because it compares an fp32 forward with a bf16 one against a 0.02 tolerance; the folder is the one it wrote.

GPU loading with openjev-sdk 0.5.0. The SDK runs a layout self-test on load in bf16 against a tolerance meant for fp32. This model trips it on an A100 (ValueError: ... uses the packed layout, which this model or transformers version does not support) and passes it on an L4 and on any CPU. Until the SDK fix lands, use device='cpu', or the transformers path above on the GPU, or run the self-test on an fp32 copy of the trunk before load_policy (a 20-line shim; nothing else in the SDK changes).

Fine-tune on your own retrieval logs (question, the passages your retriever returned, which ones helped), then re-pick the keep thresholds on your calibration split. my_rows.py in this repository builds the dataset from one JSONL line per question ({"id", "question", "passages": [...], "relevant": [indices], "group"}), prepared as the training rows were and split by group with a calibration split large enough to choose points on:

python my_rows.py my_rows.jsonl --out my_rows            # 50 / 10 / 20 / 20 split by group
openjev tod train my_rows --method encoder --model ragzilla --out my_ragzilla \
  --loss focal --focal-gamma 2 --focal-alpha 0.5 --brier 0.1 --learning-rate 1e-5 --max-steps 25
openjev tod calibrate my_ragzilla my_rows          # re-fit the temperature, re-choose tau/rho on your calibration split
openjev tod evaluate my_ragzilla my_rows           # survival and tokens kept at every point on your test split

Run end to end on a stand-in for your logs (TREC Deep Learning 2019, 2020 and 2023, 179 questions never trained on, every passage human-judged, lists of 10, split 1,399 train / 589 calibration / 511 test lists by question), answer survival and tokens kept at rho_99: as shipped 94.5% at 65.1%; with the points re-chosen on the stand-in calibration split 99.2% at 86.2%; after 25 steps of continued training, points re-chosen, 99.4% at 75.2%. A fixed top-k needs 85.1%, 100% and 100% of the tokens for the same survival.

Training

• Base model: answerdotai/ModernBERT-large (28 layers, hidden 1024, 395M parameters), full fine-tune, plus a Linear-GELU-LayerNorm-Linear head.

• Sources (share of training rows):

  source	the question	the candidate passages	label 1 means	share of train
  MS MARCO	Bing query	the ~10 Bing passages	Bing's selected passage, or contains the answer	29.9%
  Natural Questions	Google query	the Wikipedia page's paragraphs, cut to 5-10	the annotated long answer	17.8%
  2WikiMultihopQA	multi-hop question	the 10 context paragraphs	supporting paragraph	11.9%
  HotpotQA (distractor)	multi-hop question	the 10 context paragraphs	supporting paragraph	11.7%
  FEVER (pages)	claim	top-10 Wikipedia pages by dense retrieval	a page the evidence cites	9.3%
  FEVER (sentences)	claim	sentences of the evidence pages, cut to 5-10	an evidence sentence	8.6%
  SQuAD 2	question	its paragraph + same-article paragraphs	its paragraph, if answerable	7.0%
  MuSiQue	multi-hop question	20 paragraphs cut to 5-10	supporting paragraph	2.7%
  QASPER, SciFact, SciTail, WikiQA	paper / claim / hypothesis questions	paper paragraphs, abstracts, premises	cited evidence	under 0.5% each

• Label rule: these datasets mark only the passage the annotator used, and other passages in the list often contain the same answer. A 0-labelled passage containing a specific answer (2+ tokens, or one non-numeric token of 4+ characters) becomes 1; one containing only a short or numeric answer is removed from the row. Support labels are never demoted. This flipped 214k labels in MS MARCO, 81k in NQ, 61k in HotpotQA, 27k in SQuAD 2.

• Row rules: duplicate passages merged, empty ones dropped, a passage over 4,000 characters cut if negative and the row dropped if positive, rows with fewer than 2 passages dropped. Where a source has more than 10 candidates, every positive is kept plus the hardest negatives by BM25, to a count of 5 to 10 per question.

• Splits (by group, so a passage the model trained on never appears in validation, calibration or test): train 585,506 / validation 85,659 / calibration 81,796 / test 84,874. All 8,196 head-to-head and external evaluation questions were removed from every split.

• Loss: focal loss per passage (gamma 2, alpha 0.5) plus 0.1 x Brier score; no class weight. About 1 in 7 candidates is positive; with plain BCE the model learns to say "no" and the rare positives are drowned out.

• Run: one launch, one epoch, 2,602 steps on 8 x A100-40GB at 2.25 s per step (98 minutes of steps), 32,768 useful tokens per GPU per step, lr 3e-5 trunk / 1e-4 head, AdamW, weight decay 0.01, 5% warmup then linear decay, fp32 weights with bf16 autocast. The trunk is saved in bf16 and the head in fp32.

• Command (openjev-sdk 0.5.0):

  torchrun --standalone --nproc_per_node 8 -m openjev tod train data/ragzilla --method encoder \
    --model answerdotai/ModernBERT-large --out runs/ragzilla --layout packed --head mlp \
    --loss focal --focal-gamma 2 --focal-alpha 0.5 --brier 0.1 --epochs 1 --batch-tokens 32768 \
    --learning-rate 3e-5 --head-learning-rate 1e-4 --eval-every 1000 --checkpoint-every 1000
  

Evaluation

Operating-point values come from the calibration split only. Tokens on the rebuilt retrieval lists are tiktoken cl100k_base; in-distribution and external tokens are ModernBERT tokenizer counts. In every table, answer survival is the share of the questions with an answer-bearing passage that keep at least one after filtering, and tokens kept is the share of candidate passage tokens kept, over every question of the set.

External sets, never trained on

TREC DL 2019 / 2020 / 2023: the 10 judged passages BM25 ranks highest within the NIST-judged pool. TREC-COVID: BM25 top-10, 94% judged. An answer-bearing passage here is one judged top grade (grade 3 on TREC DL, grade 2 on TREC-COVID). Four stored operating points; rho_99 is the shipped default.

set	questions with a top-grade passage	all questions	tau_98 survival	tau_98 tokens kept	tau_99 survival	tau_99 tokens kept	rho_98 survival	rho_98 tokens kept	rho_99 survival	rho_99 tokens kept
TREC DL 2019	22	43	100%	60.4%	100%	69.6%	95.5%	47.2%	100%	59.1%
TREC DL 2020	38	54	100%	59.2%	100%	69.3%	100%	50.0%	100%	61.9%
TREC DL 2023	39	82	94.9%	51.7%	94.9%	63.5%	94.9%	46.5%	94.9%	56.7%
TREC-COVID	47	50	97.9%	76.6%	97.9%	85.8%	100%	66.0%	100%	81.1%

In-distribution test, per source

The held-out test split, 84,874 questions, by source. Four stored operating points; rho_99 is the shipped default.

source	questions with a relevant passage	all questions	tau_98 survival	tau_98 tokens kept	tau_99 survival	tau_99 tokens kept	rho_98 survival	rho_98 tokens kept	rho_99 survival	rho_99 tokens kept
all sources	64,948	84,874	98.1%	37.2%	99.1%	44.5%	98.2%	31.1%	99.1%	36.8%
2WikiMultihopQA	9,913	9,913	100%	44.8%	100%	46.1%	100%	41.6%	100%	42.4%
FEVER (sentences)	6,779	6,779	99.7%	49.2%	99.8%	56.8%	99.3%	32.6%	99.6%	37.8%
FEVER (pages)	8,471	8,582	99.5%	17.1%	99.7%	21.2%	99.4%	12.4%	99.6%	13.5%
HotpotQA	9,853	9,853	100%	46.4%	100%	53.7%	100%	28.7%	100%	33.9%
MS MARCO	15,435	25,106	94.4%	46.3%	97.4%	57.4%	94.7%	45.9%	97.6%	57.0%
MuSiQue	2,300	2,300	100%	55.2%	100%	61.0%	99.9%	40.0%	100%	44.6%
NQ	7,354	14,940	96.6%	50.6%	98.2%	61.1%	96.7%	48.5%	98.2%	59.3%
QASPER	372	420	98.9%	83.1%	99.7%	91.2%	98.4%	59.3%	99.2%	72.7%
SciFact	63	83	98.4%	30.2%	98.4%	37.0%	100%	32.7%	100%	44.0%
SciTail	147	147	100%	73.7%	100%	83.9%	99.3%	48.2%	100%	54.8%
SQuAD2	4,139	6,465	98.9%	34.8%	99.6%	45.5%	99.6%	23.2%	99.9%	30.0%
WikiQA	122	286	95.1%	32.7%	96.7%	41.5%	96.7%	31.4%	98.4%	39.0%

Rebuilt retrieval lists, per corpus

bge-large-en-v1.5 dense + BM25 fused top-10, 7,714 queries (7,310 with gold in the top-10), never trained on. At the shipped point rho_99.

corpus	queries with gold in top-10	all queries	answer survival	tokens kept	passages kept per query
FEVER	4,911	5,000	99.3%	14.1%	1.37
NQ	1,862	1,993	97.9%	53.9%	5.18
SciFact	67	73	97.0%	35.1%	3.45
FiQA	470	648	93.8%	71.7%	6.57
pooled	7,310	7,714	98.6%	22.7%	2.81

Choosing an operating point

Every stored point on five sets: the calibration split (81,796 questions, where the values were chosen), the in-distribution test split (84,874), the rebuilt lists pooled (7,714 queries), FiQA alone (648) and TREC-COVID (50). Any point can be named in the call (point='rho_98').

point	threshold value	calibration survival	calibration tokens kept	test survival	test tokens kept	rebuilt lists survival	rebuilt lists tokens kept	FiQA survival	FiQA tokens kept	TREC-COVID survival	TREC-COVID tokens kept
tau_95	0.0874	95.0%	29.3%	95.2%	28.4%	93.4%	18.6%	64.7%	30.7%	93.6%	63.5%
tau_98	0.0376	98.0%	38.2%	98.1%	37.2%	97.1%	25.2%	83.0%	54.2%	97.9%	76.6%
tau_99	0.0195	99.0%	45.4%	99.1%	44.5%	98.3%	31.4%	88.9%	71.2%	97.9%	85.8%
rho_95	0.4897	95.0%	24.2%	95.3%	23.4%	94.2%	14.5%	75.5%	32.6%	97.9%	44.7%
rho_98	0.2484	98.0%	32.4%	98.2%	31.1%	97.1%	18.8%	86.2%	55.5%	100%	66.0%
rho_99 (shipped)	0.1410	99.0%	38.4%	99.1%	36.8%	98.6%	22.7%	93.8%	71.7%	100%	81.1%

At matched survival rho uses fewer tokens than tau on calibration, on the rebuilt lists and on every external set.

Answer quality

Two public RAG benchmarks, the same generator (gemini-3.8-flash, temperature 0) and the benchmark's own prompt three times per question: A gets every document, B gets only the documents Ragzilla keeps at rho_99, C gets none. Scored with RAGAS 0.4.3 (gemini-3.8-flash as judge, gemini-embedding-001 for the embedding terms). RGB: 300 English questions x noise ratio 0 to 0.8 = 1,500 lists of 5 documents, every document tagged answer-bearing or not. CRAG: task 1 dev, 500 questions, the baseline's top-10 sentence chunks.

benchmark	lists	documents per list	documents kept per list	context tokens, all documents	context tokens, Ragzilla-kept	answer_correctness, all documents	answer_correctness, Ragzilla-kept	answer_correctness, no documents	faithfulness, all documents	faithfulness, Ragzilla-kept	answer_relevancy, all documents	answer_relevancy, Ragzilla-kept
RGB	1,500	5	3.98	1,304	1,046	0.732	0.736	0.131	0.974	0.968	0.934	0.932
CRAG	500	10	4.40	557	335	0.404	0.397	0.286	0.330	0.314	0.207	0.192

Per list, Ragzilla-kept against all documents, on answer_correctness (a win or loss is a difference of at least 0.05):

benchmark	lists	wins	ties	losses	mean change	95% bootstrap interval	benchmark's own scorer, all documents	benchmark's own scorer, Ragzilla-kept
RGB	1,500	125	1,272	103	+0.004	-0.001 to +0.008	string accuracy 0.958	string accuracy 0.957
CRAG	500	24	447	29	-0.007	-0.026 to +0.012	three-way score 0.134	three-way score 0.126

On RGB the filter keeps 95.4% of the answer-bearing documents and 57.8% of the others.

On RGB the filter leaves answer quality unchanged at every noise ratio while cutting the context to 80% of its tokens (66% at the highest noise); 2 of 1,500 lists went from right to wrong on RGB's own scorer because every answer-bearing document was dropped. On CRAG the context drops to 60% at the cost of 10 correct answers of 119 (7 become "I don't know", 3 become wrong); the losses sit in the lists cut to 1 to 3 chunks.

Latency

Measured on 480 held-out questions, each a question plus its top-10 passages: 240 from the in-distribution test split (mean 1,150 ModernBERT tokens per question, 9.15 passages) and 240 from the rebuilt retrieval lists, 60 per corpus (2,673 tokens, 10 passages). The policy is loaded once (4.7 s on the GPU, 3.5 s on the CPU), 32 warm-up questions are excluded, the point is rho_99. Hardware: one A100-40GB of an a2-highgpu-8g (the box has 48 physical cores; the CPU rows use 12 threads). Batch 1 is one decide call per question. Packed is decide_many over the set. Both modes kept the same passages on 479 of 480 questions.

device	trunk dtype	question set	questions	batch 1, median seconds per question	batch 1, p90 seconds per question	batch 1, questions per second	packed, questions per second	peak memory
A100-40GB	bf16	test split	240	0.030	0.041	28.2	51.0	2.2 GB GPU, 3.4 GB RSS
A100-40GB	bf16	rebuilt retrieval lists	240	0.054	0.095	17.0	19.1	2.2 GB GPU, 3.4 GB RSS
CPU, 12 threads	fp32	test split	240	0.749	2.13	0.74	0.57	5.3 GB RSS
CPU, 12 threads	fp32	rebuilt retrieval lists	240	3.13	7.30	0.27	0.21	5.3 GB RSS

On the CPU the bf16 weights are upcast to fp32. Time tracks tokens: the rebuilt lists are 2.3x longer than the test questions.

Limitations

• FiQA is the weak spot. Financial QA, a domain never trained on: 93.8% survival at the shipped point, below the 99% target. Re-choose the operating point on your own calibration data before relying on it in a new domain (the fine-tune check above shows what that buys).
• TREC DL 2023: 94.9% survival (37 of 39) at every operating point. TREC DL queries and passages are MS MARCO, which is 29.9% of training rows.
• Small n. SciFact: 73 queries, one question about 1.5 points. External sets: one question is 2.1–4.5 points. TREC-COVID loses one question at the tau points.
• English only, and passages over 4,000 characters are cut.
• Serve it with the code it ships with. Under the SDK's packed layout every layer attends across the whole question-plus-passage block; ModernBERT's 128-token local window is not applied. load_policy (openjev-sdk 0.5.0) and the bundled modeling_openjev.py both compute exactly what was trained. A stock ModernBERT pipeline does not.
• Data licences restrict use; see below.

Licence

The base model, answerdotai/ModernBERT-large, is Apache-2.0, which allows a fine-tune to be released under other terms. The training data is what restricts this model:

source	licence
MS MARCO (29.9% of training rows)	MS MARCO terms: non-commercial research use only
SciFact	CC-BY-NC-2.0 (claims and annotations); abstracts from S2ORC (ODC-By-1.0)
WikiQA	Microsoft Research Data License Agreement (research use); sentences are Wikipedia text (CC-BY-SA)
SciTail	no dataset licence declared (AI2)
FEVER, NQ	CC-BY-SA-3.0
HotpotQA, SQuAD2	CC-BY-SA-4.0
MuSiQue, QASPER	CC-BY-4.0
2WikiMultihopQA	Apache-2.0

Because MS MARCO allows only non-commercial research use (and SciFact is NC, and WikiQA is research-use), the weights are released under CC-BY-NC-4.0: non-commercial use only. For commercial use, retrain without MS MARCO, SciFact and WikiQA (and check SciTail) with the command above.

Files

One folder, two loaders. The SDK layout (load_policy, openjev tod train --model, openjev tod calibrate) is the root plus trunk/; the transformers layout (AutoModel.from_pretrained(..., trust_remote_code=True)) is config.json, model.safetensors, modeling_openjev.py and the tokenizer files at the root. The weights are the same bf16 trunk and fp32 head in both; the transformers file carries both in one safetensors.

file	bytes	loader
spec.json (the decision spec: one kind, relevant, candidates under passages)	228	SDK
tod.json (training settings, the three probes, 2,602 steps in 5,884 s)	1,905	SDK
encoder.json (packed layout, mlp head, candidate-span-mean pooling, 8,192 tokens, 4,000 characters)	263	SDK
calibration.json (temperature 0.4066)	63	SDK
operating_points.json (tau / rho at 0.95, 0.98, 0.99; default rho_99)	542	SDK
head.safetensors (Linear-GELU-LayerNorm-Linear head, fp32)	4,211,140	SDK
trunk/model.safetensors (ModernBERT-large, bf16)	789,580,328	SDK
trunk/config.json, trunk/tokenizer.json, trunk/tokenizer_config.json	2,044 + 3,583,228 + 379	SDK
config.json (OpenjevScorerConfig: base config, head, layout, temperature, operating points)	3,693	transformers
model.safetensors (trunk bf16 + head fp32)	793,793,748	transformers
modeling_openjev.py (the SDK's packing code and the OpenjevScorer wrapper)	11,669	transformers
tokenizer.json, tokenizer_config.json	3,583,228 + 379	transformers
my_rows.py (your JSONL rows -> an openjev dataset for train / calibrate / evaluate)	about 4 KB	SDK

1,594,790,859 bytes in all (1.49 GB); no training checkpoints.