Skip to content

Commit

Permalink
Add test cases for TransformersTokenCounter
Browse files Browse the repository at this point in the history
  • Loading branch information
PhilipMay committed Jan 4, 2024
1 parent af91d27 commit 1774967
Showing 1 changed file with 20 additions and 0 deletions.
20 changes: 20 additions & 0 deletions tests/test_transformers.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,29 @@
# which is available at https://opensource.org/licenses/MIT


import pytest
from hypothesis import given, settings
from hypothesis.strategies import text

from mltb2.transformers import TransformersTokenCounter


@pytest.fixture(scope="module")
def deepset_gbert_base_token_counter():
return TransformersTokenCounter("deepset/gbert-base")


@settings(max_examples=1000, deadline=None)
@given(text=text())
def test_TransformersTokenCounter_hypothesis( # noqa: N802
text: str, deepset_gbert_base_token_counter: TransformersTokenCounter
):
token_count = deepset_gbert_base_token_counter(text)

assert isinstance(token_count, int)
assert token_count >= 0


def test_TransformersTokenCounter_call_string(): # noqa: N802
transformers_token_counter = TransformersTokenCounter("deepset/gbert-base")
token_count = transformers_token_counter("Das ist ein Text.")
Expand Down

0 comments on commit 1774967

Please sign in to comment.