Merge pull request #101 from Snuffy2/Linting/Formatting-of-utils-folder

Linting/Formatting of utils folder

Author: cubist38

app/utils/__init__.py

@@ -1 +1 @@
-
+"""Utility modules for the mlx-openai-server."""

app/utils/dill.py

@@ -15,69 +15,81 @@
 # limitations under the License.
 """Extends `dill` to support pickling more types and produce more consistent dumps."""
 
-import sys
+from collections.abc import Callable, Iterable
 from io import BytesIO
+import sys
 from types import FunctionType
-from typing import Any, Dict, List, Union
+from typing import TYPE_CHECKING, Any, BinaryIO, ClassVar
 
 import dill
 import xxhash
 
+if TYPE_CHECKING:
+    import spacy
+    import torch
+
 
 class Hasher:
     """Hasher that accepts python objects as inputs."""
 
-    dispatch: Dict = {}
+    dispatch: ClassVar[dict[Any, Any]] = {}
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.m = xxhash.xxh64()
 
     @classmethod
-    def hash_bytes(cls, value: Union[bytes, List[bytes]]) -> str:
+    def hash_bytes(cls, value: bytes | list[bytes]) -> str:
+        """Hash bytes or list of bytes using xxhash."""
         value = [value] if isinstance(value, bytes) else value
         m = xxhash.xxh64()
         for x in value:
             m.update(x)
         return m.hexdigest()
 
     @classmethod
-    def hash(cls, value: Any) -> str:
+    def hash(cls, value: object) -> str:
+        """Hash a Python object by pickling it first."""
         return cls.hash_bytes(dumps(value))
 
-    def update(self, value: Any) -> None:
+    def update(self, value: object) -> None:
+        """Update the hasher with a Python object."""
         header_for_update = f"=={type(value)}=="
         value_for_update = self.hash(value)
         self.m.update(header_for_update.encode("utf8"))
         self.m.update(value_for_update.encode("utf-8"))
 
     def hexdigest(self) -> str:
+        """Return the hexadecimal digest of the hash."""
         return self.m.hexdigest()
 
 
 class Pickler(dill.Pickler):
-    dispatch = dill._dill.MetaCatchingDict(dill.Pickler.dispatch.copy())
+    """Custom Pickler that extends dill with additional type support."""
+
+    dispatch = dill._dill.MetaCatchingDict(dill.Pickler.dispatch.copy())  # noqa: SLF001
     _legacy_no_dict_keys_sorting = False
 
-    def save(self, obj, save_persistent_id=True):
+    def save(self, obj: Any, save_persistent_id: bool = True) -> None:
+        """Save an object to the pickle stream with custom handling."""
         obj_type = type(obj)
         if obj_type not in self.dispatch:
             if "regex" in sys.modules:
-                import regex  # type: ignore
+                import regex  # noqa: PLC0415
 
                 if obj_type is regex.Pattern:
                     pklregister(obj_type)(_save_regexPattern)
             if "spacy" in sys.modules:
-                import spacy  # type: ignore
+                import spacy  # noqa: PLC0415
 
                 if issubclass(obj_type, spacy.Language):
                     pklregister(obj_type)(_save_spacyLanguage)
             if "tiktoken" in sys.modules:
-                import tiktoken  # type: ignore
+                import tiktoken  # noqa: PLC0415
 
                 if obj_type is tiktoken.Encoding:
                     pklregister(obj_type)(_save_tiktokenEncoding)
             if "torch" in sys.modules:
-                import torch  # type: ignore
+                import torch  # noqa: PLC0415
 
                 if issubclass(obj_type, torch.Tensor):
                     pklregister(obj_type)(_save_torchTensor)
@@ -89,7 +101,7 @@ def save(self, obj, save_persistent_id=True):
                 if issubclass(obj_type, torch.nn.Module):
                     obj = getattr(obj, "_orig_mod", obj)
             if "transformers" in sys.modules:
-                import transformers  # type: ignore
+                import transformers  # noqa: PLC0415
 
                 if issubclass(obj_type, transformers.PreTrainedTokenizerBase):
                     pklregister(obj_type)(_save_transformersPreTrainedTokenizerBase)
@@ -99,90 +111,92 @@ def save(self, obj, save_persistent_id=True):
             obj = getattr(obj, "_torchdynamo_orig_callable", obj)
         dill.Pickler.save(self, obj, save_persistent_id=save_persistent_id)
 
-    def _batch_setitems(self, items):
+    def _batch_setitems(self, items: Iterable[tuple[Any, Any]]) -> None:
         if self._legacy_no_dict_keys_sorting:
-            return super()._batch_setitems(items)
+            super()._batch_setitems(items)
+            return
         # Ignore the order of keys in a dict
         try:
             # Faster, but fails for unorderable elements
-            items = sorted(items)
+            sorted_items = sorted(items)
         except Exception:  # TypeError, decimal.InvalidOperation, etc.
-            items = sorted(items, key=lambda x: Hasher.hash(x[0]))
-        dill.Pickler._batch_setitems(self, items)
+            sorted_items = sorted(items, key=lambda x: Hasher.hash(x[0]))
+        super()._batch_setitems(sorted_items)
 
-    def memoize(self, obj):
+    def memoize(self, obj: Any) -> None:
+        """Memoize an object, skipping strings to avoid id issues."""
         # Don't memoize strings since two identical strings can have different Python ids
         if type(obj) is not str:  # noqa: E721
             dill.Pickler.memoize(self, obj)
 
 
-def pklregister(t):
+def pklregister(t: Any) -> Callable[[Any], Any]:
     """Register a custom reducer for the type."""
 
-    def proxy(func):
+    def proxy(func: Any) -> Any:
         Pickler.dispatch[t] = func
         return func
 
     return proxy
 
 
-def dump(obj, file):
+def dump(obj: Any, file: BinaryIO) -> None:
     """Pickle an object to a file."""
     Pickler(file, recurse=True).dump(obj)
 
 
-def dumps(obj):
+def dumps(obj: Any) -> bytes:
     """Pickle an object to a string."""
     file = BytesIO()
     dump(obj, file)
     return file.getvalue()
 
 
-def log(pickler, msg):
-    pass
+def log(pickler: Any, msg: Any) -> None:
+    """Log a message from the pickler (no-op)."""
 
 
-def _save_regexPattern(pickler, obj):
-    import regex  # type: ignore
+def _save_regexPattern(pickler: Any, obj: Any) -> None:
+    import regex  # noqa: PLC0415
 
     log(pickler, f"Re: {obj}")
     args = (obj.pattern, obj.flags)
     pickler.save_reduce(regex.compile, args, obj=obj)
     log(pickler, "# Re")
 
 
-def _save_tiktokenEncoding(pickler, obj):
-    import tiktoken  # type: ignore
+def _save_tiktokenEncoding(pickler: Any, obj: Any) -> None:
+    import tiktoken  # noqa: PLC0415
 
     log(pickler, f"Enc: {obj}")
-    args = (obj.name, obj._pat_str, obj._mergeable_ranks, obj._special_tokens)
+    args = (obj.name, obj._pat_str, obj._mergeable_ranks, obj._special_tokens)  # noqa: SLF001
     pickler.save_reduce(tiktoken.Encoding, args, obj=obj)
     log(pickler, "# Enc")
 
 
-def _save_torchTensor(pickler, obj):
-    import torch  # type: ignore
+def _save_torchTensor(pickler: Any, obj: Any) -> None:
+    import torch  # noqa: PLC0415
 
     # `torch.from_numpy` is not picklable in `torch>=1.11.0`
-    def create_torchTensor(np_array, dtype=None):
+    def create_torchTensor(np_array: Any, dtype: Any = None) -> "torch.Tensor":
         tensor = torch.from_numpy(np_array)
         if dtype:
             tensor = tensor.type(dtype)
         return tensor
 
     log(pickler, f"To: {obj}")
     if obj.dtype == torch.bfloat16:
-        args = (obj.detach().to(torch.float).cpu().numpy(), torch.bfloat16)
+        args: tuple[Any, ...] = (obj.detach().to(torch.float).cpu().numpy(), torch.bfloat16)
     else:
         args = (obj.detach().cpu().numpy(),)
     pickler.save_reduce(create_torchTensor, args, obj=obj)
     log(pickler, "# To")
 
 
-def _save_torchGenerator(pickler, obj):
-    import torch  # type: ignore
+def _save_torchGenerator(pickler: Any, obj: Any) -> None:
+    import torch  # noqa: PLC0415
 
-    def create_torchGenerator(state):
+    def create_torchGenerator(state: Any) -> "torch.Generator":
         generator = torch.Generator()
         generator.set_state(state)
         return generator
@@ -193,10 +207,10 @@ def create_torchGenerator(state):
     log(pickler, "# Ge")
 
 
-def _save_spacyLanguage(pickler, obj):
-    import spacy  # type: ignore
+def _save_spacyLanguage(pickler: Any, obj: Any) -> None:
+    import spacy  # noqa: PLC0415
 
-    def create_spacyLanguage(config, bytes):
+    def create_spacyLanguage(config: Any, bytes: Any) -> "spacy.Language":
         lang_cls = spacy.util.get_lang_class(config["nlp"]["lang"])
         lang_inst = lang_cls.from_config(config)
         return lang_inst.from_bytes(bytes)
@@ -207,11 +221,11 @@ def create_spacyLanguage(config, bytes):
     log(pickler, "# Sp")
 
 
-def _save_transformersPreTrainedTokenizerBase(pickler, obj):
+def _save_transformersPreTrainedTokenizerBase(pickler: Any, obj: Any) -> None:
     log(pickler, f"Tok: {obj}")
     # Ignore the `cache` attribute
     state = obj.__dict__
     if "cache" in state and isinstance(state["cache"], dict):
         state["cache"] = {}
     pickler.save_reduce(type(obj), (), state=state, obj=obj)
-    log(pickler, "# Tok")
+    log(pickler, "# Tok")

app/utils/errors.py

@@ -1,18 +1,43 @@
+"""Utilities for creating error responses."""
+
 from http import HTTPStatus
-from typing import Union
+
 
 def create_error_response(
     message: str,
     err_type: str = "internal_error",
-    status_code: Union[int, HTTPStatus] = HTTPStatus.INTERNAL_SERVER_ERROR,
-    param: str = None,
-    code: str = None
-):
+    status_code: int | HTTPStatus = HTTPStatus.INTERNAL_SERVER_ERROR,
+    param: str | None = None,
+    code: str | None = None,
+) -> dict[str, object]:
+    """
+    Create a standardized error response dictionary.
+
+    Parameters
+    ----------
+    message : str
+        The error message to include in the response.
+    err_type : str, optional
+        The type of error, by default "internal_error".
+    status_code : int or HTTPStatus, optional
+        The HTTP status code, by default HTTPStatus.INTERNAL_SERVER_ERROR.
+    param : str or None, optional
+        The parameter that caused the error, by default None.
+    code : str or None, optional
+        The error code, by default None.
+
+    Returns
+    -------
+    dict[str, object]
+        A dictionary containing the error response structure.
+    """
     return {
         "error": {
             "message": message,
             "type": err_type,
             "param": param,
-            "code": str(code or (status_code.value if isinstance(status_code, HTTPStatus) else status_code))
+            "code": str(
+                code or (status_code.value if isinstance(status_code, HTTPStatus) else status_code)
+            ),
         }
-    }
+    }

app/utils/outlines_transformer_tokenizer.py

@@ -1,15 +1,18 @@
+"""Custom tokenizer class for outlines integration."""
+
 from outlines.models.transformers import TransformerTokenizer
 
 from .dill import Hasher
 
 
 class OutlinesTransformerTokenizer(TransformerTokenizer):
     """
-    Update the outlines TransformerTokenizer to use our own Hasher class, so that we don't need the datasets dependency
+    Update the outlines TransformerTokenizer to use our own Hasher class, so that we don't need the datasets dependency.
 
     This class and the external dependency can be removed when the following import is deleted
     https://github.com/dottxt-ai/outlines/blob/69418d/outlines/models/transformers.py#L117
     """
 
-    def __hash__(self):
+    def __hash__(self) -> int:
+        """Return a hash based on the tokenizer using custom hasher."""
         return hash(Hasher.hash(self.tokenizer))