Initial commit

.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

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

model_0/config.json

@@ -0,0 +1,50 @@
+{
+  "_name_or_path": "Alibaba-NLP/gte-multilingual-base",
+  "architectures": [
+    "NewModel"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "auto_map": {
+    "AutoConfig": "Alibaba-NLP/new-impl--configuration.NewConfig",
+    "AutoModel": "Alibaba-NLP/new-impl--modeling.NewModel",
+    "AutoModelForMaskedLM": "Alibaba-NLP/new-impl--modeling.NewForMaskedLM",
+    "AutoModelForMultipleChoice": "Alibaba-NLP/new-impl--modeling.NewForMultipleChoice",
+    "AutoModelForQuestionAnswering": "Alibaba-NLP/new-impl--modeling.NewForQuestionAnswering",
+    "AutoModelForSequenceClassification": "Alibaba-NLP/new-impl--modeling.NewForSequenceClassification",
+    "AutoModelForTokenClassification": "Alibaba-NLP/new-impl--modeling.NewForTokenClassification"
+  },
+  "classifier_dropout": 0.0,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "id2label": {
+    "0": "LABEL_0"
+  },
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "label2id": {
+    "LABEL_0": 0
+  },
+  "layer_norm_eps": 1e-12,
+  "layer_norm_type": "layer_norm",
+  "logn_attention_clip1": false,
+  "logn_attention_scale": false,
+  "max_position_embeddings": 8192,
+  "model_type": "new",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pack_qkv": true,
+  "pad_token_id": 1,
+  "position_embedding_type": "rope",
+  "rope_scaling": {
+    "factor": 8.0,
+    "type": "ntk"
+  },
+  "rope_theta": 20000,
+  "torch_dtype": "float32",
+  "transformers_version": "4.46.3",
+  "type_vocab_size": 1,
+  "unpad_inputs": false,
+  "use_memory_efficient_attention": false,
+  "vocab_size": 250048
+}

model_0/model.safetensors

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

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": 32768,
+  "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:440cc6f580c0ef00c2c9433010145fecab69c3637322d674915ef3f120ac23e3
+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_mgte_arctic_s.py

@@ -0,0 +1,177 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import BertModel, PreTrainedModel, BertConfig, PretrainedConfig, XLMRobertaTokenizerFast, \
+AutoModel, PreTrainedTokenizerFast, BertTokenizer, PreTrainedTokenizer, PreTrainedTokenizerBase, AutoTokenizer, XLMRobertaTokenizer
+from typing import *
+
+
+class ConcatModelConfig(PretrainedConfig):
+    model_type = "mgte-arctic-s"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+
+class ConcatModel(PreTrainedModel):
+    config_class = ConcatModelConfig
+
+    def __init__(self, models):
+        super().__init__(ConcatModelConfig())
+        self.models = models
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        token_type_ids: torch.Tensor = None,
+        **kwargs
+    ) -> torch.Tensor:
+        embeddings = []
+        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)
+
+    # @classmethod
+    # def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+    #     print("ConcatModel from_pretrained")
+    #     print(pretrained_model_name_or_path)
+    #     x = super().from_pretrained(pretrained_model_name_or_path)
+    #     print(x)
+    #     models = []
+    #     return cls(models)
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, model_id="model_0", **kwargs):
+        # Identify the subfolder for the model
+        model_path = f"{pretrained_model_name_or_path}/{model_id}"
+        print("---- model_path:", model_path)
+
+        # Load model configuration
+        config = cls.config_class.from_pretrained(model_path, **kwargs)
+        print("---- config:", config)
+
+        # Initialize the model
+        model = cls(config, **kwargs)
+
+        # Load weights
+        state_dict = torch.load(f"{model_path}/pytorch_model.bin", map_location="cpu")
+        model.load_state_dict(state_dict)
+
+        return model
+    
+    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()
+        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)
+
+    @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

@@ -0,0 +1,21 @@
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry]
+name = "mgte-arctic-s"
+version = "1.0.0"
+description = "Upload ConcatModels."
+authors = [
+    "Michael Dinzinger"
+]
+homepage = "https://www.fim.uni-passau.de"
+repository = "https://www.fim.uni-passau.de"
+readme = "README.md"
+license = "MIT"
+package-mode = false
+
+[tool.poetry.dependencies]
+python = ">=3.10,<3.12"
+transformers = "4.46.3"
+torch = "2.5.0"

save_safetensors.py

@@ -0,0 +1,21 @@
+from transformers import AutoModel, BertTokenizer, XLMRobertaTokenizerFast
+from modeling_mgte_arctic_s import ConcatModel, ConcatTokenizer
+
+
+models = [
+    AutoModel.from_pretrained("Alibaba-NLP/gte-multilingual-base", trust_remote_code=True),
+    AutoModel.from_pretrained("Snowflake/snowflake-arctic-embed-s", trust_remote_code=True)
+]
+model = ConcatModel(models)
+
+tokenizers = [
+    XLMRobertaTokenizerFast.from_pretrained("Alibaba-NLP/gte-multilingual-base"),
+    BertTokenizer.from_pretrained("Snowflake/snowflake-arctic-embed-s")
+]
+tokenizer = ConcatTokenizer(tokenizers)
+
+output_path = 'temp'
+model.save_pretrained(output_path)
+tokenizer.save_pretrained(output_path)
+
+print(f'Model saved as {output_path}')