Upload model

config.json

@@ -0,0 +1,36 @@
+{
+  "_name_or_path": "buffer-embedding-002",
+  "apply_residual_connection_post_layernorm": false,
+  "architectures": [
+    "DualModel"
+  ],
+  "attention_dropout": 0.0,
+  "attention_softmax_in_fp32": true,
+  "auto_map": {
+    "AutoModel": "embedding_model.DualModel"
+  },
+  "bias_dropout_fusion": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "hidden_dropout": 0.0,
+  "hidden_size": 1536,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "masked_softmax_fusion": true,
+  "model_type": "bloom",
+  "n_head": 16,
+  "n_inner": null,
+  "n_layer": 24,
+  "offset_alibi": 100,
+  "pad_token_id": 3,
+  "pretraining_tp": 1,
+  "seq_length": 2048,
+  "skip_bias_add": true,
+  "skip_bias_add_qkv": false,
+  "slow_but_exact": false,
+  "torch_dtype": "float16",
+  "transformers_version": "4.31.0",
+  "unk_token_id": 0,
+  "use_cache": true,
+  "vocab_size": 46145
+}

configuration.py

@@ -0,0 +1,55 @@
+from transformers.utils import logging
+from transformers.configuration_utils import PretrainedConfig
+
+
+logger = logging.get_logger(__name__)
+
+INTERNLM_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
+
+
+class BufferEmbeddingConfig(PretrainedConfig):
+    model_type = "buffer_embedding"
+    _auto_class = "AutoConfig"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    attribute_map = {
+        "num_hidden_layers": "n_layer",
+        "num_attention_heads": "n_head",
+    }
+    def __init__(
+        self,
+        vocab_size=250880,
+        hidden_size=64,
+        n_layer=2,
+        n_head=8,
+        layer_norm_epsilon=1e-5,
+        initializer_range=0.02,
+        use_cache=True,
+        bos_token_id=1,
+        eos_token_id=2,
+        apply_residual_connection_post_layernorm=False,
+        hidden_dropout=0.0,
+        attention_dropout=0.0,
+        pretraining_tp=1,  # TP rank used when training with megatron
+        slow_but_exact=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        # Backward compatibility with n_embed kwarg
+        n_embed = kwargs.pop("n_embed", None)
+        self.hidden_size = hidden_size if n_embed is None else n_embed
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_range = initializer_range
+        self.use_cache = use_cache
+        self.pretraining_tp = pretraining_tp
+        self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm
+        self.hidden_dropout = hidden_dropout
+        self.attention_dropout = attention_dropout
+
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.slow_but_exact = slow_but_exact
+
+        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+

embedding_model.py

@@ -0,0 +1,29 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import BloomForCausalLM, PreTrainedModel
+from .configuration import BufferEmbeddingConfig
+
+
+class DualModel(PreTrainedModel):
+    config_class = BufferEmbeddingConfig
+    _auto_class = "AutoModel"
+    def __init__(self, config):
+        super(DualModel, self).__init__(config)
+        self.model = BloomForCausalLM(config)#.from_pretrained('Langboat/bloom-800m-zh')
+        self.classifier = nn.Linear(1536, 1536)
+        self.hidden = nn.Sequential(nn.Linear(1536, 1536),
+                                    nn.Tanh())
+    def forward(self,
+                input_ids,
+                token_type_ids=None,
+                position_ids_ids=None,
+                attention_mask=None,
+                labels=None
+               ):
+        attention_mask = torch.ne(input_ids, 3)  # size: batch_size, max_len
+
+        y = self.model(input_ids, attention_mask=attention_mask, output_hidden_states=True)
+        embedding = (y.hidden_states[-1]*attention_mask.unsqueeze(-1)).sum(1)/attention_mask.sum(1).unsqueeze(-1)
+        embedding = self.classifier(self.hidden(embedding))
+        return F.normalize(embedding, p=2, dim=-1)

pytorch_model.bin

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