Ruff linting

defenseunicorns · May 10, 2024 · d0c707d · d0c707d
1 parent 4ce550d
commit d0c707d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/packages/llama-cpp-python/main.py b/packages/llama-cpp-python/main.py
@@ -35,6 +35,6 @@ def generate(
             yield res["choices"][0]["text"]  # type: ignore
 
     async def count_tokens(self, raw_text: str):
-        string_bytes: bytes = bytes(raw_text, 'utf-8')
+        string_bytes: bytes = bytes(raw_text, "utf-8")
         tokens: list[int] = self.llm.tokenize(string_bytes)
         return len(tokens)
diff --git a/packages/vllm/src/main.py b/packages/vllm/src/main.py
@@ -461,9 +461,12 @@ async def ChatCompleteStream(
         logging.info("ChatCompleteStream END:\n---")
 
     async def count_tokens(self, raw_text: str):
-        tokens: list[int] | list[str] = (await self.engine.get_tokenizer()).tokenize(raw_text)
+        tokens: list[int] | list[str] = (await self.engine.get_tokenizer()).tokenize(
+            raw_text
+        )
         return len(tokens)
 
+
 async def main():
     logging.basicConfig(level=logging.INFO)
     await serve(Model())