feat: chunkable transformer and last index pooling

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,76 @@
+---
+license: apache-2.0
+library_name: sentence-transformers
+tags:
+- transformers
+- sentence-transformers
+- vertical-chunking
+---
+
+# chunkable-sentence-transformer
+
+Custom [Sentence Transformers](https://www.sbert.net/) 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
+
+- **`ChunkableTransformer`:** extends the Sentence Transformers `Transformer` module with a `vertical_chunk_size` parameter. When set, input sequences are split into chunks that are each processed through the full model depth, carrying the recurrent states across chunks instead of materializing the entire sequence in memory at once.
+- **`LastIndexPooling`:** pools the embedding from the last token regardless of padding, which allows us to only retain outputs of the final chunk when using left padding with chunked inference.
+
+## 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:
+
+```json
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "dynatrace-oss/chunkable-sentence-transformer--models.ChunkableTransformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_LastIndexPooling",
+    "type": "dynatrace-oss/chunkable-sentence-transformer--models.LastIndexPooling"
+  }
+]
+```
+
+Constant-memory inference is then available via the `vertical_chunk_size` encode parameter:
+
+```python
+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)
+```
+
+## Requirements
+
+```bash
+pip install sentence-transformers
+```
+
+## Models
+
+This code was created for the following embedding models:
+
+- [dynatrace-oss/llama-embed-mamba2-7b](https://huggingface.co/dynatrace-oss/llama-embed-mamba2-7b)
+- [dynatrace-oss/llama-embed-mamba2-1.3b](https://huggingface.co/dynatrace-oss/llama-embed-mamba2-1.3b)
+
+## Open Source Integration Roadmap
+
+Our goal is to integrate all necessary changes to simplify the adoption of vertically chunked inference for other models:
+
+> [!Note]
+> ⚪ Planned | 🟡 In Progress | 🟢 Integrated
+
+- ⚪ **sentence-transformers:** Last index pooling
+- ⚪ **sentence-transformers:** Native vertical chunking support for transformers
+
+*This list will be updated as integration progresses.*
+
+## License
+
+Apache-2.0

models.py

@@ -0,0 +1,185 @@
+from sentence_transformers.models import Transformer, Module
+from transformers.utils.import_utils import is_peft_available
+import torch
+
+
+class ChunkableTransformer(Transformer):
+    config_keys = Transformer.config_keys + ["output_all_token_embeddings"]
+
+    def __init__(self, *args, output_all_token_embeddings=True, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.cache_param_name = (
+            "past_key_values"
+            if "past_key_values" in self.model_forward_params
+            else "cache_params"
+        )
+        self.model_forward_params = self.model_forward_params - {
+            "use_cache",
+            "past_key_values",
+            "cache_params",
+        }
+        self.forward_kwargs = set(self.forward_kwargs) | {
+            "vertical_chunk_size",
+            "output_all_token_embeddings",
+        }
+        self.output_all_token_embeddings = output_all_token_embeddings
+
+    def __repr__(self) -> str:
+        return f"ChunkableTransformer({dict(self.get_config_dict(), architecture=self.auto_model.__class__.__name__)})"
+
+    def forward(
+        self,
+        features: dict[str, torch.Tensor],
+        vertical_chunk_size: int | None = None,
+        output_all_token_embeddings: bool | None = None,
+        **kwargs,
+    ) -> dict[str, torch.Tensor]:
+        """
+        Forward pass through the transformer model.
+
+        This method processes the input features through the underlying transformers model
+        and returns the token embeddings along with any other relevant outputs.
+
+        Notes:
+            - Only passes arguments that are expected by the underlying transformer model
+
+        Args:
+            features (dict[str, torch.Tensor]): Input features dictionary containing at least
+                'input_ids' and 'attention_mask'. May also contain other tensors required by
+                the underlying transformer model.
+            vertical_chunk_size (int | None): Size of vertical chunks for processing. If None, no chunking is applied.
+            output_all_token_embeddings (bool | None): Whether to return all token embeddings or only the token embeddings of the last vertical chunk. If None, uses the default value from the model configuration.
+            **kwargs: Additional keyword arguments to pass to the underlying transformer model.
+
+        Returns:
+            dict[str, torch.Tensor]: Updated features dictionary containing the input features, plus:
+                - 'token_embeddings': Token-level embeddings from the transformer model
+                - 'attention_mask': Possibly modified attention mask if using PeftModel with prompt learning
+                - 'all_layer_embeddings': If the model outputs hidden states, contains embeddings from all layers
+        """
+        trans_features = {
+            key: value
+            for key, value in features.items()
+            if key in self.model_forward_params
+        }
+
+        input_length = features.get("input_ids", features.get("inputs_embeds")).size(1)
+        output_all_token_embeddings = (
+            output_all_token_embeddings
+            if output_all_token_embeddings is not None
+            else self.output_all_token_embeddings
+        )
+        output_hidden_states = trans_features.get(
+            "output_hidden_states", self.auto_model.config.output_hidden_states
+        )
+        use_vertical_chunking = (
+            vertical_chunk_size is not None and input_length > vertical_chunk_size
+        )
+        all_token_embedding_chunks = None
+        all_hidden_states = None
+
+        if use_vertical_chunking:
+            cache_params = None
+            chunked_features = {
+                k: v.split(vertical_chunk_size, dim=1)
+                for k, v in trans_features.items()
+                if k != "attention_mask"
+            }
+
+            if "attention_mask" in trans_features:
+                chunked_features["attention_mask"] = (
+                    trans_features["attention_mask"][:, :chunk_index]
+                    for chunk_index in range(
+                        vertical_chunk_size,
+                        input_length + vertical_chunk_size,
+                        vertical_chunk_size,
+                    )
+                )
+
+            # pivot from dict of lists to list of dicts
+            chunked_features = map(
+                dict,
+                [
+                    zip(chunked_features, values, strict=True)
+                    for values in zip(*chunked_features.values(), strict=True)
+                ],
+            )
+
+            if output_all_token_embeddings:
+                all_token_embedding_chunks = ()
+
+                if output_hidden_states:
+                    all_hidden_states = ()
+
+            for chunk_features in chunked_features:
+                outputs = self.auto_model(
+                    **chunk_features,
+                    **kwargs,
+                    use_cache=True,
+                    cache_params=cache_params,
+                    return_dict=True,
+                )
+                cache_params = getattr(outputs, self.cache_param_name)
+
+                if all_token_embedding_chunks is not None:
+                    all_token_embedding_chunks += (outputs[0],)
+                if all_hidden_states is not None and "hidden_states" in outputs:
+                    all_hidden_states += (outputs["hidden_states"],)
+
+            if output_all_token_embeddings:
+                token_embeddings = torch.cat(all_token_embedding_chunks, dim=1)
+            else:
+                token_embeddings = outputs[0]
+        else:
+            outputs = self.auto_model(**trans_features, **kwargs, return_dict=True)
+
+        if all_token_embedding_chunks is not None:
+            features["token_embeddings"] = torch.cat(all_token_embedding_chunks, dim=1)
+        else:
+            features["token_embeddings"] = outputs[0]
+
+        if all_hidden_states is not None and len(all_hidden_states) > 0:
+            features["all_layer_embeddings"] = torch.cat(all_hidden_states, dim=1)
+        elif self.auto_model.config.output_hidden_states and "hidden_states" in outputs:
+            features["all_layer_embeddings"] = outputs.get("hidden_states", None)
+
+        # If the AutoModel is wrapped with a PeftModelForFeatureExtraction, then it may have added virtual tokens
+        # We need to extend the attention mask to include these virtual tokens, or the pooling will fail
+        if is_peft_available():
+            from peft import PeftModelForFeatureExtraction
+
+            if (
+                isinstance(self.auto_model, PeftModelForFeatureExtraction)
+                and self.auto_model.active_peft_config.is_prompt_learning
+            ):
+                batch_size = token_embeddings.size(0)
+                attention_mask = features["attention_mask"]
+                prefix_attention_mask = torch.ones(
+                    batch_size,
+                    self.auto_model.active_peft_config.num_virtual_tokens,
+                    device=attention_mask.device,
+                )
+                features["attention_mask"] = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        return features
+
+
+class LastIndexPooling(Module):
+    """Pools the embedding of the last token in the sequence.
+
+    Note: This module ignores ``attention_mask`` and always returns the embedding
+    at index ``-1``. It therefore requires inputs to be left-padded so that
+    the final non-pad token aligns with the last position. Combined with
+    ``ChunkableTransformer``, this makes it safe to discard all but the final
+    vertical chunk's outputs during chunked inference.
+    """
+
+    def __repr__(self) -> str:
+        return f"LastIndexPooling({self.get_config_dict()})"
+
+    def save(self, *args, **kwargs) -> None:
+        return
+
+    def forward(self, features: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        features["sentence_embedding"] = features["token_embeddings"][:, -1]
+        return features