chunkable-sentence-transformer

Custom Sentence Transformers modules that add support for vertically chunked inference to encode arbitrarily long texts with constant memory by processing fixed-size vertical chunks through all model layers sequentially.

What this repository provides

• ExtendedTransformer: extends the Sentence Transformers Transformer module with explicit use_cache and cache arguments. The model's cache (past_key_values / cache_params) is returned as features["cache"] whenever the model produces one, so recurrent states can carry across chunks.
• ExtendedPooling: extends the Sentence Transformers Pooling module so padded inputs are pooled from per-sequence accumulators. The accumulators are its cache, so the pooled embedding can be updated chunk by chunk. A single call without a cache gives the same result as upstream Pooling. Supports the upstream pooling modes (mean, mean_sqrt_len_tokens, weightedmean, max, cls, lasttoken).
• ChunkedModule: wraps a sequence of modules and, when vertical_chunk_size is passed to encode, splits the input into chunks, runs each chunk through the full stack while carrying each module's cache forward, and returns the last chunk's sentence embedding.

Modules inside a ChunkedModule declare use_cache and cache explicitly (**kwargs does not count) and return their cache as features["cache"]. Modules that declare neither are treated as stateless, like Normalize.

Per-token embeddings are not reassembled by default, since that would bring back the full-length memory that chunking avoids. Pass return_token_embeddings=True to the ChunkedModule's forward to get them back.

Chunked inference is not supported for flattened/packed inputs (cu_seq_lens_q), those fall back to the unchunked forward with a warning.

Usage

This repository is designed to be referenced directly from Hugging Face model configs via modules.json, so that models can be loaded with trust_remote_code=True without any local installation.

Constant-memory inference is then available via the vertical_chunk_size encode parameter:

from sentence_transformers import SentenceTransformer

model = SentenceTransformer(
    "dynatrace-oss/llama-embed-mamba2-7b", trust_remote_code=True
)
embeddings = model.encode(["Your long document text here..."], vertical_chunk_size=512)

A checkpoint uses a ChunkedModule around the transformer and pooling, followed by the rest of the pipeline in modules.json:

[
  {"idx": 0, "name": "0", "path": "", "type": "dynatrace-oss/chunkable-sentence-transformer--modules.ChunkedModule"},
  {"idx": 1, "name": "1", "path": "1_Normalize", "type": "sentence_transformers.models.Normalize"}
]

The chunked part lists the transformer and the pooling in modules_chunked.json:

[
  {"idx": 0, "name": "0", "path": "", "type": "dynatrace-oss/chunkable-sentence-transformer--modules.ExtendedTransformer"},
  {"idx": 1, "name": "1", "path": "1_Pooling", "type": "dynatrace-oss/chunkable-sentence-transformer--modules.ExtendedPooling"}
]

Requirements

pip install sentence-transformers==6.2.0

Models

This code was created for the following embedding models:

• dynatrace-oss/llama-embed-mamba2-7b
• dynatrace-oss/llama-embed-mamba2-1.3b

License

Apache-2.0