|
| 1 | +import { PreTrainedTokenizer, ModernBertModel, ModernBertForMaskedLM, ModernBertForSequenceClassification, ModernBertForTokenClassification } from "../../../src/transformers.js"; |
| 2 | + |
| 3 | +import { MAX_MODEL_LOAD_TIME, MAX_TEST_EXECUTION_TIME, MAX_MODEL_DISPOSE_TIME, DEFAULT_MODEL_OPTIONS } from "../../init.js"; |
| 4 | + |
| 5 | +export default () => { |
| 6 | + describe("ModernBertModel", () => { |
| 7 | + const model_id = "hf-internal-testing/tiny-random-ModernBertModel"; |
| 8 | + |
| 9 | + /** @type {ModernBertModel} */ |
| 10 | + let model; |
| 11 | + /** @type {PreTrainedTokenizer} */ |
| 12 | + let tokenizer; |
| 13 | + beforeAll(async () => { |
| 14 | + model = await ModernBertModel.from_pretrained(model_id, DEFAULT_MODEL_OPTIONS); |
| 15 | + tokenizer = await PreTrainedTokenizer.from_pretrained(model_id); |
| 16 | + }, MAX_MODEL_LOAD_TIME); |
| 17 | + |
| 18 | + it( |
| 19 | + "batch_size=1", |
| 20 | + async () => { |
| 21 | + const inputs = tokenizer("hello"); |
| 22 | + const { last_hidden_state } = await model(inputs); |
| 23 | + expect(last_hidden_state.dims).toEqual([1, 3, 32]); |
| 24 | + expect(last_hidden_state.mean().item()).toBeCloseTo(-0.08922556787729263, 5); |
| 25 | + }, |
| 26 | + MAX_TEST_EXECUTION_TIME, |
| 27 | + ); |
| 28 | + |
| 29 | + it( |
| 30 | + "batch_size>1", |
| 31 | + async () => { |
| 32 | + const inputs = tokenizer(["hello", "hello world"], { padding: true }); |
| 33 | + const { last_hidden_state } = await model(inputs); |
| 34 | + expect(last_hidden_state.dims).toEqual([2, 4, 32]); |
| 35 | + expect(last_hidden_state.mean().item()).toBeCloseTo(0.048988230526447296, 5); |
| 36 | + }, |
| 37 | + MAX_TEST_EXECUTION_TIME, |
| 38 | + ); |
| 39 | + |
| 40 | + it( |
| 41 | + "sequence_length > local_attention_window", |
| 42 | + async () => { |
| 43 | + const text = "The sun cast long shadows across the weathered cobblestones as Thomas made his way through the ancient city. The evening air carried whispers of autumn, rustling through the golden leaves that danced and swirled around his feet. His thoughts wandered to the events that had brought him here, to this moment, in this forgotten corner of the world. The old buildings loomed above him, their facades telling stories of centuries past. Windows reflected the dying light of day, creating a kaleidoscope of amber and rose that painted the narrow streets. The distant sound of church bells echoed through the maze of alleyways, marking time's steady march forward. In his pocket, he fingered the small brass key that had belonged to his grandfather. Its weight seemed to grow heavier with each step, a tangible reminder of the promise he had made. The mystery of its purpose had consumed his thoughts for weeks, leading him through archives and dusty libraries, through conversations with local historians and elderly residents who remembered the old days. As the evening deepened into dusk, streetlamps flickered to life one by one, creating pools of warm light that guided his way. The smell of wood smoke and distant cooking fires drifted through the air, reminding him of childhood evenings spent by the hearth, listening to his grandfather's tales of hidden treasures and secret passages. His footsteps echoed against the stone walls, a rhythmic accompaniment to his journey. Each step brought him closer to his destination, though uncertainty still clouded his mind about what he might find. The old map in his other pocket, creased and worn from constant consultation, had led him this far. The street ahead narrowed, and the buildings seemed to lean in closer, their upper stories nearly touching above his head. The air grew cooler in this shadowed passage, and his breath formed small clouds in front of him. Something about this place felt different, charged with possibility and ancient secrets. He walked down the [MASK]"; |
| 44 | + const inputs = tokenizer(text); |
| 45 | + const { last_hidden_state } = await model(inputs); |
| 46 | + expect(last_hidden_state.dims).toEqual([1, 397, 32]); |
| 47 | + expect(last_hidden_state.mean().item()).toBeCloseTo(-0.06889555603265762, 5); |
| 48 | + }, |
| 49 | + MAX_TEST_EXECUTION_TIME, |
| 50 | + ); |
| 51 | + |
| 52 | + afterAll(async () => { |
| 53 | + await model?.dispose(); |
| 54 | + }, MAX_MODEL_DISPOSE_TIME); |
| 55 | + }); |
| 56 | + |
| 57 | + describe("ModernBertForMaskedLM", () => { |
| 58 | + const model_id = "hf-internal-testing/tiny-random-ModernBertForMaskedLM"; |
| 59 | + |
| 60 | + const texts = ["The goal of life is [MASK].", "Paris is the [MASK] of France."]; |
| 61 | + |
| 62 | + /** @type {ModernBertForMaskedLM} */ |
| 63 | + let model; |
| 64 | + /** @type {PreTrainedTokenizer} */ |
| 65 | + let tokenizer; |
| 66 | + beforeAll(async () => { |
| 67 | + model = await ModernBertForMaskedLM.from_pretrained(model_id, DEFAULT_MODEL_OPTIONS); |
| 68 | + tokenizer = await PreTrainedTokenizer.from_pretrained(model_id); |
| 69 | + }, MAX_MODEL_LOAD_TIME); |
| 70 | + |
| 71 | + it( |
| 72 | + "batch_size=1", |
| 73 | + async () => { |
| 74 | + const inputs = tokenizer(texts[0]); |
| 75 | + const { logits } = await model(inputs); |
| 76 | + expect(logits.dims).toEqual([1, 9, 50368]); |
| 77 | + expect(logits.mean().item()).toBeCloseTo(0.0053214821964502335, 5); |
| 78 | + }, |
| 79 | + MAX_TEST_EXECUTION_TIME, |
| 80 | + ); |
| 81 | + |
| 82 | + it( |
| 83 | + "batch_size>1", |
| 84 | + async () => { |
| 85 | + const inputs = tokenizer(texts, { padding: true }); |
| 86 | + const { logits } = await model(inputs); |
| 87 | + expect(logits.dims).toEqual([2, 9, 50368]); |
| 88 | + expect(logits.mean().item()).toBeCloseTo(0.009154772385954857, 5); |
| 89 | + }, |
| 90 | + MAX_TEST_EXECUTION_TIME, |
| 91 | + ); |
| 92 | + |
| 93 | + afterAll(async () => { |
| 94 | + await model?.dispose(); |
| 95 | + }, MAX_MODEL_DISPOSE_TIME); |
| 96 | + }); |
| 97 | + |
| 98 | + describe("ModernBertForSequenceClassification", () => { |
| 99 | + const model_id = "hf-internal-testing/tiny-random-ModernBertForSequenceClassification"; |
| 100 | + |
| 101 | + /** @type {ModernBertForSequenceClassification} */ |
| 102 | + let model; |
| 103 | + /** @type {PreTrainedTokenizer} */ |
| 104 | + let tokenizer; |
| 105 | + beforeAll(async () => { |
| 106 | + model = await ModernBertForSequenceClassification.from_pretrained(model_id, DEFAULT_MODEL_OPTIONS); |
| 107 | + tokenizer = await PreTrainedTokenizer.from_pretrained(model_id); |
| 108 | + }, MAX_MODEL_LOAD_TIME); |
| 109 | + |
| 110 | + it( |
| 111 | + "batch_size=1", |
| 112 | + async () => { |
| 113 | + const inputs = tokenizer("hello"); |
| 114 | + const { logits } = await model(inputs); |
| 115 | + const target = [[-0.7050137519836426, 2.343430519104004]]; |
| 116 | + expect(logits.dims).toEqual([1, 2]); |
| 117 | + expect(logits.tolist()).toBeCloseToNested(target, 5); |
| 118 | + }, |
| 119 | + MAX_TEST_EXECUTION_TIME, |
| 120 | + ); |
| 121 | + |
| 122 | + it( |
| 123 | + "batch_size>1", |
| 124 | + async () => { |
| 125 | + const inputs = tokenizer(["hello", "hello world"], { padding: true }); |
| 126 | + const { logits } = await model(inputs); |
| 127 | + const target = [ |
| 128 | + [-0.7050137519836426, 2.343430519104004], |
| 129 | + [-2.6860175132751465, 3.993380546569824], |
| 130 | + ]; |
| 131 | + expect(logits.dims).toEqual([2, 2]); |
| 132 | + expect(logits.tolist()).toBeCloseToNested(target, 5); |
| 133 | + }, |
| 134 | + MAX_TEST_EXECUTION_TIME, |
| 135 | + ); |
| 136 | + |
| 137 | + afterAll(async () => { |
| 138 | + await model?.dispose(); |
| 139 | + }, MAX_MODEL_DISPOSE_TIME); |
| 140 | + }); |
| 141 | + |
| 142 | + describe("ModernBertForTokenClassification", () => { |
| 143 | + const model_id = "hf-internal-testing/tiny-random-ModernBertForTokenClassification"; |
| 144 | + |
| 145 | + /** @type {ModernBertForTokenClassification} */ |
| 146 | + let model; |
| 147 | + /** @type {PreTrainedTokenizer} */ |
| 148 | + let tokenizer; |
| 149 | + beforeAll(async () => { |
| 150 | + model = await ModernBertForTokenClassification.from_pretrained(model_id, DEFAULT_MODEL_OPTIONS); |
| 151 | + tokenizer = await PreTrainedTokenizer.from_pretrained(model_id); |
| 152 | + }, MAX_MODEL_LOAD_TIME); |
| 153 | + |
| 154 | + it( |
| 155 | + "batch_size=1", |
| 156 | + async () => { |
| 157 | + const inputs = tokenizer("hello"); |
| 158 | + const { logits } = await model(inputs); |
| 159 | + expect(logits.dims).toEqual([1, 3, 2]); |
| 160 | + expect(logits.mean().item()).toBeCloseTo(1.0337047576904297, 5); |
| 161 | + }, |
| 162 | + MAX_TEST_EXECUTION_TIME, |
| 163 | + ); |
| 164 | + |
| 165 | + it( |
| 166 | + "batch_size>1", |
| 167 | + async () => { |
| 168 | + const inputs = tokenizer(["hello", "hello world"], { padding: true }); |
| 169 | + const { logits } = await model(inputs); |
| 170 | + expect(logits.dims).toEqual([2, 4, 2]); |
| 171 | + expect(logits.mean().item()).toBeCloseTo(-1.3397092819213867, 5); |
| 172 | + }, |
| 173 | + MAX_TEST_EXECUTION_TIME, |
| 174 | + ); |
| 175 | + |
| 176 | + afterAll(async () => { |
| 177 | + await model?.dispose(); |
| 178 | + }, MAX_MODEL_DISPOSE_TIME); |
| 179 | + }); |
| 180 | +}; |
0 commit comments