KoichiYasuoka commited on
Commit
e2acffb
·
1 Parent(s): 7bd35e6

initial release

Browse files
README.md ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ language:
3
+ - "ja"
4
+ tags:
5
+ - "japanese"
6
+ - "pos"
7
+ - "dependency-parsing"
8
+ - "modernbert"
9
+ base_model: KoichiYasuoka/modernbert-small-japanese-wikipedia-upos
10
+ datasets:
11
+ - "universal_dependencies"
12
+ license: "apache-2.0"
13
+ pipeline_tag: "token-classification"
14
+ widget:
15
+ - text: "全学年にわたって小学校の国語の教科書に挿し絵が用いられている"
16
+ ---
17
+
18
+ # modernbert-small-japanese-wikipedia-ud-square
19
+
20
+ ## Model Description
21
+
22
+ This is a ModernBERT model pretrained for POS-tagging and dependency-parsing (using `goeswith` for subwords), derived from [modernbert-small-japanese-wikipedia-upos](https://huggingface.co/KoichiYasuoka/modernbert-small-japanese-wikipedia-upos) and [UD_Japanese-GSDLUW](https://github.com/UniversalDependencies/UD_Japanese-GSDLUW).
23
+
24
+ ## How to Use
25
+
26
+ ```py
27
+ from transformers import pipeline
28
+ nlp=pipeline("universal-dependencies","KoichiYasuoka/modernbert-small-japanese-wikipedia-ud-square",trust_remote_code=True,aggregation_strategy="simple")
29
+ print(nlp("全学年にわたって小学校の国語の教科書に挿し絵が用いられている"))
30
+ ```
31
+
config.json ADDED
@@ -0,0 +1,501 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "ModernBertForTokenClassification"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "auto_map": {
8
+ "AutoConfig": "configuration_modernbert.ModernBertConfig",
9
+ "AutoModel": "modeling_modernbert.ModernBertModel",
10
+ "AutoModelForMaskedLM": "modeling_modernbert.ModernBertForMaskedLM",
11
+ "AutoModelForSequenceClassification": "modeling_modernbert.ModernBertForSequenceClassification",
12
+ "AutoModelForTokenClassification": "modeling_modernbert.ModernBertForTokenClassification"
13
+ },
14
+ "bos_token_id": 0,
15
+ "classifier_activation": "gelu",
16
+ "classifier_bias": false,
17
+ "classifier_dropout": 0.0,
18
+ "classifier_pooling": "mean",
19
+ "cls_token_id": 0,
20
+ "custom_pipelines": {
21
+ "universal-dependencies": {
22
+ "impl": "ud.UniversalDependenciesPipeline",
23
+ "pt": "AutoModelForTokenClassification"
24
+ }
25
+ },
26
+ "decoder_bias": true,
27
+ "deterministic_flash_attn": false,
28
+ "embedding_dropout": 0.0,
29
+ "eos_token_id": 2,
30
+ "global_attn_every_n_layers": 3,
31
+ "global_rope_theta": 160000.0,
32
+ "gradient_checkpointing": false,
33
+ "hidden_activation": "gelu",
34
+ "hidden_size": 256,
35
+ "id2label": {
36
+ "0": "ADJ",
37
+ "1": "ADJ.",
38
+ "2": "ADJ.|[acl]",
39
+ "3": "ADJ.|[advcl]",
40
+ "4": "ADJ.|[amod]",
41
+ "5": "ADJ.|[ccomp]",
42
+ "6": "ADJ.|[csubj:outer]",
43
+ "7": "ADJ.|[csubj]",
44
+ "8": "ADJ.|[dep]",
45
+ "9": "ADJ.|[nmod]",
46
+ "10": "ADJ.|[nsubj]",
47
+ "11": "ADJ.|[obj]",
48
+ "12": "ADJ.|[obl]",
49
+ "13": "ADJ.|[root]",
50
+ "14": "ADJ|[acl]",
51
+ "15": "ADJ|[advcl]",
52
+ "16": "ADJ|[amod]",
53
+ "17": "ADJ|[ccomp]",
54
+ "18": "ADJ|[csubj:outer]",
55
+ "19": "ADJ|[csubj]",
56
+ "20": "ADJ|[dep]",
57
+ "21": "ADJ|[nmod]",
58
+ "22": "ADJ|[nsubj]",
59
+ "23": "ADJ|[obj]",
60
+ "24": "ADJ|[obl]",
61
+ "25": "ADJ|[root]",
62
+ "26": "ADP",
63
+ "27": "ADP.",
64
+ "28": "ADP.|[case]",
65
+ "29": "ADP.|[fixed]",
66
+ "30": "ADP|[case]",
67
+ "31": "ADP|[fixed]",
68
+ "32": "ADV",
69
+ "33": "ADV.",
70
+ "34": "ADV.|[advcl]",
71
+ "35": "ADV.|[advmod]",
72
+ "36": "ADV.|[dep]",
73
+ "37": "ADV.|[obj]",
74
+ "38": "ADV.|[root]",
75
+ "39": "ADV|[advcl]",
76
+ "40": "ADV|[advmod]",
77
+ "41": "ADV|[dep]",
78
+ "42": "ADV|[obj]",
79
+ "43": "ADV|[root]",
80
+ "44": "AUX",
81
+ "45": "AUX.",
82
+ "46": "AUX.|Polarity=Neg",
83
+ "47": "AUX.|Polarity=Neg|[aux]",
84
+ "48": "AUX.|Polarity=Neg|[fixed]",
85
+ "49": "AUX.|[aux]",
86
+ "50": "AUX.|[cop]",
87
+ "51": "AUX.|[fixed]",
88
+ "52": "AUX.|[root]",
89
+ "53": "AUX|Polarity=Neg",
90
+ "54": "AUX|Polarity=Neg|[aux]",
91
+ "55": "AUX|Polarity=Neg|[fixed]",
92
+ "56": "AUX|[aux]",
93
+ "57": "AUX|[cop]",
94
+ "58": "AUX|[fixed]",
95
+ "59": "AUX|[root]",
96
+ "60": "CCONJ",
97
+ "61": "CCONJ.",
98
+ "62": "CCONJ.|[cc]",
99
+ "63": "CCONJ|[cc]",
100
+ "64": "DET",
101
+ "65": "DET.",
102
+ "66": "DET.|[det]",
103
+ "67": "DET|[det]",
104
+ "68": "INTJ",
105
+ "69": "INTJ.",
106
+ "70": "INTJ.|[discourse]",
107
+ "71": "INTJ.|[root]",
108
+ "72": "INTJ|[discourse]",
109
+ "73": "INTJ|[root]",
110
+ "74": "NOUN",
111
+ "75": "NOUN.",
112
+ "76": "NOUN.|Polarity=Neg",
113
+ "77": "NOUN.|Polarity=Neg|[obl]",
114
+ "78": "NOUN.|Polarity=Neg|[root]",
115
+ "79": "NOUN.|[acl]",
116
+ "80": "NOUN.|[advcl]",
117
+ "81": "NOUN.|[ccomp]",
118
+ "82": "NOUN.|[compound]",
119
+ "83": "NOUN.|[csubj:outer]",
120
+ "84": "NOUN.|[csubj]",
121
+ "85": "NOUN.|[nmod]",
122
+ "86": "NOUN.|[nsubj:outer]",
123
+ "87": "NOUN.|[nsubj]",
124
+ "88": "NOUN.|[obj]",
125
+ "89": "NOUN.|[obl]",
126
+ "90": "NOUN.|[root]",
127
+ "91": "NOUN|Polarity=Neg",
128
+ "92": "NOUN|Polarity=Neg|[obl]",
129
+ "93": "NOUN|Polarity=Neg|[root]",
130
+ "94": "NOUN|[acl]",
131
+ "95": "NOUN|[advcl]",
132
+ "96": "NOUN|[ccomp]",
133
+ "97": "NOUN|[compound]",
134
+ "98": "NOUN|[csubj:outer]",
135
+ "99": "NOUN|[csubj]",
136
+ "100": "NOUN|[nmod]",
137
+ "101": "NOUN|[nsubj:outer]",
138
+ "102": "NOUN|[nsubj]",
139
+ "103": "NOUN|[obj]",
140
+ "104": "NOUN|[obl]",
141
+ "105": "NOUN|[root]",
142
+ "106": "NUM",
143
+ "107": "NUM.",
144
+ "108": "NUM.|[advcl]",
145
+ "109": "NUM.|[compound]",
146
+ "110": "NUM.|[nmod]",
147
+ "111": "NUM.|[nsubj:outer]",
148
+ "112": "NUM.|[nsubj]",
149
+ "113": "NUM.|[nummod]",
150
+ "114": "NUM.|[obj]",
151
+ "115": "NUM.|[obl]",
152
+ "116": "NUM.|[root]",
153
+ "117": "NUM|[advcl]",
154
+ "118": "NUM|[compound]",
155
+ "119": "NUM|[nmod]",
156
+ "120": "NUM|[nsubj:outer]",
157
+ "121": "NUM|[nsubj]",
158
+ "122": "NUM|[nummod]",
159
+ "123": "NUM|[obj]",
160
+ "124": "NUM|[obl]",
161
+ "125": "NUM|[root]",
162
+ "126": "PART",
163
+ "127": "PART.",
164
+ "128": "PART.|[mark]",
165
+ "129": "PART|[mark]",
166
+ "130": "PRON",
167
+ "131": "PRON.",
168
+ "132": "PRON.|[acl]",
169
+ "133": "PRON.|[advcl]",
170
+ "134": "PRON.|[nmod]",
171
+ "135": "PRON.|[nsubj:outer]",
172
+ "136": "PRON.|[nsubj]",
173
+ "137": "PRON.|[obj]",
174
+ "138": "PRON.|[obl]",
175
+ "139": "PRON.|[root]",
176
+ "140": "PRON|[acl]",
177
+ "141": "PRON|[advcl]",
178
+ "142": "PRON|[nmod]",
179
+ "143": "PRON|[nsubj:outer]",
180
+ "144": "PRON|[nsubj]",
181
+ "145": "PRON|[obj]",
182
+ "146": "PRON|[obl]",
183
+ "147": "PRON|[root]",
184
+ "148": "PROPN",
185
+ "149": "PROPN.",
186
+ "150": "PROPN.|[acl]",
187
+ "151": "PROPN.|[advcl]",
188
+ "152": "PROPN.|[compound]",
189
+ "153": "PROPN.|[nmod]",
190
+ "154": "PROPN.|[nsubj:outer]",
191
+ "155": "PROPN.|[nsubj]",
192
+ "156": "PROPN.|[obj]",
193
+ "157": "PROPN.|[obl]",
194
+ "158": "PROPN.|[root]",
195
+ "159": "PROPN|[acl]",
196
+ "160": "PROPN|[advcl]",
197
+ "161": "PROPN|[compound]",
198
+ "162": "PROPN|[nmod]",
199
+ "163": "PROPN|[nsubj:outer]",
200
+ "164": "PROPN|[nsubj]",
201
+ "165": "PROPN|[obj]",
202
+ "166": "PROPN|[obl]",
203
+ "167": "PROPN|[root]",
204
+ "168": "PUNCT",
205
+ "169": "PUNCT.",
206
+ "170": "PUNCT.|[punct]",
207
+ "171": "PUNCT|[punct]",
208
+ "172": "SCONJ",
209
+ "173": "SCONJ.",
210
+ "174": "SCONJ.|[dep]",
211
+ "175": "SCONJ.|[fixed]",
212
+ "176": "SCONJ.|[mark]",
213
+ "177": "SCONJ|[dep]",
214
+ "178": "SCONJ|[fixed]",
215
+ "179": "SCONJ|[mark]",
216
+ "180": "SYM",
217
+ "181": "SYM.",
218
+ "182": "SYM.|[compound]",
219
+ "183": "SYM.|[dep]",
220
+ "184": "SYM.|[nmod]",
221
+ "185": "SYM.|[obl]",
222
+ "186": "SYM|[compound]",
223
+ "187": "SYM|[dep]",
224
+ "188": "SYM|[nmod]",
225
+ "189": "SYM|[obl]",
226
+ "190": "VERB",
227
+ "191": "VERB.",
228
+ "192": "VERB.|[acl]",
229
+ "193": "VERB.|[advcl]",
230
+ "194": "VERB.|[ccomp]",
231
+ "195": "VERB.|[compound]",
232
+ "196": "VERB.|[csubj:outer]",
233
+ "197": "VERB.|[csubj]",
234
+ "198": "VERB.|[nmod]",
235
+ "199": "VERB.|[obj]",
236
+ "200": "VERB.|[obl]",
237
+ "201": "VERB.|[root]",
238
+ "202": "VERB|[acl]",
239
+ "203": "VERB|[advcl]",
240
+ "204": "VERB|[ccomp]",
241
+ "205": "VERB|[compound]",
242
+ "206": "VERB|[csubj:outer]",
243
+ "207": "VERB|[csubj]",
244
+ "208": "VERB|[nmod]",
245
+ "209": "VERB|[obj]",
246
+ "210": "VERB|[obl]",
247
+ "211": "VERB|[root]",
248
+ "212": "X",
249
+ "213": "X.",
250
+ "214": "X.|[dep]",
251
+ "215": "X.|[goeswith]",
252
+ "216": "X.|[nmod]",
253
+ "217": "X|[dep]",
254
+ "218": "X|[nmod]"
255
+ },
256
+ "initializer_cutoff_factor": 2.0,
257
+ "initializer_range": 0.02,
258
+ "intermediate_size": 768,
259
+ "label2id": {
260
+ "ADJ": 0,
261
+ "ADJ.": 1,
262
+ "ADJ.|[acl]": 2,
263
+ "ADJ.|[advcl]": 3,
264
+ "ADJ.|[amod]": 4,
265
+ "ADJ.|[ccomp]": 5,
266
+ "ADJ.|[csubj:outer]": 6,
267
+ "ADJ.|[csubj]": 7,
268
+ "ADJ.|[dep]": 8,
269
+ "ADJ.|[nmod]": 9,
270
+ "ADJ.|[nsubj]": 10,
271
+ "ADJ.|[obj]": 11,
272
+ "ADJ.|[obl]": 12,
273
+ "ADJ.|[root]": 13,
274
+ "ADJ|[acl]": 14,
275
+ "ADJ|[advcl]": 15,
276
+ "ADJ|[amod]": 16,
277
+ "ADJ|[ccomp]": 17,
278
+ "ADJ|[csubj:outer]": 18,
279
+ "ADJ|[csubj]": 19,
280
+ "ADJ|[dep]": 20,
281
+ "ADJ|[nmod]": 21,
282
+ "ADJ|[nsubj]": 22,
283
+ "ADJ|[obj]": 23,
284
+ "ADJ|[obl]": 24,
285
+ "ADJ|[root]": 25,
286
+ "ADP": 26,
287
+ "ADP.": 27,
288
+ "ADP.|[case]": 28,
289
+ "ADP.|[fixed]": 29,
290
+ "ADP|[case]": 30,
291
+ "ADP|[fixed]": 31,
292
+ "ADV": 32,
293
+ "ADV.": 33,
294
+ "ADV.|[advcl]": 34,
295
+ "ADV.|[advmod]": 35,
296
+ "ADV.|[dep]": 36,
297
+ "ADV.|[obj]": 37,
298
+ "ADV.|[root]": 38,
299
+ "ADV|[advcl]": 39,
300
+ "ADV|[advmod]": 40,
301
+ "ADV|[dep]": 41,
302
+ "ADV|[obj]": 42,
303
+ "ADV|[root]": 43,
304
+ "AUX": 44,
305
+ "AUX.": 45,
306
+ "AUX.|Polarity=Neg": 46,
307
+ "AUX.|Polarity=Neg|[aux]": 47,
308
+ "AUX.|Polarity=Neg|[fixed]": 48,
309
+ "AUX.|[aux]": 49,
310
+ "AUX.|[cop]": 50,
311
+ "AUX.|[fixed]": 51,
312
+ "AUX.|[root]": 52,
313
+ "AUX|Polarity=Neg": 53,
314
+ "AUX|Polarity=Neg|[aux]": 54,
315
+ "AUX|Polarity=Neg|[fixed]": 55,
316
+ "AUX|[aux]": 56,
317
+ "AUX|[cop]": 57,
318
+ "AUX|[fixed]": 58,
319
+ "AUX|[root]": 59,
320
+ "CCONJ": 60,
321
+ "CCONJ.": 61,
322
+ "CCONJ.|[cc]": 62,
323
+ "CCONJ|[cc]": 63,
324
+ "DET": 64,
325
+ "DET.": 65,
326
+ "DET.|[det]": 66,
327
+ "DET|[det]": 67,
328
+ "INTJ": 68,
329
+ "INTJ.": 69,
330
+ "INTJ.|[discourse]": 70,
331
+ "INTJ.|[root]": 71,
332
+ "INTJ|[discourse]": 72,
333
+ "INTJ|[root]": 73,
334
+ "NOUN": 74,
335
+ "NOUN.": 75,
336
+ "NOUN.|Polarity=Neg": 76,
337
+ "NOUN.|Polarity=Neg|[obl]": 77,
338
+ "NOUN.|Polarity=Neg|[root]": 78,
339
+ "NOUN.|[acl]": 79,
340
+ "NOUN.|[advcl]": 80,
341
+ "NOUN.|[ccomp]": 81,
342
+ "NOUN.|[compound]": 82,
343
+ "NOUN.|[csubj:outer]": 83,
344
+ "NOUN.|[csubj]": 84,
345
+ "NOUN.|[nmod]": 85,
346
+ "NOUN.|[nsubj:outer]": 86,
347
+ "NOUN.|[nsubj]": 87,
348
+ "NOUN.|[obj]": 88,
349
+ "NOUN.|[obl]": 89,
350
+ "NOUN.|[root]": 90,
351
+ "NOUN|Polarity=Neg": 91,
352
+ "NOUN|Polarity=Neg|[obl]": 92,
353
+ "NOUN|Polarity=Neg|[root]": 93,
354
+ "NOUN|[acl]": 94,
355
+ "NOUN|[advcl]": 95,
356
+ "NOUN|[ccomp]": 96,
357
+ "NOUN|[compound]": 97,
358
+ "NOUN|[csubj:outer]": 98,
359
+ "NOUN|[csubj]": 99,
360
+ "NOUN|[nmod]": 100,
361
+ "NOUN|[nsubj:outer]": 101,
362
+ "NOUN|[nsubj]": 102,
363
+ "NOUN|[obj]": 103,
364
+ "NOUN|[obl]": 104,
365
+ "NOUN|[root]": 105,
366
+ "NUM": 106,
367
+ "NUM.": 107,
368
+ "NUM.|[advcl]": 108,
369
+ "NUM.|[compound]": 109,
370
+ "NUM.|[nmod]": 110,
371
+ "NUM.|[nsubj:outer]": 111,
372
+ "NUM.|[nsubj]": 112,
373
+ "NUM.|[nummod]": 113,
374
+ "NUM.|[obj]": 114,
375
+ "NUM.|[obl]": 115,
376
+ "NUM.|[root]": 116,
377
+ "NUM|[advcl]": 117,
378
+ "NUM|[compound]": 118,
379
+ "NUM|[nmod]": 119,
380
+ "NUM|[nsubj:outer]": 120,
381
+ "NUM|[nsubj]": 121,
382
+ "NUM|[nummod]": 122,
383
+ "NUM|[obj]": 123,
384
+ "NUM|[obl]": 124,
385
+ "NUM|[root]": 125,
386
+ "PART": 126,
387
+ "PART.": 127,
388
+ "PART.|[mark]": 128,
389
+ "PART|[mark]": 129,
390
+ "PRON": 130,
391
+ "PRON.": 131,
392
+ "PRON.|[acl]": 132,
393
+ "PRON.|[advcl]": 133,
394
+ "PRON.|[nmod]": 134,
395
+ "PRON.|[nsubj:outer]": 135,
396
+ "PRON.|[nsubj]": 136,
397
+ "PRON.|[obj]": 137,
398
+ "PRON.|[obl]": 138,
399
+ "PRON.|[root]": 139,
400
+ "PRON|[acl]": 140,
401
+ "PRON|[advcl]": 141,
402
+ "PRON|[nmod]": 142,
403
+ "PRON|[nsubj:outer]": 143,
404
+ "PRON|[nsubj]": 144,
405
+ "PRON|[obj]": 145,
406
+ "PRON|[obl]": 146,
407
+ "PRON|[root]": 147,
408
+ "PROPN": 148,
409
+ "PROPN.": 149,
410
+ "PROPN.|[acl]": 150,
411
+ "PROPN.|[advcl]": 151,
412
+ "PROPN.|[compound]": 152,
413
+ "PROPN.|[nmod]": 153,
414
+ "PROPN.|[nsubj:outer]": 154,
415
+ "PROPN.|[nsubj]": 155,
416
+ "PROPN.|[obj]": 156,
417
+ "PROPN.|[obl]": 157,
418
+ "PROPN.|[root]": 158,
419
+ "PROPN|[acl]": 159,
420
+ "PROPN|[advcl]": 160,
421
+ "PROPN|[compound]": 161,
422
+ "PROPN|[nmod]": 162,
423
+ "PROPN|[nsubj:outer]": 163,
424
+ "PROPN|[nsubj]": 164,
425
+ "PROPN|[obj]": 165,
426
+ "PROPN|[obl]": 166,
427
+ "PROPN|[root]": 167,
428
+ "PUNCT": 168,
429
+ "PUNCT.": 169,
430
+ "PUNCT.|[punct]": 170,
431
+ "PUNCT|[punct]": 171,
432
+ "SCONJ": 172,
433
+ "SCONJ.": 173,
434
+ "SCONJ.|[dep]": 174,
435
+ "SCONJ.|[fixed]": 175,
436
+ "SCONJ.|[mark]": 176,
437
+ "SCONJ|[dep]": 177,
438
+ "SCONJ|[fixed]": 178,
439
+ "SCONJ|[mark]": 179,
440
+ "SYM": 180,
441
+ "SYM.": 181,
442
+ "SYM.|[compound]": 182,
443
+ "SYM.|[dep]": 183,
444
+ "SYM.|[nmod]": 184,
445
+ "SYM.|[obl]": 185,
446
+ "SYM|[compound]": 186,
447
+ "SYM|[dep]": 187,
448
+ "SYM|[nmod]": 188,
449
+ "SYM|[obl]": 189,
450
+ "VERB": 190,
451
+ "VERB.": 191,
452
+ "VERB.|[acl]": 192,
453
+ "VERB.|[advcl]": 193,
454
+ "VERB.|[ccomp]": 194,
455
+ "VERB.|[compound]": 195,
456
+ "VERB.|[csubj:outer]": 196,
457
+ "VERB.|[csubj]": 197,
458
+ "VERB.|[nmod]": 198,
459
+ "VERB.|[obj]": 199,
460
+ "VERB.|[obl]": 200,
461
+ "VERB.|[root]": 201,
462
+ "VERB|[acl]": 202,
463
+ "VERB|[advcl]": 203,
464
+ "VERB|[ccomp]": 204,
465
+ "VERB|[compound]": 205,
466
+ "VERB|[csubj:outer]": 206,
467
+ "VERB|[csubj]": 207,
468
+ "VERB|[nmod]": 208,
469
+ "VERB|[obj]": 209,
470
+ "VERB|[obl]": 210,
471
+ "VERB|[root]": 211,
472
+ "X": 212,
473
+ "X.": 213,
474
+ "X.|[dep]": 214,
475
+ "X.|[goeswith]": 215,
476
+ "X.|[nmod]": 216,
477
+ "X|[dep]": 217,
478
+ "X|[nmod]": 218
479
+ },
480
+ "layer_norm_eps": 1e-05,
481
+ "local_attention": 128,
482
+ "local_rope_theta": 10000.0,
483
+ "max_position_embeddings": 8192,
484
+ "mlp_bias": false,
485
+ "mlp_dropout": 0.0,
486
+ "model_type": "modernbert",
487
+ "norm_bias": false,
488
+ "norm_eps": 1e-05,
489
+ "num_attention_heads": 8,
490
+ "num_hidden_layers": 16,
491
+ "pad_token_id": 1,
492
+ "position_embedding_type": "absolute",
493
+ "reference_compile": true,
494
+ "sep_token_id": 2,
495
+ "sparse_pred_ignore_index": -100,
496
+ "sparse_prediction": false,
497
+ "tokenizer_class": "DebertaV2TokenizerFast",
498
+ "torch_dtype": "float32",
499
+ "transformers_version": "4.47.1",
500
+ "vocab_size": 65000
501
+ }
configuration_modernbert.py ADDED
@@ -0,0 +1,213 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
2
+ # This file was automatically generated from src/transformers/models/modernbert/modular_modernbert.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_modernbert.py file directly. One of our CI enforces this.
6
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
7
+ # Copyright 2024 Answer.AI, LightOn, and contributors, and the HuggingFace Inc. team. All rights reserved.
8
+ #
9
+ #
10
+ # Licensed under the Apache License, Version 2.0 (the "License");
11
+ # you may not use this file except in compliance with the License.
12
+ # You may obtain a copy of the License at
13
+ #
14
+ # http://www.apache.org/licenses/LICENSE-2.0
15
+ #
16
+ # Unless required by applicable law or agreed to in writing, software
17
+ # distributed under the License is distributed on an "AS IS" BASIS,
18
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
19
+ # See the License for the specific language governing permissions and
20
+ # limitations under the License.
21
+
22
+ from typing import Literal
23
+
24
+ from transformers.configuration_utils import PretrainedConfig
25
+
26
+
27
+ class ModernBertConfig(PretrainedConfig):
28
+ r"""
29
+ This is the configuration class to store the configuration of a [`ModernBertModel`]. It is used to instantiate an ModernBert
30
+ model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
31
+ defaults will yield a similar configuration to that of the ModernBERT-base.
32
+ e.g. [answerdotai/ModernBERT-base](https://huggingface.co/answerdotai/ModernBERT-base)
33
+
34
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
35
+ documentation from [`PretrainedConfig`] for more information.
36
+
37
+ Args:
38
+ vocab_size (`int`, *optional*, defaults to 50368):
39
+ Vocabulary size of the ModernBert model. Defines the number of different tokens that can be represented by the
40
+ `inputs_ids` passed when calling [`ModernBertModel`]
41
+ hidden_size (`int`, *optional*, defaults to 768):
42
+ Dimension of the hidden representations.
43
+ intermediate_size (`int`, *optional*, defaults to 1152):
44
+ Dimension of the MLP representations.
45
+ num_hidden_layers (`int`, *optional*, defaults to 22):
46
+ Number of hidden layers in the Transformer decoder.
47
+ num_attention_heads (`int`, *optional*, defaults to 12):
48
+ Number of attention heads for each attention layer in the Transformer decoder.
49
+ hidden_activation (`str` or `function`, *optional*, defaults to `"gelu"`):
50
+ The non-linear activation function (function or string) in the decoder. Will default to `"gelu"`
51
+ if not specified.
52
+ max_position_embeddings (`int`, *optional*, defaults to 8192):
53
+ The maximum sequence length that this model might ever be used with.
54
+ initializer_range (`float`, *optional*, defaults to 0.02):
55
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
56
+ initializer_cutoff_factor (`float`, *optional*, defaults to 2.0):
57
+ The cutoff factor for the truncated_normal_initializer for initializing all weight matrices.
58
+ norm_eps (`float`, *optional*, defaults to 1e-05):
59
+ The epsilon used by the rms normalization layers.
60
+ norm_bias (`bool`, *optional*, defaults to `False`):
61
+ Whether to use bias in the normalization layers.
62
+ pad_token_id (`int`, *optional*, defaults to 50283):
63
+ Padding token id.
64
+ eos_token_id (`int`, *optional*, defaults to 50282):
65
+ End of stream token id.
66
+ bos_token_id (`int`, *optional*, defaults to 50281):
67
+ Beginning of stream token id.
68
+ cls_token_id (`int`, *optional*, defaults to 50281):
69
+ Classification token id.
70
+ sep_token_id (`int`, *optional*, defaults to 50282):
71
+ Separation token id.
72
+ global_rope_theta (`float`, *optional*, defaults to 160000.0):
73
+ The base period of the global RoPE embeddings.
74
+ attention_bias (`bool`, *optional*, defaults to `False`):
75
+ Whether to use a bias in the query, key, value and output projection layers during self-attention.
76
+ attention_dropout (`float`, *optional*, defaults to 0.0):
77
+ The dropout ratio for the attention probabilities.
78
+ global_attn_every_n_layers (`int`, *optional*, defaults to 3):
79
+ The number of layers between global attention layers.
80
+ local_attention (`int`, *optional*, defaults to 128):
81
+ The window size for local attention.
82
+ local_rope_theta (`float`, *optional*, defaults to 10000.0):
83
+ The base period of the local RoPE embeddings.
84
+ embedding_dropout (`float`, *optional*, defaults to 0.0):
85
+ The dropout ratio for the embeddings.
86
+ mlp_bias (`bool`, *optional*, defaults to `False`):
87
+ Whether to use bias in the MLP layers.
88
+ mlp_dropout (`float`, *optional*, defaults to 0.0):
89
+ The dropout ratio for the MLP layers.
90
+ decoder_bias (`bool`, *optional*, defaults to `True`):
91
+ Whether to use bias in the decoder layers.
92
+ classifier_pooling (`str`, *optional*, defaults to `"cls"`):
93
+ The pooling method for the classifier. Should be either `"cls"` or `"mean"`. In local attention layers, the
94
+ CLS token doesn't attend to all tokens on long sequences.
95
+ classifier_dropout (`float`, *optional*, defaults to 0.0):
96
+ The dropout ratio for the classifier.
97
+ classifier_bias (`bool`, *optional*, defaults to `False`):
98
+ Whether to use bias in the classifier.
99
+ classifier_activation (`str`, *optional*, defaults to `"gelu"`):
100
+ The activation function for the classifier.
101
+ deterministic_flash_attn (`bool`, *optional*, defaults to `False`):
102
+ Whether to use deterministic flash attention. If `False`, inference will be faster but not deterministic.
103
+ sparse_prediction (`bool`, *optional*, defaults to `False`):
104
+ Whether to use sparse prediction for the masked language model instead of returning the full dense logits.
105
+ sparse_pred_ignore_index (`int`, *optional*, defaults to -100):
106
+ The index to ignore for the sparse prediction.
107
+ reference_compile (`bool`, *optional*):
108
+ Whether to compile the layers of the model which were compiled during pretraining. If `None`, then parts of
109
+ the model will be compiled if 1) `triton` is installed, 2) the model is not on MPS, 3) the model is not
110
+ shared between devices, and 4) the model is not resized after initialization. If `True`, then the model may
111
+ be faster in some scenarios.
112
+
113
+ Examples:
114
+
115
+ ```python
116
+ >>> from transformers import ModernBertModel, ModernBertConfig
117
+
118
+ >>> # Initializing a ModernBert style configuration
119
+ >>> configuration = ModernBertConfig()
120
+
121
+ >>> # Initializing a model from the modernbert-base style configuration
122
+ >>> model = ModernBertModel(configuration)
123
+
124
+ >>> # Accessing the model configuration
125
+ >>> configuration = model.config
126
+ ```"""
127
+
128
+ model_type = "modernbert"
129
+ keys_to_ignore_at_inference = ["past_key_values"]
130
+
131
+ def __init__(
132
+ self,
133
+ vocab_size=50368,
134
+ hidden_size=768,
135
+ intermediate_size=1152,
136
+ num_hidden_layers=22,
137
+ num_attention_heads=12,
138
+ hidden_activation="gelu",
139
+ max_position_embeddings=8192,
140
+ initializer_range=0.02,
141
+ initializer_cutoff_factor=2.0,
142
+ norm_eps=1e-5,
143
+ norm_bias=False,
144
+ pad_token_id=50283,
145
+ eos_token_id=50282,
146
+ bos_token_id=50281,
147
+ cls_token_id=50281,
148
+ sep_token_id=50282,
149
+ global_rope_theta=160000.0,
150
+ attention_bias=False,
151
+ attention_dropout=0.0,
152
+ global_attn_every_n_layers=3,
153
+ local_attention=128,
154
+ local_rope_theta=10000.0,
155
+ embedding_dropout=0.0,
156
+ mlp_bias=False,
157
+ mlp_dropout=0.0,
158
+ decoder_bias=True,
159
+ classifier_pooling: Literal["cls", "mean"] = "cls",
160
+ classifier_dropout=0.0,
161
+ classifier_bias=False,
162
+ classifier_activation="gelu",
163
+ deterministic_flash_attn=False,
164
+ sparse_prediction=False,
165
+ sparse_pred_ignore_index=-100,
166
+ reference_compile=None,
167
+ **kwargs,
168
+ ):
169
+ super().__init__(
170
+ pad_token_id=pad_token_id,
171
+ bos_token_id=bos_token_id,
172
+ eos_token_id=eos_token_id,
173
+ cls_token_id=cls_token_id,
174
+ sep_token_id=sep_token_id,
175
+ **kwargs,
176
+ )
177
+ self.vocab_size = vocab_size
178
+ self.max_position_embeddings = max_position_embeddings
179
+ self.hidden_size = hidden_size
180
+ self.intermediate_size = intermediate_size
181
+ self.num_hidden_layers = num_hidden_layers
182
+ self.num_attention_heads = num_attention_heads
183
+ self.initializer_range = initializer_range
184
+ self.initializer_cutoff_factor = initializer_cutoff_factor
185
+ self.norm_eps = norm_eps
186
+ self.norm_bias = norm_bias
187
+ self.global_rope_theta = global_rope_theta
188
+ self.attention_bias = attention_bias
189
+ self.attention_dropout = attention_dropout
190
+ self.hidden_activation = hidden_activation
191
+ self.global_attn_every_n_layers = global_attn_every_n_layers
192
+ self.local_attention = local_attention
193
+ self.local_rope_theta = local_rope_theta
194
+ self.embedding_dropout = embedding_dropout
195
+ self.mlp_bias = mlp_bias
196
+ self.mlp_dropout = mlp_dropout
197
+ self.decoder_bias = decoder_bias
198
+ self.classifier_pooling = classifier_pooling
199
+ self.classifier_dropout = classifier_dropout
200
+ self.classifier_bias = classifier_bias
201
+ self.classifier_activation = classifier_activation
202
+ self.deterministic_flash_attn = deterministic_flash_attn
203
+ self.sparse_prediction = sparse_prediction
204
+ self.sparse_pred_ignore_index = sparse_pred_ignore_index
205
+ self.reference_compile = reference_compile
206
+
207
+ if self.classifier_pooling not in ["cls", "mean"]:
208
+ raise ValueError(
209
+ f'Invalid value for `classifier_pooling`, should be either "cls" or "mean", but is {self.classifier_pooling}.'
210
+ )
211
+
212
+
213
+ __all__ = ["ModernBertConfig"]
maker.py ADDED
@@ -0,0 +1,97 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #! /usr/bin/python3
2
+ src="KoichiYasuoka/modernbert-small-japanese-wikipedia-upos"
3
+ tgt="KoichiYasuoka/modernbert-small-japanese-wikipedia-ud-square"
4
+ url="https://github.com/UniversalDependencies/UD_Japanese-GSDLUW"
5
+ import os,numpy
6
+ d=os.path.basename(url)
7
+ os.system("test -d "+d+" || git clone --depth=1 "+url)
8
+ os.system("for F in train dev test ; do cp "+d+"/*-$F.conllu $F.conllu ; done")
9
+ class UDSquareDataset(object):
10
+ def __init__(self,conllu,tokenizer):
11
+ self.conllu=open(conllu,"r",encoding="utf-8")
12
+ self.tokenizer=tokenizer
13
+ self.seeks=[0]
14
+ label=set(["SYM.","X.","X.|[goeswith]"])
15
+ s=self.conllu.readline()
16
+ while s!="":
17
+ if s=="\n":
18
+ self.seeks.append(self.conllu.tell())
19
+ else:
20
+ w=s.split("\t")
21
+ if len(w)==10:
22
+ if w[0].isdecimal():
23
+ p=w[3]
24
+ q="" if w[5]=="_" else "|"+w[5]
25
+ r="|["+w[7]+"]"
26
+ label.add(p+q)
27
+ label.add(p+"."+q)
28
+ label.add(p+q+r)
29
+ label.add(p+"."+q+r)
30
+ s=self.conllu.readline()
31
+ self.label2id={l:i for i,l in enumerate(sorted(label))}
32
+ def __call__(*args):
33
+ lid={l:i for i,l in enumerate(sorted(set(sum([list(t.label2id) for t in args],[]))))}
34
+ for t in args:
35
+ t.label2id=lid
36
+ return lid
37
+ def __del__(self):
38
+ self.conllu.close()
39
+ __len__=lambda self:len(self.seeks)-1
40
+ def __getitem__(self,i):
41
+ self.conllu.seek(self.seeks[i])
42
+ c,t=[],[""]
43
+ while t[0]!="\n":
44
+ t=self.conllu.readline().split("\t")
45
+ if len(t)==10 and t[0].isdecimal():
46
+ c.append(t)
47
+ h={t[6] for t in c}
48
+ for t in c:
49
+ if t[6]!="0" and t[0] not in h:
50
+ t[3]+="."
51
+ v=self.tokenizer([t[1] for t in c],add_special_tokens=False)["input_ids"]
52
+ for i in range(len(v)-1,-1,-1):
53
+ if v[i]==[]:
54
+ v[i]=[self.tokenizer.unk_token_id]
55
+ if len(v[i])>1:
56
+ c[i][3]=c[i][3].replace(".","")
57
+ for j in range(1,len(v[i])):
58
+ c.insert(i+1,[c[i][0],"_","_","X.","_","_",c[i][0],"goeswith","_","_"])
59
+ y=["0"]+[t[0] for t in c]
60
+ h=[i if t[6]=="0" else y.index(t[6]) for i,t in enumerate(c,1)]
61
+ p=[t[3] if t[5]=="_" else t[3]+"|"+t[5] for t in c]
62
+ d=["|["+t[7]+"]" for t in c]
63
+ x=[not t[3].endswith(".") for t in c]
64
+ if len(x)<90:
65
+ x=[True]*len(x)
66
+ else:
67
+ w=(sum([1 for b in x if b])+1)*(len(x)+1)+1
68
+ for i in numpy.argsort([-abs(j-i-1) for i,j in enumerate(h)]):
69
+ if w+len(x)>8191:
70
+ break
71
+ if not x[i]:
72
+ x[i]=True
73
+ w+=len(x)+1
74
+ v=sum(v,[])
75
+ ids=[self.tokenizer.cls_token_id]+v+[self.tokenizer.sep_token_id]
76
+ upos=["SYM."]+p+["SYM."]
77
+ for i in range(len(v)):
78
+ if x[i]:
79
+ for j in range(len(v)):
80
+ ids.append(self.tokenizer.mask_token_id if i==j else v[j])
81
+ upos.append(p[j]+d[j] if h[j]==i+1 else p[j])
82
+ ids.append(self.tokenizer.sep_token_id)
83
+ upos.append("SYM.")
84
+ return {"input_ids":ids,"labels":[self.label2id[p] for p in upos]}
85
+ from transformers import AutoTokenizer,AutoConfig,AutoModelForTokenClassification,DataCollatorForTokenClassification,TrainingArguments,Trainer
86
+ tkz=AutoTokenizer.from_pretrained(src)
87
+ trainDS=UDSquareDataset("train.conllu",tkz)
88
+ devDS=UDSquareDataset("dev.conllu",tkz)
89
+ testDS=UDSquareDataset("test.conllu",tkz)
90
+ lid=trainDS(devDS,testDS)
91
+ cfg=AutoConfig.from_pretrained(src,num_labels=len(lid),label2id=lid,id2label={i:l for l,i in lid.items()},ignore_mismatched_sizes=True,trust_remote_code=True)
92
+ mdl=AutoModelForTokenClassification.from_pretrained(src,config=cfg,ignore_mismatched_sizes=True,trust_remote_code=True)
93
+ arg=TrainingArguments(num_train_epochs=3,per_device_train_batch_size=1,dataloader_pin_memory=False,output_dir=tgt,overwrite_output_dir=True,save_total_limit=2,learning_rate=5e-05,warmup_ratio=0.1,save_safetensors=False)
94
+ trn=Trainer(args=arg,data_collator=DataCollatorForTokenClassification(tkz),model=mdl,train_dataset=trainDS)
95
+ trn.train()
96
+ trn.save_model(tgt)
97
+ tkz.save_pretrained(tgt)
modeling_modernbert.py ADDED
@@ -0,0 +1,1351 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
2
+ # This file was automatically generated from src/transformers/models/modernbert/modular_modernbert.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_modernbert.py file directly. One of our CI enforces this.
6
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
7
+ # Copyright 2024 Answer.AI, LightOn, and contributors, and the HuggingFace Inc. team. All rights reserved.
8
+ #
9
+ #
10
+ # Licensed under the Apache License, Version 2.0 (the "License");
11
+ # you may not use this file except in compliance with the License.
12
+ # You may obtain a copy of the License at
13
+ #
14
+ # http://www.apache.org/licenses/LICENSE-2.0
15
+ #
16
+ # Unless required by applicable law or agreed to in writing, software
17
+ # distributed under the License is distributed on an "AS IS" BASIS,
18
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
19
+ # See the License for the specific language governing permissions and
20
+ # limitations under the License.
21
+
22
+ import math
23
+ from typing import Dict, Optional, Tuple, Union
24
+
25
+ import torch
26
+ import torch.nn.functional as F
27
+ from torch import nn
28
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
29
+
30
+ from transformers.activations import ACT2FN
31
+ from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
32
+ from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput, SequenceClassifierOutput, TokenClassifierOutput
33
+ from transformers.modeling_utils import PreTrainedModel
34
+ from transformers.utils import (
35
+ add_code_sample_docstrings,
36
+ add_start_docstrings,
37
+ add_start_docstrings_to_model_forward,
38
+ is_flash_attn_2_available,
39
+ logging,
40
+ )
41
+ import importlib
42
+ is_triton_available = lambda: importlib.util.find_spec("triton") is not None
43
+ from .configuration_modernbert import ModernBertConfig
44
+
45
+
46
+ if is_flash_attn_2_available():
47
+ from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func
48
+ from flash_attn.layers.rotary import RotaryEmbedding
49
+ from flash_attn.ops.triton.rotary import apply_rotary
50
+ else:
51
+ RotaryEmbedding = object
52
+
53
+ logger = logging.get_logger(__name__)
54
+
55
+ _CHECKPOINT_FOR_DOC = "answerdotai/ModernBERT-base"
56
+ _CONFIG_FOR_DOC = "ModernBertConfig"
57
+
58
+
59
+ class ApplyRotaryEmbUnpad(torch.autograd.Function):
60
+ @staticmethod
61
+ def forward(
62
+ ctx,
63
+ qkv,
64
+ cos,
65
+ sin,
66
+ cu_seqlens: Optional[torch.Tensor] = None,
67
+ max_seqlen: Optional[int] = None,
68
+ ):
69
+ # (total_nnz, 3, nheads, headdim)
70
+ qkv = qkv.contiguous()
71
+ total_nnz, _three, _nheads, headdim = qkv.shape
72
+ # We need qkv to be contiguous so that when we reshape to combine (3, nheads) dimensions,
73
+ # we get the same tensor
74
+ # qk = rearrange(qkv[:, :2], "b_s t h d -> b_s (t h) d")
75
+ qk = qkv[:, :2].view(total_nnz, -1, headdim)
76
+ apply_rotary(
77
+ qk,
78
+ cos,
79
+ sin,
80
+ seqlen_offsets=0,
81
+ cu_seqlens=cu_seqlens,
82
+ max_seqlen=max_seqlen,
83
+ interleaved=False,
84
+ inplace=True,
85
+ )
86
+
87
+ ctx.save_for_backward(cos, sin, cu_seqlens)
88
+ ctx.max_seqlen = max_seqlen
89
+ return qkv
90
+
91
+ @staticmethod
92
+ def backward(ctx, do):
93
+ cos, sin, cu_seqlens = ctx.saved_tensors
94
+ do = do.contiguous()
95
+ total_nnz, _three, _nheads, headdim = do.shape
96
+ # We need dqkv to be contiguous so that when we reshape to combine (3, nheads) dimensions,
97
+ # we get the same tensor
98
+ dqk = do[:, :2].view(total_nnz, -1, headdim)
99
+ apply_rotary(
100
+ dqk,
101
+ cos,
102
+ sin,
103
+ seqlen_offsets=0,
104
+ cu_seqlens=cu_seqlens,
105
+ max_seqlen=ctx.max_seqlen,
106
+ interleaved=False,
107
+ inplace=True,
108
+ conjugate=True,
109
+ )
110
+
111
+ return do, None, None, None, None, None, None
112
+
113
+
114
+ def apply_rotary_unpadded(
115
+ qkv,
116
+ cos,
117
+ sin,
118
+ cu_seqlens: Optional[torch.Tensor] = None,
119
+ max_seqlen: Optional[int] = None,
120
+ ):
121
+ """
122
+ Arguments:
123
+ qkv: (total_nnz, 3, nheads, headdim) - input tensor for packed QKV.
124
+ cos, sin: (seqlen_rotary, rotary_dim / 2)
125
+ interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead
126
+ of 1st half and 2nd half (GPT-NeoX style).
127
+ inplace: if True, apply rotary embedding in-place.
128
+ seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount.
129
+ Most commonly used in inference when we have KV cache.
130
+ cu_seqlens: (batch + 1,) or None
131
+ max_seqlen: int
132
+ Return:
133
+ out: (total_nnz, dim)
134
+ rotary_dim must be <= headdim
135
+ Apply rotary embedding to the first rotary_dim of x.
136
+ """
137
+ return ApplyRotaryEmbUnpad.apply(qkv, cos, sin, cu_seqlens, max_seqlen)
138
+
139
+
140
+ class ModernBertUnpaddedRotaryEmbedding(RotaryEmbedding):
141
+ """
142
+ The rotary position embeddings applied directly to unpadded sequences.
143
+ """
144
+
145
+ def __init__(
146
+ self,
147
+ dim: int,
148
+ base: float = 10000.0,
149
+ max_seqlen: Optional[int] = None,
150
+ device: Optional[torch.device] = None,
151
+ dtype: Optional[torch.dtype] = None,
152
+ ):
153
+ """
154
+ max_seqlen: if max_seqlen, device, and dtype are provided, we precompute the cos_sin_cache
155
+ up to max_seqlen. If the max_seqlen, device, or dtype during training/inference differ,
156
+ the cos_sin_cache wll be recomputed during the forward pass.
157
+ """
158
+ super().__init__(dim=dim, base=base, pos_idx_in_fp32=True, device=device, interleaved=False)
159
+ self.max_seqlen = max_seqlen
160
+
161
+ if max_seqlen is not None and device is not None and dtype is not None:
162
+ self._update_cos_sin_cache(max_seqlen, device=device, dtype=dtype)
163
+
164
+ def forward(
165
+ self,
166
+ qkv: torch.Tensor,
167
+ cu_seqlens: torch.Tensor,
168
+ max_seqlen: Optional[int] = None,
169
+ ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
170
+ """
171
+ Apply rotary embedding *inplace* to qkv.
172
+ qkv: (total_nnz, 3, nheads, headdim)
173
+ cu_seqlens: (batch + 1,) cumulative sequence lengths
174
+ max_seqlen: int max seq length in the batch
175
+ """
176
+ if max_seqlen is not None:
177
+ self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype)
178
+
179
+ qkv = apply_rotary_unpadded(
180
+ qkv,
181
+ self._cos_cached,
182
+ self._sin_cached,
183
+ cu_seqlens=cu_seqlens,
184
+ max_seqlen=max_seqlen,
185
+ )
186
+
187
+ return qkv
188
+
189
+ def extra_repr(self) -> str:
190
+ return f"dim={self.dim}, base={self.base}, scale_base={self.scale_base}"
191
+
192
+
193
+ class ModernBertEmbeddings(nn.Module):
194
+ """
195
+ Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.
196
+ """
197
+
198
+ def __init__(self, config: ModernBertConfig):
199
+ super().__init__()
200
+ self.config = config
201
+ self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
202
+ self.norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=config.norm_bias)
203
+ self.drop = nn.Dropout(config.embedding_dropout)
204
+
205
+ @torch.compile(dynamic=True)
206
+ def compiled_embeddings(self, input_ids: torch.LongTensor) -> torch.Tensor:
207
+ return self.drop(self.norm(self.tok_embeddings(input_ids)))
208
+
209
+ def forward(
210
+ self, input_ids: torch.LongTensor = None, inputs_embeds: Optional[torch.Tensor] = None
211
+ ) -> torch.Tensor:
212
+ if inputs_embeds is not None:
213
+ hidden_states = self.drop(self.norm(inputs_embeds))
214
+ else:
215
+ hidden_states = (
216
+ self.compiled_embeddings(input_ids)
217
+ if self.config.reference_compile
218
+ else self.drop(self.norm(self.tok_embeddings(input_ids)))
219
+ )
220
+ return hidden_states
221
+
222
+
223
+ class ModernBertMLP(nn.Module):
224
+ """Applies the GLU at the end of each ModernBERT layer.
225
+
226
+ Compared to the default BERT architecture, this block replaces :class:`~transformers.model.bert.modeling_bert.BertIntermediate`
227
+ and :class:`~transformers.model.bert.modeling_bert.SelfOutput` with a single module that has similar functionality.
228
+ """
229
+
230
+ def __init__(self, config: ModernBertConfig):
231
+ super().__init__()
232
+ self.config = config
233
+ self.Wi = nn.Linear(config.hidden_size, int(config.intermediate_size) * 2, bias=config.mlp_bias)
234
+ self.act = ACT2FN[config.hidden_activation]
235
+ self.drop = nn.Dropout(config.mlp_dropout)
236
+ self.Wo = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.mlp_bias)
237
+
238
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
239
+ input, gate = self.Wi(hidden_states).chunk(2, dim=-1)
240
+ return self.Wo(self.drop(self.act(input) * gate))
241
+
242
+
243
+ class ModernBertRotaryEmbedding(nn.Module):
244
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
245
+ super().__init__()
246
+
247
+ self.dim = dim
248
+ self.max_position_embeddings = max_position_embeddings
249
+ self.base = base
250
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim))
251
+ self.register_buffer("inv_freq", tensor=inv_freq, persistent=False)
252
+
253
+ @torch.no_grad()
254
+ def forward(self, x, position_ids, seq_len=None):
255
+ # x: [bs, num_attention_heads, seq_len, head_size]
256
+ self.inv_freq.to(x.device)
257
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
258
+ position_ids_expanded = position_ids[:, None, :].float()
259
+ # Force float32 since bfloat16 loses precision on long contexts
260
+ # See https://github.com/huggingface/transformers/pull/29285
261
+ device_type = x.device.type
262
+ device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
263
+ with torch.autocast(device_type=device_type, enabled=False):
264
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
265
+ emb = torch.cat((freqs, freqs), dim=-1)
266
+ cos = emb.cos()
267
+ sin = emb.sin()
268
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
269
+
270
+
271
+ def rotate_half(x):
272
+ """Rotates half the hidden dims of the input."""
273
+ x1 = x[..., : x.shape[-1] // 2]
274
+ x2 = x[..., x.shape[-1] // 2 :]
275
+ return torch.cat((-x2, x1), dim=-1)
276
+
277
+
278
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
279
+ """Applies Rotary Position Embedding to the query and key tensors.
280
+
281
+ Args:
282
+ q (`torch.Tensor`): The query tensor.
283
+ k (`torch.Tensor`): The key tensor.
284
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
285
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
286
+ position_ids (`torch.Tensor`, *optional*):
287
+ Deprecated and unused.
288
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
289
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
290
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
291
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
292
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
293
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
294
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
295
+ Returns:
296
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
297
+ """
298
+ cos = cos.unsqueeze(unsqueeze_dim)
299
+ sin = sin.unsqueeze(unsqueeze_dim)
300
+ q_embed = (q * cos) + (rotate_half(q) * sin)
301
+ k_embed = (k * cos) + (rotate_half(k) * sin)
302
+ return q_embed, k_embed
303
+
304
+
305
+ def eager_attention_forward(
306
+ module: "ModernBertAttention",
307
+ qkv: torch.Tensor,
308
+ attention_mask: torch.Tensor,
309
+ sliding_window_mask: torch.Tensor,
310
+ position_ids: Optional[torch.LongTensor],
311
+ local_attention: Tuple[int, int],
312
+ bs: int,
313
+ dim: int,
314
+ output_attentions: Optional[bool] = False,
315
+ **_kwargs,
316
+ ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
317
+ # qkv: [batch_size, seqlen, 3, nheads, headdim]
318
+ cos, sin = module.rotary_emb(qkv, position_ids=position_ids)
319
+ query, key, value = qkv.transpose(3, 1).unbind(dim=2)
320
+ # query, key, value: [batch_size, heads, seq_len, head_dim]
321
+ query, key = apply_rotary_pos_emb(query, key, cos, sin)
322
+
323
+ scale = module.head_dim**-0.5
324
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scale
325
+
326
+ if local_attention != (-1, -1):
327
+ attention_mask = sliding_window_mask
328
+
329
+ attn_weights = attn_weights + attention_mask
330
+
331
+ # upcast attention to fp32
332
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
333
+ attn_weights = nn.functional.dropout(attn_weights, p=module.attention_dropout, training=module.training)
334
+ attn_output = torch.matmul(attn_weights, value)
335
+ attn_output = attn_output.transpose(1, 2).contiguous()
336
+ attn_output = attn_output.view(bs, -1, dim)
337
+ if output_attentions:
338
+ return (attn_output, attn_weights)
339
+ return (attn_output,)
340
+
341
+
342
+ def flash_attention_forward(
343
+ module: "ModernBertAttention",
344
+ qkv: torch.Tensor,
345
+ rotary_emb: ModernBertUnpaddedRotaryEmbedding,
346
+ cu_seqlens: torch.Tensor,
347
+ max_seqlen: int,
348
+ local_attention: Tuple[int, int],
349
+ bs: int,
350
+ dim: int,
351
+ target_dtype: torch.dtype = torch.bfloat16,
352
+ **_kwargs,
353
+ ) -> Tuple[torch.Tensor]:
354
+ # (total_seqlen, 3, nheads, headdim)
355
+ qkv = rotary_emb(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
356
+
357
+ convert_dtype = qkv.dtype not in (torch.float16, torch.bfloat16)
358
+ if convert_dtype:
359
+ # FA2 implementation only supports fp16 and bf16. If FA2 is supported,
360
+ # bfloat16 must be supported as of FA2 2.5.7. (Turing GPUs not supported)
361
+ orig_dtype = qkv.dtype
362
+ qkv = qkv.to(target_dtype)
363
+
364
+ attn = flash_attn_varlen_qkvpacked_func(
365
+ qkv,
366
+ cu_seqlens=cu_seqlens,
367
+ max_seqlen=max_seqlen,
368
+ dropout_p=module.attention_dropout if module.training else 0.0,
369
+ deterministic=module.deterministic_flash_attn,
370
+ window_size=local_attention,
371
+ )
372
+ attn = attn.to(orig_dtype) # type: ignore
373
+ else:
374
+ attn = flash_attn_varlen_qkvpacked_func(
375
+ qkv,
376
+ cu_seqlens=cu_seqlens,
377
+ max_seqlen=max_seqlen,
378
+ dropout_p=module.attention_dropout if module.training else 0.0,
379
+ deterministic=module.deterministic_flash_attn,
380
+ window_size=local_attention,
381
+ )
382
+ return (attn.view(bs, dim),)
383
+
384
+
385
+ def sdpa_attention_forward(
386
+ module: "ModernBertAttention",
387
+ qkv: torch.Tensor,
388
+ attention_mask: torch.Tensor,
389
+ sliding_window_mask: torch.Tensor,
390
+ position_ids: Optional[torch.LongTensor],
391
+ local_attention: Tuple[int, int],
392
+ bs: int,
393
+ dim: int,
394
+ **_kwargs,
395
+ ) -> Tuple[torch.Tensor]:
396
+ # qkv: [batch_size, seqlen, 3, nheads, headdim]
397
+ cos, sin = module.rotary_emb(qkv, position_ids=position_ids)
398
+ query, key, value = qkv.transpose(3, 1).unbind(dim=2)
399
+ # query, key, value: [batch_size, heads, seq_len, head_dim]
400
+ query, key = apply_rotary_pos_emb(query, key, cos, sin)
401
+
402
+ if local_attention != (-1, -1):
403
+ attention_mask = sliding_window_mask
404
+
405
+ attn_output = (
406
+ F.scaled_dot_product_attention(
407
+ query,
408
+ key,
409
+ value,
410
+ dropout_p=module.attention_dropout if module.training else 0.0,
411
+ attn_mask=attention_mask,
412
+ )
413
+ .transpose(1, 2)
414
+ .contiguous()
415
+ )
416
+ attn_output = attn_output.view(bs, -1, dim)
417
+ return (attn_output,)
418
+
419
+
420
+ MODERNBERT_ATTENTION_FUNCTION = {
421
+ "flash_attention_2": flash_attention_forward,
422
+ "eager": eager_attention_forward,
423
+ "sdpa": sdpa_attention_forward,
424
+ }
425
+
426
+
427
+ class ModernBertAttention(nn.Module):
428
+ """Performs multi-headed self attention on a batch of unpadded sequences.
429
+
430
+ If Flash Attention 2 is installed, this module uses Flash Attention to improve throughput.
431
+ If Flash Attention 2 is not installed, the implementation will use PyTorch's SDPA kernel,
432
+ which requires padding and unpadding inputs, adding some overhead.
433
+
434
+ See `forward` method for additional details.
435
+ """
436
+
437
+ def __init__(self, config: ModernBertConfig, layer_id: Optional[int] = None):
438
+ super().__init__()
439
+ self.config = config
440
+ self.layer_id = layer_id
441
+
442
+ if config.hidden_size % config.num_attention_heads != 0:
443
+ raise ValueError(
444
+ f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention heads ({config.num_attention_heads})"
445
+ )
446
+
447
+ self.attention_dropout = config.attention_dropout
448
+ self.deterministic_flash_attn = config.deterministic_flash_attn
449
+ self.num_heads = config.num_attention_heads
450
+ self.head_dim = config.hidden_size // config.num_attention_heads
451
+ self.all_head_size = self.head_dim * self.num_heads
452
+ self.Wqkv = nn.Linear(config.hidden_size, 3 * self.all_head_size, bias=config.attention_bias)
453
+
454
+ if layer_id % config.global_attn_every_n_layers != 0:
455
+ self.local_attention = (config.local_attention // 2, config.local_attention // 2)
456
+ else:
457
+ self.local_attention = (-1, -1)
458
+
459
+ rope_theta = config.global_rope_theta
460
+ max_position_embeddings = config.max_position_embeddings
461
+ if self.local_attention != (-1, -1):
462
+ if config.local_rope_theta is not None:
463
+ rope_theta = config.local_rope_theta
464
+ max_position_embeddings = config.local_attention
465
+
466
+ if config._attn_implementation == "flash_attention_2":
467
+ self.rotary_emb = ModernBertUnpaddedRotaryEmbedding(
468
+ dim=self.head_dim, max_seqlen=max_position_embeddings, base=rope_theta
469
+ )
470
+ else:
471
+ self.rotary_emb = ModernBertRotaryEmbedding(
472
+ dim=self.head_dim, max_position_embeddings=max_position_embeddings, base=rope_theta
473
+ )
474
+
475
+ self.Wo = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias)
476
+ self.out_drop = nn.Dropout(config.attention_dropout) if config.attention_dropout > 0.0 else nn.Identity()
477
+ self.pruned_heads = set()
478
+
479
+ def forward(
480
+ self,
481
+ hidden_states: torch.Tensor,
482
+ output_attentions: Optional[bool] = False,
483
+ **kwargs,
484
+ ) -> torch.Tensor:
485
+ qkv = self.Wqkv(hidden_states)
486
+
487
+ bs = hidden_states.shape[0]
488
+ if self.config._attn_implementation == "flash_attention_2":
489
+ qkv = qkv.view(-1, 3, self.num_heads, self.head_dim)
490
+ else:
491
+ qkv = qkv.view(bs, -1, 3, self.num_heads, self.head_dim)
492
+
493
+ attn_outputs = MODERNBERT_ATTENTION_FUNCTION[self.config._attn_implementation](
494
+ self,
495
+ qkv=qkv,
496
+ rotary_emb=self.rotary_emb,
497
+ local_attention=self.local_attention,
498
+ bs=bs,
499
+ dim=self.all_head_size,
500
+ output_attentions=output_attentions,
501
+ **kwargs,
502
+ )
503
+ hidden_states = attn_outputs[0]
504
+ hidden_states = self.out_drop(self.Wo(hidden_states))
505
+
506
+ return (hidden_states,) + attn_outputs[1:] # add attentions if outputted
507
+
508
+
509
+ class ModernBertEncoderLayer(nn.Module):
510
+ def __init__(self, config: ModernBertConfig, layer_id: Optional[int] = None):
511
+ super().__init__()
512
+ self.config = config
513
+ if layer_id == 0:
514
+ self.attn_norm = nn.Identity()
515
+ else:
516
+ self.attn_norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=config.norm_bias)
517
+ self.attn = ModernBertAttention(config=config, layer_id=layer_id)
518
+ self.mlp_norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=config.norm_bias)
519
+ self.mlp = ModernBertMLP(config)
520
+
521
+ @torch.compile(dynamic=True)
522
+ def compiled_mlp(self, hidden_states: torch.Tensor) -> torch.Tensor:
523
+ return self.mlp(self.mlp_norm(hidden_states))
524
+
525
+ def forward(
526
+ self,
527
+ hidden_states: torch.Tensor,
528
+ attention_mask: Optional[torch.Tensor] = None,
529
+ sliding_window_mask: Optional[torch.Tensor] = None,
530
+ position_ids: Optional[torch.LongTensor] = None,
531
+ cu_seqlens: Optional[torch.Tensor] = None,
532
+ max_seqlen: Optional[int] = None,
533
+ output_attentions: Optional[bool] = False,
534
+ ) -> torch.Tensor:
535
+ attn_outputs = self.attn(
536
+ self.attn_norm(hidden_states),
537
+ attention_mask=attention_mask,
538
+ sliding_window_mask=sliding_window_mask,
539
+ position_ids=position_ids,
540
+ cu_seqlens=cu_seqlens,
541
+ max_seqlen=max_seqlen,
542
+ output_attentions=output_attentions,
543
+ )
544
+ hidden_states = hidden_states + attn_outputs[0]
545
+ mlp_output = (
546
+ self.compiled_mlp(hidden_states)
547
+ if self.config.reference_compile
548
+ else self.mlp(self.mlp_norm(hidden_states))
549
+ )
550
+ hidden_states = hidden_states + mlp_output
551
+
552
+ return (hidden_states,) + attn_outputs[1:] # add attentions if outputted
553
+
554
+
555
+ MODERNBERT_START_DOCSTRING = r"""
556
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
557
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
558
+ etc.)
559
+
560
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
561
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
562
+ and behavior.
563
+
564
+ Parameters:
565
+ config ([`ModernBertConfig`]):
566
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
567
+ load the weights associated with the model, only the configuration. Check out the
568
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
569
+ """
570
+
571
+
572
+ @add_start_docstrings(
573
+ "The bare ModernBert Model outputting raw hidden-states without any specific head on top.",
574
+ MODERNBERT_START_DOCSTRING,
575
+ )
576
+ class ModernBertPreTrainedModel(PreTrainedModel):
577
+ config_class = ModernBertConfig
578
+ base_model_prefix = "model"
579
+ supports_gradient_checkpointing = True
580
+ _no_split_modules = ["ModernBertEmbeddings", "ModernBertEncoderLayer"]
581
+ _supports_flash_attn_2 = True
582
+ _supports_sdpa = True
583
+ _supports_flex_attn = False
584
+
585
+ def _init_weights(self, module: nn.Module):
586
+ cutoff_factor = self.config.initializer_cutoff_factor
587
+ if cutoff_factor is None:
588
+ cutoff_factor = 3
589
+
590
+ def init_weight(module: nn.Module, std: float):
591
+ nn.init.trunc_normal_(
592
+ module.weight,
593
+ mean=0.0,
594
+ std=std,
595
+ a=-cutoff_factor * std,
596
+ b=cutoff_factor * std,
597
+ )
598
+
599
+ if isinstance(module, nn.Linear):
600
+ if module.bias is not None:
601
+ nn.init.zeros_(module.bias)
602
+
603
+ stds = {
604
+ "in": self.config.initializer_range,
605
+ "out": self.config.initializer_range / math.sqrt(2.0 * self.config.num_hidden_layers),
606
+ "embedding": self.config.initializer_range,
607
+ "final_out": self.config.hidden_size**-0.5,
608
+ }
609
+
610
+ if isinstance(module, ModernBertEmbeddings):
611
+ init_weight(module.tok_embeddings, stds["embedding"])
612
+ elif isinstance(module, ModernBertMLP):
613
+ init_weight(module.Wi, stds["in"])
614
+ init_weight(module.Wo, stds["out"])
615
+ elif isinstance(module, ModernBertAttention):
616
+ init_weight(module.Wqkv, stds["in"])
617
+ init_weight(module.Wo, stds["out"])
618
+ elif isinstance(module, ModernBertPredictionHead):
619
+ init_weight(module.dense, stds["out"])
620
+ elif isinstance(module, ModernBertForMaskedLM):
621
+ init_weight(module.decoder, stds["out"])
622
+ elif isinstance(module, (ModernBertForSequenceClassification, ModernBertForTokenClassification)):
623
+ init_weight(module.classifier, stds["final_out"])
624
+
625
+ @classmethod
626
+ def _autoset_attn_implementation(
627
+ cls,
628
+ config,
629
+ use_flash_attention_2: bool = False,
630
+ torch_dtype: Optional[torch.dtype] = None,
631
+ device_map: Optional[Union[str, Dict[str, int]]] = None,
632
+ check_device_map: bool = True,
633
+ ):
634
+ # If the user didn't specify anything, try to use flash_attention_2 if available.
635
+ # Otherwise we fall back to the default SDPA -> Eager from the super() method.
636
+ if config._attn_implementation_internal is None:
637
+ config._attn_implementation_internal = "flash_attention_2"
638
+ try:
639
+ return cls._check_and_enable_flash_attn_2(
640
+ config,
641
+ torch_dtype=torch_dtype,
642
+ device_map=device_map,
643
+ hard_check_only=False,
644
+ check_device_map=check_device_map,
645
+ )
646
+ except (ValueError, ImportError):
647
+ config._attn_implementation_internal = None
648
+ return super()._autoset_attn_implementation(
649
+ config,
650
+ use_flash_attention_2=use_flash_attention_2,
651
+ torch_dtype=torch_dtype,
652
+ device_map=device_map,
653
+ check_device_map=check_device_map,
654
+ )
655
+
656
+ def _maybe_set_compile(self):
657
+ if self.config.reference_compile is False:
658
+ return
659
+
660
+ if hasattr(self, "hf_device_map") and len(self.hf_device_map) > 1:
661
+ if self.config.reference_compile:
662
+ logger.warning_once(
663
+ "If `accelerate` split the model across devices, `torch.compile` will not work. "
664
+ "Falling back to non-compiled mode."
665
+ )
666
+ self.config.reference_compile = False
667
+
668
+ if self.device.type == "mps":
669
+ if self.config.reference_compile:
670
+ logger.warning_once(
671
+ "Compiling the model with `torch.compile` and using a `torch.mps` device is not supported. "
672
+ "Falling back to non-compiled mode."
673
+ )
674
+ self.config.reference_compile = False
675
+
676
+ if self.config.reference_compile is None:
677
+ self.config.reference_compile = is_triton_available()
678
+
679
+ def resize_token_embeddings(self, *args, **kwargs):
680
+ model_embeds = super().resize_token_embeddings(*args, **kwargs)
681
+
682
+ if self.config.reference_compile in {True, None}:
683
+ if self.config.reference_compile:
684
+ logger.warning_once(
685
+ "Resizing token embeddings with `torch.compile` is not supported. Falling back to non-compiled mode."
686
+ )
687
+ self.config.reference_compile = False
688
+
689
+ return model_embeds
690
+
691
+
692
+ def _unpad_modernbert_input(
693
+ inputs: torch.Tensor,
694
+ attention_mask: torch.Tensor,
695
+ position_ids: Optional[torch.Tensor] = None,
696
+ labels: Optional[torch.Tensor] = None,
697
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, Optional[torch.Tensor], Optional[torch.Tensor]]:
698
+ """
699
+ Remove padding from input sequences.
700
+
701
+ Args:
702
+ inputs: (batch, seqlen, ...) or (batch, seqlen)
703
+ attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
704
+ position_ids: (batch, seqlen), int, position ids
705
+ labels: (batch, seqlen), int, labels
706
+
707
+ Returns:
708
+ unpadded_inputs: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask.
709
+ indices: (total_nnz)
710
+ cu_seqlens: (batch + 1), the cumulative sequence lengths
711
+ max_seqlen_in_batch: int
712
+ unpadded_position_ids: (total_nnz) or None
713
+ unpadded_labels: (total_nnz) or None
714
+ """
715
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
716
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
717
+ max_seqlen_in_batch = int(seqlens_in_batch.max().item())
718
+ cu_seqlens = torch.nn.functional.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
719
+
720
+ if inputs.dim() == 2:
721
+ unpadded_inputs = inputs.flatten()[indices]
722
+ else:
723
+ batch, seqlen, *rest = inputs.shape
724
+ shape = batch * seqlen
725
+ unpadded_inputs = inputs.view(shape, *rest)[indices]
726
+
727
+ unpadded_position_ids = position_ids.flatten()[indices] if position_ids is not None else None
728
+ unpadded_labels = labels.flatten()[indices] if labels is not None else None
729
+
730
+ return unpadded_inputs, indices, cu_seqlens, max_seqlen_in_batch, unpadded_position_ids, unpadded_labels
731
+
732
+
733
+ def _pad_modernbert_output(
734
+ inputs: torch.Tensor,
735
+ indices: torch.Tensor,
736
+ batch: int,
737
+ seqlen: int,
738
+ ) -> torch.Tensor:
739
+ """
740
+ Add padding to sequences.
741
+
742
+ Args:
743
+ inputs: (total_nnz, ...) or (total_nnz,), where total_nnz = number of tokens selected in attention_mask.
744
+ indices: (total_nnz)
745
+ batch: int, batch size
746
+ seqlen: int, max sequence length
747
+
748
+ Returns:
749
+ padded_inputs: (batch, seqlen, ...) or (batch, seqlen)
750
+ """
751
+ if inputs.dim() == 1:
752
+ output = torch.zeros(batch * seqlen, dtype=inputs.dtype, device=inputs.device)
753
+ output[indices] = inputs
754
+ padded_inputs = output.view(batch, seqlen)
755
+ else:
756
+ _, *rest = inputs.shape
757
+ output = torch.zeros(batch * seqlen, *rest, dtype=inputs.dtype, device=inputs.device)
758
+ output[indices] = inputs
759
+ padded_inputs = output.view(batch, seqlen, *rest)
760
+
761
+ return padded_inputs
762
+
763
+
764
+ MODERNBERT_INPUTS_DOCSTRING = r"""
765
+ Args:
766
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
767
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
768
+ it.
769
+
770
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
771
+ [`PreTrainedTokenizer.__call__`] for details.
772
+
773
+ [What are input IDs?](../glossary#input-ids)
774
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
775
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
776
+
777
+ - 1 for tokens that are **not masked**,
778
+ - 0 for tokens that are **masked**.
779
+
780
+ [What are attention masks?](../glossary#attention-mask)
781
+
782
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
783
+ [`PreTrainedTokenizer.__call__`] for details.
784
+
785
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
786
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
787
+ information on the default strategy.
788
+
789
+ - 1 indicates the head is **not masked**,
790
+ - 0 indicates the head is **masked**.
791
+ sliding_window_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
792
+ Mask to avoid performing attention on padding or far-away tokens. In ModernBert, only every few layers
793
+ perform global attention, while the rest perform local attention. This mask is used to avoid attending to
794
+ far-away tokens in the local attention layers.
795
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
796
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
797
+ config.n_positions - 1]`.
798
+
799
+ [What are position IDs?](../glossary#position-ids)
800
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
801
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
802
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
803
+ model's internal embedding lookup matrix.
804
+ indices (`torch.Tensor` of shape `(total_unpadded_tokens,)`, *optional*):
805
+ Indices of the non-padding tokens in the input sequence. Used for unpadding the output.
806
+ cu_seqlens (`torch.Tensor` of shape `(batch + 1,)`, *optional*):
807
+ Cumulative sequence lengths of the input sequences. Used to index the unpadded tensors.
808
+ max_seqlen (`int`, *optional*):
809
+ Maximum sequence length in the batch. Used to pad the output tensors.
810
+ batch_size (`int`, *optional*):
811
+ Batch size of the input sequences. Used to pad the output tensors.
812
+ seq_len (`int`, *optional*):
813
+ Sequence length of the input sequences. Used to pad the output tensors.
814
+ output_attentions (`bool`, *optional*):
815
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
816
+ tensors for more detail.
817
+ output_hidden_states (`bool`, *optional*):
818
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
819
+ more detail.
820
+ return_dict (`bool`, *optional*):
821
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
822
+ """
823
+
824
+
825
+ @add_start_docstrings(
826
+ "The bare ModernBert Model outputting raw hidden-states without any specific head on top.",
827
+ MODERNBERT_START_DOCSTRING,
828
+ )
829
+ class ModernBertModel(ModernBertPreTrainedModel):
830
+ def __init__(self, config: ModernBertConfig):
831
+ super().__init__(config)
832
+ self.config = config
833
+ self.embeddings = ModernBertEmbeddings(config)
834
+ self.layers = nn.ModuleList(
835
+ [ModernBertEncoderLayer(config, layer_id) for layer_id in range(config.num_hidden_layers)]
836
+ )
837
+ self.final_norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=config.norm_bias)
838
+ self.gradient_checkpointing = False
839
+ self.post_init()
840
+
841
+ def get_input_embeddings(self):
842
+ return self.embeddings.tok_embeddings
843
+
844
+ def set_input_embeddings(self, value):
845
+ self.embeddings.tok_embeddings = value
846
+
847
+ @add_start_docstrings_to_model_forward(MODERNBERT_INPUTS_DOCSTRING)
848
+ @add_code_sample_docstrings(
849
+ checkpoint=_CHECKPOINT_FOR_DOC,
850
+ output_type=BaseModelOutput,
851
+ config_class=_CONFIG_FOR_DOC,
852
+ )
853
+ def forward(
854
+ self,
855
+ input_ids: Optional[torch.LongTensor] = None,
856
+ attention_mask: Optional[torch.Tensor] = None,
857
+ sliding_window_mask: Optional[torch.Tensor] = None,
858
+ position_ids: Optional[torch.LongTensor] = None,
859
+ inputs_embeds: Optional[torch.Tensor] = None,
860
+ indices: Optional[torch.Tensor] = None,
861
+ cu_seqlens: Optional[torch.Tensor] = None,
862
+ max_seqlen: Optional[int] = None,
863
+ batch_size: Optional[int] = None,
864
+ seq_len: Optional[int] = None,
865
+ output_attentions: Optional[bool] = None,
866
+ output_hidden_states: Optional[bool] = None,
867
+ return_dict: Optional[bool] = None,
868
+ ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutput]:
869
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
870
+ output_hidden_states = (
871
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
872
+ )
873
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
874
+
875
+ if (input_ids is None) ^ (inputs_embeds is not None):
876
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
877
+
878
+ all_hidden_states = () if output_hidden_states else None
879
+ all_self_attentions = () if output_attentions else None
880
+
881
+ self._maybe_set_compile()
882
+
883
+ if input_ids is not None:
884
+ self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
885
+
886
+ if batch_size is None and seq_len is None:
887
+ if inputs_embeds is not None:
888
+ batch_size, seq_len = inputs_embeds.shape[:2]
889
+ else:
890
+ batch_size, seq_len = input_ids.shape[:2]
891
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
892
+
893
+ if attention_mask is None:
894
+ attention_mask = torch.ones((batch_size, seq_len), device=device, dtype=torch.bool)
895
+
896
+ repad = False
897
+ if self.config._attn_implementation == "flash_attention_2":
898
+ if indices is None and cu_seqlens is None and max_seqlen is None:
899
+ repad = True
900
+ if inputs_embeds is None:
901
+ with torch.no_grad():
902
+ input_ids, indices, cu_seqlens, max_seqlen, *_ = _unpad_modernbert_input(
903
+ inputs=input_ids, attention_mask=attention_mask
904
+ )
905
+ else:
906
+ inputs_embeds, indices, cu_seqlens, max_seqlen, *_ = _unpad_modernbert_input(
907
+ inputs=inputs_embeds, attention_mask=attention_mask
908
+ )
909
+ else:
910
+ if position_ids is None:
911
+ position_ids = torch.arange(seq_len, device=device).unsqueeze(0)
912
+
913
+ attention_mask, sliding_window_mask = self._update_attention_mask(
914
+ attention_mask, output_attentions=output_attentions
915
+ )
916
+
917
+ hidden_states = self.embeddings(input_ids=input_ids, inputs_embeds=inputs_embeds)
918
+
919
+ for encoder_layer in self.layers:
920
+ if output_hidden_states:
921
+ all_hidden_states = all_hidden_states + (hidden_states,)
922
+
923
+ if self.gradient_checkpointing and self.training:
924
+ layer_outputs = self._gradient_checkpointing_func(
925
+ encoder_layer.__call__,
926
+ hidden_states,
927
+ attention_mask,
928
+ sliding_window_mask,
929
+ position_ids,
930
+ cu_seqlens,
931
+ max_seqlen,
932
+ output_attentions,
933
+ )
934
+ else:
935
+ layer_outputs = encoder_layer(
936
+ hidden_states,
937
+ attention_mask=attention_mask,
938
+ sliding_window_mask=sliding_window_mask,
939
+ position_ids=position_ids,
940
+ cu_seqlens=cu_seqlens,
941
+ max_seqlen=max_seqlen,
942
+ output_attentions=output_attentions,
943
+ )
944
+ hidden_states = layer_outputs[0]
945
+ if output_attentions and len(layer_outputs) > 1:
946
+ all_self_attentions = all_self_attentions + (layer_outputs[1],)
947
+
948
+ if output_hidden_states:
949
+ all_hidden_states = all_hidden_states + (hidden_states,)
950
+
951
+ hidden_states = self.final_norm(hidden_states)
952
+
953
+ if repad:
954
+ hidden_states = _pad_modernbert_output(
955
+ inputs=hidden_states, indices=indices, batch=batch_size, seqlen=seq_len
956
+ )
957
+ if all_hidden_states is not None:
958
+ all_hidden_states = tuple(
959
+ _pad_modernbert_output(inputs=hs, indices=indices, batch=batch_size, seqlen=seq_len)
960
+ for hs in all_hidden_states
961
+ )
962
+
963
+ if not return_dict:
964
+ return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
965
+ return BaseModelOutput(
966
+ last_hidden_state=hidden_states,
967
+ hidden_states=all_hidden_states,
968
+ attentions=all_self_attentions,
969
+ )
970
+
971
+ def _update_attention_mask(self, attention_mask: torch.Tensor, output_attentions: bool) -> torch.Tensor:
972
+ if output_attentions:
973
+ if self.config._attn_implementation == "sdpa":
974
+ logger.warning_once(
975
+ "Outputting attentions is only supported with the 'eager' attention implementation, "
976
+ 'not with "sdpa". Falling back to `attn_implementation="eager"`.'
977
+ )
978
+ self.config._attn_implementation = "eager"
979
+ elif self.config._attn_implementation != "eager":
980
+ logger.warning_once(
981
+ "Outputting attentions is only supported with the eager attention implementation, "
982
+ f'not with {self.config._attn_implementation}. Consider setting `attn_implementation="eager"`.'
983
+ " Setting `output_attentions=False`."
984
+ )
985
+
986
+ global_attention_mask = _prepare_4d_attention_mask(attention_mask, self.dtype)
987
+
988
+ # Create position indices
989
+ rows = torch.arange(global_attention_mask.shape[2]).unsqueeze(0)
990
+ # Calculate distance between positions
991
+ distance = torch.abs(rows - rows.T)
992
+
993
+ # Create sliding window mask (1 for positions within window, 0 outside)
994
+ window_mask = (
995
+ (distance <= self.config.local_attention // 2).unsqueeze(0).unsqueeze(0).to(attention_mask.device)
996
+ )
997
+ # Combine with existing mask
998
+ sliding_window_mask = global_attention_mask.masked_fill(window_mask.logical_not(), torch.finfo(self.dtype).min)
999
+
1000
+ return global_attention_mask, sliding_window_mask
1001
+
1002
+
1003
+ class ModernBertPredictionHead(nn.Module):
1004
+ def __init__(self, config: ModernBertConfig):
1005
+ super().__init__()
1006
+ self.config = config
1007
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size, config.classifier_bias)
1008
+ self.act = ACT2FN[config.classifier_activation]
1009
+ self.norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=config.norm_bias)
1010
+
1011
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
1012
+ return self.norm(self.act(self.dense(hidden_states)))
1013
+
1014
+
1015
+ @add_start_docstrings(
1016
+ "The ModernBert Model with a decoder head on top that is used for masked language modeling.",
1017
+ MODERNBERT_START_DOCSTRING,
1018
+ )
1019
+ class ModernBertForMaskedLM(ModernBertPreTrainedModel):
1020
+ _tied_weights_keys = ["decoder.weight"]
1021
+
1022
+ def __init__(self, config: ModernBertConfig):
1023
+ super().__init__(config)
1024
+ self.config = config
1025
+ self.model = ModernBertModel(config)
1026
+ self.head = ModernBertPredictionHead(config)
1027
+ self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=config.decoder_bias)
1028
+
1029
+ self.sparse_prediction = self.config.sparse_prediction
1030
+ self.sparse_pred_ignore_index = self.config.sparse_pred_ignore_index
1031
+
1032
+ # Initialize weights and apply final processing
1033
+ self.post_init()
1034
+
1035
+ def get_output_embeddings(self):
1036
+ return self.decoder
1037
+
1038
+ def set_output_embeddings(self, new_embeddings: nn.Linear):
1039
+ self.decoder = new_embeddings
1040
+
1041
+ @torch.compile(dynamic=True)
1042
+ def compiled_head(self, output: torch.Tensor) -> torch.Tensor:
1043
+ return self.decoder(self.head(output))
1044
+
1045
+ @add_start_docstrings_to_model_forward(MODERNBERT_INPUTS_DOCSTRING)
1046
+ @add_code_sample_docstrings(
1047
+ checkpoint=_CHECKPOINT_FOR_DOC,
1048
+ output_type=MaskedLMOutput,
1049
+ config_class=_CONFIG_FOR_DOC,
1050
+ )
1051
+ def forward(
1052
+ self,
1053
+ input_ids: Optional[torch.LongTensor] = None,
1054
+ attention_mask: Optional[torch.Tensor] = None,
1055
+ sliding_window_mask: Optional[torch.Tensor] = None,
1056
+ position_ids: Optional[torch.Tensor] = None,
1057
+ inputs_embeds: Optional[torch.Tensor] = None,
1058
+ labels: Optional[torch.Tensor] = None,
1059
+ indices: Optional[torch.Tensor] = None,
1060
+ cu_seqlens: Optional[torch.Tensor] = None,
1061
+ max_seqlen: Optional[int] = None,
1062
+ batch_size: Optional[int] = None,
1063
+ seq_len: Optional[int] = None,
1064
+ output_attentions: Optional[bool] = None,
1065
+ output_hidden_states: Optional[bool] = None,
1066
+ return_dict: Optional[bool] = None,
1067
+ **kwargs,
1068
+ ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
1069
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1070
+ self._maybe_set_compile()
1071
+
1072
+ if self.config._attn_implementation == "flash_attention_2":
1073
+ if indices is None and cu_seqlens is None and max_seqlen is None:
1074
+ if batch_size is None and seq_len is None:
1075
+ if inputs_embeds is not None:
1076
+ batch_size, seq_len = inputs_embeds.shape[:2]
1077
+ else:
1078
+ batch_size, seq_len = input_ids.shape[:2]
1079
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
1080
+
1081
+ if attention_mask is None:
1082
+ attention_mask = torch.ones((batch_size, seq_len), device=device, dtype=torch.bool)
1083
+
1084
+ if inputs_embeds is None:
1085
+ with torch.no_grad():
1086
+ input_ids, indices, cu_seqlens, max_seqlen, position_ids, labels = _unpad_modernbert_input(
1087
+ inputs=input_ids, attention_mask=attention_mask, position_ids=position_ids, labels=labels
1088
+ )
1089
+ else:
1090
+ inputs_embeds, indices, cu_seqlens, max_seqlen, position_ids, labels = _unpad_modernbert_input(
1091
+ inputs=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, labels=labels
1092
+ )
1093
+
1094
+ outputs = self.model(
1095
+ input_ids=input_ids,
1096
+ attention_mask=attention_mask,
1097
+ sliding_window_mask=sliding_window_mask,
1098
+ position_ids=position_ids,
1099
+ inputs_embeds=inputs_embeds,
1100
+ indices=indices,
1101
+ cu_seqlens=cu_seqlens,
1102
+ max_seqlen=max_seqlen,
1103
+ batch_size=batch_size,
1104
+ seq_len=seq_len,
1105
+ output_attentions=output_attentions,
1106
+ output_hidden_states=output_hidden_states,
1107
+ return_dict=return_dict,
1108
+ )
1109
+ last_hidden_state = outputs[0]
1110
+
1111
+ if self.sparse_prediction and labels is not None:
1112
+ # flatten labels and output first
1113
+ labels = labels.view(-1)
1114
+ last_hidden_state = last_hidden_state.view(labels.shape[0], -1)
1115
+
1116
+ # then filter out the non-masked tokens
1117
+ mask_tokens = labels != self.sparse_pred_ignore_index
1118
+ last_hidden_state = last_hidden_state[mask_tokens]
1119
+ labels = labels[mask_tokens]
1120
+
1121
+ logits = (
1122
+ self.compiled_head(last_hidden_state)
1123
+ if self.config.reference_compile
1124
+ else self.decoder(self.head(last_hidden_state))
1125
+ )
1126
+
1127
+ loss = None
1128
+ if labels is not None:
1129
+ loss = self.loss_function(logits, labels, vocab_size=self.config.vocab_size)
1130
+
1131
+ if self.config._attn_implementation == "flash_attention_2":
1132
+ with torch.no_grad():
1133
+ logits = _pad_modernbert_output(inputs=logits, indices=indices, batch=batch_size, seqlen=seq_len)
1134
+ if not return_dict:
1135
+ output = (logits,)
1136
+ return ((loss,) + output) if loss is not None else output
1137
+
1138
+ return MaskedLMOutput(
1139
+ loss=loss,
1140
+ logits=logits,
1141
+ hidden_states=outputs.hidden_states,
1142
+ attentions=outputs.attentions,
1143
+ )
1144
+
1145
+
1146
+ @add_start_docstrings(
1147
+ "The ModernBert Model with a sequence classification head on top that performs pooling.",
1148
+ MODERNBERT_START_DOCSTRING,
1149
+ )
1150
+ class ModernBertForSequenceClassification(ModernBertPreTrainedModel):
1151
+ def __init__(self, config: ModernBertConfig):
1152
+ super().__init__(config)
1153
+ self.num_labels = config.num_labels
1154
+ self.config = config
1155
+
1156
+ self.model = ModernBertModel(config)
1157
+ self.head = ModernBertPredictionHead(config)
1158
+ self.drop = torch.nn.Dropout(config.classifier_dropout)
1159
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
1160
+
1161
+ # Initialize weights and apply final processing
1162
+ self.post_init()
1163
+
1164
+ @add_start_docstrings_to_model_forward(MODERNBERT_INPUTS_DOCSTRING)
1165
+ @add_code_sample_docstrings(
1166
+ checkpoint=_CHECKPOINT_FOR_DOC,
1167
+ output_type=SequenceClassifierOutput,
1168
+ config_class=_CONFIG_FOR_DOC,
1169
+ )
1170
+ def forward(
1171
+ self,
1172
+ input_ids: Optional[torch.LongTensor] = None,
1173
+ attention_mask: Optional[torch.Tensor] = None,
1174
+ sliding_window_mask: Optional[torch.Tensor] = None,
1175
+ position_ids: Optional[torch.Tensor] = None,
1176
+ inputs_embeds: Optional[torch.Tensor] = None,
1177
+ labels: Optional[torch.Tensor] = None,
1178
+ indices: Optional[torch.Tensor] = None,
1179
+ cu_seqlens: Optional[torch.Tensor] = None,
1180
+ max_seqlen: Optional[int] = None,
1181
+ batch_size: Optional[int] = None,
1182
+ seq_len: Optional[int] = None,
1183
+ output_attentions: Optional[bool] = None,
1184
+ output_hidden_states: Optional[bool] = None,
1185
+ return_dict: Optional[bool] = None,
1186
+ **kwargs,
1187
+ ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
1188
+ r"""
1189
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1190
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1191
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1192
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1193
+ """
1194
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1195
+ self._maybe_set_compile()
1196
+
1197
+ outputs = self.model(
1198
+ input_ids=input_ids,
1199
+ attention_mask=attention_mask,
1200
+ sliding_window_mask=sliding_window_mask,
1201
+ position_ids=position_ids,
1202
+ inputs_embeds=inputs_embeds,
1203
+ indices=indices,
1204
+ cu_seqlens=cu_seqlens,
1205
+ max_seqlen=max_seqlen,
1206
+ batch_size=batch_size,
1207
+ seq_len=seq_len,
1208
+ output_attentions=output_attentions,
1209
+ output_hidden_states=output_hidden_states,
1210
+ return_dict=return_dict,
1211
+ )
1212
+ last_hidden_state = outputs[0]
1213
+
1214
+ if self.config.classifier_pooling == "cls":
1215
+ last_hidden_state = last_hidden_state[:, 0]
1216
+ elif self.config.classifier_pooling == "mean":
1217
+ last_hidden_state = (last_hidden_state * attention_mask.unsqueeze(-1)).sum(dim=1) / attention_mask.sum(
1218
+ dim=1, keepdim=True
1219
+ )
1220
+
1221
+ pooled_output = self.head(last_hidden_state)
1222
+ pooled_output = self.drop(pooled_output)
1223
+ logits = self.classifier(pooled_output)
1224
+
1225
+ loss = None
1226
+ if labels is not None:
1227
+ if self.config.problem_type is None:
1228
+ if self.num_labels == 1:
1229
+ self.config.problem_type = "regression"
1230
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1231
+ self.config.problem_type = "single_label_classification"
1232
+ else:
1233
+ self.config.problem_type = "multi_label_classification"
1234
+
1235
+ if self.config.problem_type == "regression":
1236
+ loss_fct = MSELoss()
1237
+ if self.num_labels == 1:
1238
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
1239
+ else:
1240
+ loss = loss_fct(logits, labels)
1241
+ elif self.config.problem_type == "single_label_classification":
1242
+ loss_fct = CrossEntropyLoss()
1243
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
1244
+ elif self.config.problem_type == "multi_label_classification":
1245
+ loss_fct = BCEWithLogitsLoss()
1246
+ loss = loss_fct(logits, labels)
1247
+
1248
+ if not return_dict:
1249
+ output = (logits,)
1250
+ return ((loss,) + output) if loss is not None else output
1251
+
1252
+ return SequenceClassifierOutput(
1253
+ loss=loss,
1254
+ logits=logits,
1255
+ hidden_states=outputs.hidden_states,
1256
+ attentions=outputs.attentions,
1257
+ )
1258
+
1259
+
1260
+ @add_start_docstrings(
1261
+ "The ModernBert Model with a token classification head on top, e.g. for Named Entity Recognition (NER) tasks.",
1262
+ MODERNBERT_START_DOCSTRING,
1263
+ )
1264
+ class ModernBertForTokenClassification(ModernBertPreTrainedModel):
1265
+ def __init__(self, config: ModernBertConfig):
1266
+ super().__init__(config)
1267
+ self.num_labels = config.num_labels
1268
+
1269
+ self.model = ModernBertModel(config)
1270
+ self.head = ModernBertPredictionHead(config)
1271
+ self.drop = torch.nn.Dropout(config.classifier_dropout)
1272
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
1273
+
1274
+ # Initialize weights and apply final processing
1275
+ self.post_init()
1276
+
1277
+ @add_start_docstrings_to_model_forward(MODERNBERT_INPUTS_DOCSTRING)
1278
+ @add_code_sample_docstrings(
1279
+ checkpoint=_CHECKPOINT_FOR_DOC,
1280
+ output_type=TokenClassifierOutput,
1281
+ config_class=_CONFIG_FOR_DOC,
1282
+ )
1283
+ def forward(
1284
+ self,
1285
+ input_ids: Optional[torch.LongTensor] = None,
1286
+ attention_mask: Optional[torch.Tensor] = None,
1287
+ sliding_window_mask: Optional[torch.Tensor] = None,
1288
+ position_ids: Optional[torch.Tensor] = None,
1289
+ inputs_embeds: Optional[torch.Tensor] = None,
1290
+ labels: Optional[torch.Tensor] = None,
1291
+ indices: Optional[torch.Tensor] = None,
1292
+ cu_seqlens: Optional[torch.Tensor] = None,
1293
+ max_seqlen: Optional[int] = None,
1294
+ batch_size: Optional[int] = None,
1295
+ seq_len: Optional[int] = None,
1296
+ output_attentions: Optional[bool] = None,
1297
+ output_hidden_states: Optional[bool] = None,
1298
+ return_dict: Optional[bool] = None,
1299
+ ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:
1300
+ r"""
1301
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1302
+ Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
1303
+ """
1304
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1305
+ self._maybe_set_compile()
1306
+
1307
+ outputs = self.model(
1308
+ input_ids=input_ids,
1309
+ attention_mask=attention_mask,
1310
+ sliding_window_mask=sliding_window_mask,
1311
+ position_ids=position_ids,
1312
+ inputs_embeds=inputs_embeds,
1313
+ indices=indices,
1314
+ cu_seqlens=cu_seqlens,
1315
+ max_seqlen=max_seqlen,
1316
+ batch_size=batch_size,
1317
+ seq_len=seq_len,
1318
+ output_attentions=output_attentions,
1319
+ output_hidden_states=output_hidden_states,
1320
+ return_dict=return_dict,
1321
+ )
1322
+ last_hidden_state = outputs[0]
1323
+
1324
+ last_hidden_state = self.head(last_hidden_state)
1325
+ last_hidden_state = self.drop(last_hidden_state)
1326
+ logits = self.classifier(last_hidden_state)
1327
+
1328
+ loss = None
1329
+ if labels is not None:
1330
+ loss_fct = CrossEntropyLoss()
1331
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
1332
+
1333
+ if not return_dict:
1334
+ output = (logits,) + outputs[1:]
1335
+ return ((loss,) + output) if loss is not None else output
1336
+
1337
+ return TokenClassifierOutput(
1338
+ loss=loss,
1339
+ logits=logits,
1340
+ hidden_states=outputs.hidden_states,
1341
+ attentions=outputs.attentions,
1342
+ )
1343
+
1344
+
1345
+ __all__ = [
1346
+ "ModernBertModel",
1347
+ "ModernBertPreTrainedModel",
1348
+ "ModernBertForMaskedLM",
1349
+ "ModernBertForSequenceClassification",
1350
+ "ModernBertForTokenClassification",
1351
+ ]
pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:efe2082c2825cec0a29fe7b584ef7f631f597e82cf6d3efae7fe46fa900041aa
3
+ size 121641506
special_tokens_map.json ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "[CLS]",
4
+ "lstrip": false,
5
+ "normalized": false,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "cls_token": {
10
+ "content": "[CLS]",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "eos_token": {
17
+ "content": "[SEP]",
18
+ "lstrip": false,
19
+ "normalized": false,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ },
23
+ "mask_token": {
24
+ "content": "[MASK]",
25
+ "lstrip": false,
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": "[SEP]",
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
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "added_tokens_decoder": {
3
+ "0": {
4
+ "content": "[CLS]",
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": "[SEP]",
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
+ "4": {
36
+ "content": "[MASK]",
37
+ "lstrip": false,
38
+ "normalized": false,
39
+ "rstrip": false,
40
+ "single_word": false,
41
+ "special": true
42
+ }
43
+ },
44
+ "bos_token": "[CLS]",
45
+ "clean_up_tokenization_spaces": false,
46
+ "cls_token": "[CLS]",
47
+ "do_lower_case": false,
48
+ "eos_token": "[SEP]",
49
+ "extra_special_tokens": {},
50
+ "keep_accents": true,
51
+ "mask_token": "[MASK]",
52
+ "model_input_names": [
53
+ "input_ids",
54
+ "attention_mask"
55
+ ],
56
+ "model_max_length": 1000000000000000019884624838656,
57
+ "pad_token": "[PAD]",
58
+ "sep_token": "[SEP]",
59
+ "split_by_punct": true,
60
+ "tokenizer_class": "DebertaV2TokenizerFast",
61
+ "unk_token": "[UNK]"
62
+ }
training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:bd3079607bf4955436100ad0ac0e404efb15b7865f1228bdab677a4123d0d337
3
+ size 5368
ud.py ADDED
@@ -0,0 +1,112 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy
2
+ from transformers import TokenClassificationPipeline
3
+
4
+ class UniversalDependenciesPipeline(TokenClassificationPipeline):
5
+ def __init__(self,**kwargs):
6
+ super().__init__(**kwargs)
7
+ x=self.model.config.label2id
8
+ self.root=numpy.full((len(x)),-numpy.inf)
9
+ self.arc=numpy.full((len(x)),-numpy.inf)
10
+ for k,v in x.items():
11
+ if k.endswith("|[root]"):
12
+ self.root[v]=0
13
+ elif k.endswith("]"):
14
+ self.arc[v]=0
15
+ def _forward(self,model_inputs):
16
+ import torch
17
+ v=model_inputs["input_ids"][0].tolist()
18
+ if len(v)<91:
19
+ x=[True]*(len(v)-2)
20
+ else:
21
+ with torch.no_grad():
22
+ e=self.model(input_ids=torch.tensor([v]).to(self.device))
23
+ m=e.logits[0].cpu().numpy()
24
+ e=numpy.exp(m-numpy.max(m,axis=-1,keepdims=True))
25
+ z=e/e.sum(axis=-1,keepdims=True)
26
+ k=numpy.argmax(m,axis=1).tolist()
27
+ x=[not self.model.config.id2label[p].split("|")[0].endswith(".") for p in k[1:-1]]
28
+ w=(sum([1 for b in x if b])+1)*(len(x)+1)+1
29
+ for i in numpy.argsort([z[i+1,k[i+1]] for i in range(len(x))]):
30
+ if w+len(x)>8191:
31
+ break
32
+ if not x[i]:
33
+ x[i]=True
34
+ w+=len(x)+1
35
+ ids=list(v)
36
+ for i in range(len(x)):
37
+ if x[i]:
38
+ ids+=v[1:i+1]+[self.tokenizer.mask_token_id]+v[i+2:]
39
+ with torch.no_grad():
40
+ e=self.model(input_ids=torch.tensor([ids]).to(self.device))
41
+ return {"logits":e.logits,"thin_out":x,**model_inputs}
42
+ def check_model_type(self,supported_models):
43
+ pass
44
+ def postprocess(self,model_outputs,**kwargs):
45
+ if "logits" not in model_outputs:
46
+ return "".join(self.postprocess(x,**kwargs) for x in model_outputs)
47
+ m=model_outputs["logits"][0].cpu().numpy()
48
+ x=model_outputs["thin_out"]
49
+ e=numpy.full((len(x),len(x),m.shape[-1]),m.min())
50
+ k=len(x)+2
51
+ for i in range(len(x)):
52
+ if x[i]:
53
+ for j in range(len(x)):
54
+ if i==j:
55
+ e[i,i]=m[k]+self.root
56
+ else:
57
+ e[i,j]=m[k]+self.arc
58
+ k+=1
59
+ k+=1
60
+ g=self.model.config.label2id["X.|[goeswith]"]
61
+ m,r=numpy.max(e,axis=2),numpy.tri(e.shape[0])
62
+ for i in range(e.shape[0]):
63
+ for j in range(i+2,e.shape[1]):
64
+ r[i,j]=1
65
+ if numpy.argmax(e[i,j-1])==g:
66
+ if numpy.argmax(m[:,j-1])==i:
67
+ r[i,j]=r[i,j-1]
68
+ e[:,:,g]+=numpy.where(r==0,0,-numpy.inf)
69
+ m,p=numpy.max(e,axis=2),numpy.argmax(e,axis=2)
70
+ h=self.chu_liu_edmonds(m)
71
+ z=[i for i,j in enumerate(h) if i==j]
72
+ if len(z)>1:
73
+ k,h=z[numpy.argmax(m[z,z])],numpy.min(m)-numpy.max(m)
74
+ m[:,z]+=[[0 if j in z and (i!=j or i==k) else h for i in z] for j in range(m.shape[0])]
75
+ h=self.chu_liu_edmonds(m)
76
+ v=[(s,e) for s,e in model_outputs["offset_mapping"][0].tolist() if s<e]
77
+ q=[self.model.config.id2label[p[j,i]].split("|") for i,j in enumerate(h)]
78
+ if "aggregation_strategy" in kwargs and kwargs["aggregation_strategy"]!="none":
79
+ for i,j in reversed(list(enumerate(q[1:],1))):
80
+ if j[-1]=="[goeswith]" and set([t[-1] for t in q[h[i]+1:i+1]])=={"[goeswith]"}:
81
+ h=[b if i>b else b-1 for a,b in enumerate(h) if i!=a]
82
+ v[i-1]=(v[i-1][0],v.pop(i)[1])
83
+ q.pop(i)
84
+ elif v[i-1][1]>v[i][0]:
85
+ h=[b if i>b else b-1 for a,b in enumerate(h) if i!=a]
86
+ v[i-1]=(v[i-1][0],v.pop(i)[1])
87
+ q.pop(i)
88
+ t=model_outputs["sentence"].replace("\n"," ")
89
+ u="# text = "+t+"\n"
90
+ for i,(s,e) in enumerate(v):
91
+ u+="\t".join([str(i+1),t[s:e],"_",q[i][0].replace(".",""),"_","_" if len(q[i])<3 else "|".join(q[i][1:-1]),str(0 if h[i]==i else h[i]+1),q[i][-1][1:-1],"_","_" if i+1<len(v) and e<v[i+1][0] else "SpaceAfter=No"])+"\n"
92
+ return u+"\n"
93
+ def chu_liu_edmonds(self,matrix):
94
+ import numpy
95
+ h=numpy.argmax(matrix,axis=0)
96
+ x=[-1 if i==j else j for i,j in enumerate(h)]
97
+ for b in [lambda x,i,j:-1 if i not in x else x[i],lambda x,i,j:-1 if j<0 else x[j]]:
98
+ y=[]
99
+ while x!=y:
100
+ y=list(x)
101
+ for i,j in enumerate(x):
102
+ x[i]=b(x,i,j)
103
+ if max(x)<0:
104
+ return h
105
+ y,x=[i for i,j in enumerate(x) if j==max(x)],[i for i,j in enumerate(x) if j<max(x)]
106
+ z=matrix-numpy.max(matrix,axis=0)
107
+ m=numpy.block([[z[x,:][:,x],numpy.max(z[x,:][:,y],axis=1).reshape(len(x),1)],[numpy.max(z[y,:][:,x],axis=0),numpy.max(z[y,y])]])
108
+ k=[j if i==len(x) else x[j] if j<len(x) else y[numpy.argmax(z[y,x[i]])] for i,j in enumerate(self.chu_liu_edmonds(m))]
109
+ h=[j if i in y else k[x.index(i)] for i,j in enumerate(h)]
110
+ i=y[numpy.argmax(z[x[k[-1]],y] if k[-1]<len(x) else z[y,y])]
111
+ h[i]=x[k[-1]] if k[-1]<len(x) else i
112
+ return h