Longformer chunk pruner for Q-RAG

A Longformer-base model with a trained linear keep/drop head. Given a question and a Q-RAG retrieval chain, it assigns each retrieved chunk a p_keep; the default rule retains chunks with p_keep >= 0.5. The full backbone and head are stored in model.safetensors. This is a custom chunk-pruning architecture, not a standard Transformers sequence-classification model; use the loader below.

Usage

Requires Python, PyTorch, transformers, huggingface_hub, and safetensors.

import sys
from huggingface_hub import snapshot_download

model_dir = snapshot_download("ArtTolston/longformer-base-pruner-hotpotqa_musique")
sys.path.insert(0, model_dir)
from inference import load_pruner, prune

model, tokenizer = load_pruner(model_dir, device="cpu")  # or "cuda"
question = "Which city is the university located in?"
chunks = ["First retrieved paragraph.", "Second retrieved paragraph."]
result = prune(model, tokenizer, question, chunks, threshold=0.5)
print(result["keep_probabilities"], result["kept_indices"])

Chunks must be in Q-RAG retrieval order and formatted exactly as provided to the retriever. The helper reproduces training assembly: [CLS] question [SEP] chunk_1 [SEP] ... [SEP] chunk_H, with no final separator; question body up to 128 tokens, each chunk body up to 512 tokens, total at most 4096 tokens, and global attention on [CLS] and question tokens. A chunk that cannot fit in the token budget has probability None and is excluded. During training Q-RAG produced up to six greedy, distinct chunks without using gold labels; this model does not perform retrieval.

Training and model selection

• Base: allenai/longformer-base-4096, revision 301e6a42cb0d9976a6d6a26a079fef81c18aa895; full backbone plus a per-chunk mean-pooling/linear classification head, 148,660,225 parameters.
• Data: Q-RAG/Hotpotqa_and_Musique, dataset revision 869188e85245c8ff5107f11b2d9ce24fd2bd953a: 90,447 HotpotQA and 19,938 MuSiQue train examples.
• Frozen retriever: Q-RAG/qrag-ft-gte-on-hotpotqa_musique, checkpoint SHA-256 ff2ab1db095fe05f0f854672a224e9dff6ff0a9a8ecda5cf35cdf88a94d37c56.
• Objective: masked binary cross-entropy, gold supporting chunks = keep, other retrieved chunks = drop. Three epochs; learning rate 2e-5, AdamW weight decay 0.01, seed 42. Training resumed from step 500 with batch size 8, accumulation 1; FP32.
• Best checkpoint: optimizer step 38,900, selected by highest pruned supporting-fact F1 at threshold 0.5 among periodic validation samples. This file contains model weights only; it is not an optimizer checkpoint.

Validation at the selected step

On the 300-question resampled combined HotpotQA + MuSiQue eval subset used to select this checkpoint: supporting-fact F1 0.9546, all-gold-support coverage 0.8867, precision 0.9666, recall 0.9492. The unpruned Q-RAG chain on the same questions had F1 0.5149 and coverage 0.9800. Thus pruning improved F1 on this sample but reduced all-gold coverage. ROC-AUC 0.9948 is a micro score over visible retrieved chunks only.

These are sample-based checkpoint-selection metrics, not a full-evaluation benchmark. Checkpoints were selected repeatedly on samples from the same eval pool; this score is not an independent held-out test. Coverage means every gold supporting chunk was retained, not answer exact match. Full per-threshold metrics and provenance are in training_metadata.json.