BjarneNPO commited on
Commit
abc7987
·
verified ·
1 Parent(s): 1e0e681

Add new SentenceTransformer model

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
@@ -0,0 +1,535 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0
  <br>2. Kinder kommen nicht von nach
1
  <br>3. Absage kann bei Portalstatus nicht erstellt werden.</code> | <code>Ticket
2
  <br>Userin gebeten sich an den Support zu wenden, da der Fehler liegt.</code> |
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ tags:
3
+ - sentence-transformers
4
+ - sentence-similarity
5
+ - feature-extraction
6
+ - dense
7
+ - generated_from_trainer
8
+ - dataset_size:19964
9
+ - loss:MultipleNegativesRankingLoss
10
+ base_model: Snowflake/snowflake-arctic-embed-m-v2.0
11
+ widget:
12
+ - source_sentence: 'Kollegin hat Probleme mit dem Login zu '
13
+ sentences:
14
+ - Alle genannten Kinder gab es in kitaplus. Bei einem musste nur eine neue BI angelegt
15
+ werden, bei den anderen muss der Vertrag in einer anderen Kita rückgängig gemacht
16
+ werden, damit es in kitaplus in dieser Einrichtung aus der Liste der Absagen genommen
17
+ werden kann.
18
+ - Der Bereich ist aktuell noch nicht sichtbar.
19
+ - muss mit dem Rentamt geklärt werden
20
+ - source_sentence: Benutzer möchte einen Kollegen nur für die Dokumentenbibliothek
21
+ anlegen.
22
+ sentences:
23
+ - Rücksprache mit Entwickler.
24
+ - Sie muss den Regler auf Anzahl stellen
25
+ - Zusammen die Rolle gewählt und dort dann in den individuellen Rechten alles auf
26
+ lesend bzw. ausblenden gestellt, außer die Bibliothek.
27
+ - source_sentence: Ist es richtig so, dass Mitarbeiter, wenn sie nach einer gewissen
28
+ Zeit wieder in die Einrichtung kommen, erneut angelegt werden müssen?
29
+ sentences:
30
+ - Userin an den Träger verwiesen, dieser kann bei ihr ein neues Passwort setzen.
31
+ - Ja, das ist korrekt so.
32
+ - Userin muss erst rechts über das 3-Punkte-menü die "Anmeldedaten zusammenführen".
33
+ Danach muss man in den angelegten BI die Gruppenform des Anmeldeportals angeben.
34
+ - source_sentence: Userin kann die Öffnungszeiten der Einrichtung nicht bearbeiten.
35
+ sentences:
36
+ - informiert, dass es keinen Testzugang gibt, aber Handbücher und Hilfen in zur
37
+ Verfügung stehen, wenn die Schnittstelle eingerichtet wurde.
38
+ - Bereits bekannt, die Kollegen sind schon dabei den Fehler zu beheben.
39
+ - Userin darf dies mit der Rolle nicht.
40
+ - source_sentence: fragt wie der Stand zu dem aktuellen Problem ist
41
+ sentences:
42
+ - Userin muss sich an die Bistums IT wenden.
43
+ - In Klärung mit der Kollegin - Das Problem liegt leider an deren Betreiber. Die
44
+ sind aber informiert und arbeiten bereits daran
45
+ - findet diese in der Übersicht der Gruppen.
46
+ pipeline_tag: sentence-similarity
47
+ library_name: sentence-transformers
48
+ metrics:
49
+ - cosine_accuracy@1
50
+ - cosine_accuracy@3
51
+ - cosine_accuracy@5
52
+ - cosine_accuracy@10
53
+ - cosine_precision@1
54
+ - cosine_precision@3
55
+ - cosine_precision@5
56
+ - cosine_precision@10
57
+ - cosine_recall@1
58
+ - cosine_recall@3
59
+ - cosine_recall@5
60
+ - cosine_recall@10
61
+ - cosine_ndcg@10
62
+ - cosine_mrr@10
63
+ - cosine_map@100
64
+ model-index:
65
+ - name: SentenceTransformer based on Snowflake/snowflake-arctic-embed-m-v2.0
66
+ results:
67
+ - task:
68
+ type: information-retrieval
69
+ name: Information Retrieval
70
+ dataset:
71
+ name: Snowflake/snowflake arctic embed m v2.0
72
+ type: Snowflake/snowflake-arctic-embed-m-v2.0
73
+ metrics:
74
+ - type: cosine_accuracy@1
75
+ value: 0.3277310924369748
76
+ name: Cosine Accuracy@1
77
+ - type: cosine_accuracy@3
78
+ value: 0.48739495798319327
79
+ name: Cosine Accuracy@3
80
+ - type: cosine_accuracy@5
81
+ value: 0.5462184873949579
82
+ name: Cosine Accuracy@5
83
+ - type: cosine_accuracy@10
84
+ value: 0.6470588235294118
85
+ name: Cosine Accuracy@10
86
+ - type: cosine_precision@1
87
+ value: 0.3277310924369748
88
+ name: Cosine Precision@1
89
+ - type: cosine_precision@3
90
+ value: 0.21848739495798317
91
+ name: Cosine Precision@3
92
+ - type: cosine_precision@5
93
+ value: 0.2302521008403361
94
+ name: Cosine Precision@5
95
+ - type: cosine_precision@10
96
+ value: 0.21428571428571433
97
+ name: Cosine Precision@10
98
+ - type: cosine_recall@1
99
+ value: 0.006002400960384153
100
+ name: Cosine Recall@1
101
+ - type: cosine_recall@3
102
+ value: 0.02267573696145124
103
+ name: Cosine Recall@3
104
+ - type: cosine_recall@5
105
+ value: 0.03478057889822595
106
+ name: Cosine Recall@5
107
+ - type: cosine_recall@10
108
+ value: 0.07412965186074429
109
+ name: Cosine Recall@10
110
+ - type: cosine_ndcg@10
111
+ value: 0.2371544724261345
112
+ name: Cosine Ndcg@10
113
+ - type: cosine_mrr@10
114
+ value: 0.4209383753501402
115
+ name: Cosine Mrr@10
116
+ - type: cosine_map@100
117
+ value: 0.08885586509032745
118
+ name: Cosine Map@100
119
+ ---
120
+
121
+ # SentenceTransformer based on Snowflake/snowflake-arctic-embed-m-v2.0
122
+
123
+ This is a [sentence-transformers](https://www.SBERT.net) model finetuned from [Snowflake/snowflake-arctic-embed-m-v2.0](https://huggingface.co/Snowflake/snowflake-arctic-embed-m-v2.0) on the train dataset. It maps sentences & paragraphs to a 768-dimensional dense vector space and can be used for semantic textual similarity, semantic search, paraphrase mining, text classification, clustering, and more.
124
+
125
+ ## Model Details
126
+
127
+ ### Model Description
128
+ - **Model Type:** Sentence Transformer
129
+ - **Base model:** [Snowflake/snowflake-arctic-embed-m-v2.0](https://huggingface.co/Snowflake/snowflake-arctic-embed-m-v2.0) <!-- at revision 95c2741480856aa9666782eb4afe11959938017f -->
130
+ - **Maximum Sequence Length:** 8192 tokens
131
+ - **Output Dimensionality:** 768 dimensions
132
+ - **Similarity Function:** Cosine Similarity
133
+ - **Training Dataset:**
134
+ - train
135
+ <!-- - **Language:** Unknown -->
136
+ <!-- - **License:** Unknown -->
137
+
138
+ ### Model Sources
139
+
140
+ - **Documentation:** [Sentence Transformers Documentation](https://sbert.net)
141
+ - **Repository:** [Sentence Transformers on GitHub](https://github.com/UKPLab/sentence-transformers)
142
+ - **Hugging Face:** [Sentence Transformers on Hugging Face](https://huggingface.co/models?library=sentence-transformers)
143
+
144
+ ### Full Model Architecture
145
+
146
+ ```
147
+ SentenceTransformer(
148
+ (0): Transformer({'max_seq_length': 8192, 'do_lower_case': False, 'architecture': 'GteModel'})
149
+ (1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_cls_token': True, 'pooling_mode_mean_tokens': False, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
150
+ (2): Normalize()
151
+ )
152
+ ```
153
+
154
+ ## Usage
155
+
156
+ ### Direct Usage (Sentence Transformers)
157
+
158
+ First install the Sentence Transformers library:
159
+
160
+ ```bash
161
+ pip install -U sentence-transformers
162
+ ```
163
+
164
+ Then you can load this model and run inference.
165
+ ```python
166
+ from sentence_transformers import SentenceTransformer
167
+
168
+ # Download from the 🤗 Hub
169
+ model = SentenceTransformer("BjarneNPO/finetune_21_08_2025_11_16_02")
170
+ # Run inference
171
+ queries = [
172
+ "fragt wie der Stand zu dem aktuellen Problem ist",
173
+ ]
174
+ documents = [
175
+ 'In Klärung mit der Kollegin - Das Problem liegt leider an deren Betreiber. Die sind aber informiert und arbeiten bereits daran',
176
+ 'findet diese in der Übersicht der Gruppen.',
177
+ 'Userin muss sich an die Bistums IT wenden.',
178
+ ]
179
+ query_embeddings = model.encode_query(queries)
180
+ document_embeddings = model.encode_document(documents)
181
+ print(query_embeddings.shape, document_embeddings.shape)
182
+ # [1, 768] [3, 768]
183
+
184
+ # Get the similarity scores for the embeddings
185
+ similarities = model.similarity(query_embeddings, document_embeddings)
186
+ print(similarities)
187
+ # tensor([[0.2668, 0.0872, 0.0995]])
188
+ ```
189
+
190
+ <!--
191
+ ### Direct Usage (Transformers)
192
+
193
+ <details><summary>Click to see the direct usage in Transformers</summary>
194
+
195
+ </details>
196
+ -->
197
+
198
+ <!--
199
+ ### Downstream Usage (Sentence Transformers)
200
+
201
+ You can finetune this model on your own dataset.
202
+
203
+ <details><summary>Click to expand</summary>
204
+
205
+ </details>
206
+ -->
207
+
208
+ <!--
209
+ ### Out-of-Scope Use
210
+
211
+ *List how the model may foreseeably be misused and address what users ought not to do with the model.*
212
+ -->
213
+
214
+ ## Evaluation
215
+
216
+ ### Metrics
217
+
218
+ #### Information Retrieval
219
+
220
+ * Dataset: `Snowflake/snowflake-arctic-embed-m-v2.0`
221
+ * Evaluated with <code>scripts.InformationRetrievalEvaluatorCustom.InformationRetrievalEvaluatorCustom</code> with these parameters:
222
+ ```json
223
+ {
224
+ "query_prompt_name": "query",
225
+ "corpus_prompt_name": "query"
226
+ }
227
+ ```
228
+
229
+ | Metric | Value |
230
+ |:--------------------|:-----------|
231
+ | cosine_accuracy@1 | 0.3277 |
232
+ | cosine_accuracy@3 | 0.4874 |
233
+ | cosine_accuracy@5 | 0.5462 |
234
+ | cosine_accuracy@10 | 0.6471 |
235
+ | cosine_precision@1 | 0.3277 |
236
+ | cosine_precision@3 | 0.2185 |
237
+ | cosine_precision@5 | 0.2303 |
238
+ | cosine_precision@10 | 0.2143 |
239
+ | cosine_recall@1 | 0.006 |
240
+ | cosine_recall@3 | 0.0227 |
241
+ | cosine_recall@5 | 0.0348 |
242
+ | cosine_recall@10 | 0.0741 |
243
+ | **cosine_ndcg@10** | **0.2372** |
244
+ | cosine_mrr@10 | 0.4209 |
245
+ | cosine_map@100 | 0.0889 |
246
+
247
+ <!--
248
+ ## Bias, Risks and Limitations
249
+
250
+ *What are the known or foreseeable issues stemming from this model? You could also flag here known failure cases or weaknesses of the model.*
251
+ -->
252
+
253
+ <!--
254
+ ### Recommendations
255
+
256
+ *What are recommendations with respect to the foreseeable issues? For example, filtering explicit content.*
257
+ -->
258
+
259
+ ## Training Details
260
+
261
+ ### Training Dataset
262
+
263
+ #### train
264
+
265
+ * Dataset: train
266
+ * Size: 19,964 training samples
267
+ * Columns: <code>query</code> and <code>answer</code>
268
+ * Approximate statistics based on the first 1000 samples:
269
+ | | query | answer |
270
+ |:--------|:-----------------------------------------------------------------------------------|:-----------------------------------------------------------------------------------|
271
+ | type | string | string |
272
+ | details | <ul><li>min: 4 tokens</li><li>mean: 27.77 tokens</li><li>max: 615 tokens</li></ul> | <ul><li>min: 3 tokens</li><li>mean: 22.87 tokens</li><li>max: 151 tokens</li></ul> |
273
+ * Samples:
274
+ | query | answer |
275
+ |:------------------------------------------------------------------------------------------------------------------------------------------------------------|:--------------------------------------------------------------------------------------------|
276
+ | <code>Wie kann man die Jahresurlaubsübersicht exportieren?</code> | <code>über das 3 Punkte Menü rechts oben. Mitarbeiter auswählen und exportieren</code> |
277
+ | <code>1. Vertragsabschlüsse werden nicht übertragen
278
  <br>2. Kinder kommen nicht von nach
279
  <br>3. Absage kann bei Portalstatus nicht erstellt werden.</code> | <code>Ticket
280
  <br>Userin gebeten sich an den Support zu wenden, da der Fehler liegt.</code> |
281
+ | <code>Wird im Anmeldeportal nicht gefunden.</code> | <code>Die Schnittstelle war noch nicht aktiviert und Profil ebenfalls nicht.</code> |
282
+ * Loss: [<code>MultipleNegativesRankingLoss</code>](https://sbert.net/docs/package_reference/sentence_transformer/losses.html#multiplenegativesrankingloss) with these parameters:
283
+ ```json
284
+ {
285
+ "scale": 20.0,
286
+ "similarity_fct": "cos_sim",
287
+ "gather_across_devices": false
288
+ }
289
+ ```
290
+
291
+ ### Training Hyperparameters
292
+ #### Non-Default Hyperparameters
293
+
294
+ - `eval_strategy`: epoch
295
+ - `per_device_train_batch_size`: 16
296
+ - `per_device_eval_batch_size`: 16
297
+ - `gradient_accumulation_steps`: 8
298
+ - `learning_rate`: 2e-05
299
+ - `lr_scheduler_type`: cosine
300
+ - `warmup_ratio`: 0.1
301
+ - `bf16`: True
302
+ - `tf32`: False
303
+ - `load_best_model_at_end`: True
304
+ - `batch_sampler`: no_duplicates
305
+
306
+ #### All Hyperparameters
307
+ <details><summary>Click to expand</summary>
308
+
309
+ - `overwrite_output_dir`: False
310
+ - `do_predict`: False
311
+ - `eval_strategy`: epoch
312
+ - `prediction_loss_only`: True
313
+ - `per_device_train_batch_size`: 16
314
+ - `per_device_eval_batch_size`: 16
315
+ - `per_gpu_train_batch_size`: None
316
+ - `per_gpu_eval_batch_size`: None
317
+ - `gradient_accumulation_steps`: 8
318
+ - `eval_accumulation_steps`: None
319
+ - `torch_empty_cache_steps`: None
320
+ - `learning_rate`: 2e-05
321
+ - `weight_decay`: 0.0
322
+ - `adam_beta1`: 0.9
323
+ - `adam_beta2`: 0.999
324
+ - `adam_epsilon`: 1e-08
325
+ - `max_grad_norm`: 1.0
326
+ - `num_train_epochs`: 3
327
+ - `max_steps`: -1
328
+ - `lr_scheduler_type`: cosine
329
+ - `lr_scheduler_kwargs`: {}
330
+ - `warmup_ratio`: 0.1
331
+ - `warmup_steps`: 0
332
+ - `log_level`: passive
333
+ - `log_level_replica`: warning
334
+ - `log_on_each_node`: True
335
+ - `logging_nan_inf_filter`: True
336
+ - `save_safetensors`: True
337
+ - `save_on_each_node`: False
338
+ - `save_only_model`: False
339
+ - `restore_callback_states_from_checkpoint`: False
340
+ - `no_cuda`: False
341
+ - `use_cpu`: False
342
+ - `use_mps_device`: False
343
+ - `seed`: 42
344
+ - `data_seed`: None
345
+ - `jit_mode_eval`: False
346
+ - `use_ipex`: False
347
+ - `bf16`: True
348
+ - `fp16`: False
349
+ - `fp16_opt_level`: O1
350
+ - `half_precision_backend`: auto
351
+ - `bf16_full_eval`: False
352
+ - `fp16_full_eval`: False
353
+ - `tf32`: False
354
+ - `local_rank`: 0
355
+ - `ddp_backend`: None
356
+ - `tpu_num_cores`: None
357
+ - `tpu_metrics_debug`: False
358
+ - `debug`: []
359
+ - `dataloader_drop_last`: False
360
+ - `dataloader_num_workers`: 0
361
+ - `dataloader_prefetch_factor`: None
362
+ - `past_index`: -1
363
+ - `disable_tqdm`: False
364
+ - `remove_unused_columns`: True
365
+ - `label_names`: None
366
+ - `load_best_model_at_end`: True
367
+ - `ignore_data_skip`: False
368
+ - `fsdp`: []
369
+ - `fsdp_min_num_params`: 0
370
+ - `fsdp_config`: {'min_num_params': 0, 'xla': False, 'xla_fsdp_v2': False, 'xla_fsdp_grad_ckpt': False}
371
+ - `fsdp_transformer_layer_cls_to_wrap`: None
372
+ - `accelerator_config`: {'split_batches': False, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True, 'non_blocking': False, 'gradient_accumulation_kwargs': None}
373
+ - `deepspeed`: None
374
+ - `label_smoothing_factor`: 0.0
375
+ - `optim`: adamw_torch_fused
376
+ - `optim_args`: None
377
+ - `adafactor`: False
378
+ - `group_by_length`: False
379
+ - `length_column_name`: length
380
+ - `ddp_find_unused_parameters`: None
381
+ - `ddp_bucket_cap_mb`: None
382
+ - `ddp_broadcast_buffers`: False
383
+ - `dataloader_pin_memory`: True
384
+ - `dataloader_persistent_workers`: False
385
+ - `skip_memory_metrics`: True
386
+ - `use_legacy_prediction_loop`: False
387
+ - `push_to_hub`: False
388
+ - `resume_from_checkpoint`: None
389
+ - `hub_model_id`: None
390
+ - `hub_strategy`: every_save
391
+ - `hub_private_repo`: None
392
+ - `hub_always_push`: False
393
+ - `hub_revision`: None
394
+ - `gradient_checkpointing`: False
395
+ - `gradient_checkpointing_kwargs`: None
396
+ - `include_inputs_for_metrics`: False
397
+ - `include_for_metrics`: []
398
+ - `eval_do_concat_batches`: True
399
+ - `fp16_backend`: auto
400
+ - `push_to_hub_model_id`: None
401
+ - `push_to_hub_organization`: None
402
+ - `mp_parameters`:
403
+ - `auto_find_batch_size`: False
404
+ - `full_determinism`: False
405
+ - `torchdynamo`: None
406
+ - `ray_scope`: last
407
+ - `ddp_timeout`: 1800
408
+ - `torch_compile`: False
409
+ - `torch_compile_backend`: None
410
+ - `torch_compile_mode`: None
411
+ - `include_tokens_per_second`: False
412
+ - `include_num_input_tokens_seen`: False
413
+ - `neftune_noise_alpha`: None
414
+ - `optim_target_modules`: None
415
+ - `batch_eval_metrics`: False
416
+ - `eval_on_start`: False
417
+ - `use_liger_kernel`: False
418
+ - `liger_kernel_config`: None
419
+ - `eval_use_gather_object`: False
420
+ - `average_tokens_across_devices`: False
421
+ - `prompts`: None
422
+ - `batch_sampler`: no_duplicates
423
+ - `multi_dataset_batch_sampler`: proportional
424
+ - `router_mapping`: {}
425
+ - `learning_rate_mapping`: {}
426
+
427
+ </details>
428
+
429
+ ### Training Logs
430
+ | Epoch | Step | Training Loss | Snowflake/snowflake-arctic-embed-m-v2.0_cosine_ndcg@10 |
431
+ |:-------:|:-------:|:-------------:|:------------------------------------------------------:|
432
+ | 0.0641 | 10 | 2.1446 | - |
433
+ | 0.1282 | 20 | 2.1454 | - |
434
+ | 0.1923 | 30 | 1.8973 | - |
435
+ | 0.2564 | 40 | 1.7238 | - |
436
+ | 0.3205 | 50 | 1.7305 | - |
437
+ | 0.3846 | 60 | 1.5496 | - |
438
+ | 0.4487 | 70 | 1.4557 | - |
439
+ | 0.5128 | 80 | 1.406 | - |
440
+ | 0.5769 | 90 | 1.3067 | - |
441
+ | 0.6410 | 100 | 1.2727 | - |
442
+ | 0.7051 | 110 | 1.215 | - |
443
+ | 0.7692 | 120 | 1.1902 | - |
444
+ | 0.8333 | 130 | 1.218 | - |
445
+ | 0.8974 | 140 | 1.1271 | - |
446
+ | 0.9615 | 150 | 1.0909 | - |
447
+ | **1.0** | **156** | **-** | **0.2589** |
448
+ | 1.0256 | 160 | 1.0408 | - |
449
+ | 1.0897 | 170 | 1.0112 | - |
450
+ | 1.1538 | 180 | 1.0683 | - |
451
+ | 1.2179 | 190 | 0.9405 | - |
452
+ | 1.2821 | 200 | 0.933 | - |
453
+ | 1.3462 | 210 | 0.9533 | - |
454
+ | 1.4103 | 220 | 0.9144 | - |
455
+ | 1.4744 | 230 | 0.8618 | - |
456
+ | 1.5385 | 240 | 0.8624 | - |
457
+ | 1.6026 | 250 | 0.8649 | - |
458
+ | 1.6667 | 260 | 0.8646 | - |
459
+ | 1.7308 | 270 | 0.8307 | - |
460
+ | 1.7949 | 280 | 0.8522 | - |
461
+ | 1.8590 | 290 | 0.8566 | - |
462
+ | 1.9231 | 300 | 0.8389 | - |
463
+ | 1.9872 | 310 | 0.806 | - |
464
+ | 2.0 | 312 | - | 0.2421 |
465
+ | 2.0513 | 320 | 0.7134 | - |
466
+ | 2.1154 | 330 | 0.7545 | - |
467
+ | 2.1795 | 340 | 0.8033 | - |
468
+ | 2.2436 | 350 | 0.7402 | - |
469
+ | 2.3077 | 360 | 0.7876 | - |
470
+ | 2.3718 | 370 | 0.7185 | - |
471
+ | 2.4359 | 380 | 0.7391 | - |
472
+ | 2.5 | 390 | 0.7362 | - |
473
+ | 2.5641 | 400 | 0.7276 | - |
474
+ | 2.6282 | 410 | 0.6953 | - |
475
+ | 2.6923 | 420 | 0.6853 | - |
476
+ | 2.7564 | 430 | 0.7337 | - |
477
+ | 2.8205 | 440 | 0.7704 | - |
478
+ | 2.8846 | 450 | 0.6922 | - |
479
+ | 2.9487 | 460 | 0.7153 | - |
480
+ | 3.0 | 468 | - | 0.2372 |
481
+
482
+ * The bold row denotes the saved checkpoint.
483
+
484
+ ### Framework Versions
485
+ - Python: 3.10.11
486
+ - Sentence Transformers: 5.1.0
487
+ - Transformers: 4.55.2
488
+ - PyTorch: 2.8.0+cu129
489
+ - Accelerate: 1.10.0
490
+ - Datasets: 3.6.0
491
+ - Tokenizers: 0.21.4
492
+
493
+ ## Citation
494
+
495
+ ### BibTeX
496
+
497
+ #### Sentence Transformers
498
+ ```bibtex
499
+ @inproceedings{reimers-2019-sentence-bert,
500
+ title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
501
+ author = "Reimers, Nils and Gurevych, Iryna",
502
+ booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
503
+ month = "11",
504
+ year = "2019",
505
+ publisher = "Association for Computational Linguistics",
506
+ url = "https://arxiv.org/abs/1908.10084",
507
+ }
508
+ ```
509
+
510
+ #### MultipleNegativesRankingLoss
511
+ ```bibtex
512
+ @misc{henderson2017efficient,
513
+ title={Efficient Natural Language Response Suggestion for Smart Reply},
514
+ author={Matthew Henderson and Rami Al-Rfou and Brian Strope and Yun-hsuan Sung and Laszlo Lukacs and Ruiqi Guo and Sanjiv Kumar and Balint Miklos and Ray Kurzweil},
515
+ year={2017},
516
+ eprint={1705.00652},
517
+ archivePrefix={arXiv},
518
+ primaryClass={cs.CL}
519
+ }
520
+ ```
521
+
522
+ <!--
523
+ ## Glossary
524
+
525
+ *Clearly define terms in order to be accessible across audiences.*
526
+ -->
527
+
528
+ <!--
529
+ ## Model Card Authors
530
+
531
+ *Lists the people who create the model card, providing recognition and accountability for the detailed work that goes into its construction.*
532
+ -->
533
+
534
+ <!--
535
+ ## Model Card Contact
536
+
537
+ *Provides a way for people who have updates to the Model Card, suggestions, or questions, to contact the Model Card authors.*
538
+ -->
config.json ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ "matryoshka_dimensions": [
21
+ 256
22
+ ],
23
+ "max_position_embeddings": 8192,
24
+ "model_type": "gte",
25
+ "num_attention_heads": 12,
26
+ "num_hidden_layers": 12,
27
+ "pack_qkv": true,
28
+ "pad_token_id": 1,
29
+ "position_embedding_type": "rope",
30
+ "rope_scaling": null,
31
+ "rope_theta": 160000,
32
+ "torch_dtype": "float32",
33
+ "transformers_version": "4.55.2",
34
+ "type_vocab_size": 1,
35
+ "unpad_inputs": "true",
36
+ "use_memory_efficient_attention": "true",
37
+ "vocab_size": 250048
38
+ }
config_sentence_transformers.json ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "__version__": {
3
+ "sentence_transformers": "5.1.0",
4
+ "transformers": "4.55.2",
5
+ "pytorch": "2.8.0+cu129"
6
+ },
7
+ "prompts": {
8
+ "query": "query: ",
9
+ "document": ""
10
+ },
11
+ "default_prompt_name": null,
12
+ "model_type": "SentenceTransformer",
13
+ "similarity_fn_name": "cosine"
14
+ }
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:f3322d2653bf62dd9324b7c0d644fe567a1b68493401e201362927295fab0543
3
+ size 1221487872
modeling_hf_alibaba_nlp_gte.py ADDED
@@ -0,0 +1,967 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ try:
39
+ import xformers.ops as xops
40
+ except ImportError as e:
41
+ xops = None
42
+
43
+ from .configuration_hf_alibaba_nlp_gte import GteConfig
44
+
45
+
46
+ logger = logging.get_logger(__name__)
47
+
48
+
49
+ # Adapted from https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/bert_padding.py
50
+ # Which was adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py
51
+ class IndexFirstAxis(torch.autograd.Function):
52
+ @staticmethod
53
+ def forward(ctx, input, indices):
54
+ ctx.save_for_backward(indices)
55
+ assert input.ndim >= 2
56
+ ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
57
+ second_dim = other_shape.numel()
58
+ # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
59
+ # return input[indices]
60
+ # return torch.gather(
61
+ # rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim)
62
+ # ).reshape(-1, *other_shape)
63
+ return torch.gather(
64
+ input.view(ctx.first_axis_dim, second_dim),
65
+ 0,
66
+ indices.unsqueeze(-1).expand(indices.size(0), second_dim)
67
+ ).reshape(-1, *other_shape)
68
+
69
+ @staticmethod
70
+ def backward(ctx, grad_output):
71
+ (indices,) = ctx.saved_tensors
72
+ assert grad_output.ndim >= 2
73
+ other_shape = grad_output.shape[1:]
74
+ # grad_output = rearrange(grad_output, "b ... -> b (...)")
75
+ grad_output = grad_output.view(grad_output.size(0), other_shape.numel())
76
+ grad_input = torch.zeros(
77
+ [ctx.first_axis_dim, grad_output.shape[1]],
78
+ device=grad_output.device,
79
+ dtype=grad_output.dtype,
80
+ )
81
+ # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
82
+ # grad_input[indices] = grad_output
83
+ # grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output)
84
+ grad_input.scatter_(
85
+ 0, indices.unsqueeze(-1).expand(indices.size(0), grad_output.size(1)), grad_output
86
+ )
87
+ return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
88
+
89
+
90
+ index_first_axis = IndexFirstAxis.apply
91
+
92
+
93
+ def unpad_input(hidden_states, attention_mask=None, indices=None):
94
+ """
95
+ Arguments:
96
+ hidden_states: (batch, seqlen, ...)
97
+ attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
98
+ indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence.
99
+ Return:
100
+ hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
101
+ """
102
+ if indices is None:
103
+ assert attention_mask is not None
104
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
105
+
106
+ # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the
107
+ # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim
108
+ # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to
109
+ # index with integer indices. Moreover, torch's index is a bit slower than it needs to be,
110
+ # so we write custom forward and backward to make it a bit faster.
111
+ hidden_states = hidden_states.view(-1, *hidden_states.shape[2:])
112
+ return index_first_axis(hidden_states, indices)
113
+
114
+
115
+ class IndexPutFirstAxis(torch.autograd.Function):
116
+ @staticmethod
117
+ def forward(
118
+ ctx,
119
+ values: torch.Tensor,
120
+ indices: torch.Tensor,
121
+ first_axis_dim
122
+ ) -> torch.Tensor:
123
+ ctx.save_for_backward(indices)
124
+ assert indices.ndim == 1
125
+ assert values.ndim >= 2
126
+ output = torch.zeros(
127
+ first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype
128
+ )
129
+ output[indices] = values
130
+ return output
131
+
132
+ @staticmethod
133
+ def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None, None]:
134
+ indices, = ctx.saved_tensors
135
+ grad_values = grad_output[indices]
136
+ return grad_values, None, None
137
+
138
+
139
+ index_put_first_axis = IndexPutFirstAxis.apply
140
+
141
+
142
+ def pad_input(inputs: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int) -> torch.Tensor:
143
+ """Add padding to sequences.
144
+
145
+ Arguments:
146
+ inputs: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
147
+ indices: (total_nnz), `indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()`
148
+ batch: int batch_size
149
+ seqlen: int max sequence length
150
+
151
+ Returns:
152
+ inputs: (batch, seqlen, ...)
153
+ """
154
+ output = index_put_first_axis(inputs, indices, batch * seqlen)
155
+ return output.view(batch, seqlen, *inputs.shape[1:])
156
+
157
+
158
+ def rotate_half(x):
159
+ """Rotates half the hidden dims of the input."""
160
+ x1 = x[..., : x.shape[-1] // 2]
161
+ x2 = x[..., x.shape[-1] // 2 :]
162
+ return torch.cat((-x2, x1), dim=-1)
163
+
164
+
165
+ def apply_rotary_pos_emb(q, k, cos, sin):
166
+ """Applies Rotary Position Embedding to the query and key tensors.
167
+
168
+ Args:
169
+ q (`torch.Tensor`): The query tensor.
170
+ k (`torch.Tensor`): The key tensor.
171
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
172
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
173
+ Returns:
174
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
175
+ """
176
+ cos, sin = cos.to(q.dtype), sin.to(q.dtype)
177
+ q_embed = (q * cos) + (rotate_half(q) * sin)
178
+ k_embed = (k * cos) + (rotate_half(k) * sin)
179
+ return q_embed, k_embed
180
+
181
+
182
+ class RotaryEmbedding(torch.nn.Module):
183
+ def __init__(self, dim, max_position_embeddings=512, base=10000.0, device=None):
184
+ super().__init__()
185
+
186
+ self.dim = dim
187
+ self.max_position_embeddings = max_position_embeddings
188
+ self.base = base
189
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
190
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
191
+
192
+ # Build here to make `torch.jit.trace` work.
193
+ self._set_cos_sin_cache(
194
+ seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
195
+ )
196
+
197
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
198
+ self.max_seq_len_cached = seq_len
199
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32)
200
+
201
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
202
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
203
+ emb = torch.cat((freqs, freqs), dim=-1)
204
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
205
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
206
+
207
+ def forward(self, x, seq_len=None):
208
+ # x: [bs, num_attention_heads, seq_len, head_size]
209
+ if seq_len > self.max_seq_len_cached:
210
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
211
+
212
+ return (
213
+ self.cos_cached[:seq_len, ...].to(dtype=x.dtype),
214
+ self.sin_cached[:seq_len, ...].to(dtype=x.dtype),
215
+ )
216
+
217
+
218
+ class NTKScalingRotaryEmbedding(RotaryEmbedding):
219
+ """RotaryEmbedding extended with fixed and mixed NTK scaling. https://kexue.fm/archives/9706 """
220
+
221
+ def __init__(self, dim, max_position_embeddings=512, base=10000, device=None, scaling_factor=1.0, mixed_b=None):
222
+ self.scaling_factor = scaling_factor
223
+ self.mixed_b = mixed_b
224
+ super().__init__(dim, max_position_embeddings, base, device)
225
+ max_position_embeddings = max_position_embeddings * self.scaling_factor
226
+ self._set_cos_sin_cache(max_position_embeddings, self.inv_freq.device, torch.get_default_dtype())
227
+
228
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
229
+ self.max_seq_len_cached = seq_len
230
+
231
+ if seq_len > self.max_position_embeddings:
232
+ base = self.base * (self.scaling_factor if self.mixed_b is None else 1)
233
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
234
+
235
+ if self.mixed_b is None:
236
+ inv_freq = inv_freq / self.scaling_factor ** (2 / self.dim) # (6)
237
+ else:
238
+ a = torch.tensor(self.scaling_factor).log() / (self.dim / 2) ** self.mixed_b # (13)
239
+ lambda_1_m = (a * torch.arange(1, self.dim // 2 + 1).float().to(device) ** self.mixed_b).exp() # (12)
240
+ inv_freq = inv_freq / lambda_1_m # (10)
241
+
242
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
243
+
244
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32)
245
+
246
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
247
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
248
+ emb = torch.cat((freqs, freqs), dim=-1)
249
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
250
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
251
+
252
+
253
+ class RMSNorm(nn.Module):
254
+ def __init__(self, hidden_size, eps=1e-6):
255
+ """
256
+ RMSNorm is equivalent to T5LayerNorm
257
+ """
258
+ super().__init__()
259
+ self.weight = nn.Parameter(torch.ones(hidden_size))
260
+ self.variance_epsilon = eps
261
+
262
+ def forward(self, hidden_states):
263
+ input_dtype = hidden_states.dtype
264
+ hidden_states = hidden_states.to(torch.float32)
265
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
266
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
267
+ return self.weight * hidden_states.to(input_dtype)
268
+
269
+
270
+ LAYER_NORM = {
271
+ 'layer_norm': nn.LayerNorm,
272
+ 'rms_norm': RMSNorm
273
+ }
274
+
275
+
276
+ class GteEmbeddings(nn.Module):
277
+ """
278
+ Embedding and Unpadding.
279
+ """
280
+
281
+ def __init__(self, config: GteConfig):
282
+ super().__init__()
283
+ self.padding_idx = config.pad_token_id
284
+ self.word_embeddings = nn.Embedding(
285
+ config.vocab_size, config.hidden_size, padding_idx=self.padding_idx
286
+ )
287
+
288
+ self.position_embedding_type = config.position_embedding_type
289
+ if self.position_embedding_type == 'absolute':
290
+ self.position_embeddings = nn.Embedding(
291
+ config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
292
+ )
293
+ elif self.position_embedding_type == 'rope':
294
+ self._init_rope(config)
295
+ else:
296
+ raise ValueError
297
+
298
+ self.type_vocab_size = config.type_vocab_size
299
+ if self.type_vocab_size > 0:
300
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
301
+
302
+ # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
303
+ # any TensorFlow checkpoint file
304
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
305
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
306
+ # position_ids is contiguous in memory and excluded when serialized
307
+ self.register_buffer(
308
+ "position_ids", torch.arange(config.max_position_embeddings), persistent=False
309
+ )
310
+
311
+ def _init_rope(self, config):
312
+ kwargs = dict(
313
+ dim=int(config.hidden_size / config.num_attention_heads),
314
+ max_position_embeddings=config.max_position_embeddings,
315
+ base=config.rope_theta
316
+ )
317
+ if config.rope_scaling is None:
318
+ self.rotary_emb = RotaryEmbedding(**kwargs)
319
+ else:
320
+ kwargs.update(scaling_factor=config.rope_scaling["factor"])
321
+ scaling_type = config.rope_scaling["type"]
322
+ if scaling_type == 'ntk':
323
+ kwargs.update(mixed_b=config.rope_scaling.get('mixed_b', None))
324
+ self.rotary_emb = NTKScalingRotaryEmbedding(**kwargs)
325
+ # elif scaling_type == "linear":
326
+ # self.rotary_emb = LinearScalingRotaryEmbedding(**kwargs)
327
+ # elif scaling_type == "dynamic":
328
+ # self.rotary_emb = DynamicNTKScalingRotaryEmbedding(**kwargs)
329
+ else:
330
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
331
+
332
+ def forward(
333
+ self,
334
+ unpad_inputs: bool,
335
+ input_ids: Optional[torch.Tensor] = None,
336
+ attention_mask: Optional[torch.Tensor] = None,
337
+ length: Optional[List[int]] = None,
338
+ token_type_ids: Optional[torch.Tensor] = None,
339
+ position_ids: Optional[torch.Tensor] = None,
340
+ inputs_embeds: Optional[torch.Tensor] = None,
341
+ ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Tuple], Optional[List[int]]]:
342
+ """
343
+ """
344
+ if inputs_embeds is None:
345
+ device, input_shape = input_ids.device, input_ids.shape
346
+ else:
347
+ device, input_shape = inputs_embeds.device, inputs_embeds.shape[:2]
348
+ batch_size, seq_length = input_shape
349
+
350
+ # Set attention_mask if it's None
351
+ if attention_mask is None:
352
+ attention_mask = torch.ones(input_shape, device=device)
353
+ if length is not None:
354
+ for i, l in enumerate(length):
355
+ attention_mask[i, l:] = 0
356
+
357
+ # Set attention_mask_bool for unpadding
358
+ if unpad_inputs:
359
+ attention_mask_bool = attention_mask.bool()
360
+ if length is None:
361
+ length = attention_mask.sum(-1).tolist()
362
+
363
+ # Get word embeddings
364
+ if inputs_embeds is None:
365
+ if unpad_inputs:
366
+ input_ids = input_ids[attention_mask_bool].unsqueeze(0)
367
+ inputs_embeds = self.word_embeddings(input_ids)
368
+ else:
369
+ if unpad_inputs:
370
+ inputs_embeds = inputs_embeds[attention_mask_bool].unsqueeze(0)
371
+ embeddings = inputs_embeds
372
+
373
+ # Set and unpad position_ids
374
+ if position_ids is None:
375
+ if seq_length > self.position_ids.size(0):
376
+ self.register_buffer(
377
+ "position_ids", torch.arange(seq_length, device=embeddings.device), persistent=False
378
+ )
379
+ if unpad_inputs:
380
+ # [1, cumsum_seq_len]
381
+ position_ids = torch.cat([self.position_ids[:l] for l in length]).unsqueeze(0)
382
+ else:
383
+ # [bs, seq_len]
384
+ position_ids = self.position_ids[:seq_length].expand(batch_size, -1)
385
+ elif unpad_inputs:
386
+ position_ids = position_ids[attention_mask_bool].unsqueeze(0) # [1, cumsum_seq_len]
387
+
388
+ # Compute rotary embedding
389
+ if self.position_embedding_type == 'rope':
390
+ rope_cos, rope_sin = self.rotary_emb(inputs_embeds, seq_len=seq_length)
391
+ rope_cos = rope_cos[position_ids].unsqueeze(2) # [bs, seq_len, 1, dim]
392
+ rope_sin = rope_sin[position_ids].unsqueeze(2) # [bs, seq_len, 1, dim]
393
+ rope_embeds = rope_cos, rope_sin
394
+ else:
395
+ rope_embeds = None
396
+
397
+ if self.type_vocab_size > 0:
398
+ if token_type_ids is None:
399
+ token_type_ids = position_ids.mul(0)
400
+ else:
401
+ if self.type_vocab_size < 2:
402
+ token_type_ids.mul_(0)
403
+ if unpad_inputs:
404
+ token_type_ids = token_type_ids[attention_mask_bool].unsqueeze(0)
405
+
406
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
407
+ embeddings = embeddings + token_type_embeddings
408
+
409
+ # BERT position
410
+ if self.position_embedding_type == "absolute":
411
+ position_embeddings = self.position_embeddings(position_ids)
412
+ embeddings = embeddings + position_embeddings
413
+
414
+ embeddings = self.LayerNorm(embeddings)
415
+ embeddings = self.dropout(embeddings)
416
+
417
+ return embeddings, attention_mask, rope_embeds, length
418
+
419
+
420
+ class GteAttention(nn.Module):
421
+ def __init__(self, config: GteConfig, pack_qkv=None, use_memory_efficient_attention=None):
422
+ super().__init__()
423
+ self.config = config
424
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
425
+ raise ValueError(
426
+ f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
427
+ f"heads ({config.num_attention_heads})"
428
+ )
429
+
430
+ self.hidden_size = config.hidden_size
431
+ self.num_attention_heads = config.num_attention_heads
432
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
433
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
434
+
435
+ if pack_qkv is None:
436
+ pack_qkv = config.pack_qkv
437
+ self.pack_qkv = pack_qkv
438
+
439
+ if self.pack_qkv:
440
+ self.qkv_proj = nn.Linear(config.hidden_size, self.all_head_size * 3, bias=True)
441
+ else:
442
+ self.q_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
443
+ self.k_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
444
+ self.v_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
445
+
446
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
447
+ self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=True)
448
+
449
+ if use_memory_efficient_attention is None:
450
+ use_memory_efficient_attention = self.config.use_memory_efficient_attention
451
+ self.use_memory_efficient_attention = use_memory_efficient_attention
452
+ self.memory_efficient_attention = None if xops is None else xops.memory_efficient_attention
453
+ if self.use_memory_efficient_attention:
454
+ assert self.memory_efficient_attention is not None, 'please install xformers'
455
+
456
+ def forward(
457
+ self,
458
+ hidden_states: torch.Tensor,
459
+ attention_bias: torch.FloatTensor,
460
+ rope_embeds: Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] = None,
461
+ padding_inputs: Optional[Tuple] = None, # indices, batch, seqlen
462
+ attention_scale: Optional[torch.FloatTensor] = None,
463
+ head_mask: Optional[torch.FloatTensor] = None,
464
+ output_attentions: Optional[bool] = False,
465
+ qkv_inputs: Optional[Tuple] = None, # For RetroMAE
466
+ ) -> Tuple[torch.Tensor, ...]:
467
+ shape_hd = (self.num_attention_heads, self.attention_head_size)
468
+ # qkv
469
+ if self.pack_qkv and qkv_inputs is None:
470
+ qkv_pack = self.qkv_proj(hidden_states).split(self.all_head_size, dim=-1)
471
+ else:
472
+ if qkv_inputs is None:
473
+ qkv_inputs = (hidden_states, hidden_states, hidden_states)
474
+ qkv_pack = [
475
+ getattr(self, n + '_proj')(s) for s, n in zip(qkv_inputs, 'qkv')
476
+ ]
477
+ query_states, key_states, value_states = [t.view(t.shape[:-1] + shape_hd) for t in qkv_pack]
478
+
479
+ if self.config.position_embedding_type == 'rope':
480
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, *rope_embeds)
481
+
482
+ dtype = query_states.dtype
483
+
484
+ if self.config.logn_attention_scale and attention_scale is not None:
485
+ # https://kexue.fm/archives/8823
486
+ query_states = query_states * attention_scale.to(dtype)
487
+
488
+ if padding_inputs is not None:
489
+ query_states = pad_input(query_states.squeeze(), *padding_inputs)
490
+ key_states = pad_input(key_states.squeeze(), *padding_inputs)
491
+ value_states = pad_input(value_states.squeeze(), *padding_inputs)
492
+
493
+ if self.use_memory_efficient_attention:
494
+ assert self.memory_efficient_attention is not None, "xformers is not loaded"
495
+ assert output_attentions is False, "memory_efficient_attention do not output attentions"
496
+ assert head_mask is None, "Not support yet"
497
+ attention_probs = None
498
+ if torch.is_tensor(attention_bias):
499
+ attention_bias = attention_bias.to(dtype)
500
+ context_layer = self.memory_efficient_attention(
501
+ query_states,
502
+ key_states,
503
+ value_states,
504
+ attn_bias=attention_bias,
505
+ p=self.dropout.p
506
+ )
507
+ else:
508
+ if output_attentions and isinstance(self, GteSdpaAttention):
509
+ raise RuntimeError("SDPA do not output attentions")
510
+ context_layer, attention_probs = self._attention(
511
+ query_states, key_states, value_states, attention_bias, head_mask
512
+ )
513
+
514
+ if padding_inputs is not None:
515
+ context_layer = unpad_input(context_layer, indices=padding_inputs[0])
516
+
517
+ gte_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
518
+ context_layer = context_layer.view(gte_context_layer_shape)
519
+
520
+ # output proj
521
+ attn_output = self.o_proj(context_layer)
522
+
523
+ # add attentions if we output them
524
+ outputs = (attn_output, attention_probs) if output_attentions else (attn_output,)
525
+ return outputs
526
+
527
+ def _attention(self, query_states, key_states, value_states, attention_bias, head_mask):
528
+ """
529
+ Args:
530
+ q/k/v: (B, L, n_head, head_dim),
531
+ Returns:
532
+ attn_output: (B L, n_head, head_dim)
533
+ """
534
+ query_states = query_states.transpose(1, 2)
535
+ key_states = key_states.transpose(1, 2)
536
+ value_states = value_states.transpose(1, 2)
537
+ # Take the dot product between "query" and "key" to get the raw attention scores.
538
+ attention_scores = torch.matmul(query_states, key_states.transpose(-1, -2))
539
+
540
+ attention_scores = attention_scores / math.sqrt(self.attention_head_size)
541
+ if attention_bias is not None:
542
+ # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
543
+ attention_scores = attention_scores + attention_bias
544
+
545
+ # Normalize the attention scores to probabilities.
546
+ attention_probs = nn.functional.softmax(attention_scores, dim=-1)
547
+
548
+ # This is actually dropping out entire tokens to attend to, which might
549
+ # seem a bit unusual, but is taken from the original Transformer paper.
550
+ if self.dropout.p > 0:
551
+ attention_probs = self.dropout(attention_probs)
552
+
553
+ # Mask heads if we want to
554
+ if head_mask is not None:
555
+ attention_probs = attention_probs * head_mask
556
+
557
+ context_layer = torch.matmul(attention_probs, value_states)
558
+
559
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
560
+ return context_layer, attention_probs
561
+
562
+
563
+ class GteSdpaAttention(GteAttention):
564
+ """
565
+ Gte attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
566
+ `GteAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
567
+ SDPA API.
568
+ """
569
+ def __init__(self, config: GteConfig, **kwargs):
570
+ super().__init__(config, **kwargs)
571
+ # torch.backends.cuda.enable_mem_efficient_sdp(False)
572
+ # logger.warning(
573
+ # "Disable memory efficient attention kernel for `GteSdpaAttention`, you can set "
574
+ # "`use_memory_efficient_attention=True` if it expected to use."
575
+ # )
576
+
577
+ def _attention(self, query_states, key_states, value_states, attention_bias, head_mask):
578
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
579
+ query_states.transpose(1, 2),
580
+ key_states.transpose(1, 2),
581
+ value_states.transpose(1, 2),
582
+ attn_mask=attention_bias,
583
+ dropout_p=self.dropout.p if self.training else 0.0,
584
+ )
585
+ attn_output = attn_output.permute(0, 2, 1, 3).contiguous()
586
+ return attn_output, None
587
+
588
+
589
+ GTE_ATTENTION_CLASSES = {
590
+ "eager": GteAttention,
591
+ # "flash_attention_2": , # TODO
592
+ "sdpa": GteSdpaAttention,
593
+ }
594
+
595
+
596
+ class GteGatedMLP(nn.Module):
597
+ """
598
+ GLU Variants Improve Transformer.
599
+ """
600
+
601
+ def __init__(self, config: GteConfig):
602
+ super().__init__()
603
+ self.intermediate_size = config.intermediate_size
604
+ self.up_gate_proj = nn.Linear(config.hidden_size, self.intermediate_size * 2, bias=False)
605
+ self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=True)
606
+ self.act_fn = ACT2FN[config.hidden_act]
607
+ if config.hidden_dropout_prob > 0:
608
+ self.hidden_dropout = nn.Dropout(config.hidden_dropout_prob)
609
+ else:
610
+ self.hidden_dropout = None
611
+
612
+ def forward(self, hidden_states):
613
+ up_gate = self.up_gate_proj(hidden_states)
614
+ up_states, gate = torch.split(up_gate, self.intermediate_size, dim=-1)
615
+ gate = self.act_fn(gate)
616
+ gated_states = gate * up_states
617
+ if self.hidden_dropout is not None:
618
+ gated_states = self.hidden_dropout(gated_states)
619
+ down_states = self.down_proj(gated_states)
620
+ return down_states
621
+
622
+
623
+ class GteLayer(nn.Module):
624
+ def __init__(
625
+ self,
626
+ config: GteConfig,
627
+ pack_qkv=None,
628
+ use_memory_efficient_attention=None,
629
+ attn_implementation=None
630
+ ):
631
+ super().__init__()
632
+ if attn_implementation is None:
633
+ attn_implementation = config._attn_implementation
634
+ if use_memory_efficient_attention is None:
635
+ use_memory_efficient_attention = config.use_memory_efficient_attention
636
+ if use_memory_efficient_attention:
637
+ if attn_implementation != 'eager':
638
+ logger.warning_once(f"Override {attn_implementation=} to 'eager' as {use_memory_efficient_attention=}")
639
+ attn_implementation = 'eager' # Since it will be SDPA by default for torch>=2.1.1
640
+ self.attention = GTE_ATTENTION_CLASSES[attn_implementation](
641
+ config, pack_qkv=pack_qkv, use_memory_efficient_attention=use_memory_efficient_attention
642
+ )
643
+ self.mlp = GteGatedMLP(config)
644
+
645
+ ln_class = LAYER_NORM[config.layer_norm_type]
646
+ self.attn_ln = ln_class(config.hidden_size, eps=config.layer_norm_eps)
647
+ self.mlp_ln = ln_class(config.hidden_size, eps=config.layer_norm_eps)
648
+
649
+ if config.hidden_dropout_prob > 0:
650
+ self.hidden_dropout = nn.Dropout(config.hidden_dropout_prob)
651
+ else:
652
+ self.hidden_dropout = None
653
+
654
+ def forward(
655
+ self,
656
+ hidden_states: torch.Tensor,
657
+ attention_bias: torch.FloatTensor,
658
+ rope_embeds: Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] = None,
659
+ padding_inputs: Optional[Tuple] = None, # indices, batch, seqlen
660
+ attention_scale: Optional[torch.FloatTensor] = None,
661
+ subset_indices: Optional[torch.LongTensor] = None,
662
+ head_mask: Optional[torch.FloatTensor] = None,
663
+ output_attentions: Optional[bool] = False,
664
+ qkv_inputs: Optional[Tuple] = None, # For RetroMAE
665
+ ) -> Tuple[torch.Tensor, ...]:
666
+ # Multi head self attention
667
+ residual = hidden_states if qkv_inputs is None else qkv_inputs[0]
668
+ attention_outputs = self.attention(
669
+ hidden_states,
670
+ attention_bias,
671
+ rope_embeds,
672
+ padding_inputs,
673
+ attention_scale,
674
+ head_mask,
675
+ output_attentions=output_attentions,
676
+ qkv_inputs=qkv_inputs,
677
+ )
678
+ hidden_states = attention_outputs[0]
679
+ if self.hidden_dropout is not None:
680
+ hidden_states = self.hidden_dropout(hidden_states)
681
+ hidden_states = residual + hidden_states
682
+
683
+ # In pretraining, after the attention of last layer, we only need the masked tokens.
684
+ if subset_indices is not None:
685
+ hidden_states = hidden_states[subset_indices]
686
+
687
+ hidden_states = self.attn_ln(hidden_states)
688
+
689
+ # Fully Connected
690
+ residual = hidden_states
691
+ hidden_states = self.mlp(hidden_states)
692
+ if self.hidden_dropout is not None:
693
+ hidden_states = self.hidden_dropout(hidden_states)
694
+ hidden_states = residual + hidden_states
695
+ hidden_states = self.mlp_ln(hidden_states)
696
+
697
+ # add self attentions if we output attention weights
698
+ outputs = (hidden_states,) + attention_outputs[1:]
699
+ return outputs
700
+
701
+
702
+ class GteEncoder(nn.Module):
703
+ def __init__(self, config):
704
+ super().__init__()
705
+ self.config = config
706
+ self.layer = nn.ModuleList([GteLayer(config) for _ in range(config.num_hidden_layers)])
707
+ self.gradient_checkpointing = False
708
+
709
+ def forward(
710
+ self,
711
+ hidden_states: torch.Tensor,
712
+ attention_bias: Optional[torch.FloatTensor] = None,
713
+ rope_embeds: Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] = None,
714
+ padding_inputs: Optional[Tuple] = None, # indices, batch, seqlen
715
+ attention_scale: Optional[torch.FloatTensor] = None,
716
+ subset_indices: Optional[torch.LongTensor] = None,
717
+ head_mask: Optional[torch.FloatTensor] = None,
718
+ output_attentions: Optional[bool] = False,
719
+ output_hidden_states: Optional[bool] = False,
720
+ return_dict: Optional[bool] = True,
721
+ ) -> Union[Tuple[torch.Tensor], BaseModelOutput]:
722
+ all_hidden_states = () if output_hidden_states else None
723
+ all_self_attentions = () if output_attentions else None
724
+
725
+ for i, layer_module in enumerate(self.layer):
726
+ if output_hidden_states:
727
+ all_hidden_states = all_hidden_states + (hidden_states,)
728
+
729
+ if i >= len(self.layer) - 1:
730
+ layer_subset_indices = subset_indices
731
+ else:
732
+ layer_subset_indices = None
733
+
734
+ layer_head_mask = head_mask[i] if head_mask is not None else None
735
+
736
+ if self.gradient_checkpointing and self.training:
737
+ layer_outputs = self._gradient_checkpointing_func(
738
+ layer_module.__call__,
739
+ hidden_states,
740
+ attention_bias,
741
+ rope_embeds,
742
+ padding_inputs,
743
+ attention_scale,
744
+ layer_subset_indices,
745
+ layer_head_mask,
746
+ )
747
+ else:
748
+ layer_outputs = layer_module(
749
+ hidden_states,
750
+ attention_bias,
751
+ rope_embeds,
752
+ padding_inputs,
753
+ attention_scale,
754
+ layer_subset_indices,
755
+ layer_head_mask,
756
+ output_attentions,
757
+ )
758
+
759
+ hidden_states = layer_outputs[0]
760
+ if output_attentions:
761
+ all_self_attentions = all_self_attentions + (layer_outputs[1],)
762
+
763
+ if output_hidden_states:
764
+ all_hidden_states = all_hidden_states + (hidden_states,)
765
+
766
+ if not return_dict:
767
+ return tuple(
768
+ v
769
+ for v in [
770
+ hidden_states,
771
+ all_hidden_states,
772
+ all_self_attentions,
773
+ ]
774
+ if v is not None
775
+ )
776
+ return BaseModelOutput(
777
+ last_hidden_state=hidden_states,
778
+ hidden_states=all_hidden_states,
779
+ attentions=all_self_attentions,
780
+ )
781
+
782
+
783
+ # Copied from transformers.models.bert.modeling_bert.BertPooler with Bert->Gte
784
+ class GtePooler(nn.Module):
785
+ def __init__(self, config):
786
+ super().__init__()
787
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
788
+ self.activation = nn.Tanh()
789
+
790
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
791
+ # We "pool" the model by simply taking the hidden state corresponding
792
+ # to the first token.
793
+ first_token_tensor = hidden_states[:, 0]
794
+ pooled_output = self.dense(first_token_tensor)
795
+ pooled_output = self.activation(pooled_output)
796
+ return pooled_output
797
+
798
+
799
+ class GtePreTrainedModel(PreTrainedModel):
800
+ """
801
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
802
+ models.
803
+ """
804
+
805
+ config_class = GteConfig
806
+ base_model_prefix = "gte"
807
+ supports_gradient_checkpointing = True
808
+ _supports_sdpa = True
809
+
810
+ def _init_weights(self, module):
811
+ """Initialize the weights"""
812
+ if isinstance(module, nn.Linear):
813
+ # Slightly different from the TF version which uses truncated_normal for initialization
814
+ # cf https://github.com/pytorch/pytorch/pull/5617
815
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
816
+ if module.bias is not None:
817
+ module.bias.data.zero_()
818
+ elif isinstance(module, nn.Embedding):
819
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
820
+ if module.padding_idx is not None:
821
+ module.weight.data[module.padding_idx].zero_()
822
+ elif isinstance(module, nn.LayerNorm):
823
+ module.bias.data.zero_()
824
+ module.weight.data.fill_(1.0)
825
+
826
+
827
+ class GteModel(GtePreTrainedModel):
828
+ """
829
+ The bare Gte Model transformer outputting raw hidden-states without any specific head on top.
830
+ """
831
+
832
+ def __init__(self, config: GteConfig, add_pooling_layer=False):
833
+ super().__init__(config)
834
+ self.config = config
835
+
836
+ self.embeddings = GteEmbeddings(config)
837
+ self.encoder = GteEncoder(config)
838
+
839
+ self.pooler = GtePooler(config) if add_pooling_layer else None
840
+
841
+ # Initialize weights and apply final processing
842
+ self.post_init()
843
+
844
+ def get_input_embeddings(self):
845
+ return self.embeddings.word_embeddings
846
+
847
+ def set_input_embeddings(self, value):
848
+ self.embeddings.word_embeddings = value
849
+
850
+ def forward(
851
+ self,
852
+ input_ids: Optional[torch.Tensor] = None,
853
+ attention_mask: Optional[torch.Tensor] = None,
854
+ length: Optional[List[int]] = None,
855
+ subset_indices: Optional[torch.LongTensor] = None,
856
+ token_type_ids: Optional[torch.Tensor] = None,
857
+ position_ids: Optional[torch.Tensor] = None,
858
+ head_mask: Optional[torch.Tensor] = None,
859
+ inputs_embeds: Optional[torch.Tensor] = None,
860
+ output_attentions: Optional[bool] = None,
861
+ output_hidden_states: Optional[bool] = None,
862
+ return_dict: Optional[bool] = None,
863
+ unpad_inputs: Optional[bool] = None,
864
+ ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPooling]:
865
+ r"""
866
+ length (`list` of length `batch_size`, *optional*):
867
+ If is `None`, return padded `last_hidden_state`.
868
+ subset_indices ():
869
+ pass
870
+ unpad_inputs (`bool`, *optional*):
871
+ pass
872
+ """
873
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
874
+ output_hidden_states = (
875
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
876
+ )
877
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
878
+ unpad_inputs = unpad_inputs if unpad_inputs is not None else self.config.unpad_inputs
879
+ output_padded = length is None
880
+
881
+ if input_ids is not None and inputs_embeds is not None:
882
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
883
+ elif input_ids is not None:
884
+ self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
885
+ input_shape = input_ids.size()
886
+ elif inputs_embeds is not None:
887
+ input_shape = inputs_embeds.size()[:-1]
888
+ else:
889
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
890
+
891
+ # TODO: not used
892
+ # # Prepare head mask if needed
893
+ # # 1.0 in head_mask indicate we keep the head
894
+ # # attention_probs has shape bsz x n_heads x N x N
895
+ # # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
896
+ # # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
897
+ # head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
898
+
899
+ # Get embeddings, may unpad them
900
+ (embedding_output, attention_mask, rope_embeds, length) = self.embeddings(
901
+ unpad_inputs,
902
+ input_ids=input_ids,
903
+ attention_mask=attention_mask,
904
+ length=length,
905
+ token_type_ids=token_type_ids,
906
+ position_ids=position_ids,
907
+ inputs_embeds=inputs_embeds
908
+ )
909
+
910
+ batch_size, seq_length = input_shape
911
+ if unpad_inputs and self.config.use_memory_efficient_attention:
912
+ attention_bias = xops.fmha.attn_bias.BlockDiagonalMask.from_seqlens(length)
913
+ else:
914
+ # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
915
+ # ourselves in which case we just need to make it broadcastable to all heads.
916
+ attention_bias = self.get_extended_attention_mask(attention_mask, input_shape)
917
+ if self.config.use_memory_efficient_attention:
918
+ # Invalid shape for attention bias: torch.Size([48, 1, 1, 512]) (expected (48, 12, 512, 512))
919
+ attention_bias = attention_bias.expand(-1, self.config.num_attention_heads, seq_length, -1)
920
+
921
+ padding_inputs = None
922
+ if unpad_inputs and (output_padded or not self.config.use_memory_efficient_attention):
923
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
924
+ if not self.config.use_memory_efficient_attention:
925
+ padding_inputs = (indices, *input_shape)
926
+
927
+ attention_scale = None
928
+ if self.config.logn_attention_scale:
929
+ logger.warning_once("TODO: logn_attention_scale")
930
+ # # attention scale log_512(input_len)
931
+ # attention_scale = attention_mask.sum(1).log() / torch.tensor(self.config.max_position_embeddings).log()
932
+ # # inference-time logn scale need clip 1
933
+ # if self.config.logn_attention_clip1:
934
+ # attention_scale.clip_(1)
935
+ # attention_scale = attention_scale[:, None, None, None]
936
+ # else:
937
+ # attention_scale = None
938
+
939
+ encoder_outputs = self.encoder(
940
+ embedding_output,
941
+ attention_bias=attention_bias,
942
+ rope_embeds=rope_embeds,
943
+ padding_inputs=padding_inputs,
944
+ attention_scale=attention_scale,
945
+ subset_indices=subset_indices,
946
+ head_mask=head_mask,
947
+ output_attentions=output_attentions,
948
+ output_hidden_states=output_hidden_states,
949
+ return_dict=return_dict,
950
+ )
951
+ sequence_output = encoder_outputs[0]
952
+ if unpad_inputs and output_padded:
953
+ sequence_output = pad_input(
954
+ sequence_output.squeeze(), indices, batch_size, seq_length
955
+ )
956
+
957
+ pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
958
+
959
+ if not return_dict:
960
+ return (sequence_output, pooled_output) + encoder_outputs[1:]
961
+
962
+ return BaseModelOutputWithPooling(
963
+ last_hidden_state=sequence_output,
964
+ pooler_output=pooled_output,
965
+ hidden_states=encoder_outputs.hidden_states,
966
+ attentions=encoder_outputs.attentions,
967
+ )
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
+ ]
sentence_bert_config.json ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ {
2
+ "max_seq_length": 8192,
3
+ "do_lower_case": false
4
+ }
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:aa7a6ad87a7ce8fe196787355f6af7d03aee94d19c54a5eb1392ed18c8ef451a
3
+ size 17082988
tokenizer_config.json ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ "extra_special_tokens": {},
49
+ "mask_token": "<mask>",
50
+ "max_length": 512,
51
+ "model_max_length": 32768,
52
+ "pad_to_multiple_of": null,
53
+ "pad_token": "<pad>",
54
+ "pad_token_type_id": 0,
55
+ "padding_side": "right",
56
+ "sep_token": "</s>",
57
+ "stride": 0,
58
+ "tokenizer_class": "XLMRobertaTokenizerFast",
59
+ "truncation_side": "right",
60
+ "truncation_strategy": "longest_first",
61
+ "unk_token": "<unk>"
62
+ }