Merge pull request #62 from SachinVarghese/fix-checks

Lint and type check fixes

Author: k8s-ci-robot

inference_perf/datagen/base.py

@@ -14,7 +14,7 @@
 from pydantic import BaseModel
 from inference_perf.config import APIType
 from abc import ABC, abstractmethod
-from typing import Generator, Tuple, Optional, List
+from typing import Generator, Optional, List
 
 
 class CompletionData(BaseModel):
@@ -38,9 +38,11 @@ class InferenceData(BaseModel):
 
 class DataGenerator(ABC):
     """Abstract base class for data generators."""
+
     apiType: APIType
 
     """Abstract base class for data generators."""
+
     def __init__(self, apiType: APIType) -> None:
         if apiType not in self.get_supported_apis():
             raise Exception(f"Unsupported API type {apiType}")

inference_perf/datagen/hf_sharegpt_datagen.py

@@ -62,7 +62,7 @@ def get_data(self) -> Generator[InferenceData, None, None]:
                     except (KeyError, TypeError) as e:
                         print(f"Skipping invalid completion data: {e}")
                         continue
-                elif self.APIType == APIType.Chat:
+                elif self.apiType == APIType.Chat:
                     yield InferenceData(
                         type=APIType.Chat,
                         chat=ChatCompletionData(
@@ -73,4 +73,4 @@ def get_data(self) -> Generator[InferenceData, None, None]:
                         ),
                     )
                 else:
-                    raise Exception("Unsupported API type")
+                    raise Exception("Unsupported API type")

inference_perf/datagen/mock_datagen.py

@@ -15,6 +15,7 @@
 from typing import Generator, List
 from inference_perf.config import APIType
 
+
 class MockDataGenerator(DataGenerator):
     def __init__(self, apiType: APIType) -> None:
         super().__init__(apiType)
@@ -28,9 +29,6 @@ def get_data(self) -> Generator[InferenceData, None, None]:
         while True:
             i += 1
             if self.apiType == APIType.Completion:
-                yield InferenceData(
-                    data=CompletionData(prompt="text" + str(i))
-                )
+                yield InferenceData(data=CompletionData(prompt="text" + str(i)))
             else:
                 raise Exception("Unsupported API type")
-                

inference_perf/main.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 from inference_perf.loadgen import LoadGenerator
 from inference_perf.config import DataGenType
-from inference_perf.datagen import MockDataGenerator, HFShareGPTDataGenerator
+from inference_perf.datagen import DataGenerator, MockDataGenerator, HFShareGPTDataGenerator
 from inference_perf.client import ModelServerClient, vLLMModelServerClient
 from inference_perf.reportgen import ReportGenerator, MockReportGenerator
 from inference_perf.metrics import MockMetricsClient
@@ -48,9 +48,11 @@ def main_cli() -> None:
 
     # Define DataGenerator
     if config.data:
-        datagen = MockDataGenerator(config.vllm.api)
+        datagen: DataGenerator
         if config.data.type == DataGenType.ShareGPT:
             datagen = HFShareGPTDataGenerator(config.vllm.api)
+        else:
+            datagen = MockDataGenerator(config.vllm.api)
     else:
         raise Exception("data config missing")
 

inference_perf/utils/custom_tokenizer.py

@@ -17,10 +17,12 @@
 
 class CustomTokenizer:
     def __init__(self, tokenizer_id: str, token: Optional[str], trust_remote_code: Optional[bool]):
-        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, token=token, trust_remote_code=trust_remote_code)
+        self.tokenizer: PreTrainedTokenizerBase = AutoTokenizer.from_pretrained(
+            tokenizer_id, token=token, trust_remote_code=trust_remote_code
+        )
 
     def count_tokens(self, text: str) -> int:
-        if not text:
+        if text == "":
             return 0
         return len(self.tokenizer(text).input_ids)