chore(style): reduce line length and truncate compression

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

openllm-python/src/openllm/models/auto/__init__.py

@@ -3,7 +3,12 @@ import typing as t, os
 import openllm
 from openllm_core.utils import LazyModule, is_flax_available, is_tf_available, is_torch_available, is_vllm_available
 from openllm_core.config import AutoConfig as AutoConfig, CONFIG_MAPPING as CONFIG_MAPPING, CONFIG_MAPPING_NAMES as CONFIG_MAPPING_NAMES
-_import_structure: dict[str, list[str]] = {"modeling_auto": ["MODEL_MAPPING_NAMES"], "modeling_flax_auto": ["MODEL_FLAX_MAPPING_NAMES"], "modeling_tf_auto": ["MODEL_TF_MAPPING_NAMES"], "modeling_vllm_auto": ["MODEL_VLLM_MAPPING_NAMES"]}
+_import_structure: dict[str, list[str]] = {
+    "modeling_auto": ["MODEL_MAPPING_NAMES"],
+    "modeling_flax_auto": ["MODEL_FLAX_MAPPING_NAMES"],
+    "modeling_tf_auto": ["MODEL_TF_MAPPING_NAMES"],
+    "modeling_vllm_auto": ["MODEL_VLLM_MAPPING_NAMES"]
+}
 if t.TYPE_CHECKING:
   from .modeling_auto import MODEL_MAPPING_NAMES as MODEL_MAPPING_NAMES
   from .modeling_flax_auto import MODEL_FLAX_MAPPING_NAMES as MODEL_FLAX_MAPPING_NAMES

openllm-python/src/openllm/models/auto/factory.py

@@ -23,7 +23,8 @@ class BaseAutoLLMClass:
     raise EnvironmentError(f"Cannot instantiate {self.__class__.__name__} directly. Please use '{self.__class__.__name__}.Runner(model_name)' instead.")
 
   @classmethod
-  def for_model(cls, model: str, /, model_id: str | None = None, model_version: str | None = None, llm_config: openllm.LLMConfig | None = None, ensure_available: bool = False, **attrs: t.Any) -> openllm.LLM[t.Any, t.Any]:
+  def for_model(cls, model: str, /, model_id: str | None = None, model_version: str | None = None, llm_config: openllm.LLMConfig | None = None, ensure_available: bool = False,
+                **attrs: t.Any) -> openllm.LLM[t.Any, t.Any]:
     """The lower level API for creating a LLM instance.
 
     ```python
@@ -62,14 +63,18 @@ class BaseAutoLLMClass:
     llm_class: The runnable to register.
     """
     if hasattr(llm_class, "config_class") and llm_class.config_class is not config_class:
-      raise ValueError(f"The model class you are passing has a `config_class` attribute that is not consistent with the config class you passed (model has {llm_class.config_class} and you passed {config_class}. Fix one of those so they match!")
+      raise ValueError(
+          f"The model class you are passing has a `config_class` attribute that is not consistent with the config class you passed (model has {llm_class.config_class} and you passed {config_class}. Fix one of those so they match!"
+      )
     cls._model_mapping.register(config_class, llm_class)
 
   @classmethod
   def infer_class_from_name(cls, name: str) -> type[openllm.LLM[t.Any, t.Any]]:
     config_class = openllm.AutoConfig.infer_class_from_name(name)
     if config_class in cls._model_mapping: return cls._model_mapping[config_class]
-    raise ValueError(f"Unrecognized configuration class ({config_class}) for {name}. Model name should be one of {', '.join(openllm.CONFIG_MAPPING.keys())} (Registered configuration class: {', '.join([i.__name__ for i in cls._model_mapping.keys()])}).")
+    raise ValueError(
+        f"Unrecognized configuration class ({config_class}) for {name}. Model name should be one of {', '.join(openllm.CONFIG_MAPPING.keys())} (Registered configuration class: {', '.join([i.__name__ for i in cls._model_mapping.keys()])})."
+    )
 def getattribute_from_module(module: types.ModuleType, attr: t.Any) -> t.Any:
   if attr is None: return
   if isinstance(attr, tuple): return tuple(getattribute_from_module(module, a) for a in attr)
@@ -127,13 +132,23 @@ class _LazyAutoMapping(OrderedDict, ReprMixin):
     return bool(self.keys())
 
   def keys(self) -> ConfigModelKeysView:
-    return t.cast("ConfigModelKeysView", [self._load_attr_from_module(key, name) for key, name in self._config_mapping.items() if key in self._model_mapping.keys()] + list(self._extra_content.keys()))
+    return t.cast(
+        "ConfigModelKeysView", [self._load_attr_from_module(key, name) for key, name in self._config_mapping.items() if key in self._model_mapping.keys()] + list(self._extra_content.keys())
+    )
 
   def values(self) -> ConfigModelValuesView:
-    return t.cast("ConfigModelValuesView", [self._load_attr_from_module(key, name) for key, name in self._model_mapping.items() if key in self._config_mapping.keys()] + list(self._extra_content.values()))
+    return t.cast(
+        "ConfigModelValuesView", [self._load_attr_from_module(key, name) for key, name in self._model_mapping.items() if key in self._config_mapping.keys()] + list(
+            self._extra_content.values()
+        )
+    )
 
   def items(self) -> ConfigModelItemsView:
-    return t.cast("ConfigModelItemsView", [(self._load_attr_from_module(key, self._config_mapping[key]), self._load_attr_from_module(key, self._model_mapping[key])) for key in self._model_mapping.keys() if key in self._config_mapping.keys()] + list(self._extra_content.items()))
+    return t.cast(
+        "ConfigModelItemsView",
+        [(self._load_attr_from_module(key, self._config_mapping[key]),
+          self._load_attr_from_module(key, self._model_mapping[key])) for key in self._model_mapping.keys() if key in self._config_mapping.keys()] + list(self._extra_content.items())
+    )
 
   def __iter__(self) -> t.Iterator[type[openllm.LLMConfig]]:
     return iter(t.cast("SupportsIter[t.Iterator[type[openllm.LLMConfig]]]", self.keys()))

openllm-python/src/openllm/models/auto/modeling_auto.py

@@ -3,7 +3,9 @@ import typing as t
 from collections import OrderedDict
 from .factory import BaseAutoLLMClass, _LazyAutoMapping
 from openllm_core.config import CONFIG_MAPPING_NAMES
-MODEL_MAPPING_NAMES = OrderedDict([("chatglm", "ChatGLM"), ("dolly_v2", "DollyV2"), ("falcon", "Falcon"), ("flan_t5", "FlanT5"), ("gpt_neox", "GPTNeoX"), ("llama", "Llama"), ("mpt", "MPT"), ("opt", "OPT"), ("stablelm", "StableLM"), ("starcoder", "StarCoder"), ("baichuan", "Baichuan")])
+MODEL_MAPPING_NAMES = OrderedDict([("chatglm", "ChatGLM"), ("dolly_v2", "DollyV2"), ("falcon", "Falcon"), ("flan_t5", "FlanT5"), ("gpt_neox", "GPTNeoX"), ("llama", "Llama"), ("mpt", "MPT"), (
+    "opt", "OPT"
+), ("stablelm", "StableLM"), ("starcoder", "StarCoder"), ("baichuan", "Baichuan")])
 MODEL_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_MAPPING_NAMES)
 class AutoLLM(BaseAutoLLMClass):
   _model_mapping: t.ClassVar = MODEL_MAPPING

openllm-python/src/openllm/models/auto/modeling_vllm_auto.py

@@ -3,7 +3,9 @@ import typing as t
 from collections import OrderedDict
 from .factory import BaseAutoLLMClass, _LazyAutoMapping
 from openllm_core.config import CONFIG_MAPPING_NAMES
-MODEL_VLLM_MAPPING_NAMES = OrderedDict([("baichuan", "VLLMBaichuan"), ("dolly_v2", "VLLMDollyV2"), ("falcon", "VLLMFalcon"), ("gpt_neox", "VLLMGPTNeoX"), ("mpt", "VLLMMPT"), ("opt", "VLLMOPT"), ("stablelm", "VLLMStableLM"), ("starcoder", "VLLMStarCoder"), ("llama", "VLLMLlama")])
+MODEL_VLLM_MAPPING_NAMES = OrderedDict([("baichuan", "VLLMBaichuan"), ("dolly_v2", "VLLMDollyV2"), ("falcon", "VLLMFalcon"), ("gpt_neox", "VLLMGPTNeoX"), ("mpt", "VLLMMPT"), (
+    "opt", "VLLMOPT"
+), ("stablelm", "VLLMStableLM"), ("starcoder", "VLLMStarCoder"), ("llama", "VLLMLlama")])
 MODEL_VLLM_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_VLLM_MAPPING_NAMES)
 class AutoVLLM(BaseAutoLLMClass):
   _model_mapping: t.ClassVar = MODEL_VLLM_MAPPING

openllm-python/src/openllm/models/dolly_v2/modeling_dolly_v2.py

@@ -4,7 +4,7 @@ from openllm_core._typing_compat import overload
 from openllm_core.config.configuration_dolly_v2 import DEFAULT_PROMPT_TEMPLATE, END_KEY, RESPONSE_KEY, get_special_token_id
 
 if t.TYPE_CHECKING: import torch, transformers, tensorflow as tf
-else: torch, transformers, tf = openllm.utils.LazyLoader("torch", globals(), "torch"), openllm.utils.LazyLoader("transformers", globals(), "transformers"), openllm.utils.LazyLoader("tf", globals(), "tensorflow")
+else:  torch, transformers, tf = openllm.utils.LazyLoader("torch", globals(), "torch"), openllm.utils.LazyLoader("transformers", globals(), "transformers"), openllm.utils.LazyLoader("tf", globals(), "tensorflow")
 logger = logging.getLogger(__name__)
 @overload
 def get_pipeline(model: transformers.PreTrainedModel, tokenizer: transformers.PreTrainedTokenizer, _init: t.Literal[True] = True, **attrs: t.Any) -> transformers.Pipeline:
@@ -52,7 +52,12 @@ def get_pipeline(model: transformers.PreTrainedModel, tokenizer: transformers.Pr
       input_ids, attention_mask = input_tensors["input_ids"], input_tensors.get("attention_mask", None)
       if input_ids.shape[1] == 0: input_ids, attention_mask, in_b = None, None, 1
       else: in_b = input_ids.shape[0]
-      generated_sequence = self.model.generate(input_ids=input_ids.to(self.model.device) if input_ids is not None else None, attention_mask=attention_mask.to(self.model.device) if attention_mask is not None else None, pad_token_id=self.tokenizer.pad_token_id, **generate_kwargs)
+      generated_sequence = self.model.generate(
+          input_ids=input_ids.to(self.model.device) if input_ids is not None else None,
+          attention_mask=attention_mask.to(self.model.device) if attention_mask is not None else None,
+          pad_token_id=self.tokenizer.pad_token_id,
+          **generate_kwargs
+      )
       out_b = generated_sequence.shape[0]
       if self.framework == "pt": generated_sequence = generated_sequence.reshape(in_b, out_b // in_b, *generated_sequence.shape[1:])
       elif self.framework == "tf": generated_sequence = tf.reshape(generated_sequence, (in_b, out_b // in_b, *generated_sequence.shape[1:]))

openllm-python/src/openllm/models/falcon/modeling_falcon.py

@@ -12,7 +12,14 @@ class Falcon(openllm.LLM["transformers.PreTrainedModel", "transformers.PreTraine
   def generate(self, prompt: str, **attrs: t.Any) -> list[str]:
     eos_token_id, inputs = attrs.pop("eos_token_id", self.tokenizer.eos_token_id), self.tokenizer(prompt, return_tensors="pt").to(self.device)
     with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):  # type: ignore[attr-defined]
-      return self.tokenizer.batch_decode(self.model.generate(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], generation_config=self.config.model_construct_env(eos_token_id=eos_token_id, **attrs).to_generation_config()), skip_special_tokens=True)
+      return self.tokenizer.batch_decode(
+          self.model.generate(
+              input_ids=inputs["input_ids"],
+              attention_mask=inputs["attention_mask"],
+              generation_config=self.config.model_construct_env(eos_token_id=eos_token_id, **attrs).to_generation_config()
+          ),
+          skip_special_tokens=True
+      )
 
   def generate_one(self, prompt: str, stop: list[str], **preprocess_generate_kwds: t.Any) -> list[dict[t.Literal["generated_text"], str]]:
     max_new_tokens, encoded_inputs = preprocess_generate_kwds.pop("max_new_tokens", 200), self.tokenizer(prompt, return_tensors="pt").to(self.device)

openllm-python/src/openllm/models/flan_t5/modeling_flan_t5.py

@@ -7,7 +7,10 @@ class FlanT5(openllm.LLM["transformers.T5ForConditionalGeneration", "transformer
   def generate(self, prompt: str, **attrs: t.Any) -> list[str]:
     import torch
     with torch.inference_mode():
-      return self.tokenizer.batch_decode(self.model.generate(**self.tokenizer(prompt, return_tensors="pt").to(self.device), do_sample=True, generation_config=self.config.model_construct_env(**attrs).to_generation_config()), skip_special_tokens=True)
+      return self.tokenizer.batch_decode(
+          self.model.generate(**self.tokenizer(prompt, return_tensors="pt").to(self.device), do_sample=True, generation_config=self.config.model_construct_env(**attrs).to_generation_config()),
+          skip_special_tokens=True
+      )
 
   def embeddings(self, prompts: list[str]) -> openllm.LLMEmbeddings:
     import torch, torch.nn.functional as F

openllm-python/src/openllm/models/flan_t5/modeling_flax_flan_t5.py

@@ -6,11 +6,38 @@ if t.TYPE_CHECKING: import transformers
 class FlaxFlanT5(openllm.LLM["transformers.FlaxT5ForConditionalGeneration", "transformers.T5TokenizerFast"]):
   __openllm_internal__ = True
 
-  def sanitize_parameters(self, prompt: str, max_new_tokens: int | None = None, temperature: float | None = None, top_k: int | None = None, top_p: float | None = None, repetition_penalty: float | None = None, decoder_start_token_id: int | None = None, use_default_prompt_template: bool = True, **attrs: t.Any) -> tuple[str, dict[str, t.Any], dict[str, t.Any]]:
+  def sanitize_parameters(
+      self,
+      prompt: str,
+      max_new_tokens: int | None = None,
+      temperature: float | None = None,
+      top_k: int | None = None,
+      top_p: float | None = None,
+      repetition_penalty: float | None = None,
+      decoder_start_token_id: int | None = None,
+      use_default_prompt_template: bool = True,
+      **attrs: t.Any
+  ) -> tuple[str, dict[str, t.Any], dict[str, t.Any]]:
     if decoder_start_token_id is None: decoder_start_token_id = 0
-    return process_prompt(prompt, DEFAULT_PROMPT_TEMPLATE, use_default_prompt_template, **attrs), {"max_new_tokens": max_new_tokens, "temperature": temperature, "top_k": top_k, "top_p": top_p, "repetition_penalty": repetition_penalty, "decoder_start_token_id": decoder_start_token_id}, {}
+    return process_prompt(prompt, DEFAULT_PROMPT_TEMPLATE, use_default_prompt_template, **attrs), {
+        "max_new_tokens": max_new_tokens,
+        "temperature": temperature,
+        "top_k": top_k,
+        "top_p": top_p,
+        "repetition_penalty": repetition_penalty,
+        "decoder_start_token_id": decoder_start_token_id
+    }, {}
 
   def generate(self, prompt: str, **attrs: t.Any) -> list[str]:
     # NOTE: decoder_start_token_id is extracted from https://huggingface.co/google/flan-t5-small/tree/main as it is required for encoder-decoder generation.
     decoder_start_token_id = attrs.pop("decoder_start_token_id", 0)
-    return self.tokenizer.batch_decode(self.model.generate(self.tokenizer(prompt, return_tensors="np")["input_ids"], do_sample=True, generation_config=self.config.model_construct_env(**attrs).to_generation_config(), decoder_start_token_id=decoder_start_token_id).sequences, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+    return self.tokenizer.batch_decode(
+        self.model.generate(
+            self.tokenizer(prompt, return_tensors="np")["input_ids"],
+            do_sample=True,
+            generation_config=self.config.model_construct_env(**attrs).to_generation_config(),
+            decoder_start_token_id=decoder_start_token_id
+        ).sequences,
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=True
+    )

openllm-python/src/openllm/models/flan_t5/modeling_tf_flan_t5.py

@@ -5,4 +5,7 @@ class TFFlanT5(openllm.LLM["transformers.TFT5ForConditionalGeneration", "transfo
   __openllm_internal__ = True
 
   def generate(self, prompt: str, **attrs: t.Any) -> list[str]:
-    return self.tokenizer.batch_decode(self.model.generate(self.tokenizer(prompt, return_tensors="tf").input_ids, do_sample=True, generation_config=self.config.model_construct_env(**attrs).to_generation_config()), skip_special_tokens=True)
+    return self.tokenizer.batch_decode(
+        self.model.generate(self.tokenizer(prompt, return_tensors="tf").input_ids, do_sample=True, generation_config=self.config.model_construct_env(**attrs).to_generation_config()),
+        skip_special_tokens=True
+    )

openllm-python/src/openllm/models/gpt_neox/modeling_gpt_neox.py

@@ -20,4 +20,12 @@ class GPTNeoX(openllm.LLM["transformers.GPTNeoXForCausalLM", "transformers.GPTNe
   def generate(self, prompt: str, **attrs: t.Any) -> list[str]:
     import torch
     with torch.inference_mode():
-      return self.tokenizer.batch_decode(self.model.generate(self.tokenizer(prompt, return_tensors="pt").to(self.device).input_ids, do_sample=True, generation_config=self.config.model_construct_env(**attrs).to_generation_config(), pad_token_id=self.tokenizer.eos_token_id, stopping_criteria=openllm.StoppingCriteriaList([openllm.StopOnTokens()])))
+      return self.tokenizer.batch_decode(
+          self.model.generate(
+              self.tokenizer(prompt, return_tensors="pt").to(self.device).input_ids,
+              do_sample=True,
+              generation_config=self.config.model_construct_env(**attrs).to_generation_config(),
+              pad_token_id=self.tokenizer.eos_token_id,
+              stopping_criteria=openllm.StoppingCriteriaList([openllm.StopOnTokens()])
+          )
+      )

openllm-python/src/openllm/models/mpt/modeling_mpt.py

@@ -4,12 +4,17 @@ from openllm.utils import generate_labels, is_triton_available
 if t.TYPE_CHECKING: import transformers, torch
 
 logger = logging.getLogger(__name__)
-def get_mpt_config(model_id_or_path: str, max_sequence_length: int, device: torch.device | str | int | None, device_map: str | None = None, trust_remote_code: bool = True) -> transformers.PretrainedConfig:
+def get_mpt_config(
+    model_id_or_path: str, max_sequence_length: int, device: torch.device | str | int | None, device_map: str | None = None, trust_remote_code: bool = True
+) -> transformers.PretrainedConfig:
   import torch
   config = transformers.AutoConfig.from_pretrained(model_id_or_path, trust_remote_code=trust_remote_code)
   if hasattr(config, "init_device") and device_map is None and isinstance(device, (str, torch.device)): config.init_device = str(device)
   if hasattr(config, "attn_config") and is_triton_available(): config.attn_config["attn_impl"] = "triton"
-  else: logger.debug("'triton' is not available, Flash Attention will use the default Torch implementation. For faster inference, make sure to install triton with 'pip install \"git+https://github.com/openai/triton.git#egg=triton&subdirectory=python\"'")
+  else:
+    logger.debug(
+        "'triton' is not available, Flash Attention will use the default Torch implementation. For faster inference, make sure to install triton with 'pip install \"git+https://github.com/openai/triton.git#egg=triton&subdirectory=python\"'"
+    )
   # setting max_seq_len
   config.max_seq_len = max_sequence_length
   return config
@@ -46,7 +51,9 @@ class MPT(openllm.LLM["transformers.PreTrainedModel", "transformers.GPTNeoXToken
     device_map = attrs.pop("device_map", None)
     trust_remote_code = attrs.pop("trust_remote_code", True)
     config = get_mpt_config(self._bentomodel.path, self.config.max_sequence_length, self.device, device_map=device_map, trust_remote_code=trust_remote_code,)
-    model = transformers.AutoModelForCausalLM.from_pretrained(self._bentomodel.path, config=config, trust_remote_code=trust_remote_code, torch_dtype=torch_dtype, device_map=device_map, **attrs)
+    model = transformers.AutoModelForCausalLM.from_pretrained(
+        self._bentomodel.path, config=config, trust_remote_code=trust_remote_code, torch_dtype=torch_dtype, device_map=device_map, **attrs
+    )
     model.tie_weights()
     return model
 
@@ -54,7 +61,12 @@ class MPT(openllm.LLM["transformers.PreTrainedModel", "transformers.GPTNeoXToken
     import torch
     llm_config = self.config.model_construct_env(**attrs)
     inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
-    attrs = {"do_sample": False if llm_config["temperature"] == 0 else True, "eos_token_id": self.tokenizer.eos_token_id, "pad_token_id": self.tokenizer.pad_token_id, "generation_config": llm_config.to_generation_config()}
+    attrs = {
+        "do_sample": False if llm_config["temperature"] == 0 else True,
+        "eos_token_id": self.tokenizer.eos_token_id,
+        "pad_token_id": self.tokenizer.pad_token_id,
+        "generation_config": llm_config.to_generation_config()
+    }
     with torch.inference_mode():
       if torch.cuda.is_available():
         with torch.autocast("cuda", torch.float16):  # type: ignore[attr-defined]

openllm-python/src/openllm/models/opt/modeling_flax_opt.py

@@ -13,10 +13,27 @@ class FlaxOPT(openllm.LLM["transformers.TFOPTForCausalLM", "transformers.GPT2Tok
   def import_model(self, *args: t.Any, trust_remote_code: bool = False, **attrs: t.Any) -> bentoml.Model:
     config, tokenizer = transformers.AutoConfig.from_pretrained(self.model_id), transformers.AutoTokenizer.from_pretrained(self.model_id, **self.llm_parameters[-1])
     tokenizer.pad_token_id = config.pad_token_id
-    return bentoml.transformers.save_model(self.tag, transformers.FlaxAutoModelForCausalLM.from_pretrained(self.model_id, **attrs), custom_objects={"tokenizer": tokenizer}, labels=generate_labels(self))
+    return bentoml.transformers.save_model(
+        self.tag, transformers.FlaxAutoModelForCausalLM.from_pretrained(self.model_id, **attrs), custom_objects={"tokenizer": tokenizer}, labels=generate_labels(self)
+    )
 
-  def sanitize_parameters(self, prompt: str, max_new_tokens: int | None = None, temperature: float | None = None, top_k: int | None = None, num_return_sequences: int | None = None, repetition_penalty: float | None = None, use_default_prompt_template: bool = False, **attrs: t.Any) -> tuple[str, dict[str, t.Any], dict[str, t.Any]]:
-    return process_prompt(prompt, DEFAULT_PROMPT_TEMPLATE, use_default_prompt_template, **attrs), {"max_new_tokens": max_new_tokens, "temperature": temperature, "top_k": top_k, "num_return_sequences": num_return_sequences, "repetition_penalty": repetition_penalty}, {}
+  def sanitize_parameters(
+      self,
+      prompt: str,
+      max_new_tokens: int | None = None,
+      temperature: float | None = None,
+      top_k: int | None = None,
+      num_return_sequences: int | None = None,
+      repetition_penalty: float | None = None,
+      use_default_prompt_template: bool = False,
+      **attrs: t.Any
+  ) -> tuple[str, dict[str, t.Any], dict[str, t.Any]]:
+    return process_prompt(prompt, DEFAULT_PROMPT_TEMPLATE, use_default_prompt_template, **attrs), {
+        "max_new_tokens": max_new_tokens, "temperature": temperature, "top_k": top_k, "num_return_sequences": num_return_sequences, "repetition_penalty": repetition_penalty
+    }, {}
 
   def generate(self, prompt: str, **attrs: t.Any) -> list[str]:
-    return self.tokenizer.batch_decode(self.model.generate(**self.tokenizer(prompt, return_tensors="np"), do_sample=True, generation_config=self.config.model_construct_env(**attrs).to_generation_config()).sequences, skip_special_tokens=True)
+    return self.tokenizer.batch_decode(
+        self.model.generate(**self.tokenizer(prompt, return_tensors="np"), do_sample=True, generation_config=self.config.model_construct_env(**attrs).to_generation_config()).sequences,
+        skip_special_tokens=True
+    )

openllm-python/src/openllm/models/opt/modeling_opt.py

@@ -14,4 +14,7 @@ class OPT(openllm.LLM["transformers.OPTForCausalLM", "transformers.GPT2Tokenizer
   def generate(self, prompt: str, **attrs: t.Any) -> list[str]:
     import torch
     with torch.inference_mode():
-      return self.tokenizer.batch_decode(self.model.generate(**self.tokenizer(prompt, return_tensors="pt").to(self.device), do_sample=True, generation_config=self.config.model_construct_env(**attrs).to_generation_config()), skip_special_tokens=True)
+      return self.tokenizer.batch_decode(
+          self.model.generate(**self.tokenizer(prompt, return_tensors="pt").to(self.device), do_sample=True, generation_config=self.config.model_construct_env(**attrs).to_generation_config()),
+          skip_special_tokens=True
+      )

openllm-python/src/openllm/models/opt/modeling_tf_opt.py

@@ -9,7 +9,15 @@ class TFOPT(openllm.LLM["transformers.TFOPTForCausalLM", "transformers.GPT2Token
     import transformers
     config, tokenizer = transformers.AutoConfig.from_pretrained(self.model_id), transformers.AutoTokenizer.from_pretrained(self.model_id, **self.llm_parameters[-1])
     tokenizer.pad_token_id = config.pad_token_id
-    return bentoml.transformers.save_model(self.tag, transformers.TFOPTForCausalLM.from_pretrained(self.model_id, trust_remote_code=trust_remote_code, **attrs), custom_objects={"tokenizer": tokenizer}, labels=generate_labels(self))
+    return bentoml.transformers.save_model(
+        self.tag,
+        transformers.TFOPTForCausalLM.from_pretrained(self.model_id, trust_remote_code=trust_remote_code, **attrs),
+        custom_objects={"tokenizer": tokenizer},
+        labels=generate_labels(self)
+    )
 
   def generate(self, prompt: str, **attrs: t.Any) -> list[str]:
-    return self.tokenizer.batch_decode(self.model.generate(**self.tokenizer(prompt, return_tensors="tf"), do_sample=True, generation_config=self.config.model_construct_env(**attrs).to_generation_config()), skip_special_tokens=True)
+    return self.tokenizer.batch_decode(
+        self.model.generate(**self.tokenizer(prompt, return_tensors="tf"), do_sample=True, generation_config=self.config.model_construct_env(**attrs).to_generation_config()),
+        skip_special_tokens=True
+    )

openllm-python/src/openllm/models/opt/modeling_vllm_opt.py

@@ -7,5 +7,16 @@ class VLLMOPT(openllm.LLM["vllm.LLMEngine", "transformers.GPT2Tokenizer"]):
   __openllm_internal__ = True
   tokenizer_id = "local"
 
-  def sanitize_parameters(self, prompt: str, max_new_tokens: int | None = None, temperature: float | None = None, top_k: int | None = None, num_return_sequences: int | None = None, use_default_prompt_template: bool = True, **attrs: t.Any) -> tuple[str, dict[str, t.Any], dict[str, t.Any]]:
-    return process_prompt(prompt, DEFAULT_PROMPT_TEMPLATE, use_default_prompt_template, **attrs), {"max_new_tokens": max_new_tokens, "temperature": temperature, "top_k": top_k, "num_return_sequences": num_return_sequences}, {}
+  def sanitize_parameters(
+      self,
+      prompt: str,
+      max_new_tokens: int | None = None,
+      temperature: float | None = None,
+      top_k: int | None = None,
+      num_return_sequences: int | None = None,
+      use_default_prompt_template: bool = True,
+      **attrs: t.Any
+  ) -> tuple[str, dict[str, t.Any], dict[str, t.Any]]:
+    return process_prompt(prompt, DEFAULT_PROMPT_TEMPLATE, use_default_prompt_template, **attrs), {
+        "max_new_tokens": max_new_tokens, "temperature": temperature, "top_k": top_k, "num_return_sequences": num_return_sequences
+    }, {}

openllm-python/src/openllm/models/stablelm/modeling_stablelm.py

@@ -16,4 +16,15 @@ class StableLM(openllm.LLM["transformers.GPTNeoXForCausalLM", "transformers.GPTN
   def generate(self, prompt: str, **attrs: t.Any) -> list[str]:
     import torch
     with torch.inference_mode():
-      return [self.tokenizer.decode(self.model.generate(**self.tokenizer(prompt, return_tensors="pt").to(self.device), do_sample=True, generation_config=self.config.model_construct_env(**attrs).to_generation_config(), pad_token_id=self.tokenizer.eos_token_id, stopping_criteria=openllm.StoppingCriteriaList([openllm.StopOnTokens()]))[0], skip_special_tokens=True)]
+      return [
+          self.tokenizer.decode(
+              self.model.generate(
+                  **self.tokenizer(prompt, return_tensors="pt").to(self.device),
+                  do_sample=True,
+                  generation_config=self.config.model_construct_env(**attrs).to_generation_config(),
+                  pad_token_id=self.tokenizer.eos_token_id,
+                  stopping_criteria=openllm.StoppingCriteriaList([openllm.StopOnTokens()])
+              )[0],
+              skip_special_tokens=True
+          )
+      ]

openllm-python/src/openllm/models/starcoder/modeling_starcoder.py

@@ -27,7 +27,12 @@ class StarCoder(openllm.LLM["transformers.GPTBigCodeForCausalLM", "transformers.
     with torch.inference_mode():
       # eos_token_id=self.tokenizer.convert_tokens_to_ids("<|end|>"), # NOTE: this is for finetuning starcoder
       # NOTE: support fine-tuning starcoder
-      result_tensor = self.model.generate(self.tokenizer.encode(prompt, return_tensors="pt").to(self.device), do_sample=True, pad_token_id=self.tokenizer.eos_token_id, generation_config=self.config.model_construct_env(**attrs).to_generation_config())
+      result_tensor = self.model.generate(
+          self.tokenizer.encode(prompt, return_tensors="pt").to(self.device),
+          do_sample=True,
+          pad_token_id=self.tokenizer.eos_token_id,
+          generation_config=self.config.model_construct_env(**attrs).to_generation_config()
+      )
       # TODO: We will probably want to return the tokenizer here so that we can manually process this
       # return (skip_special_tokens=False, clean_up_tokenization_spaces=False))
       return self.tokenizer.batch_decode(result_tensor[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)