refactor: cleanup typing to expose correct API (#576)

Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com>

openllm-python/src/openllm/serialisation/__init__.py

@@ -1,9 +1,9 @@
-'''Serialisation utilities for OpenLLM.
+"""Serialisation utilities for OpenLLM.
 
 Currently supports transformers for PyTorch, and vLLM.
 
 Currently, GGML format is working in progress.
-'''
+"""
 from __future__ import annotations
 import importlib
 import typing as t
@@ -32,11 +32,11 @@ else:
 P = ParamSpec('P')
 
 def load_tokenizer(llm: openllm.LLM[t.Any, T], **tokenizer_attrs: t.Any) -> T:
-  '''Load the tokenizer from BentoML store.
+  """Load the tokenizer from BentoML store.
 
   By default, it will try to find the bentomodel whether it is in store..
   If model is not found, it will raises a ``bentoml.exceptions.NotFound``.
-  '''
+  """
   from .transformers._helpers import process_config
 
   config, *_ = process_config(llm.bentomodel.path, llm.trust_remote_code)

openllm-python/src/openllm/serialisation/ggml.py

@@ -1,7 +1,7 @@
-'''Serialisation related implementation for GGML-based implementation.
+"""Serialisation related implementation for GGML-based implementation.
 
 This requires ctransformers to be installed.
-'''
+"""
 from __future__ import annotations
 import typing as t
 

openllm-python/src/openllm/serialisation/transformers/__init__.py

@@ -1,4 +1,4 @@
-'''Serialisation related implementation for Transformers-based implementation.'''
+"""Serialisation related implementation for Transformers-based implementation."""
 from __future__ import annotations
 import importlib
 import logging
@@ -150,13 +150,13 @@ def import_model(llm: openllm.LLM[M, T], *decls: t.Any, trust_remote_code: bool,
     return bentomodel
 
 def get(llm: openllm.LLM[M, T], auto_import: bool = False) -> bentoml.Model:
-  '''Return an instance of ``bentoml.Model`` from given LLM instance.
+  """Return an instance of ``bentoml.Model`` from given LLM instance.
 
   By default, it will try to check the model in the local store.
   If model is not found, and ``auto_import`` is set to True, it will try to import the model from HuggingFace Hub.
 
   Otherwise, it will raises a ``bentoml.exceptions.NotFound``.
-  '''
+  """
   try:
     model = bentoml.models.get(llm.tag)
     backend = model.info.labels['backend']

openllm-python/src/openllm/serialisation/transformers/_helpers.py

@@ -26,7 +26,7 @@ def get_hash(config: transformers.PretrainedConfig) -> str:
   return _commit_hash
 
 def process_config(model_id: str, trust_remote_code: bool, **attrs: t.Any) -> tuple[transformers.PretrainedConfig, DictStrAny, DictStrAny]:
-  '''A helper function that correctly parse config and attributes for transformers.PretrainedConfig.
+  """A helper function that correctly parse config and attributes for transformers.PretrainedConfig.
 
   Args:
     model_id: Model id to pass into ``transformers.AutoConfig``.
@@ -35,7 +35,7 @@ def process_config(model_id: str, trust_remote_code: bool, **attrs: t.Any) -> tu
 
   Returns:
     A tuple of ``transformers.PretrainedConfig``, all hub attributes, and remanining attributes that can be used by the Model class.
-  '''
+  """
   config = attrs.pop('config', None)
   # this logic below is synonymous to handling `from_pretrained` attrs.
   hub_attrs = {k: attrs.pop(k) for k in HUB_ATTRS if k in attrs}