snowflake-ml-python 1.6.2__py3-none-any.whl → 1.6.3__py3-none-any.whl

snowflake/ml/model/_deploy_client/warehouse/deploy.py

@@ -1,202 +0,0 @@
-import copy
-import logging
-import posixpath
-import tempfile
-import textwrap
-from types import ModuleType
-from typing import IO, List, Optional, Tuple, TypedDict, Union
-
-from typing_extensions import Unpack
-
-from snowflake.ml._internal import env_utils, file_utils
-from snowflake.ml._internal.exceptions import (
-    error_codes,
-    exceptions as snowml_exceptions,
-)
-from snowflake.ml.model import type_hints as model_types
-from snowflake.ml.model._deploy_client.warehouse import infer_template
-from snowflake.ml.model._packager.model_meta import model_meta
-from snowflake.snowpark import session as snowpark_session, types as st
-
-logger = logging.getLogger(__name__)
-
-
-def _deploy_to_warehouse(
-    session: snowpark_session.Session,
-    *,
-    model_stage_file_path: str,
-    model_meta: model_meta.ModelMetadata,
-    udf_name: str,
-    target_method: str,
-    **kwargs: Unpack[model_types.WarehouseDeployOptions],
-) -> None:
-    """Deploy the model to warehouse as UDF.
-
-    Args:
-        session: Snowpark session.
-        model_stage_file_path: Path to the stored model zip file in the stage.
-        model_meta: Model Metadata.
-        udf_name: Name of the UDF.
-        target_method: The name of the target method to be deployed.
-        **kwargs: Options that control some features in generated udf code.
-
-    Raises:
-        SnowflakeMLException: Raised when model file name is unable to encoded using ASCII.
-        SnowflakeMLException: Raised when incompatible model.
-        SnowflakeMLException: Raised when target method does not exist in model.
-        SnowflakeMLException: Raised when confronting invalid stage location.
-
-    """
-    # TODO(SNOW-862576): Should remove check on ASCII encoding after SNOW-862576 fixed.
-    model_stage_file_name = posixpath.basename(model_stage_file_path)
-    if not file_utils._able_ascii_encode(model_stage_file_name):
-        raise snowml_exceptions.SnowflakeMLException(
-            error_code=error_codes.INVALID_ARGUMENT,
-            original_exception=ValueError(
-                f"Model file name {model_stage_file_name} cannot be encoded using ASCII. Please rename."
-            ),
-        )
-
-    relax_version = kwargs.get("relax_version", False)
-
-    if target_method not in model_meta.signatures.keys():
-        raise snowml_exceptions.SnowflakeMLException(
-            error_code=error_codes.INVALID_ARGUMENT,
-            original_exception=ValueError(f"Target method {target_method} does not exist in model."),
-        )
-
-    final_packages = _get_model_final_packages(model_meta, session, relax_version=relax_version)
-
-    stage_location = kwargs.get("permanent_udf_stage_location", None)
-    if stage_location:
-        stage_location = posixpath.normpath(stage_location.strip())
-        if not stage_location.startswith("@"):
-            raise snowml_exceptions.SnowflakeMLException(
-                error_code=error_codes.INVALID_ARGUMENT,
-                original_exception=ValueError(f"Invalid stage location {stage_location}."),
-            )
-
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False, encoding="utf-8") as f:
-        _write_UDF_py_file(f.file, model_stage_file_name=model_stage_file_name, target_method=target_method, **kwargs)
-        logger.info(f"Generated UDF file is persisted at: {f.name}")
-
-        class _UDFParams(TypedDict):
-            file_path: str
-            func_name: str
-            name: str
-            input_types: List[st.DataType]
-            return_type: st.DataType
-            imports: List[Union[str, Tuple[str, str]]]
-            packages: List[Union[str, ModuleType]]
-
-        params = _UDFParams(
-            file_path=f.name,
-            func_name="infer",
-            name=udf_name,
-            return_type=st.PandasSeriesType(st.MapType(st.StringType(), st.VariantType())),
-            input_types=[st.PandasDataFrameType([st.MapType()])],
-            imports=[model_stage_file_path],
-            packages=list(final_packages),
-        )
-        if stage_location is None:  # Temporary UDF
-            session.udf.register_from_file(**params, replace=True)
-        else:  # Permanent UDF
-            session.udf.register_from_file(
-                **params,
-                replace=kwargs.get("replace_udf", False),
-                is_permanent=True,
-                stage_location=stage_location,
-            )
-
-    logger.info(f"{udf_name} is deployed to warehouse.")
-
-
-def _write_UDF_py_file(
-    f: IO[str],
-    model_stage_file_name: str,
-    target_method: str,
-    **kwargs: Unpack[model_types.WarehouseDeployOptions],
-) -> None:
-    """Generate and write UDF python code into a file
-
-    Args:
-        f: File descriptor to write the python code.
-        model_stage_file_name: Model zip file name.
-        target_method: The name of the target method to be deployed.
-        **kwargs: Options that control some features in generated udf code.
-    """
-    udf_code = infer_template._UDF_CODE_TEMPLATE.format(
-        model_stage_file_name=model_stage_file_name,
-        _KEEP_ORDER_COL_NAME=infer_template._KEEP_ORDER_COL_NAME,
-        target_method=target_method,
-        code_dir_name=model_meta.MODEL_CODE_DIR,
-    )
-    f.write(udf_code)
-    f.flush()
-
-
-def _get_model_final_packages(
-    meta: model_meta.ModelMetadata,
-    session: snowpark_session.Session,
-    relax_version: Optional[bool] = False,
-) -> List[str]:
-    """Generate final packages list of dependency of a model to be deployed to warehouse.
-
-    Args:
-        meta: Model metadata to get dependency information.
-        session: Snowpark connection session.
-        relax_version: Whether or not relax the version restriction when fail to resolve dependencies.
-            Defaults to False.
-
-    Raises:
-        SnowflakeMLException: Raised when PIP requirements and dependencies from non-Snowflake anaconda channel found.
-        SnowflakeMLException: Raised when not all packages are available in snowflake conda channel.
-
-    Returns:
-        List of final packages string that is accepted by Snowpark register UDF call.
-    """
-
-    if (
-        any(channel.lower() not in [env_utils.DEFAULT_CHANNEL_NAME] for channel in meta.env._conda_dependencies.keys())
-        or meta.env.pip_requirements
-    ):
-        raise snowml_exceptions.SnowflakeMLException(
-            error_code=error_codes.DEPENDENCY_VERSION_ERROR,
-            original_exception=RuntimeError(
-                "PIP requirements and dependencies from non-Snowflake anaconda channel is not supported."
-            ),
-        )
-
-    if relax_version:
-        relaxed_env = copy.deepcopy(meta.env)
-        relaxed_env.relax_version()
-        required_packages = relaxed_env._conda_dependencies[env_utils.DEFAULT_CHANNEL_NAME]
-    else:
-        required_packages = meta.env._conda_dependencies[env_utils.DEFAULT_CHANNEL_NAME]
-
-    package_availability_dict = env_utils.get_matched_package_versions_in_information_schema(
-        session, required_packages, python_version=meta.env.python_version
-    )
-    no_version_available_packages = [
-        req_name for req_name, ver_list in package_availability_dict.items() if len(ver_list) < 1
-    ]
-    unavailable_packages = [req.name for req in required_packages if req.name not in package_availability_dict]
-    if no_version_available_packages or unavailable_packages:
-        relax_version_info_str = "" if relax_version else "Try to set relax_version as True in the options. "
-        required_package_str = " ".join(map(lambda x: f'"{x}"', required_packages))
-        raise snowml_exceptions.SnowflakeMLException(
-            error_code=error_codes.DEPENDENCY_VERSION_ERROR,
-            original_exception=RuntimeError(
-                textwrap.dedent(
-                    f"""
-                The model's dependencies are not available in Snowflake Anaconda Channel. {relax_version_info_str}
-                Required packages are: {required_package_str}
-                Required Python version is: {meta.env.python_version}
-                Packages that are not available are: {unavailable_packages}
-                Packages that cannot meet your requirements are: {no_version_available_packages}
-                Package availability information of those you requested is: {package_availability_dict}
-                """
-                ),
-            ),
-        )
-    return list(sorted(map(str, required_packages)))

snowflake/ml/model/_deploy_client/warehouse/infer_template.py

@@ -1,99 +0,0 @@
-_KEEP_ORDER_COL_NAME = "_ID"
-
-_UDF_CODE_TEMPLATE = """
-import fcntl
-import functools
-import inspect
-import os
-import sys
-import threading
-import zipfile
-from types import TracebackType
-from typing import Optional, Type
-
-import anyio
-import pandas as pd
-from _snowflake import vectorized
-
-
-class FileLock:
-    def __enter__(self) -> None:
-        self._lock = threading.Lock()
-        self._lock.acquire()
-        self._fd = open("/tmp/lockfile.LOCK", "w+")
-        fcntl.lockf(self._fd, fcntl.LOCK_EX)
-
-    def __exit__(
-        self, exc_type: Optional[Type[BaseException]], exc: Optional[BaseException], traceback: Optional[TracebackType]
-    ) -> None:
-        self._fd.close()
-        self._lock.release()
-
-
-# User-defined parameters
-MODEL_FILE_NAME = "{model_stage_file_name}"
-TARGET_METHOD = "{target_method}"
-MAX_BATCH_SIZE = None
-
-
-# Retrieve the model
-IMPORT_DIRECTORY_NAME = "snowflake_import_directory"
-import_dir = sys._xoptions[IMPORT_DIRECTORY_NAME]
-
-model_dir_name = os.path.splitext(MODEL_FILE_NAME)[0]
-zip_model_path = os.path.join(import_dir, MODEL_FILE_NAME)
-extracted = "/tmp/models"
-extracted_model_dir_path = os.path.join(extracted, model_dir_name)
-
-with FileLock():
-    if not os.path.isdir(extracted_model_dir_path):
-        with zipfile.ZipFile(zip_model_path, "r") as myzip:
-            myzip.extractall(extracted_model_dir_path)
-
-sys.path.insert(0, os.path.join(extracted_model_dir_path, "{code_dir_name}"))
-
-# Load the model
-try:
-    from snowflake.ml.model._packager import model_packager
-    pk = model_packager.ModelPackager(extracted_model_dir_path)
-    pk.load(as_custom_model=True)
-    assert pk.model, "model is not loaded"
-    assert pk.meta, "model metadata is not loaded"
-
-    model = pk.model
-    meta = pk.meta
-except ImportError as e:
-    if e.name and not e.name.startswith("snowflake.ml"):
-        raise e
-    # Support Legacy model
-    from snowflake.ml.model import _model
-    # Backward for <= 1.0.5
-    if hasattr(_model, "_load_model_for_deploy"):
-        model, meta = _model._load_model_for_deploy(extracted_model_dir_path)
-    else:
-        model, meta = _model._load(local_dir_path=extracted_model_dir_path, as_custom_model=True)
-
-# Determine the actual runner
-func = getattr(model, TARGET_METHOD)
-if inspect.iscoroutinefunction(func):
-    runner = functools.partial(anyio.run, func)
-else:
-    runner = functools.partial(func)
-
-# Determine preprocess parameters
-features = meta.signatures[TARGET_METHOD].inputs
-input_cols = [feature.name for feature in features]
-dtype_map = {{feature.name: feature.as_dtype() for feature in features}}
-
-
-# Actual handler
-@vectorized(input=pd.DataFrame, max_batch_size=MAX_BATCH_SIZE)
-def infer(df: pd.DataFrame) -> dict:
-    input_df = pd.json_normalize(df[0]).astype(dtype=dtype_map)
-    predictions_df = runner(input_df[input_cols])
-
-    if "{_KEEP_ORDER_COL_NAME}" in input_df.columns:
-        predictions_df["{_KEEP_ORDER_COL_NAME}"] = input_df["{_KEEP_ORDER_COL_NAME}"]
-
-    return predictions_df.to_dict("records")
-"""

snowflake/ml/model/_packager/model_handlers/llm.py

@@ -1,269 +0,0 @@
-import logging
-import os
-from typing import Dict, Optional, Type, cast, final
-
-import cloudpickle
-import pandas as pd
-from typing_extensions import TypeGuard, Unpack
-
-from snowflake.ml._internal import file_utils
-from snowflake.ml.model import custom_model, model_signature, type_hints as model_types
-from snowflake.ml.model._packager.model_env import model_env
-from snowflake.ml.model._packager.model_handlers import _base
-from snowflake.ml.model._packager.model_handlers_migrator import base_migrator
-from snowflake.ml.model._packager.model_meta import (
-    model_blob_meta,
-    model_meta as model_meta_api,
-    model_meta_schema,
-)
-from snowflake.ml.model.models import llm
-
-logger = logging.getLogger(__name__)
-
-
-@final
-class LLMHandler(_base.BaseModelHandler[llm.LLM]):
-    HANDLER_TYPE = "llm"
-    HANDLER_VERSION = "2023-12-01"
-    _MIN_SNOWPARK_ML_VERSION = "1.0.12"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
-
-    MODEL_BLOB_FILE_OR_DIR = "model"
-    LLM_META = "llm_meta"
-    IS_AUTO_SIGNATURE = True
-
-    @classmethod
-    def can_handle(
-        cls,
-        model: model_types.SupportedModelType,
-    ) -> TypeGuard[llm.LLM]:
-        return isinstance(model, llm.LLM)
-
-    @classmethod
-    def cast_model(
-        cls,
-        model: model_types.SupportedModelType,
-    ) -> llm.LLM:
-        assert isinstance(model, llm.LLM)
-        return cast(llm.LLM, model)
-
-    @classmethod
-    def save_model(
-        cls,
-        name: str,
-        model: llm.LLM,
-        model_meta: model_meta_api.ModelMetadata,
-        model_blobs_dir_path: str,
-        sample_input_data: Optional[model_types.SupportedDataType] = None,
-        is_sub_model: Optional[bool] = False,
-        **kwargs: Unpack[model_types.LLMSaveOptions],
-    ) -> None:
-        assert not is_sub_model, "LLM can not be sub-model."
-        enable_explainability = kwargs.get("enable_explainability", False)
-        if enable_explainability:
-            raise NotImplementedError("Explainability is not supported for llm model.")
-        model_blob_path = os.path.join(model_blobs_dir_path, name)
-        os.makedirs(model_blob_path, exist_ok=True)
-        model_blob_dir_path = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
-
-        sig = model_signature.ModelSignature(
-            inputs=[
-                model_signature.FeatureSpec(name="input", dtype=model_signature.DataType.STRING),
-            ],
-            outputs=[
-                model_signature.FeatureSpec(name="generated_text", dtype=model_signature.DataType.STRING),
-            ],
-        )
-        model_meta.signatures = {"infer": sig}
-        if os.path.isdir(model.model_id_or_path):
-            file_utils.copytree(model.model_id_or_path, model_blob_dir_path)
-
-        os.makedirs(model_blob_dir_path, exist_ok=True)
-        with open(
-            os.path.join(model_blob_dir_path, cls.LLM_META),
-            "wb",
-        ) as f:
-            cloudpickle.dump(model, f)
-
-        base_meta = model_blob_meta.ModelBlobMeta(
-            name=name,
-            model_type=cls.HANDLER_TYPE,
-            handler_version=cls.HANDLER_VERSION,
-            path=cls.MODEL_BLOB_FILE_OR_DIR,
-            options=model_meta_schema.LLMModelBlobOptions(
-                {
-                    "batch_size": model.max_batch_size,
-                }
-            ),
-        )
-        model_meta.models[name] = base_meta
-        model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
-
-        pkgs_requirements = [
-            model_env.ModelDependency(requirement="transformers>=4.32.1", pip_name="transformers"),
-            model_env.ModelDependency(requirement="pytorch==2.0.1", pip_name="torch"),
-        ]
-        if model.model_type == llm.SupportedLLMType.LLAMA_MODEL_TYPE.value:
-            pkgs_requirements = [
-                model_env.ModelDependency(requirement="sentencepiece", pip_name="sentencepiece"),
-                model_env.ModelDependency(requirement="protobuf", pip_name="protobuf"),
-                *pkgs_requirements,
-            ]
-        model_meta.env.include_if_absent(pkgs_requirements, check_local_version=True)
-        # Recent peft versions are only available in PYPI.
-        model_meta.env.include_if_absent_pip(["peft==0.5.0", "vllm==0.2.1.post1"])
-
-        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
-
-    @classmethod
-    def load_model(
-        cls,
-        name: str,
-        model_meta: model_meta_api.ModelMetadata,
-        model_blobs_dir_path: str,
-        **kwargs: Unpack[model_types.LLMLoadOptions],
-    ) -> llm.LLM:
-        model_blob_path = os.path.join(model_blobs_dir_path, name)
-        if not hasattr(model_meta, "models"):
-            raise ValueError("Ill model metadata found.")
-        model_blobs_metadata = model_meta.models
-        if name not in model_blobs_metadata:
-            raise ValueError(f"Blob of model {name} does not exist.")
-        model_blob_metadata = model_blobs_metadata[name]
-        model_blob_filename = model_blob_metadata.path
-        model_blob_dir_path = os.path.join(model_blob_path, model_blob_filename)
-        assert model_blob_dir_path, "It must be a directory."
-        with open(os.path.join(model_blob_dir_path, cls.LLM_META), "rb") as f:
-            m = cloudpickle.load(f)
-        assert isinstance(m, llm.LLM)
-        if m.mode == llm.LLM.Mode.LOCAL_LORA:
-            # Switch to local path
-            m.model_id_or_path = model_blob_dir_path
-        return m
-
-    @classmethod
-    def convert_as_custom_model(
-        cls,
-        raw_model: llm.LLM,
-        model_meta: model_meta_api.ModelMetadata,
-        background_data: Optional[pd.DataFrame] = None,
-        **kwargs: Unpack[model_types.LLMLoadOptions],
-    ) -> custom_model.CustomModel:
-        import gc
-        import tempfile
-
-        import torch
-        import transformers
-        import vllm
-
-        assert torch.cuda.is_available(), "LLM inference only works on GPUs."
-        device_count = torch.cuda.device_count()
-        logger.warning(f"There's total {device_count} GPUs visible to use.")
-
-        class _LLMCustomModel(custom_model.CustomModel):
-            def _memory_stats(self, msg: str) -> None:
-                logger.warning(msg)
-                logger.warning(f"Torch VRAM {torch.cuda.memory_allocated()/1024**2} MB allocated.")
-                logger.warning(f"Torch VRAM {torch.cuda.memory_reserved()/1024**2} MB reserved.")
-
-            def _prepare_for_pretrain(self) -> None:
-                hub_kwargs = {
-                    "revision": raw_model.revision,
-                    "token": raw_model.token,
-                }
-                model_dir_path = raw_model.model_id_or_path
-                tokenizer = transformers.AutoTokenizer.from_pretrained(
-                    model_dir_path,
-                    padding_side="right",
-                    use_fast=False,
-                    **hub_kwargs,
-                )
-                if not tokenizer.pad_token:
-                    tokenizer.pad_token = tokenizer.eos_token
-                tokenizer.save_pretrained(self.local_model_dir)
-                hf_model = transformers.AutoModelForCausalLM.from_pretrained(
-                    model_dir_path,
-                    device_map="auto",
-                    torch_dtype="auto",
-                    **hub_kwargs,
-                )
-                hf_model.eval()
-                hf_model.save_pretrained(self.local_model_dir)
-                logger.warning(f"Model state is saved to {self.local_model_dir}.")
-                del tokenizer
-                del hf_model
-                gc.collect()
-                torch.cuda.empty_cache()
-                self._memory_stats("After GC on model.")
-
-            def _prepare_for_lora(self) -> None:
-                self._memory_stats("Before model load & merge.")
-                import peft
-
-                hub_kwargs = {
-                    "revision": raw_model.revision,
-                    "token": raw_model.token,
-                }
-                model_dir_path = raw_model.model_id_or_path
-                peft_config = peft.PeftConfig.from_pretrained(  # type: ignore[no-untyped-call, attr-defined]
-                    model_dir_path
-                )
-                base_model_path = peft_config.base_model_name_or_path
-                tokenizer = transformers.AutoTokenizer.from_pretrained(
-                    base_model_path,
-                    padding_side="right",
-                    use_fast=False,
-                    **hub_kwargs,
-                )
-                if not tokenizer.pad_token:
-                    tokenizer.pad_token = tokenizer.eos_token
-                tokenizer.save_pretrained(self.local_model_dir)
-                logger.warning(f"Tokenizer state is saved to {self.local_model_dir}.")
-                hf_model = peft.AutoPeftModelForCausalLM.from_pretrained(  # type: ignore[attr-defined]
-                    model_dir_path,
-                    device_map="auto",
-                    torch_dtype="auto",
-                    **hub_kwargs,  # type: ignore[arg-type]
-                )
-                hf_model.eval()
-                hf_model = hf_model.merge_and_unload()
-                hf_model.save_pretrained(self.local_model_dir)
-                logger.warning(f"Merged model state is saved to {self.local_model_dir}.")
-                self._memory_stats("After model load & merge.")
-                del hf_model
-                gc.collect()
-                torch.cuda.empty_cache()
-                self._memory_stats("After GC on model.")
-
-            def __init__(self, context: custom_model.ModelContext) -> None:
-                self.local_tmp_holder = tempfile.TemporaryDirectory()
-                self.local_model_dir = self.local_tmp_holder.name
-                if raw_model.mode == llm.LLM.Mode.LOCAL_LORA:
-                    self._prepare_for_lora()
-                elif raw_model.mode == llm.LLM.Mode.REMOTE_PRETRAIN:
-                    self._prepare_for_pretrain()
-                self.sampling_params = vllm.SamplingParams(
-                    temperature=raw_model.temperature,
-                    top_p=raw_model.top_p,
-                    max_tokens=raw_model.max_tokens,
-                )
-                self._init_engine()
-
-            # This has to have same lifetime as main thread
-            # in order to avoid pre-maturely terminate ray.
-            def _init_engine(self) -> None:
-                tp_size = torch.cuda.device_count() if raw_model.enable_tp else 1
-                self.llm_engine = vllm.LLM(
-                    model=self.local_model_dir,
-                    tensor_parallel_size=tp_size,
-                )
-
-            @custom_model.inference_api
-            def infer(self, X: pd.DataFrame) -> pd.DataFrame:
-                input_data = X.to_dict("list")["input"]
-                res = self.llm_engine.generate(input_data, self.sampling_params)
-                return pd.DataFrame({"generated_text": [o.outputs[0].text for o in res]})
-
-        llm_custom = _LLMCustomModel(custom_model.ModelContext())
-
-        return llm_custom

snowflake/ml/model/_packager/model_meta/_core_requirements.py

@@ -1,11 +0,0 @@
-REQUIREMENTS = [
-    "absl-py>=0.15,<2",
-    "anyio>=3.5.0,<4",
-    "cloudpickle>=2.0.0",
-    "numpy>=1.23,<2",
-    "packaging>=20.9,<24",
-    "pandas>=1.0.0,<3",
-    "pyyaml>=6.0,<7",
-    "snowflake-snowpark-python>=1.17.0,<2",
-    "typing-extensions>=4.1.0,<5"
-]

snowflake/ml/model/deploy_platforms.py

@@ -1,6 +0,0 @@
-from enum import Enum
-
-
-class TargetPlatform(Enum):
-    WAREHOUSE = "warehouse"
-    SNOWPARK_CONTAINER_SERVICES = "SNOWPARK_CONTAINER_SERVICES"

snowflake/ml/model/models/llm.py

@@ -1,106 +0,0 @@
-import os
-from dataclasses import dataclass, field
-from enum import Enum
-from typing import Optional, Set
-
-_PEFT_CONFIG_NAME = "adapter_config.json"
-
-
-class SupportedLLMType(Enum):
-    LLAMA_MODEL_TYPE = "llama"
-    OPT_MODEL_TYPE = "opt"
-
-    @classmethod
-    def valid_values(cls) -> Set[str]:
-        return {member.value for member in cls}
-
-
-@dataclass(frozen=True)
-class LLMOptions:
-    """
-    This is the option class for LLM.
-
-    Args:
-        revision: Revision of HF model. Defaults to None.
-        token: The token to use as HTTP bearer authorization for remote files. Defaults to None.
-        max_batch_size: Max batch size allowed for single inferenced. Defaults to 1.
-    """
-
-    revision: Optional[str] = field(default=None)
-    token: Optional[str] = field(default=None)
-    max_batch_size: int = field(default=1)
-    enable_tp: bool = field(default=False)
-    # TODO(halu): Below could be per query call param instead.
-    temperature: float = field(default=0.01)
-    top_p: float = field(default=1.0)
-    max_tokens: int = field(default=100)
-
-
-class LLM:
-    class Mode(Enum):
-        LOCAL_LORA = "local_lora"
-        REMOTE_PRETRAIN = "remote_pretrain"
-
-    def __init__(
-        self,
-        model_id_or_path: str,
-        *,
-        options: Optional[LLMOptions] = None,
-    ) -> None:
-        """
-
-        Args:
-            model_id_or_path: model_id or local dir to PEFT lora weights.
-            options: Options for LLM. Defaults to be None.
-
-        Raises:
-            ValueError: When unsupported.
-        """
-        if not options:
-            options = LLMOptions()
-        hub_kwargs = {
-            "revision": options.revision,
-            "token": options.token,
-        }
-        import transformers
-
-        if os.path.isdir(model_id_or_path):
-            if not os.path.isfile(os.path.join(model_id_or_path, _PEFT_CONFIG_NAME)):
-                raise ValueError("Peft config is not found.")
-
-            import peft
-
-            peft_config = peft.PeftConfig.from_pretrained(  # type: ignore[no-untyped-call, attr-defined]
-                model_id_or_path, **hub_kwargs
-            )
-            if peft_config.peft_type != peft.PeftType.LORA:  # type: ignore[attr-defined]
-                raise ValueError("Only LORA is supported.")
-            if peft_config.task_type != peft.TaskType.CAUSAL_LM:  # type: ignore[attr-defined]
-                raise ValueError("Only CAUSAL_LM is supported.")
-            base_model = peft_config.base_model_name_or_path
-            base_config = transformers.AutoConfig.from_pretrained(base_model, **hub_kwargs)
-            assert (
-                base_config.model_type in SupportedLLMType.valid_values()
-            ), f"{base_config.model_type} is not supported."
-            self.mode = LLM.Mode.LOCAL_LORA
-            self.model_type = base_config.model_type
-        else:
-            # We support pre-train model as well
-            model_config = transformers.AutoConfig.from_pretrained(
-                model_id_or_path,
-                **hub_kwargs,
-            )
-            assert (
-                model_config.model_type in SupportedLLMType.valid_values()
-            ), f"{model_config.model_type} is not supported."
-            self.mode = LLM.Mode.REMOTE_PRETRAIN
-            self.model_type = model_config.model_type
-
-        self.model_id_or_path = model_id_or_path
-        self.token = options.token
-        self.revision = options.revision
-        self.max_batch_size = options.max_batch_size
-        self.temperature = options.temperature
-        self.top_p = options.top_p
-        self.max_tokens = options.max_tokens
-        self.enable_tp = options.enable_tp