taeshahn commited on
Commit
2eaaf70
·
verified ·
1 Parent(s): 519bba8

Add new SentenceTransformer model

Browse files
0_KananaEmbeddingWrapper/kanana2vec/config.json ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "/data/public/model_registry/huggingface/hub/models--datalama--kanana-nano-2.1b-embedding/snapshots/77c87d91735397f0416f5bf595ba48f6af97333a/0_KananaEmbeddingWrapper/kanana2vec",
3
+ "architectures": [
4
+ "Kanana2VecModel"
5
+ ],
6
+ "attention_bias": false,
7
+ "attention_dropout": 0.0,
8
+ "auto_map": {
9
+ "AutoConfig": "configuration_kanana2vec.Kanana2VecConfig",
10
+ "AutoModel": "modeling_kanana2vec.Kanana2VecModel"
11
+ },
12
+ "bos_token_id": 128000,
13
+ "eos_token_id": 128001,
14
+ "head_dim": 128,
15
+ "hidden_act": "silu",
16
+ "hidden_size": 1792,
17
+ "initializer_range": 0.02,
18
+ "intermediate_size": 8064,
19
+ "max_position_embeddings": 8192,
20
+ "mlp_bias": false,
21
+ "model_type": "kanana2vec",
22
+ "num_attention_heads": 24,
23
+ "num_hidden_layers": 32,
24
+ "num_key_value_heads": 8,
25
+ "pad_token_id": 128001,
26
+ "pretraining_tp": 1,
27
+ "rms_norm_eps": 1e-05,
28
+ "rope_scaling": null,
29
+ "rope_theta": 500000.0,
30
+ "tie_word_embeddings": true,
31
+ "torch_dtype": "float32",
32
+ "transformers_version": "4.48.3",
33
+ "use_cache": true,
34
+ "vocab_size": 128256
35
+ }
0_KananaEmbeddingWrapper/kanana2vec/configuration_kanana2vec.py ADDED
@@ -0,0 +1,194 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2024 Kakao Corp. team. All rights reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ """Kanana2Vec model configuration"""
16
+
17
+ from transformers import AutoConfig
18
+ from transformers.models.llama import LlamaConfig
19
+ from transformers.configuration_utils import PretrainedConfig
20
+ from transformers.modeling_rope_utils import rope_config_validation
21
+
22
+
23
+ class Kanana2VecConfig(PretrainedConfig):
24
+ r"""
25
+ This is the configuration class to store the configuration of a [`Kanana2VecModel`]. It is used to instantiate an LLaMA
26
+ model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
27
+ defaults will yield a similar configuration to that of the LLaMA-7B.
28
+
29
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
30
+ documentation from [`PretrainedConfig`] for more information.
31
+
32
+
33
+ Args:
34
+ vocab_size (`int`, *optional*, defaults to 32000):
35
+ Vocabulary size of the LLaMA model. Defines the number of different tokens that can be represented by the
36
+ `inputs_ids` passed when calling [`Kanana2VecModel`]
37
+ hidden_size (`int`, *optional*, defaults to 4096):
38
+ Dimension of the hidden representations.
39
+ intermediate_size (`int`, *optional*, defaults to 11008):
40
+ Dimension of the MLP representations.
41
+ num_hidden_layers (`int`, *optional*, defaults to 32):
42
+ Number of hidden layers in the Transformer decoder.
43
+ num_attention_heads (`int`, *optional*, defaults to 32):
44
+ Number of attention heads for each attention layer in the Transformer decoder.
45
+ num_key_value_heads (`int`, *optional*):
46
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
47
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
48
+ `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
49
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
50
+ by meanpooling all the original heads within that group. For more details checkout [this
51
+ paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
52
+ `num_attention_heads`.
53
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
54
+ The non-linear activation function (function or string) in the decoder.
55
+ max_position_embeddings (`int`, *optional*, defaults to 2048):
56
+ The maximum sequence length that this model might ever be used with. Llama 1 supports up to 2048 tokens,
57
+ Llama 2 up to 4096, CodeLlama up to 16384.
58
+ initializer_range (`float`, *optional*, defaults to 0.02):
59
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
60
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
61
+ The epsilon used by the rms normalization layers.
62
+ use_cache (`bool`, *optional*, defaults to `True`):
63
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
64
+ relevant if `config.is_decoder=True`.
65
+ pad_token_id (`int`, *optional*):
66
+ Padding token id.
67
+ bos_token_id (`int`, *optional*, defaults to 1):
68
+ Beginning of stream token id.
69
+ eos_token_id (`int`, *optional*, defaults to 2):
70
+ End of stream token id.
71
+ pretraining_tp (`int`, *optional*, defaults to 1):
72
+ Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
73
+ document](https://huggingface.co/docs/transformers/main/perf_train_gpu_many#tensor-parallelism) to
74
+ understand more about it. This value is necessary to ensure exact reproducibility of the pretraining
75
+ results. Please refer to [this issue](https://github.com/pytorch/pytorch/issues/76232).
76
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
77
+ Whether to tie weight embeddings
78
+ rope_theta (`float`, *optional*, defaults to 10000.0):
79
+ The base period of the RoPE embeddings.
80
+ rope_scaling (`Dict`, *optional*):
81
+ Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
82
+ and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
83
+ accordingly.
84
+ Expected contents:
85
+ `rope_type` (`str`):
86
+ The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
87
+ 'llama3'], with 'default' being the original RoPE implementation.
88
+ `factor` (`float`, *optional*):
89
+ Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
90
+ most scaling types, a `factor` of x will enable the model to handle sequences of length x *
91
+ original maximum pre-trained length.
92
+ `original_max_position_embeddings` (`int`, *optional*):
93
+ Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
94
+ pretraining.
95
+ `attention_factor` (`float`, *optional*):
96
+ Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
97
+ computation. If unspecified, it defaults to value recommended by the implementation, using the
98
+ `factor` field to infer the suggested value.
99
+ `beta_fast` (`float`, *optional*):
100
+ Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
101
+ ramp function. If unspecified, it defaults to 32.
102
+ `beta_slow` (`float`, *optional*):
103
+ Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
104
+ ramp function. If unspecified, it defaults to 1.
105
+ `short_factor` (`List[float]`, *optional*):
106
+ Only used with 'longrope'. The scaling factor to be applied to short contexts (<
107
+ `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
108
+ size divided by the number of attention heads divided by 2
109
+ `long_factor` (`List[float]`, *optional*):
110
+ Only used with 'longrope'. The scaling factor to be applied to long contexts (<
111
+ `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
112
+ size divided by the number of attention heads divided by 2
113
+ `low_freq_factor` (`float`, *optional*):
114
+ Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
115
+ `high_freq_factor` (`float`, *optional*):
116
+ Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
117
+ attention_bias (`bool`, *optional*, defaults to `False`):
118
+ Whether to use a bias in the query, key, value and output projection layers during self-attention.
119
+ attention_dropout (`float`, *optional*, defaults to 0.0):
120
+ The dropout ratio for the attention probabilities.
121
+ mlp_bias (`bool`, *optional*, defaults to `False`):
122
+ Whether to use a bias in up_proj, down_proj and gate_proj layers in the MLP layers.
123
+ head_dim (`int`, *optional*):
124
+ The attention head dimension. If None, it will default to hidden_size // num_heads"""
125
+
126
+ model_type = "kanana2vec"
127
+ keys_to_ignore_at_inference = ["past_key_values"]
128
+
129
+ def __init__(
130
+ self,
131
+ vocab_size=32000,
132
+ hidden_size=4096,
133
+ intermediate_size=11008,
134
+ num_hidden_layers=32,
135
+ num_attention_heads=32,
136
+ num_key_value_heads=None,
137
+ hidden_act="silu",
138
+ max_position_embeddings=2048,
139
+ initializer_range=0.02,
140
+ rms_norm_eps=1e-6,
141
+ use_cache=True,
142
+ pad_token_id=None,
143
+ bos_token_id=1,
144
+ eos_token_id=2,
145
+ pretraining_tp=1,
146
+ tie_word_embeddings=False,
147
+ rope_theta=10000.0,
148
+ rope_scaling=None,
149
+ attention_bias=False,
150
+ attention_dropout=0.0,
151
+ mlp_bias=False,
152
+ head_dim=None,
153
+ **kwargs,
154
+ ):
155
+ self.vocab_size = vocab_size
156
+ self.max_position_embeddings = max_position_embeddings
157
+ self.hidden_size = hidden_size
158
+ self.intermediate_size = intermediate_size
159
+ self.num_hidden_layers = num_hidden_layers
160
+ self.num_attention_heads = num_attention_heads
161
+
162
+ # for backward compatibility
163
+ if num_key_value_heads is None:
164
+ num_key_value_heads = num_attention_heads
165
+
166
+ self.num_key_value_heads = num_key_value_heads
167
+ self.hidden_act = hidden_act
168
+ self.initializer_range = initializer_range
169
+ self.rms_norm_eps = rms_norm_eps
170
+ self.pretraining_tp = pretraining_tp
171
+ self.use_cache = use_cache
172
+ self.rope_theta = rope_theta
173
+ self.rope_scaling = rope_scaling
174
+ self.attention_bias = attention_bias
175
+ self.attention_dropout = attention_dropout
176
+ self.mlp_bias = mlp_bias
177
+ self.head_dim = head_dim if head_dim is not None else self.hidden_size // self.num_attention_heads
178
+ # Validate the correctness of rotary position embeddings parameters
179
+ # BC: if there is a 'type' field, copy it it to 'rope_type'.
180
+ if self.rope_scaling is not None and "type" in self.rope_scaling:
181
+ self.rope_scaling["rope_type"] = self.rope_scaling["type"]
182
+ rope_config_validation(self)
183
+
184
+ super().__init__(
185
+ pad_token_id=pad_token_id,
186
+ bos_token_id=bos_token_id,
187
+ eos_token_id=eos_token_id,
188
+ tie_word_embeddings=tie_word_embeddings,
189
+ **kwargs,
190
+ )
191
+
192
+
193
+ class BiLlamaConfig(LlamaConfig):
194
+ model_type = "billama"
0_KananaEmbeddingWrapper/kanana2vec/model-00001-of-00002.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c9b71eba036018909451d8b8523c6326d852d9cae59ed9ead84175371ea4bfe2
3
+ size 4982308720
0_KananaEmbeddingWrapper/kanana2vec/model-00002-of-00002.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:7dabb33c15497d9505fb1ba107e46dc9af665b02ed3e7c8daa043f3533a17d07
3
+ size 3365641920
0_KananaEmbeddingWrapper/kanana2vec/model.safetensors.index.json ADDED
@@ -0,0 +1,297 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "metadata": {
3
+ "total_size": 8347917312
4
+ },
5
+ "weight_map": {
6
+ "model.embed_tokens.weight": "model-00001-of-00002.safetensors",
7
+ "model.layers.0.input_layernorm.weight": "model-00001-of-00002.safetensors",
8
+ "model.layers.0.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
9
+ "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
10
+ "model.layers.0.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
11
+ "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
12
+ "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
13
+ "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
14
+ "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
15
+ "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
16
+ "model.layers.1.input_layernorm.weight": "model-00001-of-00002.safetensors",
17
+ "model.layers.1.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
18
+ "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
19
+ "model.layers.1.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
20
+ "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
21
+ "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
22
+ "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
23
+ "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
24
+ "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
25
+ "model.layers.10.input_layernorm.weight": "model-00001-of-00002.safetensors",
26
+ "model.layers.10.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
27
+ "model.layers.10.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
28
+ "model.layers.10.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
29
+ "model.layers.10.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
30
+ "model.layers.10.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
31
+ "model.layers.10.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
32
+ "model.layers.10.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
33
+ "model.layers.10.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
34
+ "model.layers.11.input_layernorm.weight": "model-00001-of-00002.safetensors",
35
+ "model.layers.11.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
36
+ "model.layers.11.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
37
+ "model.layers.11.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
38
+ "model.layers.11.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
39
+ "model.layers.11.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
40
+ "model.layers.11.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
41
+ "model.layers.11.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
42
+ "model.layers.11.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
43
+ "model.layers.12.input_layernorm.weight": "model-00001-of-00002.safetensors",
44
+ "model.layers.12.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
45
+ "model.layers.12.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
46
+ "model.layers.12.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
47
+ "model.layers.12.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
48
+ "model.layers.12.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
49
+ "model.layers.12.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
50
+ "model.layers.12.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
51
+ "model.layers.12.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
52
+ "model.layers.13.input_layernorm.weight": "model-00001-of-00002.safetensors",
53
+ "model.layers.13.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
54
+ "model.layers.13.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
55
+ "model.layers.13.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
56
+ "model.layers.13.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
57
+ "model.layers.13.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
58
+ "model.layers.13.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
59
+ "model.layers.13.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
60
+ "model.layers.13.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
61
+ "model.layers.14.input_layernorm.weight": "model-00001-of-00002.safetensors",
62
+ "model.layers.14.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
63
+ "model.layers.14.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
64
+ "model.layers.14.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
65
+ "model.layers.14.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
66
+ "model.layers.14.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
67
+ "model.layers.14.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
68
+ "model.layers.14.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
69
+ "model.layers.14.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
70
+ "model.layers.15.input_layernorm.weight": "model-00001-of-00002.safetensors",
71
+ "model.layers.15.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
72
+ "model.layers.15.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
73
+ "model.layers.15.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
74
+ "model.layers.15.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
75
+ "model.layers.15.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
76
+ "model.layers.15.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
77
+ "model.layers.15.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
78
+ "model.layers.15.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
79
+ "model.layers.16.input_layernorm.weight": "model-00001-of-00002.safetensors",
80
+ "model.layers.16.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
81
+ "model.layers.16.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
82
+ "model.layers.16.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
83
+ "model.layers.16.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
84
+ "model.layers.16.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
85
+ "model.layers.16.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
86
+ "model.layers.16.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
87
+ "model.layers.16.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
88
+ "model.layers.17.input_layernorm.weight": "model-00002-of-00002.safetensors",
89
+ "model.layers.17.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
90
+ "model.layers.17.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
91
+ "model.layers.17.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
92
+ "model.layers.17.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
93
+ "model.layers.17.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
94
+ "model.layers.17.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
95
+ "model.layers.17.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
96
+ "model.layers.17.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
97
+ "model.layers.18.input_layernorm.weight": "model-00002-of-00002.safetensors",
98
+ "model.layers.18.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
99
+ "model.layers.18.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
100
+ "model.layers.18.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
101
+ "model.layers.18.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
102
+ "model.layers.18.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
103
+ "model.layers.18.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
104
+ "model.layers.18.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
105
+ "model.layers.18.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
106
+ "model.layers.19.input_layernorm.weight": "model-00002-of-00002.safetensors",
107
+ "model.layers.19.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
108
+ "model.layers.19.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
109
+ "model.layers.19.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
110
+ "model.layers.19.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
111
+ "model.layers.19.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
112
+ "model.layers.19.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
113
+ "model.layers.19.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
114
+ "model.layers.19.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
115
+ "model.layers.2.input_layernorm.weight": "model-00001-of-00002.safetensors",
116
+ "model.layers.2.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
117
+ "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
118
+ "model.layers.2.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
119
+ "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
120
+ "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
121
+ "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
122
+ "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
123
+ "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
124
+ "model.layers.20.input_layernorm.weight": "model-00002-of-00002.safetensors",
125
+ "model.layers.20.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
126
+ "model.layers.20.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
127
+ "model.layers.20.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
128
+ "model.layers.20.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
129
+ "model.layers.20.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
130
+ "model.layers.20.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
131
+ "model.layers.20.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
132
+ "model.layers.20.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
133
+ "model.layers.21.input_layernorm.weight": "model-00002-of-00002.safetensors",
134
+ "model.layers.21.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
135
+ "model.layers.21.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
136
+ "model.layers.21.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
137
+ "model.layers.21.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
138
+ "model.layers.21.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
139
+ "model.layers.21.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
140
+ "model.layers.21.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
141
+ "model.layers.21.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
142
+ "model.layers.22.input_layernorm.weight": "model-00002-of-00002.safetensors",
143
+ "model.layers.22.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
144
+ "model.layers.22.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
145
+ "model.layers.22.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
146
+ "model.layers.22.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
147
+ "model.layers.22.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
148
+ "model.layers.22.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
149
+ "model.layers.22.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
150
+ "model.layers.22.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
151
+ "model.layers.23.input_layernorm.weight": "model-00002-of-00002.safetensors",
152
+ "model.layers.23.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
153
+ "model.layers.23.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
154
+ "model.layers.23.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
155
+ "model.layers.23.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
156
+ "model.layers.23.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
157
+ "model.layers.23.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
158
+ "model.layers.23.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
159
+ "model.layers.23.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
160
+ "model.layers.24.input_layernorm.weight": "model-00002-of-00002.safetensors",
161
+ "model.layers.24.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
162
+ "model.layers.24.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
163
+ "model.layers.24.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
164
+ "model.layers.24.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
165
+ "model.layers.24.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
166
+ "model.layers.24.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
167
+ "model.layers.24.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
168
+ "model.layers.24.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
169
+ "model.layers.25.input_layernorm.weight": "model-00002-of-00002.safetensors",
170
+ "model.layers.25.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
171
+ "model.layers.25.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
172
+ "model.layers.25.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
173
+ "model.layers.25.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
174
+ "model.layers.25.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
175
+ "model.layers.25.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
176
+ "model.layers.25.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
177
+ "model.layers.25.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
178
+ "model.layers.26.input_layernorm.weight": "model-00002-of-00002.safetensors",
179
+ "model.layers.26.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
180
+ "model.layers.26.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
181
+ "model.layers.26.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
182
+ "model.layers.26.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
183
+ "model.layers.26.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
184
+ "model.layers.26.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
185
+ "model.layers.26.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
186
+ "model.layers.26.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
187
+ "model.layers.27.input_layernorm.weight": "model-00002-of-00002.safetensors",
188
+ "model.layers.27.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
189
+ "model.layers.27.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
190
+ "model.layers.27.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
191
+ "model.layers.27.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
192
+ "model.layers.27.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
193
+ "model.layers.27.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
194
+ "model.layers.27.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
195
+ "model.layers.27.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
196
+ "model.layers.28.input_layernorm.weight": "model-00002-of-00002.safetensors",
197
+ "model.layers.28.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
198
+ "model.layers.28.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
199
+ "model.layers.28.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
200
+ "model.layers.28.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
201
+ "model.layers.28.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
202
+ "model.layers.28.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
203
+ "model.layers.28.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
204
+ "model.layers.28.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
205
+ "model.layers.29.input_layernorm.weight": "model-00002-of-00002.safetensors",
206
+ "model.layers.29.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
207
+ "model.layers.29.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
208
+ "model.layers.29.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
209
+ "model.layers.29.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
210
+ "model.layers.29.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
211
+ "model.layers.29.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
212
+ "model.layers.29.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
213
+ "model.layers.29.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
214
+ "model.layers.3.input_layernorm.weight": "model-00001-of-00002.safetensors",
215
+ "model.layers.3.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
216
+ "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
217
+ "model.layers.3.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
218
+ "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
219
+ "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
220
+ "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
221
+ "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
222
+ "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
223
+ "model.layers.30.input_layernorm.weight": "model-00002-of-00002.safetensors",
224
+ "model.layers.30.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
225
+ "model.layers.30.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
226
+ "model.layers.30.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
227
+ "model.layers.30.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
228
+ "model.layers.30.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
229
+ "model.layers.30.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
230
+ "model.layers.30.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
231
+ "model.layers.30.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
232
+ "model.layers.31.input_layernorm.weight": "model-00002-of-00002.safetensors",
233
+ "model.layers.31.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
234
+ "model.layers.31.mlp.gate_proj.weight": "model-00002-of-00002.safetensors",
235
+ "model.layers.31.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
236
+ "model.layers.31.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
237
+ "model.layers.31.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
238
+ "model.layers.31.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
239
+ "model.layers.31.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
240
+ "model.layers.31.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
241
+ "model.layers.4.input_layernorm.weight": "model-00001-of-00002.safetensors",
242
+ "model.layers.4.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
243
+ "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
244
+ "model.layers.4.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
245
+ "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
246
+ "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
247
+ "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
248
+ "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
249
+ "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
250
+ "model.layers.5.input_layernorm.weight": "model-00001-of-00002.safetensors",
251
+ "model.layers.5.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
252
+ "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
253
+ "model.layers.5.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
254
+ "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
255
+ "model.layers.5.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
256
+ "model.layers.5.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
257
+ "model.layers.5.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
258
+ "model.layers.5.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
259
+ "model.layers.6.input_layernorm.weight": "model-00001-of-00002.safetensors",
260
+ "model.layers.6.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
261
+ "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
262
+ "model.layers.6.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
263
+ "model.layers.6.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
264
+ "model.layers.6.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
265
+ "model.layers.6.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
266
+ "model.layers.6.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
267
+ "model.layers.6.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
268
+ "model.layers.7.input_layernorm.weight": "model-00001-of-00002.safetensors",
269
+ "model.layers.7.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
270
+ "model.layers.7.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
271
+ "model.layers.7.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
272
+ "model.layers.7.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
273
+ "model.layers.7.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
274
+ "model.layers.7.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
275
+ "model.layers.7.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
276
+ "model.layers.7.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
277
+ "model.layers.8.input_layernorm.weight": "model-00001-of-00002.safetensors",
278
+ "model.layers.8.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
279
+ "model.layers.8.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
280
+ "model.layers.8.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
281
+ "model.layers.8.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
282
+ "model.layers.8.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
283
+ "model.layers.8.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
284
+ "model.layers.8.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
285
+ "model.layers.8.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
286
+ "model.layers.9.input_layernorm.weight": "model-00001-of-00002.safetensors",
287
+ "model.layers.9.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
288
+ "model.layers.9.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
289
+ "model.layers.9.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
290
+ "model.layers.9.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
291
+ "model.layers.9.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
292
+ "model.layers.9.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
293
+ "model.layers.9.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
294
+ "model.layers.9.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
295
+ "model.norm.weight": "model-00002-of-00002.safetensors"
296
+ }
297
+ }
0_KananaEmbeddingWrapper/kanana_embedding_config.json ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ {
2
+ "model_name_or_path": "/data/public/model_registry/huggingface/hub/models--datalama--kanana-nano-2.1b-embedding/snapshots/77c87d91735397f0416f5bf595ba48f6af97333a/0_KananaEmbeddingWrapper/kanana2vec",
3
+ "trust_remote_code": true,
4
+ "device": "cpu",
5
+ "hidden_size": 1792,
6
+ "max_seq_length": 8192
7
+ }
0_KananaEmbeddingWrapper/kanana_embedding_wrapper.py ADDED
@@ -0,0 +1,182 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import json
3
+ import torch
4
+ from torch import nn
5
+ import torch.nn.functional as F
6
+ from typing import List, Dict
7
+
8
+ from transformers import AutoModel
9
+
10
+ os.environ["TOKENIZERS_PARALLELISM"] = "false"
11
+
12
+ class KananaEmbeddingWrapper(nn.Module):
13
+ def __init__(self, model_name_or_path: str, trust_remote_code=True, device: str = "cpu", max_seq_length:int=None):
14
+ """
15
+ Initialize the KananaEmbeddingWrapper.
16
+
17
+ Args:
18
+ model_name_or_path: Path or name of the pretrained model
19
+ trust_remote_code: Whether to trust remote code when loading the model
20
+ device: Device to load the model on (e.g., 'cpu', 'cuda')
21
+ """
22
+ super(KananaEmbeddingWrapper, self).__init__()
23
+ self.model_name_or_path = model_name_or_path
24
+ self.trust_remote_code = trust_remote_code
25
+ self.device = device
26
+ self.kanana2vec = AutoModel.from_pretrained(
27
+ model_name_or_path, trust_remote_code=trust_remote_code
28
+ ).to(self.device)
29
+
30
+ self.max_seq_length = max_seq_length if max_seq_length is not None else self.kanana2vec.config.max_position_embeddings
31
+
32
+ def get_sentence_embedding_dimension(self) -> int:
33
+ """
34
+ Returns the dimension of the sentence embeddings.
35
+
36
+ Returns:
37
+ Dimensionality of the sentence embeddings
38
+ """
39
+ return self.kanana2vec.config.hidden_size
40
+
41
+ def get_max_seq_length(self) -> int:
42
+ """
43
+ Returns the maximum sequence length this module can process.
44
+
45
+ Returns:
46
+ Maximum sequence length
47
+ """
48
+ return self.max_seq_length
49
+
50
+ def tokenize(self, texts: List[str]) -> Dict[str, torch.Tensor]:
51
+ """
52
+ Tokenize input texts.
53
+
54
+ Args:
55
+ texts: List of input texts to tokenize
56
+
57
+ Returns:
58
+ Dictionary containing tokenized inputs
59
+ """
60
+ return self.kanana2vec.tokenizer(
61
+ texts,
62
+ padding=True,
63
+ return_token_type_ids=False,
64
+ return_tensors="pt",
65
+ truncation=True,
66
+ max_length=self.max_seq_length
67
+ ).to(self.device)
68
+
69
+ def forward(self, features: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
70
+ """
71
+ Forward pass of the module.
72
+
73
+ Args:
74
+ features: Dictionary with inputs including 'input_ids', 'attention_mask', etc.
75
+
76
+ Returns:
77
+ Dictionary with updated features including 'sentence_embedding'
78
+ """
79
+ # Extract only the required features for the model
80
+ model_inputs = self._extract_model_inputs(features)
81
+
82
+ # Create pool mask considering prompt length if available
83
+ model_inputs["pool_mask"] = self._create_pool_mask(features)
84
+
85
+ # Get embeddings from the model and normalize
86
+ embedding = self.kanana2vec.forward(**model_inputs).embedding
87
+ normalized_embedding = F.normalize(embedding, p=2, dim=1)
88
+
89
+ # Update features with sentence embedding
90
+ features['sentence_embedding'] = normalized_embedding
91
+ return features
92
+
93
+ def _extract_model_inputs(self, features: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
94
+ """
95
+ Extract only the inputs needed for the model.
96
+
97
+ Args:
98
+ features: Complete feature dictionary
99
+
100
+ Returns:
101
+ Dictionary with only the required keys for the model
102
+ """
103
+ return {k: v for k, v in features.items() if k in ['input_ids', 'attention_mask']}
104
+
105
+ def _create_pool_mask(self, features: Dict[str, torch.Tensor]) -> torch.Tensor:
106
+ """
107
+ Create a pool mask based on attention mask and prompt length.
108
+
109
+ Args:
110
+ features: Feature dictionary containing attention_mask and optionally prompt_length
111
+
112
+ Returns:
113
+ Pool mask tensor
114
+ """
115
+ pool_mask = features['attention_mask'].clone()
116
+ if "prompt_length" in features:
117
+ pool_mask[:, :features['prompt_length']] = 0
118
+ return pool_mask
119
+
120
+ def get_config_dict(self) -> Dict:
121
+ """
122
+ Returns a dictionary with the module's configuration.
123
+
124
+ Returns:
125
+ Dictionary with module configuration
126
+ """
127
+ return {
128
+ "model_name_or_path": self.model_name_or_path,
129
+ "trust_remote_code": self.trust_remote_code,
130
+ "device": self.device,
131
+ "hidden_size": self.get_sentence_embedding_dimension(),
132
+ "max_seq_length": self.get_max_seq_length()
133
+ }
134
+
135
+ def save(self, save_dir: str) -> None:
136
+ """
137
+ Saves the module's configuration and model to the specified directory.
138
+
139
+ Args:
140
+ save_dir: Directory to save the module configuration
141
+ """
142
+ os.makedirs(save_dir, exist_ok=True)
143
+
144
+ # Save model configuration
145
+ config_path = os.path.join(save_dir, "kanana_embedding_config.json")
146
+ with open(config_path, 'w', encoding='utf-8') as f:
147
+ json.dump(self.get_config_dict(), f, ensure_ascii=False, indent=2)
148
+
149
+ # Save the underlying model
150
+ model_save_path = os.path.join(save_dir, "kanana2vec")
151
+ self.kanana2vec.save_pretrained(model_save_path)
152
+
153
+ print(f"KananaEmbeddingWrapper model saved to {save_dir}")
154
+
155
+ @staticmethod
156
+ def load(load_dir: str, device: str = "cpu") -> 'KananaEmbeddingWrapper':
157
+ """
158
+ Loads a KananaEmbeddingWrapper model from the specified directory.
159
+
160
+ Args:
161
+ load_dir: Directory containing the saved module
162
+ device: Device to load the model on
163
+
164
+ Returns:
165
+ Initialized KananaEmbeddingWrapper
166
+ """
167
+ # Load configuration
168
+ config_path = os.path.join(load_dir, "kanana_embedding_config.json")
169
+ with open(config_path, 'r', encoding='utf-8') as f:
170
+ config = json.load(f)
171
+
172
+ # Use the saved model path
173
+ model_load_path = os.path.join(load_dir, "kanana2vec")
174
+
175
+ # Create instance with saved configuration
176
+ instance = KananaEmbeddingWrapper(
177
+ model_name_or_path=model_load_path,
178
+ trust_remote_code=config.get("trust_remote_code", True),
179
+ device=device # Use the provided device or default
180
+ )
181
+
182
+ return instance
README.md ADDED
@@ -0,0 +1,157 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ language:
3
+ - en
4
+ - ko
5
+ tags:
6
+ - sentence-transformers
7
+ - sentence-similarity
8
+ - feature-extraction
9
+ pipeline_tag: sentence-similarity
10
+ library_name: sentence-transformers
11
+ model_id: datalama/kanana-nano-2.1b-embedding
12
+ repo: datalama/kanana-nano-2.1b-embedding
13
+ developers: datalama
14
+ license: cc-by-nc-4.0
15
+ ---
16
+
17
+ # Sentence-Transformers Compatible Kanana-Nano-2.1b-Embedding
18
+
19
+ This repository contains a sentence-transformers compatible version of the [Kanana-Nano-2.1b-Embedding](https://huggingface.co/kakaocorp/kanana-nano-2.1b-embedding) model developed by Kakao.
20
+
21
+ For detailed information about the model architecture, training methodology, and comprehensive performance benchmarks, please refer to the [original model repository](https://huggingface.co/kakaocorp/kanana-nano-2.1b-embedding) and the [Kanana technical report](https://arxiv.org/abs/2502.18934).
22
+
23
+ ## Key Adaptations
24
+
25
+ This version has been modified to work seamlessly with the sentence-transformers library with the following changes:
26
+
27
+ * Implemented `KananaEmbeddingWrapper` module to enable loading via SentenceTransformer
28
+ * Added L2 normalization within the `KananaEmbeddingWrapper`'s forward method
29
+ * max_seq_length is fixed with 8192.
30
+ * Embed the query prompt related parts into the model. You can encode the query with `query_name`.
31
+
32
+ ## Usage
33
+
34
+ ### Installation
35
+
36
+ ```bash
37
+ pip install sentence-transformers
38
+ ```
39
+
40
+ ### Basic Usage
41
+
42
+ ```python
43
+ from sentence_transformers import SentenceTransformer
44
+
45
+ # Load the model
46
+ model = SentenceTransformer("datalama/kanana-nano-2.1b-embedding", device="cpu", trust_remote_code=True)
47
+
48
+ # Encode sentences
49
+ sentences = [
50
+ "이 문장은 한국어로 작성되었습니다.",
51
+ "This sentence is written in English."
52
+ ]
53
+
54
+ embeddings = model.encode(sentences)
55
+ ```
56
+
57
+ ### Advanced Usage with Query/Passage Format
58
+
59
+ * You can use `prompt_name` or `prompt`.
60
+
61
+ ```python
62
+ import numpy as np
63
+ from sentence_transformers import SentenceTransformer
64
+
65
+ model = SentenceTransformer("datalama/kanana-nano-2.1b-embedding", device="cpu", trust_remote_code=True)
66
+
67
+ # For retrieval tasks
68
+ instruction = "Given a question, retrieve passages that answer the question"
69
+ queries = [
70
+ "are judo throws allowed in wrestling?",
71
+ "how to become a radiology technician in michigan?",
72
+ ]
73
+
74
+
75
+ # You can encode query by prompt_name with predefiend prompt template.
76
+ embedding_a = model.encode(queries, prompt_name="query")
77
+
78
+ # You can directly encode the query with prompt.
79
+ prompt_template = """Instruct: {instruction}\nQuery: """
80
+ embedding_b = model.encode(queries, prompt=prompt_template.format(instruction=instruction))
81
+
82
+ # compare input.
83
+ np.allclose(embedding_a, embedding_b)
84
+ # True
85
+ ```
86
+
87
+ * Compare embedding with original code.
88
+
89
+ ```python
90
+ import torch.nn.functional as F
91
+ import numpy as np
92
+ from transformers import AutoModel
93
+ from sentence_transformers import SentenceTransformer
94
+
95
+ # For retrieval tasks
96
+ instruction = "Given a question, retrieve passages that answer the question"
97
+ queries = [
98
+ "are judo throws allowed in wrestling?",
99
+ "how to become a radiology technician in michigan?",
100
+ ]
101
+
102
+ passages = [
103
+ "Since you're reading this, you are probably someone from a judo background or someone who is just wondering how judo techniques can be applied under wrestling rules. So without further ado, let's get to the question. Are Judo throws allowed in wrestling? Yes, judo throws are allowed in freestyle and folkstyle wrestling. You only need to be careful to follow the slam rules when executing judo throws. In wrestling, a slam is lifting and returning an opponent to the mat with unnecessary force.",
104
+ "Below are the basic steps to becoming a radiologic technologist in Michigan:Earn a high school diploma. As with most careers in health care, a high school education is the first step to finding entry-level employment. Taking classes in math and science, such as anatomy, biology, chemistry, physiology, and physics, can help prepare students for their college studies and future careers.Earn an associate degree. Entry-level radiologic positions typically require at least an Associate of Applied Science. Before enrolling in one of these degree programs, students should make sure it has been properly accredited by the Joint Review Committee on Education in Radiologic Technology (JRCERT).Get licensed or certified in the state of Michigan.",
105
+ ]
106
+
107
+ # compare originaml model and this model.
108
+ model_a = AutoModel.from_pretrained("kakaocorp/kanana-nano-2.1b-embedding",trust_remote_code=True,).to("cpu")
109
+ model_b = SentenceTransformer("datalama/kanana-nano-2.1b-embedding", device="cpu", trust_remote_code=True)
110
+
111
+ # original encoding method.
112
+ max_length = 512
113
+ query_embeddings = model_a.encode(queries, instruction=instruction, max_length=max_length)
114
+ passage_embeddings = model_a.encode(passages, instruction="", max_length=max_length)
115
+
116
+ query_embeddings = F.normalize(query_embeddings, p=2, dim=1)
117
+ passage_embeddings = F.normalize(passage_embeddings, p=2, dim=1)
118
+
119
+ scores_a = (query_embeddings @ passage_embeddings.T) * 100
120
+
121
+ # sentence_transformers compatible encoding method.
122
+ query_embeddings = model_b.encode(queries, prompt_name="query")
123
+ passage_embeddings = model_b.encode(passages)
124
+
125
+ scores_b = (query_embeddings @ passage_embeddings.T) * 100
126
+
127
+ # compare embedding
128
+ np.allclose(scores_a.cpu().numpy(), scores_b)
129
+ # True
130
+ ```
131
+
132
+ Note: Unlike the original model, you don't need to manually perform L2 normalization as this is handled by the `KananaEmbeddingWrapper` module during the forward pass.
133
+
134
+ ## License
135
+
136
+ This model is licensed under [CC-BY-NC-4.0](https://spdx.org/licenses/CC-BY-NC-4.0).
137
+
138
+ ## Citation
139
+
140
+ If you use this model, please cite the original work:
141
+
142
+ ```
143
+ @misc{kananallmteam2025kananacomputeefficientbilinguallanguage,
144
+ title={Kanana: Compute-efficient Bilingual Language Models},
145
+ author={Kanana LLM Team and Yunju Bak and Hojin Lee and Minho Ryu and Jiyeon Ham and Seungjae Jung and Daniel Wontae Nam and Taegyeong Eo and Donghun Lee and Doohae Jung and Boseop Kim and Nayeon Kim and Jaesun Park and Hyunho Kim and Hyunwoong Ko and Changmin Lee and Kyoung-Woon On and Seulye Baeg and Junrae Cho and Sunghee Jung and Jieun Kang and EungGyun Kim and Eunhwa Kim and Byeongil Ko and Daniel Lee and Minchul Lee and Miok Lee and Shinbok Lee and Gaeun Seo},
146
+ year={2025},
147
+ eprint={2502.18934},
148
+ archivePrefix={arXiv},
149
+ primaryClass={cs.CL},
150
+ url={https://arxiv.org/abs/2502.18934},
151
+ }
152
+ ```
153
+
154
+ ## Acknowledgements
155
+
156
+ - Original model developed by the Kanana LLM Team at Kakao
157
+ - Adaptation to sentence-transformers format by datalama
config_sentence_transformers.json ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "__version__": {
3
+ "sentence_transformers": "3.4.1",
4
+ "transformers": "4.48.3",
5
+ "pytorch": "2.5.1+cu124"
6
+ },
7
+ "prompts": {
8
+ "query": "Instruct: Given a question, retrieve passages that answer the question\nQuery: "
9
+ },
10
+ "default_prompt_name": null,
11
+ "similarity_fn_name": "cosine"
12
+ }
modules.json ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ [
2
+ {
3
+ "idx": 0,
4
+ "name": "0",
5
+ "path": "0_KananaEmbeddingWrapper",
6
+ "type": "kanana_embedding_wrapper.KananaEmbeddingWrapper"
7
+ }
8
+ ]