upload last-token model

.gitattributes

@@ -0,0 +1 @@
+*.safetensors filter=lfs diff=lfs merge=lfs -text

config.json

@@ -0,0 +1,76 @@
+{
+  "architectures": [
+    "JinaEmbeddingsC1Model"
+  ],
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "modeling_jina_embeddings_c1.JinaEmbeddingsC1Model"
+  },
+  "bos_token_id": 151643,
+  "eos_token_id": 151643,
+  "hidden_act": "silu",
+  "hidden_size": 896,
+  "initializer_range": 0.02,
+  "intermediate_size": 4864,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "matryoshka_dims": [
+    64,
+    128,
+    256,
+    512,
+    896
+  ],
+  "max_position_embeddings": 32768,
+  "max_window_layers": 24,
+  "model_type": "qwen2",
+  "num_attention_heads": 14,
+  "num_hidden_layers": 24,
+  "num_key_value_heads": 2,
+  "prompt_names": [
+    "query",
+    "passage"
+  ],
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000.0,
+  "sliding_window": null,
+  "task_names": [
+    "nl2code",
+    "qa",
+    "code2code",
+    "code2nl",
+    "code2completion"
+  ],
+  "tie_word_embeddings": true,
+  "tokenizer_class": "Qwen2TokenizerFast",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.53.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e4be672e203959409ac76ded4387f324d17539f3456cab628bd66bd5d0439675
+size 988096088

modeling_jina_embeddings_c1.py

@@ -0,0 +1,180 @@
+from typing import List, Union
+
+import torch
+import numpy as np
+
+from transformers.utils import is_flash_attn_2_available
+from transformers.models.qwen2 import Qwen2Model
+from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+
+
+INSTRUCTION_CONFIG = {
+    "nl2code": {
+        "query": "Find the most relevant code snippet given the following query:\n",
+        "passage": "Candidate code snippet:\n"
+    },
+    "qa": {
+        "query": "Find the most relevant answer given the following question:\n",
+        "passage": "Candidate answer:\n"
+    },
+    "code2code": {
+        "query": "Find an equivalent code snippet given the following code snippet:\n",
+        "passage": "Candidate code snippet:\n"
+    },
+    "code2nl": {
+        "query": "Find the most relevant comment given the following code snippet:\n",
+        "passage": "Candidate comment:\n"
+    },
+    "code2completion": {
+        "query": "Find the most relevant completion given the following start of code snippet:\n",
+        "passage": "Candidate completion:\n"
+    }
+}
+
+
+def batch(iterable, n=1):
+    items = len(iterable)
+    for ndx in range(0, items, n):
+        yield iterable[ndx : min(ndx + n, items)]
+
+
+def mean_pooling(model_output, attention_mask):
+    token_embeddings = model_output[0]
+    input_mask_expanded = (
+        attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+    )
+    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
+        input_mask_expanded.sum(1), min=1e-9
+    )
+
+
+def last_token_pooling(model_output, attention_mask):
+    token_embeddings = model_output[0]
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return token_embeddings[:, -1]
+    else:
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = token_embeddings.shape[0]
+        return token_embeddings[torch.arange(batch_size, device=token_embeddings.device), sequence_lengths].float()
+
+
+class JinaEmbeddingsC1Model(Qwen2Model):
+    def __init__(self, config: Qwen2Config):
+        Qwen2Model.__init__(self, config)
+        self.instructions = INSTRUCTION_CONFIG
+    
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: torch.Tensor,
+        **kwargs
+    ) -> List[torch.Tensor]:
+        """
+        Forward pass through the model.
+        """
+        batch_model_output = super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            **kwargs
+        )
+        batch_sentence_embeddings = last_token_pooling(
+            batch_model_output, attention_mask
+        )
+        return batch_sentence_embeddings
+
+
+    def encode(
+        self,
+        sentences: List[str],
+        batch_size: int = 32,
+        max_length: int = 32768,
+        task: str = "nl2code",
+        prompt_name: str = "query",
+        return_numpy: bool = False,
+        truncate_dim: int = 896,
+    ) -> Union[np.ndarray, List[torch.Tensor]]:
+        """
+        Encodes a list of texts into embeddings.
+        Args:
+            sentences: list of text strings to encode
+            batch_size: Number of texts to process at once
+            max_length: Maximum token length for text processing
+            task: Type of retrieval task ('nl2code', 'qa', or 'code2code')
+            prompt_name: Type of text being encoded ('query' or 'passage')
+            return_numpy: Whether to return numpy arrays instead of torch tensors
+            truncate_dim: Dimension to truncate embeddings to (64, 128, 256, 512, or 896)
+        Returns:
+            List of text embeddings as tensors or numpy arrays
+        """
+        assert task in self.config.task_names, \
+            f"Invalid task: {task}. Must be one of {self.config.task_names}."
+        assert prompt_name in self.config.prompt_names, \
+            f"Invalid prompt name: {prompt_name}. Must be one of {self.config.prompt_names}."
+        assert truncate_dim in self.config.matryoshka_dims, \
+            f"Invalid embedding dimension: {truncate_dim}. Must be one of {self.config.matryoshka_dims}."
+
+        instruction = self.instructions[task][prompt_name]
+        sentences = [f'{instruction}{sentence}' for sentence in sentences]
+        embeddings = []
+
+        self.eval()
+
+        with torch.inference_mode():
+            for batch_of_sentences in batch(sentences, n=batch_size):
+                batch_encoded_input = self.tokenizer(
+                    batch_of_sentences,
+                    padding=True,
+                    truncation=True,
+                    return_tensors="pt",
+                    max_length=max_length
+                ).to(self.device)
+
+                batch_sentence_embeddings = self(
+                    **batch_encoded_input,
+                    output_attentions=False,
+                    return_dict=True,
+                    max_length=max_length
+                )
+
+                batch_sentence_embeddings = batch_sentence_embeddings[:, :truncate_dim]
+                batch_sentence_embeddings = torch.nn.functional.normalize(
+                    batch_sentence_embeddings, p=2, dim=-1
+                ).to("cpu")
+
+                embeddings.append(batch_sentence_embeddings)
+
+        if return_numpy:
+             return np.concatenate([b.numpy() for b in embeddings], axis=0)
+        return [t for b in embeddings for t in torch.unbind(b, dim=0)]        
+        
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path,
+        *args,
+        **kwargs,
+    ):
+        """
+        Loads a pretrained model.
+        """
+        if "torch_dtype" not in kwargs:
+            kwargs["torch_dtype"] = "auto"
+        
+        if "attn_implementation" not in kwargs:
+            kwargs["attn_implementation"] = "flash_attention_2" if is_flash_attn_2_available() else "sdpa"
+        
+        model = super().from_pretrained(
+            pretrained_model_name_or_path, *args, **kwargs
+        )
+
+        model.tokenizer = Qwen2TokenizerFast.from_pretrained(
+            pretrained_model_name_or_path,
+            trust_remote_code=True
+        )
+
+        return model
+