Skip to content

Commit

Permalink
Added unit test for checking truncation parameter
Browse files Browse the repository at this point in the history
Signed-off-by: yerzhaisang <tasqali1697@gmail.com>
  • Loading branch information
Yerzhaisang committed Jul 16, 2023
1 parent 64696fe commit 8a55e7b
Showing 1 changed file with 27 additions and 0 deletions.
27 changes: 27 additions & 0 deletions tests/ml_models/test_sentencetransformermodel_pytest.py
Original file line number Diff line number Diff line change
Expand Up @@ -452,5 +452,32 @@ def test_missing_fields_in_config_json():
clean_test_folder(TEST_FOLDER)


def test_truncation_parameter():
model_id = "sentence-transformers/msmarco-distilbert-base-tas-b"

clean_test_folder(TEST_FOLDER)
test_model10 = SentenceTransformerModel(
folder_path=TEST_FOLDER,
model_id=model_id,
)

test_model10.save_as_pt(model_id=model_id, sentences=["today is sunny"])

tokenizer_json_file_path = os.path.join(TEST_FOLDER, "tokenizer.json")
try:
with open(tokenizer_json_file_path, "r") as json_file:
tokenizer_json = json.load(json_file)
except Exception as exec:
assert (
False
), f"Creating tokenizer.json file for tracing raised an exception {exec}"

assert tokenizer_json[
"truncation"
], "truncation parameter in tokenizer.json is null"

clean_test_folder(TEST_FOLDER)


clean_test_folder(TEST_FOLDER)
clean_test_folder(TESTDATA_UNZIP_FOLDER)

0 comments on commit 8a55e7b

Please sign in to comment.