snippet-extraction

Query-Aware Snippet Extraction

A query-aware extractive context compressor for RAG. The model fine-tunes answerdotai/ModernBERT-base and adds a linear head that scores each document token for relevance to a paired query. Token scores are pooled over document units such as sentences, tables, and code blocks to decide which context to keep.

The full data generation, training, inference, and serving implementation is available in query-aware-snippets.

Intended use

Use this checkpoint to reduce retrieved documents before passing them to an LLM. Input is a (query, document) pair with a maximum sequence length of 8192 tokens. The production runtime handles unit segmentation, long-document windowing, score pooling, and thresholding.

This checkpoint is not a generative model and does not produce summaries. It selects spans from the source document, so output can still contain incorrect, unsafe, or sensitive source content.

Validation metrics

The selected checkpoint is epoch 3 of run 9. On its held-out validation split:

The split was generated by this project's labeling pipeline and is not a standard public benchmark, so these numbers should not be compared directly with unrelated compression datasets.

Loading

The checkpoint includes the fine-tuned encoder and classification head in model.pt. Use the included loader after downloading the repository:

import torch
from huggingface_hub import snapshot_download

model_dir = snapshot_download("gziz/snippet-extraction")

import sys
sys.path.insert(0, model_dir)
from load_model import load_model

model, tokenizer = load_model(model_dir)
inputs = tokenizer(
    "What is retrieval-augmented generation?",
    "Retrieval-augmented generation gives an LLM external context.",
    return_tensors="pt",
)
with torch.no_grad():
    token_logits = model(**inputs)
    token_probabilities = token_logits.sigmoid()

For sentence segmentation, long-document support, and the calibrated dual-threshold decision rule, use snippets_runtime from the source repository with this downloaded checkpoint directory.

Training configuration

• Base model: answerdotai/ModernBERT-base
• Maximum sequence length: 8192
• Precision: bfloat16
• Optimizer learning rate: 4e-5
• Effective batch size: 32
• Seed: 42
• Selection metric: unit F1 at token threshold 0.5 and unit threshold 0.3