Upload new script

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* 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
+model_0/tokenizer.json filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,4 +1,4 @@
 .env
 poetry.lock
 **/__pycache__/
-model/
+temp/

model_0/config.json

@@ -0,0 +1,67 @@
+{
+  "_name_or_path": "jinaai/jina-embeddings-v3",
+  "architectures": [
+    "XLMRobertaLoRA"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "auto_map": {
+    "AutoConfig": "jinaai/xlm-roberta-flash-implementation--configuration_xlm_roberta.XLMRobertaFlashConfig",
+    "AutoModel": "jinaai/xlm-roberta-flash-implementation--modeling_lora.XLMRobertaLoRA",
+    "AutoModelForMaskedLM": "jinaai/xlm-roberta-flash-implementation--modeling_xlm_roberta.XLMRobertaForMaskedLM",
+    "AutoModelForPreTraining": "jinaai/xlm-roberta-flash-implementation--modeling_xlm_roberta.XLMRobertaForPreTraining"
+  },
+  "bos_token_id": 0,
+  "classifier_dropout": null,
+  "emb_pooler": null,
+  "eos_token_id": 2,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-05,
+  "load_trained_adapters": true,
+  "lora_adaptations": [
+    "retrieval.query",
+    "retrieval.passage",
+    "separation",
+    "classification",
+    "text-matching"
+  ],
+  "lora_alpha": 1,
+  "lora_dropout_p": 0.0,
+  "lora_main_params_trainable": false,
+  "lora_rank": 4,
+  "matryoshka_dimensions": [
+    32,
+    64,
+    128,
+    256,
+    512,
+    768,
+    1024
+  ],
+  "max_position_embeddings": 8194,
+  "model_type": "xlm-roberta",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 24,
+  "output_past": true,
+  "pad_token_id": 1,
+  "position_embedding_type": "rotary",
+  "rotary_emb_base": 20000.0,
+  "task_instructions": {
+    "classification": "",
+    "retrieval.passage": "Represent the document for retrieval: ",
+    "retrieval.query": "Represent the query for retrieving evidence documents: ",
+    "separation": "",
+    "text-matching": ""
+  },
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.46.3",
+  "truncate_dim": null,
+  "type_vocab_size": 1,
+  "use_cache": true,
+  "use_flash_attn": true,
+  "use_reentrant": false,
+  "vocab_size": 250002
+}

model_0/model.safetensors

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

model_0/special_tokens_map.json

@@ -0,0 +1,51 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "cls_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "<mask>",
+    "lstrip": true,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

model_0/tokenizer.json

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

model_0/tokenizer_config.json

@@ -0,0 +1,54 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "250001": {
+      "content": "<mask>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "mask_token": "<mask>",
+  "model_max_length": 8194,
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "tokenizer_class": "XLMRobertaTokenizer",
+  "unk_token": "<unk>"
+}

model_1/config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "Snowflake/snowflake-arctic-embed-s",
+  "architectures": [
+    "BertModel"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "classifier_dropout": null,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 384,
+  "initializer_range": 0.02,
+  "intermediate_size": 1536,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.46.3",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 30522
+}

model_1/model.safetensors

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

model_1/special_tokens_map.json

@@ -0,0 +1,37 @@
+{
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

model_1/tokenizer_config.json

@@ -0,0 +1,64 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "101": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "102": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "103": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "do_basic_tokenize": true,
+  "do_lower_case": true,
+  "mask_token": "[MASK]",
+  "max_length": 512,
+  "model_max_length": 512,
+  "never_split": null,
+  "pad_to_multiple_of": null,
+  "pad_token": "[PAD]",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "sep_token": "[SEP]",
+  "stride": 0,
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "[UNK]"
+}

modeling_jina_v3_arctic_s.py

@@ -1,7 +1,9 @@
+import os
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import BertModel, PreTrainedModel, BertConfig, PretrainedConfig, AutoModel
+from transformers import BertModel, PreTrainedModel, BertConfig, PretrainedConfig, XLMRobertaTokenizerFast, \
+AutoModel, PreTrainedTokenizerFast, BertTokenizer, PreTrainedTokenizer, PreTrainedTokenizerBase, AutoTokenizer, XLMRobertaTokenizer
 from typing import *
 
 
@@ -12,48 +14,12 @@ class ConcatModelConfig(PretrainedConfig):
         super().__init__(**kwargs)
 
 
-# See https://huggingface.co/Marqo/marqo-chimera-arctic-bge-m
 class ConcatModel(PreTrainedModel):
     config_class = ConcatModelConfig
 
-    def __init__(self, config: ConcatModelConfig):
-        super().__init__(config)
-        bert_config_1 = BertConfig(
-            vocab_size=30522,
-            hidden_size=384,
-            num_hidden_layers=12,
-            num_attention_heads=12,
-            intermediate_size=1536,
-            hidden_act="gelu",
-            hidden_dropout_prob=0.1,
-            attention_probs_dropout_prob=0.1,
-            max_position_embeddings=512,
-            type_vocab_size=2,
-            initializer_range=0.02,
-            layer_norm_eps=1e-12,
-        )
-
-        bert_config_2 = BertConfig(
-            vocab_size=30522,
-            hidden_size=384,
-            num_hidden_layers=12,
-            num_attention_heads=12,
-            intermediate_size=1536,
-            hidden_act="gelu",
-            hidden_dropout_prob=0.1,
-            attention_probs_dropout_prob=0.1,
-            max_position_embeddings=512,
-            type_vocab_size=2,
-            initializer_range=0.02,
-            layer_norm_eps=1e-12,
-        )
-
-        self.model = nn.ModuleDict(
-            {
-                "model_0": BertModel(bert_config_1),
-                "model_1": BertModel(bert_config_2),
-            }
-        )
+    def __init__(self, models):
+        super().__init__(ConcatModelConfig())
+        self.models = models
 
     def forward(
         self,
@@ -63,31 +29,129 @@ class ConcatModel(PreTrainedModel):
         **kwargs
     ) -> torch.Tensor:
         embeddings = []
-        for _, model in self.model.items():
-            model_output = model(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                token_type_ids=token_type_ids,
-            )
+        for i, model in enumerate(self.models):
+            if i == 0:
+                model_output = model(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    token_type_ids=token_type_ids,
+                )
+            else:
+                model_output = model(
+                    input_ids=kwargs["input_ids_" + str(i)],
+                    attention_mask=kwargs["attention_mask_" + str(i)],
+                    token_type_ids=kwargs.get("token_type_ids_" + str(i)),
+                )
             pooled_output = model_output[0][:, 0]
             pooled_output = F.normalize(pooled_output, p=2, dim=-1)
             embeddings.append(pooled_output)
 
         return torch.cat(embeddings, dim=-1)
+    
+    def save_pretrained(self, save_directory):
+        for i, model in enumerate(self.models):
+            path = os.path.join(save_directory, f"model_{i}")
+            model.save_pretrained(path)
 
-    def load_weights_from_automodels(
-        self, in_models: List[str], has_pooling_layer: List[bool]
-    ):
-        model_list = []
-        for i, model_name in enumerate(in_models):
-            model = AutoModel.from_pretrained(
-                model_name,
-                add_pooling_layer=has_pooling_layer[i],
-                trust_remote_code=True,
-            )
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        models = [
+            AutoModel.from_pretrained(f"{pretrained_model_name_or_path}/model_0", trust_remote_code=True),
+            AutoModel.from_pretrained(f"{pretrained_model_name_or_path}/model_1", trust_remote_code=True)
+        ]
+        return cls(models)
+    
+    def __repr__(self):
+        s = "ConcatModel with models:"
+        for i, model in enumerate(self.models):
+            s += f"\nModel {i}: {model}"
+        return s
+    
+    def eval(self):
+        for model in self.models:
             model.eval()
-            model_list.append(model)
+        return self
+    
+    def cuda(self):
+        for i, model in enumerate(self.models):
+            self.models[i] = model.cuda()
+        return self
+
+class ConcatTokenizer(PreTrainedTokenizer):
+    """
+    A custom tokenizer to handle multiple tokenizers for concatenated models.
+    This tokenizer will delegate tokenization to the underlying individual tokenizers.
+    """
+    
+    def __init__(self, tokenizers, **kwargs):
+        self.tokenizers = tokenizers
+
+    def tokenize(self, text: str, **kwargs):
+        """
+        Tokenizes text using all tokenizers.
+        """
+        return [tokenizer.tokenize(text, **kwargs) for tokenizer in self.tokenizers]
+    
+    def __call__(self, text, **kwargs):
+        """
+        Tokenize and encode input text using all tokenizers.
+        Returns combined inputs.
+        """
+        combined_inputs = {}
+        for i, tokenizer in enumerate(self.tokenizers):
+            encoded = tokenizer(text, **kwargs)
+            # Prefix the keys to distinguish between tokenizers
+            for key, value in encoded.items():
+                _key = key
+                if i > 0:
+                    _key = f"{key}_{i}"
+                combined_inputs[_key] = value
+        
+        return combined_inputs
+    
+    def batch_encode_plus(self, batch_text_or_text_pairs, **kwargs):
+        """
+        Handles batch tokenization for all tokenizers.
+        """
+        combined_inputs = {}
+        for i, tokenizer in enumerate(self.tokenizers):
+            encoded_batch = tokenizer.batch_encode_plus(batch_text_or_text_pairs, **kwargs)
+            for key, value in encoded_batch.items():
+                _key = key
+                if i > 0:
+                    _key = f"{key}_{i}"
+                combined_inputs[_key] = value
+
+        return combined_inputs
+
+    def decode(self, token_ids, **kwargs):
+        """
+        Decode tokens using the first tokenizer (or specific one, if required).
+        """
+        # Choose the primary tokenizer for decoding (default: model_0)
+        return self.tokenizers[0].decode(token_ids, **kwargs)
+    
+    def save_pretrained(self, save_directory):
+        """
+        Save the tokenizers to the specified directory.
+        """
+        for i, tokenizer in enumerate(self.tokenizers):
+            path = os.path.join(save_directory, f"model_{i}")
+            tokenizer.save_pretrained(path)
 
-        self.model = nn.ModuleDict(
-            {f"model_{i}": model for i, model in enumerate(model_list)}
-        )
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        """
+        Load the tokenizers from the specified directory.
+        """
+        tokenizers = [
+            XLMRobertaTokenizerFast.from_pretrained(f"{pretrained_model_name_or_path}/model_0"),
+            BertTokenizer.from_pretrained(f"{pretrained_model_name_or_path}/model_1")
+        ]
+        return cls(tokenizers)
+    
+    def __repr__(self):
+        s = "ConcatTokenizer with tokenizers:"
+        for i, tokenizer in enumerate(self.tokenizers):
+            s += f"\nTokenizer {i}: {tokenizer}"
+        return s

pyproject.toml

@@ -22,4 +22,7 @@ torch = "2.5.0"
 einops = "0.8.0"
 numpy = "1.26.4"
 setuptools = "75.6.0"
-wheel = "0.45.1"
+wheel = "0.45.1"
+tiktoken = "0.8.0"
+protobuf = "5.29.0"
+sentencepiece = "0.2.0"

save_safetensors.py

@@ -1,16 +1,20 @@
-from transformers import BertTokenizer
-from modeling_jina_v3_arctic_s import ConcatModel, ConcatModelConfig
+from transformers import AutoModel, BertTokenizer, XLMRobertaTokenizerFast
+from temp.modeling_jina_v3_arctic_s import ConcatModel, ConcatTokenizer
 
-config = ConcatModelConfig()
-model = ConcatModel(config)
-model.load_weights_from_automodels(
-    in_models=['jinaai/jina-embeddings-v3', 'Snowflake/snowflake-arctic-embed-s'],
-    has_pooling_layer=[True, True]
-)
 
-tokenizer = BertTokenizer(vocab_file='vocab.txt')
+models = [
+    AutoModel.from_pretrained("jinaai/jina-embeddings-v3", trust_remote_code=True),
+    AutoModel.from_pretrained("Snowflake/snowflake-arctic-embed-s", trust_remote_code=True)
+]
+model = ConcatModel(models)
 
-output_path = 'model'
+tokenizers = [
+    XLMRobertaTokenizerFast.from_pretrained("jinaai/jina-embeddings-v3"),
+    BertTokenizer.from_pretrained("Snowflake/snowflake-arctic-embed-s")
+]
+tokenizer = ConcatTokenizer(tokenizers)
+
+output_path = 'temp'
 model.save_pretrained(output_path)
 tokenizer.save_pretrained(output_path)