Обновление весов модели

.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
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

config.json

@@ -1,5 +1,5 @@
 {
-  "_name_or_path": "ai-sage/Giga-Embeddings-instruct",
+  "_name_or_path": "/home/jovyan/ekolodin/gigachat-embeddings/ckpt/multitask_prenorm_lr2e-5/checkpoint-6591",
   "add_eos": true,
   "add_pad_token": true,
   "architectures": [
@@ -10,7 +10,7 @@
     "AutoModel": "modeling_gigarembed.GigarEmbedModel"
   },
   "hidden_size": 2048,
-  "is_mask_instruction": false,
+  "is_mask_instruction": true,
   "latent_attention_config": {
     "cross_dim_head": 2048,
     "hidden_dim": 2048,
@@ -21,7 +21,8 @@
   "model_type": "gigarembed",
   "padding_side": "right",
   "text_config": {
-    "_name_or_path": "ai-sage/Giga-Embeddings-instruct",
+    "_attn_implementation_autoset": false,
+    "_name_or_path": "/home/jovyan/ekolodin/models/qiwiembed2.5_3b_pretrain/",
     "activation_checkpoint_layers_num": null,
     "add_cross_attention": false,
     "architectures": [
@@ -78,7 +79,7 @@
     "num_attention_heads": 16,
     "num_beam_groups": 1,
     "num_beams": 1,
-    "num_hidden_layers": 27,
+    "num_hidden_layers": 36,
     "num_key_value_heads": 2,
     "num_return_sequences": 1,
     "output_attentions": false,
@@ -122,5 +123,5 @@
     "vocab_size": 128256
   },
   "torch_dtype": "float32",
-  "transformers_version": "4.40.0.dev0"
+  "transformers_version": "4.46.3"
 }

config_sentence_transformers.json

@@ -1,9 +1,10 @@
 {
   "__version__": {
-    "sentence_transformers": "2.2.2",
-    "transformers": "4.40.0.dev0",
-    "pytorch": "2.0.1+cu118"
+    "sentence_transformers": "3.3.1",
+    "transformers": "4.46.3",
+    "pytorch": "2.1.1+cu121"
   },
   "prompts": {},
-  "default_prompt_name": null
-}
+  "default_prompt_name": null,
+  "similarity_fn_name": "cosine"
+}

configuration_gigarembed.py

@@ -76,7 +76,6 @@ class LatentAttentionConfig(PretrainedConfig):
         self.cross_dim_head = cross_dim_head
         self._attn_implementation = "eager"
 
-
 class BidirectionalLlamaConfig(LlamaConfig):
     model_type = BIDIR_LLAMA_TYPE
     keys_to_ignore_at_inference = ["past_key_values"]

model-00001-of-00003.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9f5bdcf6ab584d8e3fa1e53ad3061a203125be81e043b11e7c867e04670e8aa7
-size 4913926592
+oid sha256:f32eedfe2127f8e9507427af1d796c6547df4c8e5795e4ea8b3a22a96e782292
+size 4930720644

model-00002-of-00003.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:29f27ca086d141ee606037884f8bb93176a61032230ba1ca92ba0c2fe7615cb2
+oid sha256:edc8c0c52613a2712e8c65b3d8b4249b6e99622c695ee5aec698ca37a5a556d3
 size 4932780264

model-00003-of-00003.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:60c4fbeddfae2f00234d2d9e6646cd7a11914454eaf46e278f85bfb088d1417f
-size 270557856
+oid sha256:8cb8d3f0fb5c526162adc18efcc9e1c13d07088602775e742037a5e53d1531b9
+size 3045246736

model.safetensors.index.json

@@ -1,9 +1,26 @@
 {
   "metadata": {
-    "total_size": 10117234688
+    "total_size": 12908707844
   },
   "weight_map": {
+    "latent_attention_model.cross_attend_blocks.0.fn.to_kv.weight": "model-00001-of-00003.safetensors",
+    "latent_attention_model.cross_attend_blocks.0.fn.to_out.weight": "model-00001-of-00003.safetensors",
+    "latent_attention_model.cross_attend_blocks.0.fn.to_q.weight": "model-00001-of-00003.safetensors",
+    "latent_attention_model.cross_attend_blocks.0.norm.bias": "model-00001-of-00003.safetensors",
+    "latent_attention_model.cross_attend_blocks.0.norm.weight": "model-00001-of-00003.safetensors",
+    "latent_attention_model.cross_attend_blocks.0.norm_context.bias": "model-00001-of-00003.safetensors",
+    "latent_attention_model.cross_attend_blocks.0.norm_context.weight": "model-00001-of-00003.safetensors",
+    "latent_attention_model.cross_attend_blocks.1.fn.net.0.bias": "model-00001-of-00003.safetensors",
+    "latent_attention_model.cross_attend_blocks.1.fn.net.0.weight": "model-00001-of-00003.safetensors",
+    "latent_attention_model.cross_attend_blocks.1.fn.net.2.bias": "model-00001-of-00003.safetensors",
+    "latent_attention_model.cross_attend_blocks.1.fn.net.2.weight": "model-00001-of-00003.safetensors",
+    "latent_attention_model.cross_attend_blocks.1.norm.bias": "model-00001-of-00003.safetensors",
+    "latent_attention_model.cross_attend_blocks.1.norm.weight": "model-00001-of-00003.safetensors",
     "latent_attention_model.latents": "model-00001-of-00003.safetensors",
+    "latent_attention_model.w_lexical.bias": "model-00001-of-00003.safetensors",
+    "latent_attention_model.w_lexical.weight": "model-00001-of-00003.safetensors",
+    "latent_attention_model.w_multi_vector.bias": "model-00001-of-00003.safetensors",
+    "latent_attention_model.w_multi_vector.weight": "model-00001-of-00003.safetensors",
     "model.embed_tokens.weight": "model-00001-of-00003.safetensors",
     "model.layers.0.input_layernorm.weight": "model-00001-of-00003.safetensors",
     "model.layers.0.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
@@ -185,6 +202,33 @@
     "model.layers.26.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
     "model.layers.26.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
     "model.layers.26.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
     "model.layers.3.input_layernorm.weight": "model-00001-of-00003.safetensors",
     "model.layers.3.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
     "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
@@ -194,6 +238,60 @@
     "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
     "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
     "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.30.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.32.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.33.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.34.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.35.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.35.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.35.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.35.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.35.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.35.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.35.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.35.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.35.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
     "model.layers.4.input_layernorm.weight": "model-00001-of-00003.safetensors",
     "model.layers.4.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
     "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",

modeling_gigarembed.py

@@ -3,6 +3,8 @@ import torch
 import os
 import json
 import numpy as np
+import torch.nn.functional as F
+
 from functools import partial
 from contextlib import nullcontext
 from transformers import AutoModel, PreTrainedTokenizerFast, BatchEncoding, DataCollatorWithPadding
@@ -227,15 +229,22 @@ def input_transform_func(
 class PreNorm(torch.nn.Module):
     def __init__(self, dim, fn, context_dim = None):
         super().__init__()
-        # TODO remove this layer, we don't use it
+        self.fn = fn
+        self.norm = torch.nn.LayerNorm(dim)
+        self.norm_context = torch.nn.LayerNorm(context_dim) if exists(context_dim) else None
 
     def forward(self, x, **kwargs):
-        return x
+        x = self.norm(x)
+        if exists(self.norm_context):
+            context = kwargs['context']
+            normed_context = self.norm_context(context)
+            kwargs.update(context = normed_context)
+        return self.fn(x, **kwargs)
 
 class GEGLU(torch.nn.Module):
     def forward(self, x):
         x, gates = x.chunk(2, dim = -1)
-        return x * torch.nn.functional.gelu(gates)
+        return x * F.gelu(gates)
 
 class FeedForward(torch.nn.Module):
     def __init__(self, dim, mult = 4):
@@ -275,17 +284,8 @@ class Attention(torch.nn.Module):
         k, v = self.to_kv(context).chunk(2, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h = h), (q, k, v))
 
-        attn_weights = torch.matmul(q, k.transpose(-1, -2)) / self.scale
-
-        mask_value = torch.finfo(attn_weights.dtype).min
-        mask_value = torch.full([], mask_value, dtype=attn_weights.dtype).to(attn_weights.device)
-
-        padding_mask = mask[:, :, None].repeat(self.heads, 1, 1).bool()
-
-        attn_weights = torch.where(padding_mask, attn_weights, mask_value)
-        attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1)
-
-        out = torch.matmul(attn_weights, v)
+        with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_mem_efficient=True):
+            out = torch.nn.functional.scaled_dot_product_attention(q, k, v)
         out = rearrange(out, '(b h) n d -> b n (h d)', h = h)
         return self.to_out(out)
 
@@ -304,31 +304,35 @@ class LatentAttentionModel(PreTrainedModel):
                     context_dim = dim),
             PreNorm(latent_dim, FeedForward(latent_dim)),
         ])
-        self.output_normalize = config.output_normalize
+
+        self.w_lexical = torch.nn.Linear(latent_dim, 1)
+        self.w_multi_vector = torch.nn.Linear(latent_dim, latent_dim)
+
+        # self.output_normalize = config.output_normalize
         self.register_parameter("latents", torch.nn.Parameter(torch.randn(num_latents, latent_dim)))
+        self._attn_implementation = "eager"
 
     def forward(self, hiddens, attention_mask: torch.Tensor=None):
         # cross-attention block
         cross_attn, cross_ff = self.cross_attend_blocks
         b, *_, device = *hiddens.shape, hiddens.device
         x = repeat(self.latents, 'n d -> b n d', b = b)
-        hiddens = cross_attn(hiddens, context=x, mask=attention_mask) + hiddens
-        hiddens = cross_ff(hiddens) + hiddens
+        output = cross_attn(hiddens, context=x, mask=attention_mask) + hiddens
+        output = cross_ff(output) + output
         if attention_mask != None:
-            s = torch.sum(hiddens * attention_mask.unsqueeze(-1).float(), dim=1)
-            d = attention_mask.sum(dim=1, keepdim=True).float()
-            hiddens = s / d
-            if self.output_normalize:
-                hiddens = torch.nn.functional.normalize(hiddens, p=2, dim=-1)
-        return hiddens
-    
+            s = torch.sum(output * attention_mask.unsqueeze(-1), dim=1)
+            d = attention_mask.sum(dim=1, keepdim=True)
+            output = s / d
+            output = F.normalize(output, p=2, dim=-1)
+        return output
+
 class GigarEmbedModel(PreTrainedModel):
     config_class = GigarEmbedConfig
     _no_split_modules = ["LlamaDecoderLayer", "LatentAttentionModel"]
     
     def __init__(self, config: GigarEmbedConfig):
         super().__init__(config)
-        self.latent_attention_model = AutoModel.from_config(config.latent_attention_config).float()
+        self.latent_attention_model = AutoModel.from_config(config.latent_attention_config)
         self.model = AutoModel.from_config(
             config.text_config,
         ) if config.text_config is not None else None
@@ -339,12 +343,6 @@ class GigarEmbedModel(PreTrainedModel):
         self.mask_type = config.mask_type
         if config.add_pad_token and self.tokenizer is not None:
             self.add_pad_token()
-            
-        self.latent_attention_model.apply(self._init_weights)
-
-    def _init_weights(self, module):
-        if isinstance(module, torch.nn.Linear):
-            torch.nn.init.xavier_normal_(module.weight)
 
     def add_pad_token(self):
         self.tokenizer.pad_token_id = 0
@@ -360,7 +358,7 @@ class GigarEmbedModel(PreTrainedModel):
             # Mask out the instruction tokens for mean-pooling
             attention_mask[:, :instruction_lens] = 0
         features: GigarEmbedFeatures = {
-            'input_ids': batch_dict['input_ids'],
+            'input_ids': torch.tensor(batch_dict.get('input_ids').to(batch_dict.get('input_ids')).long()),
             'attention_mask': batch_dict['attention_mask'],
             'pool_mask': attention_mask,
         }
@@ -410,12 +408,14 @@ class GigarEmbedModel(PreTrainedModel):
     def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, pool_mask: Optional[torch.Tensor]=None, 
                 return_dict: bool=True, **kwargs):
         kwargs.pop('token_type_ids', None)
-        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
-        
-        embeds = self.latent_attention_model(
-            outputs.last_hidden_state,
-            pool_mask,
-        )
+
+        with torch.autocast('cuda', dtype=torch.bfloat16):
+            outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
+
+            if pool_mask is None: pool_mask = attention_mask.clone()
+
+            embeds = self.latent_attention_model(outputs.last_hidden_state, pool_mask)
+
         if not return_dict:
             return (embeds,)
         return {"sentence_embeddings": embeds}

modules.json

@@ -10,11 +10,5 @@
     "name": "1",
     "path": "1_Pooling",
     "type": "sentence_transformers.models.Pooling"
-  },
-  {
-    "idx": 2,
-    "name": "2",
-    "path": "2_Normalize",
-    "type": "sentence_transformers.models.Normalize"
   }
-]
+]

sentence_bert_config.json

@@ -1,4 +1,4 @@
 {
-  "max_seq_length": 4096,
+  "max_seq_length": null,
   "do_lower_case": false
-}
+}

special_tokens_map.json

@@ -1,5 +1,37 @@
 {
-  "bos_token": "<s>",
-  "eos_token": "</s>",
-  "unk_token": "<unk>"
+  "bos_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
+  },
+  "pad_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
 }

tokenizer_config.json

@@ -2076,8 +2076,16 @@
   "bos_token": "<s>",
   "clean_up_tokenization_spaces": true,
   "eos_token": "</s>",
+  "max_length": 512,
   "model_max_length": 1000000000000000019884624838656,
+  "pad_to_multiple_of": null,
+  "pad_token": "<unk>",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "sep_token": "<unk>",
+  "stride": 0,
   "tokenizer_class": "PreTrainedTokenizerFast",
-  "unk_token": "<unk>",
-  "pad_token": "<unk>"
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "<unk>"
 }