diff --git a/tests/entrypoints/test_server_tokenize_truncate.py b/tests/entrypoints/test_server_tokenize_truncate.py index 252d20c74..ad5c9469e 100644 --- a/tests/entrypoints/test_server_tokenize_truncate.py +++ b/tests/entrypoints/test_server_tokenize_truncate.py @@ -17,6 +17,7 @@ # The tokenizer was tested using the following model: MODEL_NAME = "facebook/opt-125m" + @pytest.fixture(scope="module") def server(): ray.init() @@ -31,6 +32,7 @@ def server(): yield server_runner ray.shutdown() + # Fixture to create a gRPC stub for the GenerationService @pytest.fixture(scope="module") def grpc_stub(): @@ -39,6 +41,7 @@ def grpc_stub(): yield stub channel.close() + # Test cases @pytest.mark.parametrize("test_case", [ { @@ -48,15 +51,37 @@ def grpc_stub(): "return_offsets": True, }, "response": { - "tokenCount": 7, + "tokenCount": + 7, "offsets": [ - {"start": 0, "end": 0}, - {"start": 0, "end": 3}, - {"start": 3, "end": 8}, - {"start": 8, "end": 13}, - {"start": 13, "end": 19}, - {"start": 19, "end": 22}, - {"start": 22, "end": 30}, + { + "start": 0, + "end": 0 + }, + { + "start": 0, + "end": 3 + }, + { + "start": 3, + "end": 8 + }, + { + "start": 8, + "end": 13 + }, + { + "start": 13, + "end": 19 + }, + { + "start": 19, + "end": 22 + }, + { + "start": 22, + "end": 30 + }, ], }, }, @@ -68,24 +93,39 @@ def grpc_stub(): "return_offsets": True, }, "response": { - "tokenCount": 7, - "tokens": [ - "", - "The", - "Ġvery", - "Ġlong", - "Ġstory", - "Ġis", - "Ġwritten" - ], + "tokenCount": + 7, + "tokens": + ["", "The", "Ġvery", "Ġlong", "Ġstory", "Ġis", "Ġwritten"], "offsets": [ - {"start": 0, "end": 0}, - {"start": 0, "end": 3}, - {"start": 3, "end": 8}, - {"start": 8, "end": 13}, - {"start": 13, "end": 19}, - {"start": 19, "end": 22}, - {"start": 22, "end": 30}, + { + "start": 0, + "end": 0 + }, + { + "start": 0, + "end": 3 + }, + { + "start": 3, + "end": 8 + }, + { + "start": 8, + "end": 13 + }, + { + "start": 13, + "end": 19 + }, + { + "start": 19, + "end": 22 + }, + { + "start": 22, + "end": 30 + }, ], }, }, @@ -97,7 +137,8 @@ def grpc_stub(): "truncate_input_tokens": 10, }, "response": { - "tokenCount": 10, + "tokenCount": + 10, "tokens": [ "Ġvery", "Ġlong", @@ -113,7 +154,8 @@ def grpc_stub(): }, }, { - "name": "Tokenize, trunc and offset for a request with no text message", + "name": + "Tokenize, trunc and offset for a request with no text message", "request": { "text": "", "return_offsets": True, @@ -122,15 +164,13 @@ def grpc_stub(): }, "response": { "tokenCount": 1, - "tokens": [ - "" - ], + "tokens": [""], }, }, { "name": "A request without text ('') and parameters", "request": { - "text" : "" + "text": "" }, "response": { "tokenCount": 1 @@ -139,7 +179,7 @@ def grpc_stub(): { "name": "A request without text (None) and parameters", "request": { - "text" : None + "text": None }, "response": { "tokenCount": 1 @@ -158,8 +198,7 @@ def test_tokenization(server, grpc_stub, test_case): requests=[pb2.TokenizeRequest(text=text)], return_tokens=request.get('return_tokens', False), return_offsets=request.get('return_offsets', False), - truncate_input_tokens=truncate_input_tokens - ) + truncate_input_tokens=truncate_input_tokens) try: responses = grpc_stub.Tokenize(batch).responses @@ -189,5 +228,6 @@ def test_tokenization(server, grpc_stub, test_case): print("Test case passed: ", test_case["name"]) + if __name__ == "__main__": pytest.main([__file__])