noname002 commited on
Commit
5d96034
·
verified ·
1 Parent(s): 6d0e380

Upload folder using huggingface_hub

Browse files
1_Pooling/config.json ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ {
2
+ "embedding_dimension": 768,
3
+ "pooling_mode": "lasttoken",
4
+ "include_prompt": true
5
+ }
README.md ADDED
@@ -0,0 +1,265 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ pipeline_tag: sentence-similarity
3
+ tags:
4
+ - gguf
5
+ - embedding
6
+ - eurobert
7
+ - llama-cpp
8
+ - jina-embeddings-v5
9
+ - feature-extraction
10
+ - mteb
11
+ - vllm
12
+ - sentence-transformers
13
+ language:
14
+ - multilingual
15
+ base_model: jinaai/jina-embeddings-v5-text-nano
16
+ base_model_relation: quantized
17
+ inference: false
18
+ license: cc-by-nc-4.0
19
+ library_name: llama.cpp
20
+ ---
21
+ <br><br>
22
+
23
+ <p align="center">
24
+ <img src="https://huggingface.co/datasets/jinaai/documentation-images/resolve/main/logo.webp" alt="Jina AI: Your Search Foundation, Supercharged!" width="150px">
25
+ </p>
26
+
27
+ ### **jina-embeddings-v5-text**: Task-Targeted Embedding Distillation
28
+
29
+ [Elastic Inference Service](https://www.elastic.co/docs/explore-analyze/elastic-inference/eis) | [ArXiv](https://arxiv.org/abs/2602.15547) | [Release Note](https://jina.ai/news/jina-embeddings-v5-text-distilling-4b-quality-into-sub-1b-multilingual-embeddings) | [Blog](https://www.elastic.co/search-labs/blog/jina-embeddings-v5-text)
30
+
31
+ ### Model Overview
32
+
33
+ <p align="center">
34
+ <img src="https://jina-ai-gmbh.ghost.io/content/images/2026/02/v5_architecture_1771470917.png" alt="jina-embeddings-v5-text Architecture" width="600px">
35
+ </p>
36
+ `jina-embeddings-v5-text-nano-retrieval` is a compact, high-performance text embedding model designed for information retrieval.
37
+
38
+ It is part of the **jina-embeddings-v5-text** model family, which also includes [jina-embeddings-v5-text-small](https://huggingface.co/jinaai/jina-embeddings-v5-text-small), for better performance at a bigger size.
39
+
40
+ Trained using a novel approach that combines distillation with task-specific contrastive losses, `jina-embeddings-v5-text-nano-retrieval` outperforms existing state-of-the-art models of similar size across diverse embedding benchmarks.
41
+ | Feature | Value |
42
+ | --- | --- |
43
+ | Parameters | 239M |
44
+ | Supported Tasks | `retrieval` |
45
+ | Max Sequence Length | 8192 |
46
+ | Embedding Dimension | 768 |
47
+ | Matryoshka Dimensions | 32, 64, 128, 256, 512, 768 |
48
+ | Pooling Strategy | Last-token pooling |
49
+ | Base Model | jinaai/jina-embeddings-v5-text-nano |
50
+
51
+ ![image](https://cdn-uploads.huggingface.co/production/uploads/6476ff2699a5ce743ccea3fc/SJw9j09PkErQ0v9P052S9.png)
52
+
53
+
54
+ ### Training and Evaluation
55
+
56
+ For training details and evaluation results, see our [technical report](https://arxiv.org/abs/2602.15547).
57
+
58
+ ### Usage
59
+
60
+ <details>
61
+ <summary>Requirements</a></summary>
62
+
63
+ The following Python packages are required:
64
+
65
+ - `transformers>=5.1.0`
66
+ - `torch>=2.8.0`
67
+ - `peft>=0.15.2`
68
+ - `vllm==0.15.1`
69
+
70
+ ### Optional / Recommended
71
+ - **flash-attention**: Installing [flash-attention](https://github.com/Dao-AILab/flash-attention) is recommended for improved inference speed and efficiency, but not mandatory.
72
+ - **sentence-transformers**: If you want to use the model via the `sentence-transformers` interface, install this package as well.
73
+
74
+ </details>
75
+
76
+ <details open>
77
+ <summary>via <a href="https://www.elastic.co/docs/explore-analyze/elastic-inference/eis">Elastic Inference Service</a></summary>
78
+
79
+ The fastest way to use v5-text in production. Elastic Inference Service (EIS) provides managed embedding inference with built-in scaling, so you can generate embeddings directly within your Elastic deployment.
80
+
81
+ ```bash
82
+ PUT _inference/text_embedding/jina-v5
83
+ {
84
+ "service": "elastic",
85
+ "service_settings": {
86
+ "model_id": "jina-embeddings-v5-text-nano"
87
+ }
88
+ }
89
+ ```
90
+
91
+ See the [Elastic Inference Service documentation](https://www.elastic.co/docs/explore-analyze/elastic-inference/eis) for setup details.
92
+
93
+ </details>
94
+
95
+
96
+ <details>
97
+ <summary>via <a href="https://sbert.net/">sentence-transformers</a></summary>
98
+
99
+ ```python
100
+ from sentence_transformers import SentenceTransformer
101
+ import torch
102
+
103
+ model = SentenceTransformer(
104
+ "jinaai/jina-embeddings-v5-text-nano-retrieval",
105
+ trust_remote_code=True,
106
+ model_kwargs={"dtype": torch.bfloat16}, # Recommended for GPUs
107
+ config_kwargs={"_attn_implementation": "flash_attention_2"}, # Recommended but optional
108
+ )
109
+ # Optional: set truncate_dim in encode() to control embedding size
110
+
111
+ query = "Which planet is known as the Red Planet?"
112
+ documents = [
113
+ "Venus is often called Earth's twin because of its similar size and proximity.",
114
+ "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
115
+ "Jupiter, the largest planet in our solar system, has a prominent red spot.",
116
+ "Saturn, famous for its rings, is sometimes mistaken for the Red Planet.",
117
+ ]
118
+
119
+ # Encode query and documents
120
+ query_embeddings = model.encode(sentences=query, prompt_name="query")
121
+ document_embeddings = model.encode(sentences=documents, prompt_name="document")
122
+ print(query_embeddings.shape, document_embeddings.shape)
123
+ # (768,) (4, 768)
124
+
125
+ similarity = model.similarity(query_embeddings, document_embeddings)
126
+ print(similarity)
127
+ # tensor([[0.5013, 0.7914, 0.6133, 0.5736]])
128
+ ```
129
+
130
+ </details>
131
+
132
+ <details>
133
+ <summary>via <a href="https://github.com/vllm-project/vllm">vLLM</a></summary>
134
+
135
+ ```python
136
+ from vllm import LLM
137
+ from vllm.config.pooler import PoolerConfig
138
+
139
+ # Initialize model
140
+ name = "jinaai/jina-embeddings-v5-text-nano-retrieval"
141
+ model = LLM(
142
+ model=name,
143
+ dtype="float16",
144
+ runner="pooling",
145
+ trust_remote_code=True,
146
+ pooler_config=PoolerConfig(seq_pooling_type="LAST", normalize=True),
147
+ )
148
+
149
+ # Create text prompts
150
+ query = "Overview of climate change impacts on coastal cities"
151
+ query_prompt = f"Query: {query}"
152
+
153
+ document = "The impacts of climate change on coastal cities are significant.."
154
+ document_prompt = f"Document: {document}"
155
+
156
+ # Encode all prompts
157
+ prompts = [query_prompt, document_prompt]
158
+ outputs = model.encode(prompts, pooling_task="embed")
159
+
160
+ ```
161
+
162
+ </details>
163
+
164
+ <details>
165
+ <summary>via <a href="https://github.com/ggml-org/llama.cpp">llama.cpp (GGUF)</a></summary>
166
+
167
+ Since our nano model is based on `jinaai/jina-embeddings-v5-text-nano`, which is not yet supported by llama.cpp, we provide our own branch of [llama.cpp](https://github.com/jina-ai/llama.cpp/tree/feat-jina-v5-text), which implements the necessary changes to support it for now.
168
+
169
+ To start the OpenAI API compatible HTTP server, run with the respective model version:
170
+
171
+ ```
172
+ llama-server \
173
+ -hf jinaai/jina-embeddings-v5-text-nano-retrieval:F16 \
174
+ --embedding \
175
+ --pooling last \
176
+ --batch-size 8192 \
177
+ --ubatch-size 8192 \
178
+ --ctx-size 8192
179
+ ```
180
+
181
+ Client:
182
+
183
+ ```
184
+ curl -X POST "http://127.0.0.1:8080/v1/embeddings" \
185
+ -H "Content-Type: application/json" \
186
+ -d '{
187
+ "input": [
188
+ "Query: A beautiful sunset over the beach",
189
+ "Query: Un beau coucher de soleil sur la plage",
190
+ "Document: 海滩上美丽的日落",
191
+ "Document: 浜辺に沈む美しい夕日",
192
+ "Document: Golden sunlight melts into the horizon, painting waves in warm amber and rose, while the sky whispers goodnight to the quiet, endless sea."
193
+ ]
194
+ }'
195
+ ```
196
+
197
+ Note: For the retrieval variant, add `Query: ` or `Document: ` prefix in front of your input as shown above.
198
+
199
+ </details>
200
+
201
+ <details>
202
+ <summary> via <a href="https://huggingface.co/docs/optimum/index">Optimum (ONNX)</a></summary>
203
+
204
+ You can run the ONNX-optimized version of the model locally using Hugging Face's `optimum` library. Make sure you have the required dependencies installed (e.g., `pip install optimum[onnxruntime] transformers torch`):
205
+
206
+ ```python
207
+ from optimum.onnxruntime import ORTModelForFeatureExtraction
208
+ from transformers import AutoTokenizer
209
+ import torch
210
+
211
+ model_id = "jinaai/jina-embeddings-v5-text-nano-retrieval"
212
+
213
+ # 1. Load tokenizer and ONNX model
214
+ # We specify the subfolder 'onnx' where the weights are located
215
+ tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
216
+ model = ORTModelForFeatureExtraction.from_pretrained(
217
+ model_id,
218
+ subfolder="onnx",
219
+ file_name="model.onnx",
220
+ provider="CPUExecutionProvider", # Or "CUDAExecutionProvider" for GPU
221
+ trust_remote_code=True,
222
+ )
223
+
224
+ # 2. Prepare input
225
+ texts = ["Query: How do I use Jina ONNX models?", "Document: Information about semantic matching."]
226
+ inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
227
+
228
+
229
+ # 4. Inference
230
+ with torch.no_grad():
231
+ outputs = model(**inputs)
232
+
233
+ # 5. Pooling (Crucial for Jina-v5)
234
+ # Jina-v5 uses LAST-TOKEN pooling.
235
+ # We take the hidden state of the last non-padding token.
236
+ last_hidden_state = outputs.last_hidden_state
237
+ # Find the indices of the last token (usually the end of the sequence)
238
+ sequence_lengths = inputs.attention_mask.sum(dim=1) - 1
239
+ embeddings = last_hidden_state[torch.arange(last_hidden_state.size(0)), sequence_lengths]
240
+
241
+ print('embeddings shape:', embeddings.shape)
242
+ print('embeddings:', embeddings)
243
+ ```
244
+
245
+ </details>
246
+
247
+ ### License
248
+
249
+ The model is licensed under CC BY-NC 4.0. For commercial use, please [contact us](mailto:sales@jina.ai).
250
+
251
+ ### Citation
252
+
253
+ If you find `jina-embeddings-v5-text-nano-retrieval` useful in your research, please cite the following paper:
254
+
255
+ ```
256
+ @misc{akram2026jinaembeddingsv5texttasktargetedembeddingdistillation,
257
+ title={jina-embeddings-v5-text: Task-Targeted Embedding Distillation},
258
+ author={Mohammad Kalim Akram and Saba Sturua and Nastia Havriushenko and Quentin Herreros and Michael Günther and Maximilian Werk and Han Xiao},
259
+ year={2026},
260
+ eprint={2602.15547},
261
+ archivePrefix={arXiv},
262
+ primaryClass={cs.CL},
263
+ url={https://arxiv.org/abs/2602.15547},
264
+ }
265
+ ```
config.json ADDED
@@ -0,0 +1,81 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "EuroBertModel"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "auto_map": {
8
+ "AutoConfig": "configuration_eurobert.EuroBertConfig",
9
+ "AutoModel": "modeling_eurobert.EuroBertModel",
10
+ "AutoModelForMaskedLM": "modeling_eurobert.EuroBertForMaskedLM"
11
+ },
12
+ "bos_token": "<|begin_of_text|>",
13
+ "bos_token_id": 33717,
14
+ "classifier_pooling": "late",
15
+ "clf_pooling": "late",
16
+ "dtype": "bfloat16",
17
+ "eos_token": "<|end_of_text|>",
18
+ "eos_token_id": 33718,
19
+ "head_dim": 64,
20
+ "hidden_act": "silu",
21
+ "hidden_dropout": 0.0,
22
+ "hidden_size": 768,
23
+ "initializer_range": 0.02,
24
+ "intermediate_size": 3072,
25
+ "is_decoder": false,
26
+ "mask_token": "<|mask|>",
27
+ "mask_token_id": 33719,
28
+ "max_position_embeddings": 8192,
29
+ "mlp_bias": false,
30
+ "model_type": "eurobert",
31
+ "num_attention_heads": 12,
32
+ "num_hidden_layers": 12,
33
+ "num_key_value_heads": 12,
34
+ "pad_token": "<|end_of_text|>",
35
+ "pad_token_id": 33718,
36
+ "pretraining_tp": 1,
37
+ "rms_norm_eps": 1e-05,
38
+ "rope_parameters": {
39
+ "factor": 1.0,
40
+ "rope_theta": 1000000,
41
+ "rope_type": "default"
42
+ },
43
+ "task_names": [
44
+ "retrieval",
45
+ "text-matching",
46
+ "clustering",
47
+ "classification"
48
+ ],
49
+ "tie_word_embeddings": false,
50
+ "transformers.js_config": {
51
+ "use_external_data_format": true
52
+ },
53
+ "transformers_version": "5.5.4",
54
+ "use_cache": false,
55
+ "vocab_size": 34048,
56
+ "vocabtrimmer": {
57
+ "mining_config": {
58
+ "dataset": [
59
+ "minhnguyent546/mc4-vi"
60
+ ],
61
+ "dataset_column": [
62
+ "text",
63
+ "pos"
64
+ ],
65
+ "dataset_split": "validation",
66
+ "min_frequency": 20,
67
+ "subset": null,
68
+ "target_vocab_size": null
69
+ },
70
+ "stats": {
71
+ "compression_rate_embedding": 26.546906187624753,
72
+ "compression_rate_full": 65.83411003963923,
73
+ "parameter_size_embedding/raw": 98500608,
74
+ "parameter_size_embedding/trimmed": 26148864,
75
+ "parameter_size_full/raw": 211766016,
76
+ "parameter_size_full/trimmed": 139414272,
77
+ "vocab_size/raw": 128256,
78
+ "vocab_size/trimmed": 34048
79
+ }
80
+ }
81
+ }
config_sentence_transformers.json ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "__version__": {
3
+ "pytorch": "2.9.1+cu128",
4
+ "sentence_transformers": "5.4.1",
5
+ "transformers": "5.5.4"
6
+ },
7
+ "default_prompt_name": null,
8
+ "model_type": "SentenceTransformer",
9
+ "prompts": {
10
+ "document": "Document: ",
11
+ "query": "Query: "
12
+ },
13
+ "similarity_fn_name": "cosine"
14
+ }
configuration_eurobert.py ADDED
@@ -0,0 +1,216 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
2
+ # This file was automatically generated from src/transformers/models/eurobert/modular_eurobert.py.
3
+ # Do NOT edit this file manually as any edits will be overwritten by the generation of
4
+ # the file from the modular. If any change should be done, please apply the change to the
5
+ # modular_eurobert.py file directly. One of our CI enforces this.
6
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
7
+ # coding=utf-8
8
+ # Copyright 2025 Nicolas Boizard, Duarte M. Alves, Hippolyte Gisserot-Boukhlef and the EuroBert team. All rights reserved.
9
+ #
10
+ #
11
+ # Licensed under the Apache License, Version 2.0 (the "License");
12
+ # you may not use this file except in compliance with the License.
13
+ # You may obtain a copy of the License at
14
+ #
15
+ # http://www.apache.org/licenses/LICENSE-2.0
16
+ #
17
+ # Unless required by applicable law or agreed to in writing, software
18
+ # distributed under the License is distributed on an "AS IS" BASIS,
19
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20
+ # See the License for the specific language governing permissions and
21
+ # limitations under the License.
22
+
23
+ from transformers.utils import logging
24
+ from transformers.models.llama import LlamaConfig
25
+
26
+
27
+ logger = logging.get_logger(__name__)
28
+
29
+
30
+ class EuroBertConfig(LlamaConfig):
31
+ r"""
32
+ This is the configuration class to store the configuration of a [`EuroBertModel`]. It is used to instantiate an EuroBert
33
+ model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
34
+ defaults will yield a similar configuration to that of the EuroBERT-210m.
35
+
36
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
37
+ documentation from [`PretrainedConfig`] for more information.
38
+
39
+
40
+ Args:
41
+ vocab_size (`int`, *optional*, defaults to 128256):
42
+ Vocabulary size of the EuroBert model. Defines the number of different tokens that can be represented by the
43
+ `inputs_ids` passed when calling [`EuroBertModel`]
44
+ hidden_size (`int`, *optional*, defaults to 768):
45
+ Dimensionality of the encoder layers and the pooler layer.
46
+ intermediate_size (`int`, *optional*, defaults to 3072):
47
+ Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
48
+ num_hidden_layers (`int`, *optional*, defaults to 12):
49
+ Number of hidden layers in the Transformer encoder.
50
+ num_attention_heads (`int`, *optional*, defaults to 12):
51
+ Number of attention heads for each attention layer in the Transformer encoder.
52
+ num_key_value_heads (`int`, *optional*):
53
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
54
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
55
+ `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
56
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
57
+ by meanpooling all the original heads within that group. For more details checkout [this
58
+ paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
59
+ `num_attention_heads`.
60
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
61
+ The non-linear activation function (function or string) in the encoder and pooler.
62
+ max_position_embeddings (`int`, *optional*, defaults to 8192):
63
+ The maximum sequence length that this model might ever be used with. EuroBert supports up to 8192 tokens,
64
+ EuroBert-pretrained up to 2048.
65
+ initializer_range (`float`, *optional*, defaults to 0.02):
66
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
67
+ rms_norm_eps (`float`, *optional*, defaults to 1e-05):
68
+ The epsilon used by the rms normalization layers.
69
+ bos_token_id (`int`, *optional*, defaults to 128000):
70
+ Beginning of stream token id.
71
+ eos_token_id (`int`, *optional*, defaults to 128001):
72
+ End of stream token id.
73
+ pad_token_id (`int`, *optional*, defaults to 128001):
74
+ Padding token id.
75
+ mask_token_id (`int`, *optional*, defaults to 128002):
76
+ Mask token id.
77
+ pretraining_tp (`int`, *optional*, defaults to 1):
78
+ Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
79
+ document](https://huggingface.co/docs/transformers/main/perf_train_gpu_many#tensor-parallelism) to
80
+ understand more about it. This value is necessary to ensure exact reproducibility of the pretraining
81
+ results. Please refer to [this issue](https://github.com/pytorch/pytorch/issues/76232).
82
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
83
+ Whether to tie weight embeddings
84
+ rope_theta (`float`, *optional*, defaults to 250000.0):
85
+ The base period of the RoPE embeddings. EuroBert used base period of 250000.0,
86
+ EuroBert-pretrained 10000.0.
87
+ rope_scaling (`Dict`, *optional*):
88
+ Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
89
+ and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
90
+ accordingly.
91
+ Expected contents:
92
+ `rope_type` (`str`):
93
+ The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
94
+ 'eurobert3'], with 'default' being the original RoPE implementation.
95
+ `factor` (`float`, *optional*):
96
+ Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
97
+ most scaling types, a `factor` of x will enable the model to handle sequences of length x *
98
+ original maximum pre-trained length.
99
+ `original_max_position_embeddings` (`int`, *optional*):
100
+ Used with 'dynamic', 'longrope' and 'eurobert3'. The original max position embeddings used during
101
+ pretraining.
102
+ `attention_factor` (`float`, *optional*):
103
+ Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
104
+ computation. If unspecified, it defaults to value recommended by the implementation, using the
105
+ `factor` field to infer the suggested value.
106
+ `beta_fast` (`float`, *optional*):
107
+ Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
108
+ ramp function. If unspecified, it defaults to 32.
109
+ `beta_slow` (`float`, *optional*):
110
+ Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
111
+ ramp function. If unspecified, it defaults to 1.
112
+ `short_factor` (`List[float]`, *optional*):
113
+ Only used with 'longrope'. The scaling factor to be applied to short contexts (<
114
+ `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
115
+ size divided by the number of attention heads divided by 2
116
+ `long_factor` (`List[float]`, *optional*):
117
+ Only used with 'longrope'. The scaling factor to be applied to long contexts (<
118
+ `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
119
+ size divided by the number of attention heads divided by 2
120
+ `low_freq_factor` (`float`, *optional*):
121
+ Only used with 'eurobert3'. Scaling factor applied to low frequency components of the RoPE
122
+ `high_freq_factor` (`float`, *optional*):
123
+ Only used with 'eurobert3'. Scaling factor applied to high frequency components of the RoPE
124
+ attention_bias (`bool`, *optional*, defaults to `False`):
125
+ Whether to use a bias in the query, key, value and output projection layers during self-attention.
126
+ attention_dropout (`float`, *optional*, defaults to 0.0):
127
+ The dropout ratio for the attention probabilities.
128
+ mlp_bias (`bool`, *optional*, defaults to `False`):
129
+ Whether to use a bias in up_proj, down_proj and gate_proj layers in the MLP layers.
130
+ head_dim (`int`, *optional*):
131
+ The attention head dimension. If None, it will default to hidden_size // num_attention_heads
132
+ classifier_pooling (`str`, *optional*, defaults to `"late"`):
133
+ The pooling strategy to use for the classifier. Can be one of ['bos', 'mean', 'late'].
134
+
135
+ ```python
136
+ >>> from transformers import EuroBertModel, EuroBertConfig
137
+
138
+ >>> # Initializing a EuroBert eurobert-base style configuration
139
+ >>> configuration = EuroBertConfig()
140
+
141
+ >>> # Initializing a model from the eurobert-base style configuration
142
+ >>> model = EuroBertModel(configuration)
143
+
144
+ >>> # Accessing the model configuration
145
+ >>> configuration = model.config
146
+ ```"""
147
+
148
+ model_type = "eurobert"
149
+
150
+ def __init__(
151
+ self,
152
+ vocab_size=128256,
153
+ hidden_size=768,
154
+ intermediate_size=3072,
155
+ num_hidden_layers=12,
156
+ num_attention_heads=12,
157
+ num_key_value_heads=None,
158
+ hidden_act="silu",
159
+ max_position_embeddings=8192,
160
+ initializer_range=0.02,
161
+ rms_norm_eps=1e-05,
162
+ bos_token_id=128000,
163
+ eos_token_id=128001,
164
+ pad_token_id=128001,
165
+ mask_token_id=128002,
166
+ pretraining_tp=1,
167
+ tie_word_embeddings=False,
168
+ rope_theta=250000.0,
169
+ rope_scaling=None,
170
+ attention_bias=False,
171
+ attention_dropout=0.0,
172
+ mlp_bias=False,
173
+ head_dim=None,
174
+ classifier_pooling="late",
175
+ **kwargs,
176
+ ):
177
+ # use_cache is specific to decoder models and should be set to False for encoder models
178
+ use_cache = kwargs.pop("use_cache", None)
179
+ if use_cache:
180
+ logger.warning_once(
181
+ "The `use_cache` argument to EuroBertConfig is set to `False`, as caching is never used for encoder models."
182
+ )
183
+
184
+ if num_key_value_heads is None:
185
+ num_key_value_heads = num_attention_heads
186
+
187
+ super().__init__(
188
+ vocab_size=vocab_size,
189
+ hidden_size=hidden_size,
190
+ intermediate_size=intermediate_size,
191
+ num_hidden_layers=num_hidden_layers,
192
+ num_attention_heads=num_attention_heads,
193
+ num_key_value_heads=num_key_value_heads,
194
+ hidden_act=hidden_act,
195
+ max_position_embeddings=max_position_embeddings,
196
+ initializer_range=initializer_range,
197
+ rms_norm_eps=rms_norm_eps,
198
+ use_cache=False,
199
+ bos_token_id=bos_token_id,
200
+ eos_token_id=eos_token_id,
201
+ pad_token_id=pad_token_id,
202
+ pretraining_tp=pretraining_tp,
203
+ tie_word_embeddings=tie_word_embeddings,
204
+ rope_theta=rope_theta,
205
+ rope_scaling=rope_scaling,
206
+ attention_bias=attention_bias,
207
+ attention_dropout=attention_dropout,
208
+ mlp_bias=mlp_bias,
209
+ head_dim=head_dim,
210
+ **kwargs,
211
+ )
212
+ self.mask_token_id = mask_token_id
213
+ self.clf_pooling = classifier_pooling
214
+
215
+
216
+ __all__ = ["EuroBertConfig"]
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5eabe908e7c1933301f9e283dc009e9f65705c7149404b79386d3d4e8ef83abf
3
+ size 278840176
modeling_eurobert.py ADDED
@@ -0,0 +1,1094 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
2
+ # This file was automatically generated from src/transformers/models/eurobert/modular_eurobert.py.
3
+ # Do NOT edit this file manually as any edits will be overwritten by the generation of
4
+ # the file from the modular. If any change should be done, please apply the change to the
5
+ # modular_eurobert.py file directly. One of our CI enforces this.
6
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
7
+ # coding=utf-8
8
+ # Copyright 2025 Nicolas Boizard, Duarte M. Alves, Hippolyte Gisserot-Boukhlef and the EuroBert team. All rights reserved.
9
+ #
10
+ #
11
+ # Licensed under the Apache License, Version 2.0 (the "License");
12
+ # you may not use this file except in compliance with the License.
13
+ # You may obtain a copy of the License at
14
+ #
15
+ # http://www.apache.org/licenses/LICENSE-2.0
16
+ #
17
+ # Unless required by applicable law or agreed to in writing, software
18
+ # distributed under the License is distributed on an "AS IS" BASIS,
19
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20
+ # See the License for the specific language governing permissions and
21
+ # limitations under the License.
22
+
23
+ from typing import Callable, Optional, Tuple, Union
24
+
25
+ import torch
26
+ from torch import nn
27
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
28
+
29
+ try:
30
+ from transformers import initialization as init
31
+ except ImportError: # transformers < v5
32
+ from transformers import modeling_utils as _modeling_utils
33
+ init = getattr(_modeling_utils, "init", torch.nn.init)
34
+
35
+ from transformers.activations import ACT2FN
36
+ from transformers.cache_utils import Cache, StaticCache
37
+ from transformers.modeling_attn_mask_utils import AttentionMaskConverter
38
+ from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
39
+ from transformers.modeling_outputs import (
40
+ BaseModelOutput,
41
+ BaseModelOutputWithPast,
42
+ MaskedLMOutput,
43
+ QuestionAnsweringModelOutput,
44
+ SequenceClassifierOutput,
45
+ TokenClassifierOutput,
46
+ )
47
+ from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
48
+ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
49
+ from transformers.processing_utils import Unpack
50
+ from transformers.utils import (
51
+ add_code_sample_docstrings,
52
+ add_start_docstrings,
53
+ add_start_docstrings_to_model_forward,
54
+ logging,
55
+ )
56
+
57
+ try:
58
+ from .configuration_eurobert import EuroBertConfig
59
+ except ImportError:
60
+ from configuration_eurobert import EuroBertConfig
61
+
62
+
63
+ logger = logging.get_logger(__name__)
64
+
65
+ _CHECKPOINT_FOR_DOC = "EuroBERT/EuroBERT-210m"
66
+ _CONFIG_FOR_DOC = "EuroBertConfig"
67
+
68
+
69
+ class EuroBertRMSNorm(nn.Module):
70
+ def __init__(self, hidden_size, eps=1e-5):
71
+ """
72
+ EuroBertRMSNorm is equivalent to T5LayerNorm
73
+ """
74
+ super().__init__()
75
+ self.weight = nn.Parameter(torch.ones(hidden_size))
76
+ self.variance_epsilon = eps
77
+
78
+ def forward(self, hidden_states):
79
+ input_dtype = hidden_states.dtype
80
+ hidden_states = hidden_states.to(torch.float32)
81
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
82
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
83
+ return self.weight * hidden_states.to(input_dtype)
84
+
85
+ def extra_repr(self):
86
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
87
+
88
+
89
+ def rotate_half(x):
90
+ """Rotates half the hidden dims of the input."""
91
+ x1 = x[..., : x.shape[-1] // 2]
92
+ x2 = x[..., x.shape[-1] // 2 :]
93
+ return torch.cat((-x2, x1), dim=-1)
94
+
95
+
96
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
97
+ """Applies Rotary Position Embedding to the query and key tensors.
98
+
99
+ Args:
100
+ q (`torch.Tensor`): The query tensor.
101
+ k (`torch.Tensor`): The key tensor.
102
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
103
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
104
+ position_ids (`torch.Tensor`, *optional*):
105
+ Deprecated and unused.
106
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
107
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
108
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
109
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
110
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
111
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
112
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
113
+ Returns:
114
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
115
+ """
116
+ cos = cos.unsqueeze(unsqueeze_dim)
117
+ sin = sin.unsqueeze(unsqueeze_dim)
118
+ q_embed = (q * cos) + (rotate_half(q) * sin)
119
+ k_embed = (k * cos) + (rotate_half(k) * sin)
120
+ return q_embed, k_embed
121
+
122
+
123
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
124
+ """
125
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
126
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
127
+ """
128
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
129
+ if n_rep == 1:
130
+ return hidden_states
131
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
132
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
133
+
134
+
135
+ def eager_attention_forward(
136
+ module: nn.Module,
137
+ query: torch.Tensor,
138
+ key: torch.Tensor,
139
+ value: torch.Tensor,
140
+ attention_mask: Optional[torch.Tensor],
141
+ scaling: float,
142
+ dropout: float = 0.0,
143
+ **kwargs,
144
+ ):
145
+ key_states = repeat_kv(key, module.num_key_value_groups)
146
+ value_states = repeat_kv(value, module.num_key_value_groups)
147
+
148
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
149
+ if attention_mask is not None:
150
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
151
+ attn_weights = attn_weights + causal_mask
152
+
153
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
154
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
155
+ attn_output = torch.matmul(attn_weights, value_states)
156
+ attn_output = attn_output.transpose(1, 2).contiguous()
157
+
158
+ return attn_output, attn_weights
159
+
160
+
161
+ class EuroBertAttention(nn.Module):
162
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
163
+
164
+ def __init__(self, config: EuroBertConfig, layer_idx: int):
165
+ super().__init__()
166
+ self.config = config
167
+ self.layer_idx = layer_idx
168
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
169
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
170
+ self.scaling = self.head_dim**-0.5
171
+ self.attention_dropout = config.attention_dropout
172
+ self.is_causal = False
173
+
174
+ self.q_proj = nn.Linear(
175
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
176
+ )
177
+ self.k_proj = nn.Linear(
178
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
179
+ )
180
+ self.v_proj = nn.Linear(
181
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
182
+ )
183
+ self.o_proj = nn.Linear(
184
+ config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
185
+ )
186
+
187
+ def forward(
188
+ self,
189
+ hidden_states: torch.Tensor,
190
+ position_embeddings: Tuple[torch.Tensor, torch.Tensor],
191
+ attention_mask: Optional[torch.Tensor],
192
+ **kwargs: Unpack[FlashAttentionKwargs],
193
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
194
+ input_shape = hidden_states.shape[:-1]
195
+
196
+ hidden_shape = (*input_shape, -1, self.head_dim)
197
+
198
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
199
+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
200
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
201
+
202
+ cos, sin = position_embeddings
203
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
204
+
205
+ attention_interface: Callable = eager_attention_forward
206
+ if self.config._attn_implementation != "eager":
207
+ if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):
208
+ logger.warning_once(
209
+ "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "
210
+ 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
211
+ )
212
+ else:
213
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
214
+ attn_output, attn_weights = attention_interface(
215
+ self,
216
+ query_states,
217
+ key_states,
218
+ value_states,
219
+ attention_mask,
220
+ dropout=0.0 if not self.training else self.attention_dropout,
221
+ scaling=self.scaling,
222
+ is_causal=False,
223
+ **kwargs,
224
+ )
225
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
226
+ attn_output = self.o_proj(attn_output)
227
+ return attn_output, attn_weights
228
+
229
+
230
+ EUROBERT_START_DOCSTRING = r"""
231
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
232
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
233
+ etc.)
234
+
235
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
236
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
237
+ and behavior.
238
+
239
+ Parameters:
240
+ config ([`EuroBertConfig`]):
241
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
242
+ load the weights associated with the model, only the configuration. Check out the
243
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
244
+ """
245
+
246
+
247
+ @add_start_docstrings(
248
+ "The bare EuroBERT Model outputting raw hidden-states without any specific head on top.",
249
+ EUROBERT_START_DOCSTRING,
250
+ )
251
+ class EuroBertPreTrainedModel(PreTrainedModel):
252
+ config_class = EuroBertConfig
253
+ base_model_prefix = "model"
254
+ supports_gradient_checkpointing = True
255
+ _no_split_modules = ["EuroBertDecoderLayer"]
256
+ _skip_keys_device_placement = ["past_key_values"]
257
+ _supports_flash_attn_2 = True
258
+ _supports_sdpa = True
259
+ _supports_flex_attn = True
260
+ _supports_cache_class = True
261
+ _supports_quantized_cache = True
262
+ _supports_static_cache = True
263
+ _supports_attention_backend = True
264
+
265
+ def _init_weights(self, module):
266
+ std = self.config.initializer_range
267
+ if isinstance(module, nn.Linear):
268
+ init.normal_(module.weight, mean=0.0, std=std)
269
+ if module.bias is not None:
270
+ init.zeros_(module.bias)
271
+ elif isinstance(module, nn.Embedding):
272
+ init.normal_(module.weight, mean=0.0, std=std)
273
+ if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False):
274
+ init.zeros_(module.weight[module.padding_idx])
275
+ elif isinstance(module, EuroBertRotaryEmbedding):
276
+ rope_fn = ROPE_INIT_FUNCTIONS[module.rope_type]
277
+ buffer_value, _ = rope_fn(module.config, device=module.inv_freq.device)
278
+ if hasattr(init, "copy_"):
279
+ init.copy_(module.inv_freq, buffer_value)
280
+ init.copy_(module.original_inv_freq, buffer_value)
281
+ else:
282
+ module.inv_freq.copy_(buffer_value)
283
+ module.original_inv_freq.copy_(buffer_value)
284
+
285
+
286
+ class EuroBertRotaryEmbedding(nn.Module):
287
+ def __init__(self, config: EuroBertConfig, device=None):
288
+ super().__init__()
289
+ # BC: "rope_type" was originally "type"
290
+ if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
291
+ self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
292
+ else:
293
+ self.rope_type = "default"
294
+ if self.rope_type == "default":
295
+ self.rope_type = "linear"
296
+ # Ensure rope_scaling is set up with factor=1.0 for linear (no scaling, equivalent to default)
297
+ if not hasattr(config, "rope_scaling") or config.rope_scaling is None:
298
+ config.rope_scaling = {"rope_type": "linear", "factor": 1.0}
299
+ elif "factor" not in config.rope_scaling:
300
+ config.rope_scaling["factor"] = 1.0
301
+ self.max_seq_len_cached = config.max_position_embeddings
302
+ self.original_max_seq_len = config.max_position_embeddings
303
+
304
+ self.config = config
305
+ self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
306
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
307
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
308
+ self.original_inv_freq = self.inv_freq
309
+
310
+ def _dynamic_frequency_update(self, position_ids, device):
311
+ """
312
+ dynamic RoPE layers should recompute `inv_freq` in the following situations:
313
+ 1 - growing beyond the cached sequence length (allow scaling)
314
+ 2 - the current sequence length is in the original scale (avoid losing precision with small sequences)
315
+ """
316
+ seq_len = torch.max(position_ids) + 1
317
+ if seq_len > self.max_seq_len_cached: # growth
318
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, seq_len=seq_len)
319
+ self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation
320
+ self.max_seq_len_cached = seq_len
321
+
322
+ if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset
323
+ # This .to() is needed if the model has been moved to a device after being initialized (because
324
+ # the buffer is automatically moved, but not the original copy)
325
+ self.original_inv_freq = self.original_inv_freq.to(device)
326
+ self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
327
+ self.max_seq_len_cached = self.original_max_seq_len
328
+
329
+ @torch.no_grad()
330
+ def forward(self, x, position_ids):
331
+ if "dynamic" in self.rope_type:
332
+ self._dynamic_frequency_update(position_ids, device=x.device)
333
+
334
+ # Core RoPE block
335
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
336
+ position_ids_expanded = position_ids[:, None, :].float()
337
+ # Force float32 (see https://github.com/huggingface/transformers/pull/29285)
338
+ device_type = x.device.type
339
+ device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
340
+ with torch.autocast(device_type=device_type, enabled=False):
341
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
342
+ emb = torch.cat((freqs, freqs), dim=-1)
343
+ cos = emb.cos()
344
+ sin = emb.sin()
345
+
346
+ # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention
347
+ cos = cos * self.attention_scaling
348
+ sin = sin * self.attention_scaling
349
+
350
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
351
+
352
+
353
+ class EuroBertMLP(nn.Module):
354
+ def __init__(self, config):
355
+ super().__init__()
356
+ self.config = config
357
+ self.hidden_size = config.hidden_size
358
+ self.intermediate_size = config.intermediate_size
359
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
360
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
361
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
362
+ self.act_fn = ACT2FN[config.hidden_act]
363
+
364
+ def forward(self, x):
365
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
366
+ return down_proj
367
+
368
+
369
+ class EuroBertDecoderLayer(nn.Module):
370
+ def __init__(self, config: EuroBertConfig, layer_idx: int):
371
+ super().__init__()
372
+ self.hidden_size = config.hidden_size
373
+
374
+ self.self_attn = EuroBertAttention(config=config, layer_idx=layer_idx)
375
+
376
+ self.mlp = EuroBertMLP(config)
377
+ self.input_layernorm = EuroBertRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
378
+ self.post_attention_layernorm = EuroBertRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
379
+
380
+ def forward(
381
+ self,
382
+ hidden_states: torch.Tensor,
383
+ attention_mask: Optional[torch.Tensor] = None,
384
+ position_ids: Optional[torch.LongTensor] = None,
385
+ past_key_value: Optional[Cache] = None,
386
+ output_attentions: Optional[bool] = False,
387
+ use_cache: Optional[bool] = False,
388
+ cache_position: Optional[torch.LongTensor] = None,
389
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
390
+ **kwargs: Unpack[FlashAttentionKwargs],
391
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
392
+ residual = hidden_states
393
+
394
+ hidden_states = self.input_layernorm(hidden_states)
395
+
396
+ # Self Attention
397
+ hidden_states, self_attn_weights = self.self_attn(
398
+ hidden_states=hidden_states,
399
+ attention_mask=attention_mask,
400
+ position_ids=position_ids,
401
+ past_key_value=past_key_value,
402
+ output_attentions=output_attentions,
403
+ use_cache=use_cache,
404
+ cache_position=cache_position,
405
+ position_embeddings=position_embeddings,
406
+ **kwargs,
407
+ )
408
+
409
+ hidden_states = residual + hidden_states
410
+
411
+ # Fully Connected
412
+ residual = hidden_states
413
+ hidden_states = self.post_attention_layernorm(hidden_states)
414
+ hidden_states = self.mlp(hidden_states)
415
+ hidden_states = residual + hidden_states
416
+
417
+ outputs = (hidden_states,)
418
+ if output_attentions:
419
+ outputs += (self_attn_weights,)
420
+
421
+ return outputs
422
+
423
+
424
+ EUROBERT_INPUTS_DOCSTRING = r"""
425
+ Args:
426
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
427
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
428
+ it.
429
+
430
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
431
+ [`PreTrainedTokenizer.__call__`] for details.
432
+
433
+ [What are input IDs?](../glossary#input-ids)
434
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
435
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
436
+
437
+ - 1 for tokens that are **not masked**,
438
+ - 0 for tokens that are **masked**.
439
+
440
+ [What are attention masks?](../glossary#attention-mask)
441
+
442
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
443
+ [`PreTrainedTokenizer.__call__`] for details.
444
+
445
+ If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
446
+ `past_key_values`).
447
+
448
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
449
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
450
+ information on the default strategy.
451
+
452
+ - 1 indicates the head is **not masked**,
453
+ - 0 indicates the head is **masked**.
454
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
455
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
456
+ config.n_positions - 1]`.
457
+
458
+ [What are position IDs?](../glossary#position-ids)
459
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
460
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
461
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
462
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
463
+
464
+ Two formats are allowed:
465
+ - a [`~cache_utils.Cache`] instance, see our
466
+ [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache);
467
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
468
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
469
+ cache format.
470
+
471
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
472
+ legacy cache format will be returned.
473
+
474
+ If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
475
+ have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
476
+ of shape `(batch_size, sequence_length)`.
477
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
478
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
479
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
480
+ model's internal embedding lookup matrix.
481
+ use_cache (`bool`, *optional*):
482
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
483
+ `past_key_values`).
484
+ output_attentions (`bool`, *optional*):
485
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
486
+ tensors for more detail.
487
+ output_hidden_states (`bool`, *optional*):
488
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
489
+ more detail.
490
+ return_dict (`bool`, *optional*):
491
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
492
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
493
+ Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,
494
+ this tensor is not affected by padding. It is used to update the cache in the correct position and to infer
495
+ the complete sequence length.
496
+ """
497
+
498
+
499
+ @add_start_docstrings(
500
+ "The bare EuroBert Model outputting raw hidden-states without any specific head on top.",
501
+ EUROBERT_START_DOCSTRING,
502
+ )
503
+ class EuroBertModel(EuroBertPreTrainedModel):
504
+ """
505
+ Transformer encoder consisting of *config.num_hidden_layers* layers. Each layer is a [`EuroBertDecoderLayer`]
506
+
507
+ Args:
508
+ config: EuroBertConfig
509
+ """
510
+
511
+ def __init__(self, config: EuroBertConfig):
512
+ super().__init__(config)
513
+ self.padding_idx = config.pad_token_id
514
+ self.vocab_size = config.vocab_size
515
+
516
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
517
+ self.layers = nn.ModuleList(
518
+ [EuroBertDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
519
+ )
520
+ self.norm = EuroBertRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
521
+ self.rotary_emb = EuroBertRotaryEmbedding(config=config)
522
+ self.gradient_checkpointing = False
523
+ self.mask_converter = AttentionMaskConverter(is_causal=False)
524
+
525
+ # Initialize weights and apply final processing
526
+ self.post_init()
527
+
528
+ def get_input_embeddings(self):
529
+ return self.embed_tokens
530
+
531
+ def set_input_embeddings(self, value):
532
+ self.embed_tokens = value
533
+
534
+ @add_start_docstrings_to_model_forward(EUROBERT_INPUTS_DOCSTRING)
535
+ @add_code_sample_docstrings(
536
+ checkpoint=_CHECKPOINT_FOR_DOC,
537
+ output_type=BaseModelOutput,
538
+ config_class=_CONFIG_FOR_DOC,
539
+ )
540
+ def forward(
541
+ self,
542
+ input_ids: torch.LongTensor = None,
543
+ attention_mask: Optional[torch.Tensor] = None,
544
+ position_ids: Optional[torch.LongTensor] = None,
545
+ inputs_embeds: Optional[torch.FloatTensor] = None,
546
+ output_attentions: Optional[bool] = None,
547
+ output_hidden_states: Optional[bool] = None,
548
+ return_dict: Optional[bool] = None,
549
+ **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
550
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
551
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
552
+ output_hidden_states = (
553
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
554
+ )
555
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
556
+
557
+ if (input_ids is None) ^ (inputs_embeds is not None):
558
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
559
+
560
+ if inputs_embeds is None:
561
+ inputs_embeds = self.embed_tokens(input_ids)
562
+
563
+ if attention_mask is not None and self.config._attn_implementation != "flash_attention_2":
564
+ mask = self.mask_converter.to_4d(attention_mask, attention_mask.shape[1], inputs_embeds.dtype)
565
+ else:
566
+ mask = attention_mask
567
+
568
+ hidden_states = inputs_embeds
569
+
570
+ # create position embeddings to be shared across the encoder layers
571
+ if position_ids is None:
572
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0)
573
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
574
+
575
+ # encoder layers
576
+ all_hidden_states = () if output_hidden_states else None
577
+ all_self_attns = () if output_attentions else None
578
+
579
+ for encoder_layer in self.layers[: self.config.num_hidden_layers]:
580
+ if output_hidden_states:
581
+ all_hidden_states += (hidden_states,)
582
+
583
+ if self.gradient_checkpointing and self.training:
584
+ layer_outputs = self._gradient_checkpointing_func(
585
+ encoder_layer.__call__,
586
+ hidden_states,
587
+ mask,
588
+ position_ids,
589
+ None,
590
+ output_attentions,
591
+ False,
592
+ None,
593
+ position_embeddings,
594
+ )
595
+ else:
596
+ layer_outputs = encoder_layer(
597
+ hidden_states,
598
+ attention_mask=mask,
599
+ position_ids=position_ids,
600
+ output_attentions=output_attentions,
601
+ position_embeddings=position_embeddings,
602
+ **flash_attn_kwargs,
603
+ )
604
+
605
+ hidden_states = layer_outputs[0]
606
+
607
+ if output_attentions:
608
+ all_self_attns += (layer_outputs[1],)
609
+
610
+ hidden_states = self.norm(hidden_states)
611
+
612
+ # add hidden states from the last encoder layer
613
+ if output_hidden_states:
614
+ all_hidden_states += (hidden_states,)
615
+
616
+ output = BaseModelOutput(
617
+ last_hidden_state=hidden_states,
618
+ hidden_states=all_hidden_states,
619
+ attentions=all_self_attns,
620
+ )
621
+ return output if return_dict else output.to_tuple()
622
+
623
+ def _update_causal_mask(
624
+ self,
625
+ attention_mask: torch.Tensor,
626
+ input_tensor: torch.Tensor,
627
+ cache_position: torch.Tensor,
628
+ past_key_values: Cache,
629
+ output_attentions: bool,
630
+ ):
631
+ if self.config._attn_implementation == "flash_attention_2":
632
+ if attention_mask is not None and (attention_mask == 0.0).any():
633
+ return attention_mask
634
+ return None
635
+
636
+ # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
637
+ # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
638
+ # to infer the attention mask.
639
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
640
+ using_static_cache = isinstance(past_key_values, StaticCache)
641
+
642
+ # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
643
+ if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
644
+ if AttentionMaskConverter._ignore_causal_mask_sdpa(
645
+ attention_mask,
646
+ inputs_embeds=input_tensor,
647
+ past_key_values_length=past_seen_tokens,
648
+ is_training=self.training,
649
+ ):
650
+ return None
651
+
652
+ dtype, device = input_tensor.dtype, input_tensor.device
653
+ sequence_length = input_tensor.shape[1]
654
+ if using_static_cache:
655
+ target_length = past_key_values.get_max_cache_shape()
656
+ else:
657
+ target_length = (
658
+ attention_mask.shape[-1]
659
+ if isinstance(attention_mask, torch.Tensor)
660
+ else past_seen_tokens + sequence_length + 1
661
+ )
662
+
663
+ # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
664
+ causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(
665
+ attention_mask,
666
+ sequence_length=sequence_length,
667
+ target_length=target_length,
668
+ dtype=dtype,
669
+ device=device,
670
+ cache_position=cache_position,
671
+ batch_size=input_tensor.shape[0],
672
+ )
673
+
674
+ if (
675
+ self.config._attn_implementation == "sdpa"
676
+ and attention_mask is not None
677
+ and attention_mask.device.type in ["cuda", "xpu"]
678
+ and not output_attentions
679
+ ):
680
+ # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
681
+ # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
682
+ # Details: https://github.com/pytorch/pytorch/issues/110213
683
+ min_dtype = torch.finfo(dtype).min
684
+ causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
685
+
686
+ return causal_mask
687
+
688
+ @staticmethod
689
+ def _prepare_4d_causal_attention_mask_with_cache_position(
690
+ attention_mask: torch.Tensor,
691
+ sequence_length: int,
692
+ target_length: int,
693
+ dtype: torch.dtype,
694
+ device: torch.device,
695
+ cache_position: torch.Tensor,
696
+ batch_size: int,
697
+ **kwargs,
698
+ ):
699
+ """
700
+ Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
701
+ `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
702
+
703
+ Args:
704
+ attention_mask (`torch.Tensor`):
705
+ A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape
706
+ `(batch_size, 1, query_length, key_value_length)`.
707
+ sequence_length (`int`):
708
+ The sequence length being processed.
709
+ target_length (`int`):
710
+ The target length: when generating with static cache, the mask should be as long as the static cache,
711
+ to account for the 0 padding, the part of the cache that is not filled yet.
712
+ dtype (`torch.dtype`):
713
+ The dtype to use for the 4D attention mask.
714
+ device (`torch.device`):
715
+ The device to plcae the 4D attention mask on.
716
+ cache_position (`torch.Tensor`):
717
+ Indices depicting the position of the input sequence tokens in the sequence.
718
+ batch_size (`torch.Tensor`):
719
+ Batch size.
720
+ """
721
+ if attention_mask is not None and attention_mask.dim() == 4:
722
+ # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
723
+ causal_mask = attention_mask
724
+ else:
725
+ min_dtype = torch.finfo(dtype).min
726
+ causal_mask = torch.full(
727
+ (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
728
+ )
729
+ if sequence_length != 1:
730
+ causal_mask = torch.triu(causal_mask, diagonal=1)
731
+ causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
732
+ causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
733
+ if attention_mask is not None:
734
+ causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
735
+ mask_length = attention_mask.shape[-1]
736
+ padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
737
+ causal_mask.device
738
+ )
739
+ padding_mask = padding_mask == 0
740
+ causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
741
+ padding_mask, min_dtype
742
+ )
743
+
744
+ return causal_mask
745
+
746
+
747
+ @add_start_docstrings(
748
+ "The EuroBert Model with a decoder head on top that is used for masked language modeling.",
749
+ EUROBERT_START_DOCSTRING,
750
+ )
751
+ class EuroBertForMaskedLM(EuroBertPreTrainedModel):
752
+ def __init__(self, config: EuroBertConfig):
753
+ super().__init__(config)
754
+ self.model = EuroBertModel(config)
755
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, config.mlp_bias)
756
+ self.post_init()
757
+
758
+ @add_start_docstrings_to_model_forward(EUROBERT_INPUTS_DOCSTRING)
759
+ @add_code_sample_docstrings(
760
+ checkpoint=_CHECKPOINT_FOR_DOC,
761
+ output_type=BaseModelOutput,
762
+ config_class=_CONFIG_FOR_DOC,
763
+ )
764
+ def forward(
765
+ self,
766
+ input_ids: Optional[torch.LongTensor] = None,
767
+ attention_mask: Optional[torch.Tensor] = None,
768
+ position_ids: Optional[torch.LongTensor] = None,
769
+ inputs_embeds: Optional[torch.FloatTensor] = None,
770
+ labels: Optional[torch.LongTensor] = None,
771
+ output_attentions: Optional[bool] = None,
772
+ output_hidden_states: Optional[bool] = None,
773
+ return_dict: Optional[bool] = None,
774
+ ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
775
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
776
+
777
+ encoder_output = self.model(
778
+ input_ids,
779
+ attention_mask=attention_mask,
780
+ position_ids=position_ids,
781
+ inputs_embeds=inputs_embeds,
782
+ output_attentions=output_attentions,
783
+ output_hidden_states=output_hidden_states,
784
+ return_dict=return_dict,
785
+ )
786
+
787
+ prediction_scores = self.lm_head(encoder_output[0])
788
+ masked_lm_loss = None
789
+ if labels is not None:
790
+ labels = labels.to(prediction_scores.device)
791
+ masked_lm_loss = self.loss_function(prediction_scores, labels, vocab_size=self.config.vocab_size)
792
+
793
+ if not return_dict:
794
+ output = (prediction_scores,) + encoder_output[1:]
795
+ return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
796
+
797
+ return MaskedLMOutput(
798
+ loss=masked_lm_loss,
799
+ logits=prediction_scores,
800
+ hidden_states=encoder_output.hidden_states,
801
+ attentions=encoder_output.attentions,
802
+ )
803
+
804
+
805
+ @add_start_docstrings(
806
+ "The EuroBert Model with a sequence classification head on top that performs pooling.",
807
+ EUROBERT_START_DOCSTRING,
808
+ )
809
+ class EuroBertForSequenceClassification(EuroBertPreTrainedModel):
810
+ def __init__(self, config: EuroBertConfig):
811
+ super().__init__(config)
812
+ self.num_labels = config.num_labels
813
+ self.clf_pooling = config.clf_pooling
814
+
815
+ self.model = EuroBertModel(config)
816
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
817
+ self.activation = nn.GELU()
818
+ self.classifier = nn.Linear(config.hidden_size, self.num_labels)
819
+ self.post_init()
820
+
821
+ @add_start_docstrings_to_model_forward(EUROBERT_INPUTS_DOCSTRING)
822
+ @add_code_sample_docstrings(
823
+ checkpoint=_CHECKPOINT_FOR_DOC,
824
+ output_type=BaseModelOutput,
825
+ config_class=_CONFIG_FOR_DOC,
826
+ )
827
+ def forward(
828
+ self,
829
+ input_ids: Optional[torch.LongTensor] = None,
830
+ attention_mask: Optional[torch.Tensor] = None,
831
+ position_ids: Optional[torch.LongTensor] = None,
832
+ inputs_embeds: Optional[torch.FloatTensor] = None,
833
+ labels: Optional[torch.LongTensor] = None,
834
+ output_attentions: Optional[bool] = None,
835
+ output_hidden_states: Optional[bool] = None,
836
+ return_dict: Optional[bool] = None,
837
+ ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
838
+ r"""
839
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
840
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
841
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
842
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
843
+ """
844
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
845
+
846
+ encoder_output = self.model(
847
+ input_ids,
848
+ attention_mask=attention_mask,
849
+ position_ids=position_ids,
850
+ inputs_embeds=inputs_embeds,
851
+ output_attentions=output_attentions,
852
+ output_hidden_states=output_hidden_states,
853
+ return_dict=return_dict,
854
+ )
855
+ last_hidden_state = encoder_output[0]
856
+
857
+ if self.clf_pooling in ["bos", "mean"]:
858
+ if self.clf_pooling == "bos":
859
+ pooled_output = last_hidden_state[:, 0]
860
+
861
+ elif self.clf_pooling == "mean":
862
+ if attention_mask is None:
863
+ pooled_output = last_hidden_state.mean(dim=1)
864
+ else:
865
+ pooled_output = (last_hidden_state * attention_mask.unsqueeze(-1)).sum(dim=1)
866
+ pooled_output /= attention_mask.sum(dim=1, keepdim=True)
867
+
868
+ pooled_output = self.dense(pooled_output)
869
+ pooled_output = self.activation(pooled_output)
870
+ logits = self.classifier(pooled_output)
871
+
872
+ elif self.clf_pooling == "late":
873
+ x = self.dense(last_hidden_state)
874
+ x = self.activation(x)
875
+ logits = self.classifier(x)
876
+ if attention_mask is None:
877
+ logits = logits.mean(dim=1)
878
+ else:
879
+ logits = (logits * attention_mask.unsqueeze(-1)).sum(dim=1)
880
+ logits /= attention_mask.sum(dim=1, keepdim=True)
881
+
882
+ loss = None
883
+ if labels is not None:
884
+ labels = labels.to(logits.device)
885
+ if self.config.problem_type is None:
886
+ if self.num_labels == 1:
887
+ self.config.problem_type = "regression"
888
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
889
+ self.config.problem_type = "single_label_classification"
890
+ else:
891
+ self.config.problem_type = "multi_label_classification"
892
+
893
+ if self.config.problem_type == "regression":
894
+ loss_fct = MSELoss()
895
+ if self.num_labels == 1:
896
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
897
+ else:
898
+ loss = loss_fct(logits, labels)
899
+ elif self.config.problem_type == "single_label_classification":
900
+ loss_fct = CrossEntropyLoss()
901
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
902
+ elif self.config.problem_type == "multi_label_classification":
903
+ loss_fct = BCEWithLogitsLoss()
904
+ loss = loss_fct(logits, labels)
905
+
906
+ if not return_dict:
907
+ output = (logits,) + encoder_output[1:]
908
+ return ((loss,) + output) if loss is not None else output
909
+
910
+ return SequenceClassifierOutput(
911
+ loss=loss,
912
+ logits=logits,
913
+ hidden_states=encoder_output.hidden_states,
914
+ attentions=encoder_output.attentions,
915
+ )
916
+
917
+
918
+ @add_start_docstrings(
919
+ """
920
+ The EuroBert Model with a token classification head on top (a linear layer on top of the hidden-states
921
+ output) e.g. for Named-Entity-Recognition (NER) tasks."
922
+ """,
923
+ EUROBERT_START_DOCSTRING,
924
+ )
925
+ class EuroBertForTokenClassification(EuroBertPreTrainedModel):
926
+ def __init__(self, config: EuroBertConfig):
927
+ super().__init__(config)
928
+ self.num_labels = config.num_labels
929
+ self.model = EuroBertModel(config)
930
+
931
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
932
+ self.post_init()
933
+
934
+ def get_input_embeddings(self):
935
+ return self.model.embed_tokens
936
+
937
+ def set_input_embeddings(self, value):
938
+ self.model.embed_tokens = value
939
+
940
+ @add_start_docstrings_to_model_forward(EUROBERT_INPUTS_DOCSTRING)
941
+ def forward(
942
+ self,
943
+ input_ids: Optional[torch.LongTensor] = None,
944
+ attention_mask: Optional[torch.Tensor] = None,
945
+ position_ids: Optional[torch.LongTensor] = None,
946
+ inputs_embeds: Optional[torch.FloatTensor] = None,
947
+ labels: Optional[torch.LongTensor] = None,
948
+ use_cache: Optional[bool] = None,
949
+ output_attentions: Optional[bool] = None,
950
+ output_hidden_states: Optional[bool] = None,
951
+ return_dict: Optional[bool] = None,
952
+ ) -> Union[Tuple, TokenClassifierOutput]:
953
+ r"""
954
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
955
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
956
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
957
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
958
+ """
959
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
960
+
961
+ outputs = self.model(
962
+ input_ids,
963
+ attention_mask=attention_mask,
964
+ position_ids=position_ids,
965
+ inputs_embeds=inputs_embeds,
966
+ use_cache=use_cache,
967
+ output_attentions=output_attentions,
968
+ output_hidden_states=output_hidden_states,
969
+ return_dict=return_dict,
970
+ )
971
+ sequence_output = outputs[0]
972
+ logits = self.classifier(sequence_output)
973
+
974
+ loss = None
975
+ if labels is not None:
976
+ loss_fct = CrossEntropyLoss()
977
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
978
+
979
+ if not return_dict:
980
+ output = (logits,) + outputs[2:]
981
+ return ((loss,) + output) if loss is not None else output
982
+
983
+ return TokenClassifierOutput(
984
+ loss=loss,
985
+ logits=logits,
986
+ hidden_states=outputs.hidden_states,
987
+ attentions=outputs.attentions,
988
+ )
989
+
990
+
991
+ @add_start_docstrings(
992
+ """
993
+ The EuroBert Model with a span classification head on top for extractive question-answering tasks
994
+ like SQuAD (a linear layers on top of the hidden-states output to compute span start logits
995
+ and span end logits).
996
+ """,
997
+ EUROBERT_START_DOCSTRING,
998
+ )
999
+ class EuroBertForQuestionAnswering(EuroBertPreTrainedModel):
1000
+ def __init__(self, config: EuroBertConfig):
1001
+ super().__init__(config)
1002
+ self.num_labels = config.num_labels
1003
+ self.model = EuroBertModel(config)
1004
+
1005
+ self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
1006
+ self.post_init()
1007
+
1008
+ def get_input_embeddings(self):
1009
+ return self.model.embed_tokens
1010
+
1011
+ def set_input_embeddings(self, value):
1012
+ self.model.embed_tokens = value
1013
+
1014
+ @add_start_docstrings_to_model_forward(EUROBERT_INPUTS_DOCSTRING)
1015
+ def forward(
1016
+ self,
1017
+ input_ids: Optional[torch.Tensor] = None,
1018
+ attention_mask: Optional[torch.Tensor] = None,
1019
+ position_ids: Optional[torch.Tensor] = None,
1020
+ inputs_embeds: Optional[torch.Tensor] = None,
1021
+ use_cache: Optional[bool] = None,
1022
+ start_positions: Optional[torch.Tensor] = None,
1023
+ end_positions: Optional[torch.Tensor] = None,
1024
+ output_attentions: Optional[bool] = None,
1025
+ output_hidden_states: Optional[bool] = None,
1026
+ return_dict: Optional[bool] = None,
1027
+ ) -> Union[Tuple[torch.Tensor], QuestionAnsweringModelOutput]:
1028
+ r"""
1029
+ start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1030
+ Labels for position (index) of the start of the labelled span for computing the token classification loss.
1031
+ Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
1032
+ are not taken into account for computing the loss.
1033
+ end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1034
+ Labels for position (index) of the end of the labelled span for computing the token classification loss.
1035
+ Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
1036
+ are not taken into account for computing the loss.
1037
+ """
1038
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1039
+
1040
+ outputs = self.model(
1041
+ input_ids,
1042
+ attention_mask=attention_mask,
1043
+ position_ids=position_ids,
1044
+ inputs_embeds=inputs_embeds,
1045
+ use_cache=use_cache,
1046
+ output_attentions=output_attentions,
1047
+ output_hidden_states=output_hidden_states,
1048
+ return_dict=return_dict,
1049
+ )
1050
+ sequence_output = outputs[0]
1051
+
1052
+ logits = self.qa_outputs(sequence_output)
1053
+ start_logits, end_logits = logits.split(1, dim=-1)
1054
+ start_logits = start_logits.squeeze(-1).contiguous()
1055
+ end_logits = end_logits.squeeze(-1).contiguous()
1056
+
1057
+ total_loss = None
1058
+ if start_positions is not None and end_positions is not None:
1059
+ # If we are on multi-GPU, split add a dimension
1060
+ if len(start_positions.size()) > 1:
1061
+ start_positions = start_positions.squeeze(-1)
1062
+ if len(end_positions.size()) > 1:
1063
+ end_positions = end_positions.squeeze(-1)
1064
+ # sometimes the start/end positions are outside our model inputs, we ignore these terms
1065
+ ignored_index = start_logits.size(1)
1066
+ start_positions = start_positions.clamp(0, ignored_index)
1067
+ end_positions = end_positions.clamp(0, ignored_index)
1068
+
1069
+ loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
1070
+ start_loss = loss_fct(start_logits, start_positions)
1071
+ end_loss = loss_fct(end_logits, end_positions)
1072
+ total_loss = (start_loss + end_loss) / 2
1073
+
1074
+ if not return_dict:
1075
+ output = (start_logits, end_logits) + outputs[2:]
1076
+ return ((total_loss,) + output) if total_loss is not None else output
1077
+
1078
+ return QuestionAnsweringModelOutput(
1079
+ loss=total_loss,
1080
+ start_logits=start_logits,
1081
+ end_logits=end_logits,
1082
+ hidden_states=outputs.hidden_states,
1083
+ attentions=outputs.attentions,
1084
+ )
1085
+
1086
+
1087
+ __all__ = [
1088
+ "EuroBertPreTrainedModel",
1089
+ "EuroBertModel",
1090
+ "EuroBertForMaskedLM",
1091
+ "EuroBertForSequenceClassification",
1092
+ "EuroBertForTokenClassification",
1093
+ "EuroBertForQuestionAnswering",
1094
+ ]
modules.json ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [
2
+ {
3
+ "idx": 0,
4
+ "name": "0",
5
+ "path": "",
6
+ "type": "sentence_transformers.base.modules.transformer.Transformer"
7
+ },
8
+ {
9
+ "idx": 1,
10
+ "name": "1",
11
+ "path": "1_Pooling",
12
+ "type": "sentence_transformers.sentence_transformer.modules.pooling.Pooling"
13
+ },
14
+ {
15
+ "idx": 2,
16
+ "name": "2",
17
+ "path": "2_Normalize",
18
+ "type": "sentence_transformers.sentence_transformer.modules.normalize.Normalize"
19
+ }
20
+ ]
sentence_bert_config.json ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "transformer_task": "feature-extraction",
3
+ "modality_config": {
4
+ "text": {
5
+ "method": "forward",
6
+ "method_output_name": "last_hidden_state"
7
+ }
8
+ },
9
+ "module_output_name": "token_embeddings"
10
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "backend": "tokenizers",
3
+ "bos_token": "<|begin_of_text|>",
4
+ "clean_up_tokenization_spaces": true,
5
+ "eos_token": "<|end_of_text|>",
6
+ "is_local": true,
7
+ "mask_token": "<|mask|>",
8
+ "max_length": null,
9
+ "model_input_names": [
10
+ "input_ids",
11
+ "attention_mask"
12
+ ],
13
+ "model_max_length": 1000000000000000019884624838656,
14
+ "pad_to_multiple_of": null,
15
+ "pad_token": "<|pad|>",
16
+ "pad_token_type_id": 0,
17
+ "padding_side": "right",
18
+ "tokenizer_class": "TokenizersBackend"
19
+ }