Initial commit

.gitignore

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

config.json

@@ -0,0 +1,12 @@
+{
+  "architectures": [
+    "ConcatModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_arctic_s_bge_small.ConcatModelConfig",
+    "AutoModel": "modeling_arctic_s_bge_small.ConcatModel"
+  },
+  "model_type": "arctic-l-bge-small",
+  "torch_dtype": "float32",
+  "transformers_version": "4.42.3"
+}

model.safetensors

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

modeling_arctic_l_bge_small.py

@@ -0,0 +1,93 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import BertModel, PreTrainedModel, BertConfig, PretrainedConfig, AutoModel
+from typing import *
+
+
+class ConcatModelConfig(PretrainedConfig):
+    model_type = "arctic-l-bge-small"
+
+    def __init__(self, **kwargs):
+        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=1024,
+            num_hidden_layers=24,
+            num_attention_heads=16,
+            intermediate_size=4096,
+            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 forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        token_type_ids: torch.Tensor = None,
+        **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,
+            )
+            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 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,
+            )
+            model.eval()
+            model_list.append(model)
+
+        self.model = nn.ModuleDict(
+            {f"model_{i}": model for i, model in enumerate(model_list)}
+        )

pyproject.toml

@@ -0,0 +1,21 @@
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry]
+name = "arctic-l-bge-small"
+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.42.3"
+torch = "2.5.0"

save_safetensors.py

@@ -0,0 +1,17 @@
+from transformers import BertTokenizer
+from modeling_arctic_l_bge_small import ConcatModel, ConcatModelConfig
+
+config = ConcatModelConfig()
+model = ConcatModel(config)
+model.load_weights_from_automodels(
+    in_models=['Snowflake/snowflake-arctic-embed-l', 'BAAI/bge-small-en-v1.5'],
+    has_pooling_layer=[True, True]
+)
+
+tokenizer = BertTokenizer(vocab_file='vocab.txt')
+
+output_path = 'model'
+model.save_pretrained(output_path)
+tokenizer.save_pretrained(output_path)
+
+print(f'Model saved as {output_path}')

special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

tokenizer_config.json

@@ -0,0 +1,57 @@
+{
+  "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]",
+  "model_max_length": 1000000000000000019884624838656,
+  "never_split": null,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "unk_token": "[UNK]"
+}