Skip to content

Commit

Permalink
Add CPU test for ModernBert Model
Browse files Browse the repository at this point in the history
  • Loading branch information
pdeviTT committed Jan 30, 2025
1 parent 82af76c commit 32e2d21
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 1 deletion.
2 changes: 1 addition & 1 deletion env/core_requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ scipy>=1.8.0
tensorflow==2.13
tensorboard==2.13
tf2onnx==1.15.1
transformers==4.47.0
transformers==4.48.1
# To avoid warning during the import
requests==2.28.2
urllib3==1.26.14
Expand Down
28 changes: 28 additions & 0 deletions forge/test/models/pytorch/text/modernbert/test_modernbert.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
#
# SPDX-License-Identifier: Apache-2.0
import pytest
from transformers import AutoTokenizer, ModernBertModel

variants = [
"answerdotai/ModernBERT-base",
]


@pytest.mark.nightly
@pytest.mark.parametrize("variant", variants, ids=variants)
def test_modernbert_test_generation(variant):
tokenizer = AutoTokenizer.from_pretrained(variant)
framework_model = ModernBertModel.from_pretrained(variant) # , return_dict=False)
# model = AutoModelForMaskedLM.from_pretrained("answerdotai/ModernBERT-base")

inputs = tokenizer(
"Hello, my dog is cute",
return_tensors="pt",
max_length=150,
pad_to_max_length=True,
truncation=True,
)
outputs = framework_model(**inputs)

print("outputs:", outputs)

0 comments on commit 32e2d21

Please sign in to comment.