cnmoro commited on
Commit
541365d
·
verified ·
1 Parent(s): 108fa29

Upload 11 files

Browse files
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ tokenizer.json filter=lfs diff=lfs merge=lfs -text
1_Pooling/config.json ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "word_embedding_dimension": 768,
3
+ "pooling_mode_cls_token": true,
4
+ "pooling_mode_mean_tokens": false,
5
+ "pooling_mode_max_tokens": false,
6
+ "pooling_mode_mean_sqrt_len_tokens": false,
7
+ "pooling_mode_weightedmean_tokens": false,
8
+ "pooling_mode_lasttoken": false,
9
+ "include_prompt": true
10
+ }
README.md ADDED
The diff for this file is too large to render. See raw diff
 
config.json ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "GteModel"
4
+ ],
5
+ "attention_probs_dropout_prob": 0.0,
6
+ "auto_map": {
7
+ "AutoConfig": "configuration_hf_alibaba_nlp_gte.GteConfig",
8
+ "AutoModel": "modeling_hf_alibaba_nlp_gte.GteModel"
9
+ },
10
+ "classifier_dropout": 0.1,
11
+ "hidden_act": "gelu",
12
+ "hidden_dropout_prob": 0.1,
13
+ "hidden_size": 768,
14
+ "initializer_range": 0.02,
15
+ "intermediate_size": 3072,
16
+ "layer_norm_eps": 1e-12,
17
+ "layer_norm_type": "layer_norm",
18
+ "logn_attention_clip1": false,
19
+ "logn_attention_scale": false,
20
+ "max_position_embeddings": 8192,
21
+ "model_type": "gte",
22
+ "num_attention_heads": 12,
23
+ "num_hidden_layers": 12,
24
+ "pack_qkv": true,
25
+ "pad_token_id": 1,
26
+ "position_embedding_type": "rope",
27
+ "rope_scaling": null,
28
+ "rope_theta": 160000,
29
+ "torch_dtype": "float32",
30
+ "transformers_version": "4.39.3",
31
+ "type_vocab_size": 1,
32
+ "unpad_inputs": "true",
33
+ "use_memory_efficient_attention": "true",
34
+ "vocab_size": 250048
35
+ }
config_sentence_transformers.json ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "__version__": {
3
+ "sentence_transformers": "2.7.0.dev0",
4
+ "transformers": "4.39.3",
5
+ "pytorch": "2.1.0+cu121"
6
+ },
7
+ "prompts": {
8
+ "query": "query: "
9
+ },
10
+ "default_prompt_name": null
11
+ }
configuration_hf_alibaba_nlp_gte.py ADDED
@@ -0,0 +1,145 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2024 The GTE Team Authors and Alibaba Group.
3
+ # Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
4
+ #
5
+ # Licensed under the Apache License, Version 2.0 (the "License");
6
+ # you may not use this file except in compliance with the License.
7
+ # You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing, software
12
+ # distributed under the License is distributed on an "AS IS" BASIS,
13
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
+ # See the License for the specific language governing permissions and
15
+ # limitations under the License.
16
+ """ GTE model configuration"""
17
+ from transformers.configuration_utils import PretrainedConfig
18
+ from transformers.utils import logging
19
+
20
+ logger = logging.get_logger(__name__)
21
+
22
+
23
+ class GteConfig(PretrainedConfig):
24
+ r"""
25
+ This is the configuration class to store the configuration of a [`NewModel`] or a [`TFNewModel`]. It is used to
26
+ instantiate a NEW model according to the specified arguments, defining the model architecture. Instantiating a
27
+ configuration with the defaults will yield a similar configuration to that of the NEW
28
+ [izhx/new-base-en](https://huggingface.co/izhx/new-base-en) architecture.
29
+
30
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
31
+ documentation from [`PretrainedConfig`] for more information.
32
+
33
+
34
+ Args:
35
+ vocab_size (`int`, *optional*, defaults to 30522):
36
+ Vocabulary size of the NEW model. Defines the number of different tokens that can be represented by the
37
+ `inputs_ids` passed when calling [`NewModel`] or [`TFNewModel`].
38
+ hidden_size (`int`, *optional*, defaults to 768):
39
+ Dimensionality of the encoder layers and the pooler layer.
40
+ num_hidden_layers (`int`, *optional*, defaults to 12):
41
+ Number of hidden layers in the Transformer encoder.
42
+ num_attention_heads (`int`, *optional*, defaults to 12):
43
+ Number of attention heads for each attention layer in the Transformer encoder.
44
+ intermediate_size (`int`, *optional*, defaults to 3072):
45
+ Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
46
+ hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`):
47
+ The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
48
+ `"relu"`, `"silu"` and `"gelu_new"` are supported.
49
+ hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
50
+ The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
51
+ attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
52
+ The dropout ratio for the attention probabilities.
53
+ max_position_embeddings (`int`, *optional*, defaults to 512):
54
+ The maximum sequence length that this model might ever be used with. Typically set this to something large
55
+ just in case (e.g., 512 or 1024 or 2048).
56
+ type_vocab_size (`int`, *optional*, defaults to 2):
57
+ The vocabulary size of the `token_type_ids` passed when calling [`NewModel`] or [`TFNewModel`].
58
+ initializer_range (`float`, *optional*, defaults to 0.02):
59
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
60
+ layer_norm_eps (`float`, *optional*, defaults to 1e-12):
61
+ The epsilon used by the layer normalization layers.
62
+ position_embedding_type (`str`, *optional*, defaults to `"rope"`):
63
+ Type of position embedding. Choose one of `"absolute"`, `"rope"`.
64
+ rope_theta (`float`, *optional*, defaults to 10000.0):
65
+ The base period of the RoPE embeddings.
66
+ rope_scaling (`Dict`, *optional*):
67
+ Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
68
+ strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
69
+ `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
70
+ `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
71
+ these scaling strategies behave:
72
+ https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
73
+ experimental feature, subject to breaking API changes in future versions.
74
+ classifier_dropout (`float`, *optional*):
75
+ The dropout ratio for the classification head.
76
+
77
+ Examples:
78
+
79
+ ```python
80
+ >>> from transformers import NewConfig, NewModel
81
+
82
+ >>> # Initializing a NEW izhx/new-base-en style configuration
83
+ >>> configuration = NewConfig()
84
+
85
+ >>> # Initializing a model (with random weights) from the izhx/new-base-en style configuration
86
+ >>> model = NewModel(configuration)
87
+
88
+ >>> # Accessing the model configuration
89
+ >>> configuration = model.config
90
+ ```"""
91
+
92
+ model_type = "gte"
93
+
94
+ def __init__(
95
+ self,
96
+ vocab_size=30528,
97
+ hidden_size=768,
98
+ num_hidden_layers=12,
99
+ num_attention_heads=12,
100
+ intermediate_size=3072,
101
+ hidden_act="gelu",
102
+ hidden_dropout_prob=0.1,
103
+ attention_probs_dropout_prob=0.0,
104
+ max_position_embeddings=2048,
105
+ type_vocab_size=1,
106
+ initializer_range=0.02,
107
+ layer_norm_type='layer_norm',
108
+ layer_norm_eps=1e-12,
109
+ # pad_token_id=0,
110
+ position_embedding_type="rope",
111
+ rope_theta=10000.0,
112
+ rope_scaling=None,
113
+ classifier_dropout=None,
114
+ pack_qkv=True,
115
+ unpad_inputs=False,
116
+ use_memory_efficient_attention=False,
117
+ logn_attention_scale=False,
118
+ logn_attention_clip1=False,
119
+ **kwargs,
120
+ ):
121
+ super().__init__(**kwargs)
122
+
123
+ self.vocab_size = vocab_size
124
+ self.hidden_size = hidden_size
125
+ self.num_hidden_layers = num_hidden_layers
126
+ self.num_attention_heads = num_attention_heads
127
+ self.hidden_act = hidden_act
128
+ self.intermediate_size = intermediate_size
129
+ self.hidden_dropout_prob = hidden_dropout_prob
130
+ self.attention_probs_dropout_prob = attention_probs_dropout_prob
131
+ self.max_position_embeddings = max_position_embeddings
132
+ self.type_vocab_size = type_vocab_size
133
+ self.initializer_range = initializer_range
134
+ self.layer_norm_type = layer_norm_type
135
+ self.layer_norm_eps = layer_norm_eps
136
+ self.position_embedding_type = position_embedding_type
137
+ self.rope_theta = rope_theta
138
+ self.rope_scaling = rope_scaling
139
+ self.classifier_dropout = classifier_dropout
140
+
141
+ self.pack_qkv = pack_qkv
142
+ self.unpad_inputs = unpad_inputs
143
+ self.use_memory_efficient_attention = use_memory_efficient_attention
144
+ self.logn_attention_scale = logn_attention_scale
145
+ self.logn_attention_clip1 = logn_attention_clip1
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3d80d4727ac8759fb8624b690697c053a3d1992120111dc4a71178e608c26604
3
+ size 1221487872
modeling_hf_alibaba_nlp_gte.py ADDED
@@ -0,0 +1,936 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2024 The GTE Team Authors and Alibaba Group.
3
+ # Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
4
+ #
5
+ # Licensed under the Apache License, Version 2.0 (the "License");
6
+ # you may not use this file except in compliance with the License.
7
+ # You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing, software
12
+ # distributed under the License is distributed on an "AS IS" BASIS,
13
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
+ # See the License for the specific language governing permissions and
15
+ # limitations under the License.
16
+
17
+ import math
18
+ from dataclasses import dataclass
19
+ from typing import List, Optional, Tuple, Union
20
+
21
+ import torch
22
+ import torch.utils.checkpoint
23
+ from torch import nn
24
+
25
+ from transformers.activations import ACT2FN
26
+ from transformers.modeling_outputs import (
27
+ BaseModelOutput,
28
+ BaseModelOutputWithPooling,
29
+ MaskedLMOutput,
30
+ MultipleChoiceModelOutput,
31
+ QuestionAnsweringModelOutput,
32
+ SequenceClassifierOutput,
33
+ ModelOutput,
34
+ )
35
+ from transformers.modeling_utils import PreTrainedModel
36
+ from transformers.utils import logging
37
+
38
+ xops = None
39
+
40
+ from .configuration_hf_alibaba_nlp_gte import GteConfig
41
+
42
+
43
+ logger = logging.get_logger(__name__)
44
+
45
+
46
+ # Adapted from https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/bert_padding.py
47
+ # Which was adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py
48
+ class IndexFirstAxis(torch.autograd.Function):
49
+ @staticmethod
50
+ def forward(ctx, input, indices):
51
+ ctx.save_for_backward(indices)
52
+ assert input.ndim >= 2
53
+ ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
54
+ second_dim = other_shape.numel()
55
+ # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
56
+ # return input[indices]
57
+ # return torch.gather(
58
+ # rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim)
59
+ # ).reshape(-1, *other_shape)
60
+ return torch.gather(
61
+ input.view(ctx.first_axis_dim, second_dim),
62
+ 0,
63
+ indices.unsqueeze(-1).expand(indices.size(0), second_dim)
64
+ ).reshape(-1, *other_shape)
65
+
66
+ @staticmethod
67
+ def backward(ctx, grad_output):
68
+ (indices,) = ctx.saved_tensors
69
+ assert grad_output.ndim >= 2
70
+ other_shape = grad_output.shape[1:]
71
+ # grad_output = rearrange(grad_output, "b ... -> b (...)")
72
+ grad_output = grad_output.view(grad_output.size(0), other_shape.numel())
73
+ grad_input = torch.zeros(
74
+ [ctx.first_axis_dim, grad_output.shape[1]],
75
+ device=grad_output.device,
76
+ dtype=grad_output.dtype,
77
+ )
78
+ # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
79
+ # grad_input[indices] = grad_output
80
+ # grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output)
81
+ grad_input.scatter_(
82
+ 0, indices.unsqueeze(-1).expand(indices.size(0), grad_output.size(1)), grad_output
83
+ )
84
+ return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
85
+
86
+
87
+ index_first_axis = IndexFirstAxis.apply
88
+
89
+
90
+ def unpad_input(hidden_states, attention_mask=None, indices=None):
91
+ """
92
+ Arguments:
93
+ hidden_states: (batch, seqlen, ...)
94
+ attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
95
+ indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence.
96
+ Return:
97
+ hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
98
+ """
99
+ if indices is None:
100
+ assert attention_mask is not None
101
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
102
+
103
+ # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the
104
+ # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim
105
+ # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to
106
+ # index with integer indices. Moreover, torch's index is a bit slower than it needs to be,
107
+ # so we write custom forward and backward to make it a bit faster.
108
+ hidden_states = hidden_states.view(-1, *hidden_states.shape[2:])
109
+ return index_first_axis(hidden_states, indices)
110
+
111
+
112
+ class IndexPutFirstAxis(torch.autograd.Function):
113
+ @staticmethod
114
+ def forward(
115
+ ctx,
116
+ values: torch.Tensor,
117
+ indices: torch.Tensor,
118
+ first_axis_dim
119
+ ) -> torch.Tensor:
120
+ ctx.save_for_backward(indices)
121
+ assert indices.ndim == 1
122
+ assert values.ndim >= 2
123
+ output = torch.zeros(
124
+ first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype
125
+ )
126
+ output[indices] = values
127
+ return output
128
+
129
+ @staticmethod
130
+ def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None, None]:
131
+ indices, = ctx.saved_tensors
132
+ grad_values = grad_output[indices]
133
+ return grad_values, None, None
134
+
135
+
136
+ index_put_first_axis = IndexPutFirstAxis.apply
137
+
138
+
139
+ def pad_input(inputs: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int) -> torch.Tensor:
140
+ """Add padding to sequences.
141
+
142
+ Arguments:
143
+ inputs: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
144
+ indices: (total_nnz), `indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()`
145
+ batch: int batch_size
146
+ seqlen: int max sequence length
147
+
148
+ Returns:
149
+ inputs: (batch, seqlen, ...)
150
+ """
151
+ output = index_put_first_axis(inputs, indices, batch * seqlen)
152
+ return output.view(batch, seqlen, *inputs.shape[1:])
153
+
154
+
155
+ def rotate_half(x):
156
+ """Rotates half the hidden dims of the input."""
157
+ x1 = x[..., : x.shape[-1] // 2]
158
+ x2 = x[..., x.shape[-1] // 2 :]
159
+ return torch.cat((-x2, x1), dim=-1)
160
+
161
+
162
+ def apply_rotary_pos_emb(q, k, cos, sin):
163
+ """Applies Rotary Position Embedding to the query and key tensors.
164
+
165
+ Args:
166
+ q (`torch.Tensor`): The query tensor.
167
+ k (`torch.Tensor`): The key tensor.
168
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
169
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
170
+ Returns:
171
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
172
+ """
173
+ cos, sin = cos.to(q.dtype), sin.to(q.dtype)
174
+ q_embed = (q * cos) + (rotate_half(q) * sin)
175
+ k_embed = (k * cos) + (rotate_half(k) * sin)
176
+ return q_embed, k_embed
177
+
178
+
179
+ class RotaryEmbedding(torch.nn.Module):
180
+ def __init__(self, dim, max_position_embeddings=512, base=10000.0, device=None):
181
+ super().__init__()
182
+
183
+ self.dim = dim
184
+ self.max_position_embeddings = max_position_embeddings
185
+ self.base = base
186
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
187
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
188
+
189
+ # Build here to make `torch.jit.trace` work.
190
+ self._set_cos_sin_cache(
191
+ seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
192
+ )
193
+
194
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
195
+ self.max_seq_len_cached = seq_len
196
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32)
197
+
198
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
199
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
200
+ emb = torch.cat((freqs, freqs), dim=-1)
201
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
202
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
203
+
204
+ def forward(self, x, seq_len=None):
205
+ # x: [bs, num_attention_heads, seq_len, head_size]
206
+ if seq_len > self.max_seq_len_cached:
207
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
208
+
209
+ return (
210
+ self.cos_cached[:seq_len, ...].to(dtype=x.dtype),
211
+ self.sin_cached[:seq_len, ...].to(dtype=x.dtype),
212
+ )
213
+
214
+
215
+ class NTKScalingRotaryEmbedding(RotaryEmbedding):
216
+ """RotaryEmbedding extended with fixed and mixed NTK scaling. https://kexue.fm/archives/9706 """
217
+
218
+ def __init__(self, dim, max_position_embeddings=512, base=10000, device=None, scaling_factor=1.0, mixed_b=None):
219
+ self.scaling_factor = scaling_factor
220
+ self.mixed_b = mixed_b
221
+ super().__init__(dim, max_position_embeddings, base, device)
222
+ max_position_embeddings = max_position_embeddings * self.scaling_factor
223
+ self._set_cos_sin_cache(max_position_embeddings, self.inv_freq.device, torch.get_default_dtype())
224
+
225
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
226
+ self.max_seq_len_cached = seq_len
227
+
228
+ if seq_len > self.max_position_embeddings:
229
+ base = self.base * (self.scaling_factor if self.mixed_b is None else 1)
230
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
231
+
232
+ if self.mixed_b is None:
233
+ inv_freq = inv_freq / self.scaling_factor ** (2 / self.dim) # (6)
234
+ else:
235
+ a = torch.tensor(self.scaling_factor).log() / (self.dim / 2) ** self.mixed_b # (13)
236
+ lambda_1_m = (a * torch.arange(1, self.dim // 2 + 1).float().to(device) ** self.mixed_b).exp() # (12)
237
+ inv_freq = inv_freq / lambda_1_m # (10)
238
+
239
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
240
+
241
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32)
242
+
243
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
244
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
245
+ emb = torch.cat((freqs, freqs), dim=-1)
246
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
247
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
248
+
249
+
250
+ class RMSNorm(nn.Module):
251
+ def __init__(self, hidden_size, eps=1e-6):
252
+ """
253
+ RMSNorm is equivalent to T5LayerNorm
254
+ """
255
+ super().__init__()
256
+ self.weight = nn.Parameter(torch.ones(hidden_size))
257
+ self.variance_epsilon = eps
258
+
259
+ def forward(self, hidden_states):
260
+ input_dtype = hidden_states.dtype
261
+ hidden_states = hidden_states.to(torch.float32)
262
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
263
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
264
+ return self.weight * hidden_states.to(input_dtype)
265
+
266
+
267
+ LAYER_NORM = {
268
+ 'layer_norm': nn.LayerNorm,
269
+ 'rms_norm': RMSNorm
270
+ }
271
+
272
+
273
+ class GteEmbeddings(nn.Module):
274
+ """
275
+ Embedding and Unpadding.
276
+ """
277
+
278
+ def __init__(self, config: GteConfig):
279
+ super().__init__()
280
+ self.padding_idx = config.pad_token_id
281
+ self.word_embeddings = nn.Embedding(
282
+ config.vocab_size, config.hidden_size, padding_idx=self.padding_idx
283
+ )
284
+
285
+ self.position_embedding_type = config.position_embedding_type
286
+ if self.position_embedding_type == 'absolute':
287
+ self.position_embeddings = nn.Embedding(
288
+ config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
289
+ )
290
+ elif self.position_embedding_type == 'rope':
291
+ self._init_rope(config)
292
+ else:
293
+ raise ValueError
294
+
295
+ self.type_vocab_size = config.type_vocab_size
296
+ if self.type_vocab_size > 0:
297
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
298
+
299
+ # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
300
+ # any TensorFlow checkpoint file
301
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
302
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
303
+ # position_ids is contiguous in memory and excluded when serialized
304
+ self.register_buffer(
305
+ "position_ids", torch.arange(config.max_position_embeddings), persistent=False
306
+ )
307
+
308
+ def _init_rope(self, config):
309
+ kwargs = dict(
310
+ dim=int(config.hidden_size / config.num_attention_heads),
311
+ max_position_embeddings=config.max_position_embeddings,
312
+ base=config.rope_theta
313
+ )
314
+ if config.rope_scaling is None:
315
+ self.rotary_emb = RotaryEmbedding(**kwargs)
316
+ else:
317
+ kwargs.update(scaling_factor=config.rope_scaling["factor"])
318
+ scaling_type = config.rope_scaling["type"]
319
+ if scaling_type == 'ntk':
320
+ kwargs.update(mixed_b=config.rope_scaling.get('mixed_b', None))
321
+ self.rotary_emb = NTKScalingRotaryEmbedding(**kwargs)
322
+ # elif scaling_type == "linear":
323
+ # self.rotary_emb = LinearScalingRotaryEmbedding(**kwargs)
324
+ # elif scaling_type == "dynamic":
325
+ # self.rotary_emb = DynamicNTKScalingRotaryEmbedding(**kwargs)
326
+ else:
327
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
328
+
329
+ def forward(
330
+ self,
331
+ unpad_inputs: bool,
332
+ input_ids: Optional[torch.Tensor] = None,
333
+ attention_mask: Optional[torch.Tensor] = None,
334
+ length: Optional[List[int]] = None,
335
+ token_type_ids: Optional[torch.Tensor] = None,
336
+ position_ids: Optional[torch.Tensor] = None,
337
+ inputs_embeds: Optional[torch.Tensor] = None,
338
+ ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Tuple], Optional[List[int]]]:
339
+ """
340
+ """
341
+ if inputs_embeds is None:
342
+ device, input_shape = input_ids.device, input_ids.shape
343
+ else:
344
+ device, input_shape = inputs_embeds.device, inputs_embeds.shape[:2]
345
+ batch_size, seq_length = input_shape
346
+
347
+ # Set attention_mask if it's None
348
+ if attention_mask is None:
349
+ attention_mask = torch.ones(input_shape, device=device)
350
+ if length is not None:
351
+ for i, l in enumerate(length):
352
+ attention_mask[i, l:] = 0
353
+
354
+ # Set attention_mask_bool for unpadding
355
+ if unpad_inputs:
356
+ attention_mask_bool = attention_mask.bool()
357
+ if length is None:
358
+ length = attention_mask.sum(-1).tolist()
359
+
360
+ # Get word embeddings
361
+ if inputs_embeds is None:
362
+ if unpad_inputs:
363
+ input_ids = input_ids[attention_mask_bool].unsqueeze(0)
364
+ inputs_embeds = self.word_embeddings(input_ids)
365
+ else:
366
+ if unpad_inputs:
367
+ inputs_embeds = inputs_embeds[attention_mask_bool].unsqueeze(0)
368
+ embeddings = inputs_embeds
369
+
370
+ # Set and unpad position_ids
371
+ if position_ids is None:
372
+ if seq_length > self.position_ids.size(0):
373
+ self.register_buffer(
374
+ "position_ids", torch.arange(seq_length, device=embeddings.device), persistent=False
375
+ )
376
+ if unpad_inputs:
377
+ # [1, cumsum_seq_len]
378
+ position_ids = torch.cat([self.position_ids[:l] for l in length]).unsqueeze(0)
379
+ else:
380
+ # [bs, seq_len]
381
+ position_ids = self.position_ids[:seq_length].expand(batch_size, -1)
382
+ elif unpad_inputs:
383
+ position_ids = position_ids[attention_mask_bool].unsqueeze(0) # [1, cumsum_seq_len]
384
+
385
+ # Compute rotary embedding
386
+ if self.position_embedding_type == 'rope':
387
+ rope_cos, rope_sin = self.rotary_emb(inputs_embeds, seq_len=seq_length)
388
+ rope_cos = rope_cos[position_ids].unsqueeze(2) # [bs, seq_len, 1, dim]
389
+ rope_sin = rope_sin[position_ids].unsqueeze(2) # [bs, seq_len, 1, dim]
390
+ rope_embeds = rope_cos, rope_sin
391
+ else:
392
+ rope_embeds = None
393
+
394
+ if self.type_vocab_size > 0:
395
+ if token_type_ids is None:
396
+ token_type_ids = position_ids.mul(0)
397
+ else:
398
+ if self.type_vocab_size < 2:
399
+ token_type_ids.mul_(0)
400
+ if unpad_inputs:
401
+ token_type_ids = token_type_ids[attention_mask_bool].unsqueeze(0)
402
+
403
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
404
+ embeddings = embeddings + token_type_embeddings
405
+
406
+ # BERT position
407
+ if self.position_embedding_type == "absolute":
408
+ position_embeddings = self.position_embeddings(position_ids)
409
+ embeddings = embeddings + position_embeddings
410
+
411
+ embeddings = self.LayerNorm(embeddings)
412
+ embeddings = self.dropout(embeddings)
413
+
414
+ return embeddings, attention_mask, rope_embeds, length
415
+
416
+
417
+ class GteAttention(nn.Module):
418
+ def __init__(self, config: GteConfig, pack_qkv=None, use_memory_efficient_attention=False):
419
+ super().__init__()
420
+ self.config = config
421
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
422
+ raise ValueError(
423
+ f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
424
+ f"heads ({config.num_attention_heads})"
425
+ )
426
+
427
+ self.hidden_size = config.hidden_size
428
+ self.num_attention_heads = config.num_attention_heads
429
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
430
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
431
+
432
+ if pack_qkv is None:
433
+ pack_qkv = config.pack_qkv
434
+ self.pack_qkv = pack_qkv
435
+
436
+ if self.pack_qkv:
437
+ self.qkv_proj = nn.Linear(config.hidden_size, self.all_head_size * 3, bias=True)
438
+ else:
439
+ self.q_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
440
+ self.k_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
441
+ self.v_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
442
+
443
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
444
+ self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=True)
445
+
446
+ self.use_memory_efficient_attention = False
447
+ self.memory_efficient_attention = None if xops is None else xops.memory_efficient_attention
448
+
449
+ def forward(
450
+ self,
451
+ hidden_states: torch.Tensor,
452
+ attention_bias: torch.FloatTensor,
453
+ rope_embeds: Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] = None,
454
+ padding_inputs: Optional[Tuple] = None, # indices, batch, seqlen
455
+ attention_scale: Optional[torch.FloatTensor] = None,
456
+ head_mask: Optional[torch.FloatTensor] = None,
457
+ output_attentions: Optional[bool] = False,
458
+ qkv_inputs: Optional[Tuple] = None, # For RetroMAE
459
+ ) -> Tuple[torch.Tensor, ...]:
460
+ shape_hd = (self.num_attention_heads, self.attention_head_size)
461
+ # qkv
462
+ if self.pack_qkv and qkv_inputs is None:
463
+ qkv_pack = self.qkv_proj(hidden_states).split(self.all_head_size, dim=-1)
464
+ else:
465
+ if qkv_inputs is None:
466
+ qkv_inputs = (hidden_states, hidden_states, hidden_states)
467
+ qkv_pack = [
468
+ getattr(self, n + '_proj')(s) for s, n in zip(qkv_inputs, 'qkv')
469
+ ]
470
+ query_states, key_states, value_states = [t.view(t.shape[:-1] + shape_hd) for t in qkv_pack]
471
+
472
+ if self.config.position_embedding_type == 'rope':
473
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, *rope_embeds)
474
+
475
+ dtype = query_states.dtype
476
+
477
+ if self.config.logn_attention_scale and attention_scale is not None:
478
+ # https://kexue.fm/archives/8823
479
+ query_states = query_states * attention_scale.to(dtype)
480
+
481
+ if padding_inputs is not None:
482
+ query_states = pad_input(query_states.squeeze(), *padding_inputs)
483
+ key_states = pad_input(key_states.squeeze(), *padding_inputs)
484
+ value_states = pad_input(value_states.squeeze(), *padding_inputs)
485
+
486
+ if output_attentions and isinstance(self, GteSdpaAttention):
487
+ raise RuntimeError("SDPA do not output attentions")
488
+ context_layer, attention_probs = self._attention(
489
+ query_states, key_states, value_states, attention_bias, head_mask
490
+ )
491
+
492
+ if padding_inputs is not None:
493
+ context_layer = unpad_input(context_layer, indices=padding_inputs[0])
494
+
495
+ gte_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
496
+ context_layer = context_layer.view(gte_context_layer_shape)
497
+
498
+ # output proj
499
+ attn_output = self.o_proj(context_layer)
500
+
501
+ # add attentions if we output them
502
+ outputs = (attn_output, attention_probs) if output_attentions else (attn_output,)
503
+ return outputs
504
+
505
+ def _attention(self, query_states, key_states, value_states, attention_bias, head_mask):
506
+ """
507
+ Args:
508
+ q/k/v: (B, L, n_head, head_dim),
509
+ Returns:
510
+ attn_output: (B L, n_head, head_dim)
511
+ """
512
+ query_states = query_states.transpose(1, 2)
513
+ key_states = key_states.transpose(1, 2)
514
+ value_states = value_states.transpose(1, 2)
515
+ # Take the dot product between "query" and "key" to get the raw attention scores.
516
+ attention_scores = torch.matmul(query_states, key_states.transpose(-1, -2))
517
+
518
+ attention_scores = attention_scores / math.sqrt(self.attention_head_size)
519
+ if attention_bias is not None:
520
+ # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
521
+ attention_scores = attention_scores + attention_bias
522
+
523
+ # Normalize the attention scores to probabilities.
524
+ attention_probs = nn.functional.softmax(attention_scores, dim=-1)
525
+
526
+ # This is actually dropping out entire tokens to attend to, which might
527
+ # seem a bit unusual, but is taken from the original Transformer paper.
528
+ if self.dropout.p > 0:
529
+ attention_probs = self.dropout(attention_probs)
530
+
531
+ # Mask heads if we want to
532
+ if head_mask is not None:
533
+ attention_probs = attention_probs * head_mask
534
+
535
+ context_layer = torch.matmul(attention_probs, value_states)
536
+
537
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
538
+ return context_layer, attention_probs
539
+
540
+
541
+ class GteSdpaAttention(GteAttention):
542
+ """
543
+ Gte attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
544
+ `GteAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
545
+ SDPA API.
546
+ """
547
+ def __init__(self, config: GteConfig, **kwargs):
548
+ super().__init__(config, **kwargs)
549
+ # torch.backends.cuda.enable_mem_efficient_sdp(False)
550
+ # logger.warning(
551
+ # "Disable memory efficient attention kernel for `GteSdpaAttention`, you can set "
552
+ # "`use_memory_efficient_attention=True` if it expected to use."
553
+ # )
554
+
555
+ def _attention(self, query_states, key_states, value_states, attention_bias, head_mask):
556
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
557
+ query_states.transpose(1, 2),
558
+ key_states.transpose(1, 2),
559
+ value_states.transpose(1, 2),
560
+ attn_mask=attention_bias,
561
+ dropout_p=self.dropout.p if self.training else 0.0,
562
+ )
563
+ attn_output = attn_output.permute(0, 2, 1, 3).contiguous()
564
+ return attn_output, None
565
+
566
+
567
+ GTE_ATTENTION_CLASSES = {
568
+ "eager": GteAttention,
569
+ # "flash_attention_2": , # TODO
570
+ "sdpa": GteSdpaAttention,
571
+ }
572
+
573
+
574
+ class GteGatedMLP(nn.Module):
575
+ """
576
+ GLU Variants Improve Transformer.
577
+ """
578
+
579
+ def __init__(self, config: GteConfig):
580
+ super().__init__()
581
+ self.intermediate_size = config.intermediate_size
582
+ self.up_gate_proj = nn.Linear(config.hidden_size, self.intermediate_size * 2, bias=False)
583
+ self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=True)
584
+ self.act_fn = ACT2FN[config.hidden_act]
585
+ if config.hidden_dropout_prob > 0:
586
+ self.hidden_dropout = nn.Dropout(config.hidden_dropout_prob)
587
+ else:
588
+ self.hidden_dropout = None
589
+
590
+ def forward(self, hidden_states):
591
+ up_gate = self.up_gate_proj(hidden_states)
592
+ up_states, gate = torch.split(up_gate, self.intermediate_size, dim=-1)
593
+ gate = self.act_fn(gate)
594
+ gated_states = gate * up_states
595
+ if self.hidden_dropout is not None:
596
+ gated_states = self.hidden_dropout(gated_states)
597
+ down_states = self.down_proj(gated_states)
598
+ return down_states
599
+
600
+
601
+ class GteLayer(nn.Module):
602
+ def __init__(
603
+ self,
604
+ config: GteConfig,
605
+ pack_qkv=None,
606
+ use_memory_efficient_attention=None,
607
+ attn_implementation=None
608
+ ):
609
+ super().__init__()
610
+ if attn_implementation is None:
611
+ attn_implementation = config._attn_implementation
612
+
613
+ use_memory_efficient_attention = False
614
+
615
+ self.attention = GTE_ATTENTION_CLASSES[attn_implementation](
616
+ config, pack_qkv=pack_qkv, use_memory_efficient_attention=use_memory_efficient_attention
617
+ )
618
+ self.mlp = GteGatedMLP(config)
619
+
620
+ ln_class = LAYER_NORM[config.layer_norm_type]
621
+ self.attn_ln = ln_class(config.hidden_size, eps=config.layer_norm_eps)
622
+ self.mlp_ln = ln_class(config.hidden_size, eps=config.layer_norm_eps)
623
+
624
+ if config.hidden_dropout_prob > 0:
625
+ self.hidden_dropout = nn.Dropout(config.hidden_dropout_prob)
626
+ else:
627
+ self.hidden_dropout = None
628
+
629
+ def forward(
630
+ self,
631
+ hidden_states: torch.Tensor,
632
+ attention_bias: torch.FloatTensor,
633
+ rope_embeds: Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] = None,
634
+ padding_inputs: Optional[Tuple] = None, # indices, batch, seqlen
635
+ attention_scale: Optional[torch.FloatTensor] = None,
636
+ subset_indices: Optional[torch.LongTensor] = None,
637
+ head_mask: Optional[torch.FloatTensor] = None,
638
+ output_attentions: Optional[bool] = False,
639
+ qkv_inputs: Optional[Tuple] = None, # For RetroMAE
640
+ ) -> Tuple[torch.Tensor, ...]:
641
+ # Multi head self attention
642
+ residual = hidden_states if qkv_inputs is None else qkv_inputs[0]
643
+ attention_outputs = self.attention(
644
+ hidden_states,
645
+ attention_bias,
646
+ rope_embeds,
647
+ padding_inputs,
648
+ attention_scale,
649
+ head_mask,
650
+ output_attentions=output_attentions,
651
+ qkv_inputs=qkv_inputs,
652
+ )
653
+ hidden_states = attention_outputs[0]
654
+ if self.hidden_dropout is not None:
655
+ hidden_states = self.hidden_dropout(hidden_states)
656
+ hidden_states = residual + hidden_states
657
+
658
+ # In pretraining, after the attention of last layer, we only need the masked tokens.
659
+ if subset_indices is not None:
660
+ hidden_states = hidden_states[subset_indices]
661
+
662
+ hidden_states = self.attn_ln(hidden_states)
663
+
664
+ # Fully Connected
665
+ residual = hidden_states
666
+ hidden_states = self.mlp(hidden_states)
667
+ if self.hidden_dropout is not None:
668
+ hidden_states = self.hidden_dropout(hidden_states)
669
+ hidden_states = residual + hidden_states
670
+ hidden_states = self.mlp_ln(hidden_states)
671
+
672
+ # add self attentions if we output attention weights
673
+ outputs = (hidden_states,) + attention_outputs[1:]
674
+ return outputs
675
+
676
+
677
+ class GteEncoder(nn.Module):
678
+ def __init__(self, config):
679
+ super().__init__()
680
+ self.config = config
681
+ self.layer = nn.ModuleList([GteLayer(config) for _ in range(config.num_hidden_layers)])
682
+ self.gradient_checkpointing = False
683
+
684
+ def forward(
685
+ self,
686
+ hidden_states: torch.Tensor,
687
+ attention_bias: Optional[torch.FloatTensor] = None,
688
+ rope_embeds: Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] = None,
689
+ padding_inputs: Optional[Tuple] = None, # indices, batch, seqlen
690
+ attention_scale: Optional[torch.FloatTensor] = None,
691
+ subset_indices: Optional[torch.LongTensor] = None,
692
+ head_mask: Optional[torch.FloatTensor] = None,
693
+ output_attentions: Optional[bool] = False,
694
+ output_hidden_states: Optional[bool] = False,
695
+ return_dict: Optional[bool] = True,
696
+ ) -> Union[Tuple[torch.Tensor], BaseModelOutput]:
697
+ all_hidden_states = () if output_hidden_states else None
698
+ all_self_attentions = () if output_attentions else None
699
+
700
+ for i, layer_module in enumerate(self.layer):
701
+ if output_hidden_states:
702
+ all_hidden_states = all_hidden_states + (hidden_states,)
703
+
704
+ if i >= len(self.layer) - 1:
705
+ layer_subset_indices = subset_indices
706
+ else:
707
+ layer_subset_indices = None
708
+
709
+ layer_head_mask = head_mask[i] if head_mask is not None else None
710
+
711
+ if self.gradient_checkpointing and self.training:
712
+ layer_outputs = self._gradient_checkpointing_func(
713
+ layer_module.__call__,
714
+ hidden_states,
715
+ attention_bias,
716
+ rope_embeds,
717
+ padding_inputs,
718
+ attention_scale,
719
+ layer_subset_indices,
720
+ layer_head_mask,
721
+ )
722
+ else:
723
+ layer_outputs = layer_module(
724
+ hidden_states,
725
+ attention_bias,
726
+ rope_embeds,
727
+ padding_inputs,
728
+ attention_scale,
729
+ layer_subset_indices,
730
+ layer_head_mask,
731
+ output_attentions,
732
+ )
733
+
734
+ hidden_states = layer_outputs[0]
735
+ if output_attentions:
736
+ all_self_attentions = all_self_attentions + (layer_outputs[1],)
737
+
738
+ if output_hidden_states:
739
+ all_hidden_states = all_hidden_states + (hidden_states,)
740
+
741
+ if not return_dict:
742
+ return tuple(
743
+ v
744
+ for v in [
745
+ hidden_states,
746
+ all_hidden_states,
747
+ all_self_attentions,
748
+ ]
749
+ if v is not None
750
+ )
751
+ return BaseModelOutput(
752
+ last_hidden_state=hidden_states,
753
+ hidden_states=all_hidden_states,
754
+ attentions=all_self_attentions,
755
+ )
756
+
757
+
758
+ # Copied from transformers.models.bert.modeling_bert.BertPooler with Bert->Gte
759
+ class GtePooler(nn.Module):
760
+ def __init__(self, config):
761
+ super().__init__()
762
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
763
+ self.activation = nn.Tanh()
764
+
765
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
766
+ # We "pool" the model by simply taking the hidden state corresponding
767
+ # to the first token.
768
+ first_token_tensor = hidden_states[:, 0]
769
+ pooled_output = self.dense(first_token_tensor)
770
+ pooled_output = self.activation(pooled_output)
771
+ return pooled_output
772
+
773
+
774
+ class GtePreTrainedModel(PreTrainedModel):
775
+ """
776
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
777
+ models.
778
+ """
779
+
780
+ config_class = GteConfig
781
+ base_model_prefix = "gte"
782
+ supports_gradient_checkpointing = True
783
+ _supports_sdpa = True
784
+
785
+ def _init_weights(self, module):
786
+ """Initialize the weights"""
787
+ if isinstance(module, nn.Linear):
788
+ # Slightly different from the TF version which uses truncated_normal for initialization
789
+ # cf https://github.com/pytorch/pytorch/pull/5617
790
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
791
+ if module.bias is not None:
792
+ module.bias.data.zero_()
793
+ elif isinstance(module, nn.Embedding):
794
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
795
+ if module.padding_idx is not None:
796
+ module.weight.data[module.padding_idx].zero_()
797
+ elif isinstance(module, nn.LayerNorm):
798
+ module.bias.data.zero_()
799
+ module.weight.data.fill_(1.0)
800
+
801
+
802
+ class GteModel(GtePreTrainedModel):
803
+ """
804
+ The bare Gte Model transformer outputting raw hidden-states without any specific head on top.
805
+ """
806
+
807
+ def __init__(self, config: GteConfig, add_pooling_layer=False):
808
+ super().__init__(config)
809
+ self.config = config
810
+
811
+ self.embeddings = GteEmbeddings(config)
812
+ self.encoder = GteEncoder(config)
813
+
814
+ self.pooler = GtePooler(config) if add_pooling_layer else None
815
+
816
+ # Initialize weights and apply final processing
817
+ self.post_init()
818
+
819
+ def get_input_embeddings(self):
820
+ return self.embeddings.word_embeddings
821
+
822
+ def set_input_embeddings(self, value):
823
+ self.embeddings.word_embeddings = value
824
+
825
+ def forward(
826
+ self,
827
+ input_ids: Optional[torch.Tensor] = None,
828
+ attention_mask: Optional[torch.Tensor] = None,
829
+ length: Optional[List[int]] = None,
830
+ subset_indices: Optional[torch.LongTensor] = None,
831
+ token_type_ids: Optional[torch.Tensor] = None,
832
+ position_ids: Optional[torch.Tensor] = None,
833
+ head_mask: Optional[torch.Tensor] = None,
834
+ inputs_embeds: Optional[torch.Tensor] = None,
835
+ output_attentions: Optional[bool] = None,
836
+ output_hidden_states: Optional[bool] = None,
837
+ return_dict: Optional[bool] = None,
838
+ unpad_inputs: Optional[bool] = None,
839
+ ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPooling]:
840
+ r"""
841
+ length (`list` of length `batch_size`, *optional*):
842
+ If is `None`, return padded `last_hidden_state`.
843
+ subset_indices ():
844
+ pass
845
+ unpad_inputs (`bool`, *optional*):
846
+ pass
847
+ """
848
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
849
+ output_hidden_states = (
850
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
851
+ )
852
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
853
+ unpad_inputs = unpad_inputs if unpad_inputs is not None else self.config.unpad_inputs
854
+ output_padded = length is None
855
+
856
+ if input_ids is not None and inputs_embeds is not None:
857
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
858
+ elif input_ids is not None:
859
+ self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
860
+ input_shape = input_ids.size()
861
+ elif inputs_embeds is not None:
862
+ input_shape = inputs_embeds.size()[:-1]
863
+ else:
864
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
865
+
866
+ # TODO: not used
867
+ # # Prepare head mask if needed
868
+ # # 1.0 in head_mask indicate we keep the head
869
+ # # attention_probs has shape bsz x n_heads x N x N
870
+ # # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
871
+ # # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
872
+ # head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
873
+
874
+ # Get embeddings, may unpad them
875
+ (embedding_output, attention_mask, rope_embeds, length) = self.embeddings(
876
+ unpad_inputs,
877
+ input_ids=input_ids,
878
+ attention_mask=attention_mask,
879
+ length=length,
880
+ token_type_ids=token_type_ids,
881
+ position_ids=position_ids,
882
+ inputs_embeds=inputs_embeds
883
+ )
884
+
885
+ batch_size, seq_length = input_shape
886
+
887
+ # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
888
+ # ourselves in which case we just need to make it broadcastable to all heads.
889
+ attention_bias = self.get_extended_attention_mask(attention_mask, input_shape)
890
+
891
+ padding_inputs = None
892
+ if unpad_inputs:
893
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
894
+ padding_inputs = (indices, *input_shape)
895
+
896
+ attention_scale = None
897
+ if self.config.logn_attention_scale:
898
+ logger.warning_once("TODO: logn_attention_scale")
899
+ # # attention scale log_512(input_len)
900
+ # attention_scale = attention_mask.sum(1).log() / torch.tensor(self.config.max_position_embeddings).log()
901
+ # # inference-time logn scale need clip 1
902
+ # if self.config.logn_attention_clip1:
903
+ # attention_scale.clip_(1)
904
+ # attention_scale = attention_scale[:, None, None, None]
905
+ # else:
906
+ # attention_scale = None
907
+
908
+ encoder_outputs = self.encoder(
909
+ embedding_output,
910
+ attention_bias=attention_bias,
911
+ rope_embeds=rope_embeds,
912
+ padding_inputs=padding_inputs,
913
+ attention_scale=attention_scale,
914
+ subset_indices=subset_indices,
915
+ head_mask=head_mask,
916
+ output_attentions=output_attentions,
917
+ output_hidden_states=output_hidden_states,
918
+ return_dict=return_dict,
919
+ )
920
+ sequence_output = encoder_outputs[0]
921
+ if unpad_inputs and output_padded:
922
+ sequence_output = pad_input(
923
+ sequence_output.squeeze(), indices, batch_size, seq_length
924
+ )
925
+
926
+ pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
927
+
928
+ if not return_dict:
929
+ return (sequence_output, pooled_output) + encoder_outputs[1:]
930
+
931
+ return BaseModelOutputWithPooling(
932
+ last_hidden_state=sequence_output,
933
+ pooler_output=pooled_output,
934
+ hidden_states=encoder_outputs.hidden_states,
935
+ attentions=encoder_outputs.attentions,
936
+ )
modules.json ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [
2
+ {
3
+ "idx": 0,
4
+ "name": "0",
5
+ "path": "",
6
+ "type": "sentence_transformers.models.Transformer"
7
+ },
8
+ {
9
+ "idx": 1,
10
+ "name": "1",
11
+ "path": "1_Pooling",
12
+ "type": "sentence_transformers.models.Pooling"
13
+ },
14
+ {
15
+ "idx": 2,
16
+ "name": "2",
17
+ "path": "2_Normalize",
18
+ "type": "sentence_transformers.models.Normalize"
19
+ }
20
+ ]
special_tokens_map.json ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "<s>",
4
+ "lstrip": false,
5
+ "normalized": false,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "cls_token": {
10
+ "content": "<s>",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "eos_token": {
17
+ "content": "</s>",
18
+ "lstrip": false,
19
+ "normalized": false,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ },
23
+ "mask_token": {
24
+ "content": "<mask>",
25
+ "lstrip": true,
26
+ "normalized": false,
27
+ "rstrip": false,
28
+ "single_word": false
29
+ },
30
+ "pad_token": {
31
+ "content": "<pad>",
32
+ "lstrip": false,
33
+ "normalized": false,
34
+ "rstrip": false,
35
+ "single_word": false
36
+ },
37
+ "sep_token": {
38
+ "content": "</s>",
39
+ "lstrip": false,
40
+ "normalized": false,
41
+ "rstrip": false,
42
+ "single_word": false
43
+ },
44
+ "unk_token": {
45
+ "content": "<unk>",
46
+ "lstrip": false,
47
+ "normalized": false,
48
+ "rstrip": false,
49
+ "single_word": false
50
+ }
51
+ }
tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f1cc44ad7faaeec47241864835473fd5403f2da94673f3f764a77ebcb0a803ec
3
+ size 17083009
tokenizer_config.json ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "added_tokens_decoder": {
3
+ "0": {
4
+ "content": "<s>",
5
+ "lstrip": false,
6
+ "normalized": false,
7
+ "rstrip": false,
8
+ "single_word": false,
9
+ "special": true
10
+ },
11
+ "1": {
12
+ "content": "<pad>",
13
+ "lstrip": false,
14
+ "normalized": false,
15
+ "rstrip": false,
16
+ "single_word": false,
17
+ "special": true
18
+ },
19
+ "2": {
20
+ "content": "</s>",
21
+ "lstrip": false,
22
+ "normalized": false,
23
+ "rstrip": false,
24
+ "single_word": false,
25
+ "special": true
26
+ },
27
+ "3": {
28
+ "content": "<unk>",
29
+ "lstrip": false,
30
+ "normalized": false,
31
+ "rstrip": false,
32
+ "single_word": false,
33
+ "special": true
34
+ },
35
+ "250001": {
36
+ "content": "<mask>",
37
+ "lstrip": true,
38
+ "normalized": false,
39
+ "rstrip": false,
40
+ "single_word": false,
41
+ "special": true
42
+ }
43
+ },
44
+ "bos_token": "<s>",
45
+ "clean_up_tokenization_spaces": true,
46
+ "cls_token": "<s>",
47
+ "eos_token": "</s>",
48
+ "mask_token": "<mask>",
49
+ "max_length": 512,
50
+ "model_max_length": 32768,
51
+ "pad_to_multiple_of": null,
52
+ "pad_token": "<pad>",
53
+ "pad_token_type_id": 0,
54
+ "padding_side": "right",
55
+ "sep_token": "</s>",
56
+ "stride": 0,
57
+ "tokenizer_class": "XLMRobertaTokenizer",
58
+ "truncation_side": "right",
59
+ "truncation_strategy": "longest_first",
60
+ "unk_token": "<unk>"
61
+ }