From 2a7d2331e4dcee1450d2d81fe326b94b82c882a3 Mon Sep 17 00:00:00 2001 From: Gautier Solard Date: Tue, 30 Jul 2024 16:53:33 +0200 Subject: [PATCH] Fix mypy errors --- src/benchmark_llm_serving/backends.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/benchmark_llm_serving/backends.py b/src/benchmark_llm_serving/backends.py index 0cc344b..14a67a1 100644 --- a/src/benchmark_llm_serving/backends.py +++ b/src/benchmark_llm_serving/backends.py @@ -23,9 +23,9 @@ def get_payload(self, query_input: QueryInput, args: argparse.Namespace) -> dict Returns: dict : The payload """ - raise NotImplemented("The subclass should implement this method") + raise NotImplemented("The subclass should implement this method") # type: ignore - def get_newly_generated_text(self, json_chunk: str) -> str: + def get_newly_generated_text(self, json_chunk: dict) -> str: """Gets the newly generated text Args: @@ -34,7 +34,7 @@ def get_newly_generated_text(self, json_chunk: str) -> str: Returns: str : The newly generated text """ - raise NotImplemented("The subclass should implement this method") + raise NotImplemented("The subclass should implement this method") # type: ignore def test_chunk_validity(self, chunk: str) -> bool: """Tests if the chunk is valid or should not be considered. @@ -113,7 +113,7 @@ def get_payload(self, query_input: QueryInput, args: argparse.Namespace) -> dict "stream_options": {"include_usage": True} } - def get_newly_generated_text(self, json_chunk: str) -> str: + def get_newly_generated_text(self, json_chunk: dict) -> str: """Gets the newly generated text Args: @@ -173,7 +173,7 @@ def get_completions_headers(self) -> dict: return {"Accept": "application/json", "Content-Type": "application/json"} - def get_newly_generated_text(self, json_chunk: str) -> str: + def get_newly_generated_text(self, json_chunk: dict) -> str: """Gets the newly generated text Args: @@ -192,8 +192,9 @@ def get_newly_generated_text(self, json_chunk: str) -> str: def get_backend(backend_name: str) -> BackEnd: implemented_backends = ["mistral", "happy_vllm"] if backend_name not in implemented_backends: - raise ValueError(f"The specified backend {backend_name} is not implemented. Please use one of the following : {implemented_backends}") + raise ValueError(f"The specified backend {backend_name} is not implemented. Please use one of the following : {implemented_backends}") if backend_name == "happy_vllm": return BackendHappyVllm(backend_name, chunk_prefix="data: ", last_chunk="[DONE]") if backend_name == "mistral": - return BackEndMistral(backend_name, chunk_prefix="data: ", last_chunk="[DONE]") \ No newline at end of file + return BackEndMistral(backend_name, chunk_prefix="data: ", last_chunk="[DONE]") + return BackEnd("not_implemented") \ No newline at end of file