Skip to content

feat: sentence-level pruning (prune_async / open_provence) mirroring the reranker #1144

Description

@Goldziher

Summary

Add a sentence-level pruning capability — prune_async / prune — that mirrors the existing
reranker, using an open_provence / Provence sentence-pruning
model. Provence drops off-topic sentences within a document (and scores relevance) in a single
ModernBERT/DeBERTa pass — complementary to the reranker, which only reorders whole documents.

The headline consumer is basemind's search_documents and compress (search → rerank → prune
return), but this is a generally useful kreuzberg primitive.

Proposed API (mirror the reranker)

#[cfg(all(feature = "pruner", feature = "tokio-runtime"))]
pub async fn prune_async(
    query: String,
    documents: Vec<String>,
    config: &core::config::PrunerConfig,
) -> crate::Result<Vec<PruneResult>>;

pub struct PruneResult {
    pub index: usize,                          // position in input `documents`
    pub score: f32,                            // document relevance in [0,1]
    pub document: String,                      // original text, unmodified
    pub kept_sentence_spans: Vec<(usize, usize)>, // half-open byte ranges kept
    pub compression_ratio: f32,
    pub sentences_kept: usize,
}

pub struct PrunerConfig {
    pub model: PrunerModelType,                // Preset | Custom | Plugin
    pub retention_threshold: f32,              // keep sentences with score >= threshold
    pub batch_size: usize,
    pub show_download_progress: bool,
    pub cache_dir: Option<PathBuf>,
    pub acceleration: Option<AccelerationConfig>,
    pub max_prune_duration_secs: Option<u64>,
}

Return byte spans of kept sentences (never destructively rewrite text) so callers keep
provenance.

Implementation notes (reranker is the template)

  • Mirror crates/kreuzberg/src/reranking/{mod.rs,engine.rs} → new pruning/{mod.rs,engine.rs,segmenter.rs}
    and core/config/reranker.rscore/config/pruner.rs.
    • rerank_async template: reranking/mod.rs:926; ONNX inference: reranking/engine.rs:45;
      config: core/config/reranker.rs:21.
  • Reuse the existing hf-hub download + cross-process lock + ort session + ONNX engine cache; pair
    encoding (query, sentence) via the tokenizers stack.
  • Sentence segmentation: a splitter already exists at
    text/summarization/textrank.rs:81 (split_sentences) — adapt it to track byte offsets, or use
    the OpenProvence reference segmenter for fidelity.
  • New Cargo features pruner + pruner-presets, mirroring reranker / reranker-presets
    (same deps: hf-hub, ort, ndarray, tokenizers). No new dependencies. Add to full /
    server / no-ort-target / wasm-target aggregates as the reranker does.
  • Async dispatch via a dedicated PRUNE_SEMAPHORE + spawn_blocking, same as the reranker.

Blocker: no ONNX weights published

The open_provence checkpoints ship PyTorch .safetensors only — no ONNX on HF:

Model Repo Params License
OpenProvence xsmall hotchpotch/open-provence-reranker-xsmall-v1 ~30M MIT
OpenProvence base hotchpotch/open-provence-reranker-v1 ~130M MIT
OpenProvence large hotchpotch/open-provence-reranker-large-v1 ~310M MIT

The preset/download path runs ort on ONNX, so turnkey presets (download-on-first-use, selectable
size) require a PyTorch→ONNX export + hosting first, and the custom Provence prune+rerank head
may not export cleanly via optimum out of the box.

Suggested sequencing:

  1. Land the pruner module architecture with Custom / Plugin model paths (no presets) so the
    API and consumers can be built and tested.
  2. Export open_provence xsmall + large to ONNX, host under a controlled HF repo, and register them as
    pruner presets.

Acceptance

  • prune_async(query, docs, &PrunerConfig) returns kept-sentence byte spans + per-doc score.
  • Feature-gated pruner / pruner-presets mirroring the reranker; builds across the existing target
    matrix (WASM-safe pruner-presets).
  • Tests mirror the reranker suite (empty docs, threshold filtering, serde round-trip, segmentation).
  • Once ONNX exports exist: xsmall + large selectable as presets, downloaded on first use, absent
    by default.

Metadata

Metadata

Assignees

Labels

enhancementNew feature or request

Fields

No fields configured for Feature.

Projects

Status
Todo

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions