snowflake-ml-python 1.12.0__tar.gz → 1.13.0__tar.gz

{snowflake_ml_python-1.12.0 → snowflake_ml_python-1.13.0}/CHANGELOG.md

@@ -1,11 +1,27 @@
 # Release History
 
+## 1.13.0
+
+### Bug Fixes
+
+### Behavior Changes
+
+### New Features
+
+* Registry: Log a HuggingFace model without having to load the model in memory using
+ the `huggingface_pipeline.HuggingFacePipelineModel`. Requires `huggingface_hub` package to installed.
+ To disable downloading HuggingFace repository, provide `download_snapshot=False` while creating the
+ `huggingface_pipeline.HuggingFacePipelineModel` object.
+* Registry: Added support for XGBoost models to use `enable_categorical=True` with pandas DataFrame
+* Registry: Added support to display privatelink inference endpoint in ModelVersion list services.
+
 ## 1.12.0
 
 ### Bug Fixes
 
 * Registry: Fixed an issue where the string representation of dictionary-type output columns was being incorrectly
   created during structured output deserialization. Now, the original data type is properly preserved.
+* Registry: Fixed the inference server performance issue for wide (500+ features) and JSON inputs.
 
 ### Behavior Changes
 

{snowflake_ml_python-1.12.0/snowflake_ml_python.egg-info → snowflake_ml_python-1.13.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowflake-ml-python
-Version: 1.12.0
+Version: 1.13.0
 Summary: The machine learning client library that is used for interacting with Snowflake to build machine learning solutions.
 Author-email: "Snowflake, Inc" <support@snowflake.com>
 License: 
@@ -410,12 +410,28 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 
 # Release History
 
+## 1.13.0
+
+### Bug Fixes
+
+### Behavior Changes
+
+### New Features
+
+* Registry: Log a HuggingFace model without having to load the model in memory using
+ the `huggingface_pipeline.HuggingFacePipelineModel`. Requires `huggingface_hub` package to installed.
+ To disable downloading HuggingFace repository, provide `download_snapshot=False` while creating the
+ `huggingface_pipeline.HuggingFacePipelineModel` object.
+* Registry: Added support for XGBoost models to use `enable_categorical=True` with pandas DataFrame
+* Registry: Added support to display privatelink inference endpoint in ModelVersion list services.
+
 ## 1.12.0
 
 ### Bug Fixes
 
 * Registry: Fixed an issue where the string representation of dictionary-type output columns was being incorrectly
   created during structured output deserialization. Now, the original data type is properly preserved.
+* Registry: Fixed the inference server performance issue for wide (500+ features) and JSON inputs.
 
 ### Behavior Changes
 

{snowflake_ml_python-1.12.0 → snowflake_ml_python-1.13.0}/snowflake/ml/_internal/telemetry.py

@@ -680,8 +680,10 @@ def _extract_arg_value(field: str, func_spec: inspect.FullArgSpec, args: Any, kw
     elif func_spec.kwonlydefaults and field in func_spec.kwonlyargs:
         if field in kwargs:
             return True, kwargs[field]
-        else:
+        elif field in func_spec.kwonlydefaults:
             return True, func_spec.kwonlydefaults[field]
+        else:
+            return False, None
     else:
         return False, None
 

{snowflake_ml_python-1.12.0 → snowflake_ml_python-1.13.0}/snowflake/ml/experiment/experiment_tracking.py

@@ -7,7 +7,7 @@ from urllib.parse import quote
 import snowflake.snowpark._internal.utils as snowpark_utils
 from snowflake.ml import model as ml_model, registry
 from snowflake.ml._internal.human_readable_id import hrid_generator
-from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml._internal.utils import mixins, sql_identifier
 from snowflake.ml.experiment import (
     _entities as entities,
     _experiment_info as experiment_info,
@@ -23,7 +23,7 @@ from snowflake.snowpark import session
 DEFAULT_EXPERIMENT_NAME = sql_identifier.SqlIdentifier("DEFAULT")
 
 
-class ExperimentTracking:
+class ExperimentTracking(mixins.SerializableSessionMixin):
     """
     Class to manage experiments in Snowflake.
     """
@@ -73,12 +73,34 @@ class ExperimentTracking:
             database_name=self._database_name,
             schema_name=self._schema_name,
         )
+        self._session = session
 
         # The experiment in context
         self._experiment: Optional[entities.Experiment] = None
         # The run in context
         self._run: Optional[entities.Run] = None
 
+    def __getstate__(self) -> dict[str, Any]:
+        state = super().__getstate__()
+        # Remove unpicklable attributes
+        state["_sql_client"] = None
+        state["_registry"] = None
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        super().__setstate__(state)
+        # Restore unpicklable attributes
+        self._sql_client = sql_client.ExperimentTrackingSQLClient(
+            session=self._session,
+            database_name=self._database_name,
+            schema_name=self._schema_name,
+        )
+        self._registry = registry.Registry(
+            session=self._session,
+            database_name=self._database_name,
+            schema_name=self._schema_name,
+        )
+
     def set_experiment(
         self,
         experiment_name: str,

{snowflake_ml_python-1.12.0 → snowflake_ml_python-1.13.0}/snowflake/ml/jobs/_utils/scripts/mljob_launcher.py

@@ -236,7 +236,9 @@ def run_script(script_path: str, *script_args: Any, main_func: Optional[str] = N
 
     # Create a Snowpark session before running the script
     # Session can be retrieved from using snowflake.snowpark.context.get_active_session()
-    session = Session.builder.configs(SnowflakeLoginOptions()).create()  # noqa: F841
+    config = SnowflakeLoginOptions()
+    config["client_session_keep_alive"] = "True"
+    session = Session.builder.configs(config).create()  # noqa: F841
 
     try:
 
@@ -264,6 +266,7 @@ def run_script(script_path: str, *script_args: Any, main_func: Optional[str] = N
     finally:
         # Restore original sys.argv
         sys.argv = original_argv
+        session.close()
 
 
 def main(script_path: str, *script_args: Any, script_main_func: Optional[str] = None) -> ExecutionResult:
@@ -287,6 +290,13 @@ def main(script_path: str, *script_args: Any, script_main_func: Optional[str] =
     output_dir = os.path.dirname(result_abs_path)
     os.makedirs(output_dir, exist_ok=True)
 
+    try:
+        import ray
+
+        ray.init(address="auto")
+    except ModuleNotFoundError:
+        warnings.warn("Ray is not installed, skipping Ray initialization", ImportWarning, stacklevel=1)
+
     try:
         # Wait for minimum required instances if specified
         min_instances_str = os.environ.get(MIN_INSTANCES_ENV_VAR) or "1"

{snowflake_ml_python-1.12.0 → snowflake_ml_python-1.13.0}/snowflake/ml/jobs/job.py

@@ -50,7 +50,7 @@ class MLJob(Generic[T], SerializableSessionMixin):
     def min_instances(self) -> int:
         try:
             return int(self._container_spec["env"].get(constants.MIN_INSTANCES_ENV_VAR, 1))
-        except TypeError:
+        except (TypeError, ValueError):
             return 1
 
     @property
@@ -83,7 +83,10 @@ class MLJob(Generic[T], SerializableSessionMixin):
     def _container_spec(self) -> dict[str, Any]:
         """Get the job's main container spec."""
         containers = self._service_spec["spec"]["containers"]
-        container_spec = next(c for c in containers if c["name"] == constants.DEFAULT_CONTAINER_NAME)
+        try:
+            container_spec = next(c for c in containers if c["name"] == constants.DEFAULT_CONTAINER_NAME)
+        except StopIteration:
+            raise ValueError(f"Container '{constants.DEFAULT_CONTAINER_NAME}' not found in job {self.name}")
         return cast(dict[str, Any], container_spec)
 
     @property

{snowflake_ml_python-1.12.0 → snowflake_ml_python-1.13.0}/snowflake/ml/jobs/manager.py

@@ -58,7 +58,8 @@ def list_jobs(
         >>> from snowflake.ml.jobs import list_jobs
         >>> list_jobs(limit=5)
     """
-    session = session or get_active_session()
+
+    session = _ensure_session(session)
     try:
         df = _get_job_history_spcs(
             session,
@@ -162,7 +163,7 @@ def _get_job_history_spcs(
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
 def get_job(job_id: str, session: Optional[snowpark.Session] = None) -> jb.MLJob[Any]:
     """Retrieve a job service from the backend."""
-    session = session or get_active_session()
+    session = _ensure_session(session)
     try:
         database, schema, job_name = identifier.parse_schema_level_object_identifier(job_id)
         database = identifier.resolve_identifier(cast(str, database or session.get_current_database()))
@@ -434,8 +435,10 @@ def _submit_job(
 
     Raises:
         ValueError: If database or schema value(s) are invalid
+        RuntimeError: If schema is not specified in session context or job submission
+        snowpark.exceptions.SnowparkSQLException: if failed to upload payload
     """
-    session = session or get_active_session()
+    session = _ensure_session(session)
 
     # Check for deprecated args
     if "num_instances" in kwargs:
@@ -486,10 +489,17 @@ def _submit_job(
     stage_name = f"@{'.'.join(filter(None, stage_path_parts[:3]))}"
     stage_path = pathlib.PurePosixPath(f"{stage_name}{stage_path_parts[-1].rstrip('/')}/{job_name}")
 
-    # Upload payload
-    uploaded_payload = payload_utils.JobPayload(
-        source, entrypoint=entrypoint, pip_requirements=pip_requirements, additional_payloads=additional_payloads
-    ).upload(session, stage_path)
+    try:
+        # Upload payload
+        uploaded_payload = payload_utils.JobPayload(
+            source, entrypoint=entrypoint, pip_requirements=pip_requirements, additional_payloads=additional_payloads
+        ).upload(session, stage_path)
+    except snowpark.exceptions.SnowparkSQLException as e:
+        if e.sql_error_code == 90106:
+            raise RuntimeError(
+                "Please specify a schema, either in the session context or as a parameter in the job submission"
+            )
+        raise
 
     if feature_flags.FeatureFlags.USE_SUBMIT_JOB_V2.is_enabled():
         # Add default env vars (extracted from spec_utils.generate_service_spec)
@@ -651,3 +661,15 @@ def _do_submit_job_v2(
     actual_job_id = query_helper.run_query(session, query_template, params=params)[0][0]
 
     return get_job(actual_job_id, session=session)
+
+
+def _ensure_session(session: Optional[snowpark.Session]) -> snowpark.Session:
+    try:
+        session = session or get_active_session()
+    except snowpark.exceptions.SnowparkSessionException as e:
+        if "More than one active session" in e.message:
+            raise RuntimeError("Please specify the session as a parameter in API call")
+        if "No default Session is found" in e.message:
+            raise RuntimeError("Please create a session before API call")
+        raise
+    return session

{snowflake_ml_python-1.12.0 → snowflake_ml_python-1.13.0}/snowflake/ml/model/_client/model/batch_inference_specs.py

@@ -4,14 +4,11 @@ from pydantic import BaseModel
 
 
 class InputSpec(BaseModel):
-    input_stage_location: str
-    input_file_pattern: str = "*"
+    stage_location: str
 
 
 class OutputSpec(BaseModel):
-    output_stage_location: str
-    output_file_prefix: Optional[str] = None
-    completion_filename: str = "_SUCCESS"
+    stage_location: str
 
 
 class JobSpec(BaseModel):
@@ -25,3 +22,4 @@ class JobSpec(BaseModel):
     warehouse: Optional[str] = None
     cpu_requests: Optional[str] = None
     memory_requests: Optional[str] = None
+    replicas: Optional[int] = None

{snowflake_ml_python-1.12.0 → snowflake_ml_python-1.13.0}/snowflake/ml/model/_client/model/model_version_impl.py

@@ -590,11 +590,12 @@ class ModelVersion(lineage_node.LineageNode):
             cpu_requests=job_spec.cpu_requests,
             memory_requests=job_spec.memory_requests,
             job_name=job_name,
+            replicas=job_spec.replicas,
             # input and output
-            input_stage_location=input_spec.input_stage_location,
-            input_file_pattern=input_spec.input_file_pattern,
-            output_stage_location=output_spec.output_stage_location,
-            completion_filename=output_spec.completion_filename,
+            input_stage_location=input_spec.stage_location,
+            input_file_pattern="*",
+            output_stage_location=output_spec.stage_location,
+            completion_filename="_SUCCESS",
             # misc
             statement_params=statement_params,
         )

{snowflake_ml_python-1.12.0 → snowflake_ml_python-1.13.0}/snowflake/ml/model/_client/ops/model_ops.py

@@ -47,6 +47,7 @@ class ServiceInfo(TypedDict):
 class ModelOperator:
     INFERENCE_SERVICE_ENDPOINT_NAME = "inference"
     INGRESS_ENDPOINT_URL_SUFFIX = "snowflakecomputing.app"
+    PRIVATELINK_INGRESS_ENDPOINT_URL_SUBSTRING = "privatelink.snowflakecomputing"
 
     def __init__(
         self,
@@ -612,6 +613,30 @@ class ModelOperator:
             statement_params=statement_params,
         )
 
+    def _is_privatelink_connection(self) -> bool:
+        """Detect if the current session is using a privatelink connection."""
+        try:
+            host = self._session.connection.host
+            return ModelOperator.PRIVATELINK_INGRESS_ENDPOINT_URL_SUBSTRING in host
+        except AttributeError:
+            return False
+
+    def _extract_and_validate_ingress_url(self, res_row: "row.Row") -> Optional[str]:
+        """Extract and validate ingress URL from endpoint row."""
+        url_value = res_row[self._service_client.MODEL_INFERENCE_SERVICE_ENDPOINT_INGRESS_URL_COL_NAME]
+        if url_value is None:
+            return None
+        url_str = str(url_value)
+        return url_str if url_str.endswith(ModelOperator.INGRESS_ENDPOINT_URL_SUFFIX) else None
+
+    def _extract_and_validate_privatelink_url(self, res_row: "row.Row") -> Optional[str]:
+        """Extract and validate privatelink ingress URL from endpoint row."""
+        url_value = res_row[self._service_client.MODEL_INFERENCE_SERVICE_ENDPOINT_PRIVATELINK_INGRESS_URL_COL_NAME]
+        if url_value is None:
+            return None
+        url_str = str(url_value)
+        return url_str if ModelOperator.PRIVATELINK_INGRESS_ENDPOINT_URL_SUBSTRING in url_str else None
+
     def show_services(
         self,
         *,
@@ -644,8 +669,10 @@ class ModelOperator:
         fully_qualified_service_names = [str(service) for service in json_array if "MODEL_BUILD_" not in service]
 
         result: list[ServiceInfo] = []
+        is_privatelink_connection = self._is_privatelink_connection()
+
         for fully_qualified_service_name in fully_qualified_service_names:
-            ingress_url: Optional[str] = None
+            inference_endpoint: Optional[str] = None
             db, schema, service_name = sql_identifier.parse_fully_qualified_name(fully_qualified_service_name)
             statuses = self._service_client.get_service_container_statuses(
                 database_name=db, schema_name=schema, service_name=service_name, statement_params=statement_params
@@ -659,17 +686,23 @@ class ModelOperator:
             ):
                 if (
                     res_row[self._service_client.MODEL_INFERENCE_SERVICE_ENDPOINT_NAME_COL_NAME]
-                    == self.INFERENCE_SERVICE_ENDPOINT_NAME
-                    and res_row[self._service_client.MODEL_INFERENCE_SERVICE_ENDPOINT_INGRESS_URL_COL_NAME] is not None
+                    != self.INFERENCE_SERVICE_ENDPOINT_NAME
                 ):
-                    ingress_url = str(
-                        res_row[self._service_client.MODEL_INFERENCE_SERVICE_ENDPOINT_INGRESS_URL_COL_NAME]
-                    )
-                    if not ingress_url.endswith(ModelOperator.INGRESS_ENDPOINT_URL_SUFFIX):
-                        ingress_url = None
+                    continue
+
+                ingress_url = self._extract_and_validate_ingress_url(res_row)
+                privatelink_ingress_url = self._extract_and_validate_privatelink_url(res_row)
+
+                if is_privatelink_connection and privatelink_ingress_url is not None:
+                    inference_endpoint = privatelink_ingress_url
+                else:
+                    inference_endpoint = ingress_url
+
             result.append(
                 ServiceInfo(
-                    name=fully_qualified_service_name, status=service_status.value, inference_endpoint=ingress_url
+                    name=fully_qualified_service_name,
+                    status=service_status.value,
+                    inference_endpoint=inference_endpoint,
                 )
             )
 

{snowflake_ml_python-1.12.0 → snowflake_ml_python-1.13.0}/snowflake/ml/model/_client/ops/service_ops.py

@@ -881,6 +881,7 @@ class ServiceOperator:
         max_batch_rows: Optional[int],
         cpu_requests: Optional[str],
         memory_requests: Optional[str],
+        replicas: Optional[int],
         statement_params: Optional[dict[str, Any]] = None,
     ) -> jobs.MLJob[Any]:
         database_name = self._database_name
@@ -914,6 +915,7 @@ class ServiceOperator:
             warehouse=warehouse,
             cpu=cpu_requests,
             memory=memory_requests,
+            replicas=replicas,
         )
 
         self._model_deployment_spec.add_image_build_spec(

{snowflake_ml_python-1.12.0 → snowflake_ml_python-1.13.0}/snowflake/ml/model/_client/service/model_deployment_spec.py

@@ -207,6 +207,7 @@ class ModelDeploymentSpec:
         gpu: Optional[Union[str, int]] = None,
         num_workers: Optional[int] = None,
         max_batch_rows: Optional[int] = None,
+        replicas: Optional[int] = None,
     ) -> "ModelDeploymentSpec":
         """Add job specification to the deployment spec.
 
@@ -226,6 +227,7 @@ class ModelDeploymentSpec:
             gpu: GPU requirement.
             num_workers: Number of workers.
             max_batch_rows: Maximum batch rows for inference.
+            replicas: Number of replicas.
 
         Raises:
             ValueError: If a service spec already exists.
@@ -260,6 +262,7 @@ class ModelDeploymentSpec:
                 output_stage_location=output_stage_location,
                 completion_filename=completion_filename,
             ),
+            replicas=replicas,
             **self._inference_spec,
         )
         return self

{snowflake_ml_python-1.12.0 → snowflake_ml_python-1.13.0}/snowflake/ml/model/_client/service/model_deployment_spec_schema.py

@@ -57,6 +57,7 @@ class Job(BaseModel):
     function_name: str
     input: Input
     output: Output
+    replicas: Optional[int] = None
 
 
 class LogModelArgs(BaseModel):

{snowflake_ml_python-1.12.0 → snowflake_ml_python-1.13.0}/snowflake/ml/model/_client/sql/service.py

@@ -63,6 +63,7 @@ class ServiceStatusInfo:
 class ServiceSQLClient(_base._BaseSQLClient):
     MODEL_INFERENCE_SERVICE_ENDPOINT_NAME_COL_NAME = "name"
     MODEL_INFERENCE_SERVICE_ENDPOINT_INGRESS_URL_COL_NAME = "ingress_url"
+    MODEL_INFERENCE_SERVICE_ENDPOINT_PRIVATELINK_INGRESS_URL_COL_NAME = "privatelink_ingress_url"
     SERVICE_STATUS = "service_status"
     INSTANCE_ID = "instance_id"
     INSTANCE_STATUS = "instance_status"
@@ -255,6 +256,9 @@ class ServiceSQLClient(_base._BaseSQLClient):
             )
             .has_column(ServiceSQLClient.MODEL_INFERENCE_SERVICE_ENDPOINT_NAME_COL_NAME, allow_empty=True)
             .has_column(ServiceSQLClient.MODEL_INFERENCE_SERVICE_ENDPOINT_INGRESS_URL_COL_NAME, allow_empty=True)
+            .has_column(
+                ServiceSQLClient.MODEL_INFERENCE_SERVICE_ENDPOINT_PRIVATELINK_INGRESS_URL_COL_NAME, allow_empty=True
+            )
         )
 
         return res.validate()

{snowflake_ml_python-1.12.0 → snowflake_ml_python-1.13.0}/snowflake/ml/model/_packager/model_handlers/huggingface_pipeline.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+import shutil
 import time
 import uuid
 import warnings
@@ -88,6 +89,7 @@ class HuggingFacePipelineHandler(
     _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     MODEL_BLOB_FILE_OR_DIR = "model"
+    MODEL_PICKLE_FILE = "snowml_huggingface_pipeline.pkl"
     ADDITIONAL_CONFIG_FILE = "pipeline_config.pt"
     DEFAULT_TARGET_METHODS = ["__call__"]
     IS_AUTO_SIGNATURE = True
@@ -199,6 +201,7 @@ class HuggingFacePipelineHandler(
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
 
+        is_repo_downloaded = False
         if type_utils.LazyType("transformers.Pipeline").isinstance(model):
             save_path = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
             model.save_pretrained(  # type:ignore[attr-defined]
@@ -224,11 +227,22 @@ class HuggingFacePipelineHandler(
             ) as f:
                 cloudpickle.dump(pipeline_params, f)
         else:
+            model_blob_file_or_dir = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
+            model_blob_pickle_file = os.path.join(model_blob_file_or_dir, cls.MODEL_PICKLE_FILE)
+            os.makedirs(model_blob_file_or_dir, exist_ok=True)
             with open(
-                os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR),
+                model_blob_pickle_file,
                 "wb",
             ) as f:
                 cloudpickle.dump(model, f)
+            if model.repo_snapshot_dir:
+                logger.info("model's repo_snapshot_dir is available, copying snapshot")
+                shutil.copytree(
+                    model.repo_snapshot_dir,
+                    model_blob_file_or_dir,
+                    dirs_exist_ok=True,
+                )
+                is_repo_downloaded = True
 
         base_meta = model_blob_meta.ModelBlobMeta(
             name=name,
@@ -236,13 +250,12 @@ class HuggingFacePipelineHandler(
             handler_version=cls.HANDLER_VERSION,
             path=cls.MODEL_BLOB_FILE_OR_DIR,
             options=model_meta_schema.HuggingFacePipelineModelBlobOptions(
-                {
-                    "task": task,
-                    "batch_size": batch_size if batch_size is not None else 1,
-                    "has_tokenizer": has_tokenizer,
-                    "has_feature_extractor": has_feature_extractor,
-                    "has_image_preprocessor": has_image_preprocessor,
-                }
+                task=task,
+                batch_size=batch_size if batch_size is not None else 1,
+                has_tokenizer=has_tokenizer,
+                has_feature_extractor=has_feature_extractor,
+                has_image_preprocessor=has_image_preprocessor,
+                is_repo_downloaded=is_repo_downloaded,
             ),
         )
         model_meta.models[name] = base_meta
@@ -286,6 +299,27 @@ class HuggingFacePipelineHandler(
 
         return device_config
 
+    @staticmethod
+    def _load_pickle_model(
+        pickle_file: str,
+        **kwargs: Unpack[model_types.HuggingFaceLoadOptions],
+    ) -> huggingface_pipeline.HuggingFacePipelineModel:
+        with open(pickle_file, "rb") as f:
+            m = cloudpickle.load(f)
+        assert isinstance(m, huggingface_pipeline.HuggingFacePipelineModel)
+        torch_dtype: Optional[str] = None
+        device_config = None
+        if getattr(m, "device", None) is None and getattr(m, "device_map", None) is None:
+            device_config = HuggingFacePipelineHandler._get_device_config(**kwargs)
+            m.__dict__.update(device_config)
+
+        if getattr(m, "torch_dtype", None) is None and kwargs.get("use_gpu", False):
+            torch_dtype = "auto"
+            m.__dict__.update(torch_dtype=torch_dtype)
+        else:
+            m.__dict__.update(torch_dtype=None)
+        return m
+
     @classmethod
     def load_model(
         cls,
@@ -310,7 +344,13 @@ class HuggingFacePipelineHandler(
             raise ValueError("Missing field `batch_size` in model blob metadata for type `huggingface_pipeline`")
 
         model_blob_file_or_dir_path = os.path.join(model_blob_path, model_blob_filename)
-        if os.path.isdir(model_blob_file_or_dir_path):
+        is_repo_downloaded = model_blob_options.get("is_repo_downloaded", False)
+
+        def _create_pipeline_from_dir(
+            model_blob_file_or_dir_path: str,
+            model_blob_options: model_meta_schema.HuggingFacePipelineModelBlobOptions,
+            **kwargs: Unpack[model_types.HuggingFaceLoadOptions],
+        ) -> "transformers.Pipeline":
             import transformers
 
             additional_pipeline_params = {}
@@ -330,7 +370,7 @@ class HuggingFacePipelineHandler(
             ) as f:
                 pipeline_params = cloudpickle.load(f)
 
-            device_config = cls._get_device_config(**kwargs)
+            device_config = HuggingFacePipelineHandler._get_device_config(**kwargs)
 
             m = transformers.pipeline(
                 model_blob_options["task"],
@@ -359,18 +399,59 @@ class HuggingFacePipelineHandler(
                 m.tokenizer.chat_template = DEFAULT_CHAT_TEMPLATE
 
             m.__dict__.update(pipeline_params)
+            return m
 
+        def _create_pipeline_from_model(
+            model_blob_file_or_dir_path: str,
+            m: huggingface_pipeline.HuggingFacePipelineModel,
+            **kwargs: Unpack[model_types.HuggingFaceLoadOptions],
+        ) -> "transformers.Pipeline":
+            import transformers
+
+            return transformers.pipeline(
+                m.task,
+                model=model_blob_file_or_dir_path,
+                trust_remote_code=m.trust_remote_code,
+                torch_dtype=getattr(m, "torch_dtype", None),
+                revision=m.revision,
+                # pass device or device_map when creating the pipeline
+                **HuggingFacePipelineHandler._get_device_config(**kwargs),
+                # pass other model_kwargs to transformers.pipeline.from_pretrained method
+                **m.model_kwargs,
+            )
+
+        if os.path.isdir(model_blob_file_or_dir_path) and not is_repo_downloaded:
+            # the logged model is a transformers.Pipeline object
+            # weights of the model are saved in the directory
+            return _create_pipeline_from_dir(model_blob_file_or_dir_path, model_blob_options, **kwargs)
         else:
-            assert os.path.isfile(model_blob_file_or_dir_path)
-            with open(model_blob_file_or_dir_path, "rb") as f:
-                m = cloudpickle.load(f)
-            assert isinstance(m, huggingface_pipeline.HuggingFacePipelineModel)
-            if getattr(m, "device", None) is None and getattr(m, "device_map", None) is None:
-                m.__dict__.update(cls._get_device_config(**kwargs))
-
-            if getattr(m, "torch_dtype", None) is None and kwargs.get("use_gpu", False):
-                m.__dict__.update(torch_dtype="auto")
-        return m
+            # case 1: LEGACY logging, repo snapshot is not logged
+            if os.path.isfile(model_blob_file_or_dir_path):
+                # LEGACY logging that had model as a pickle file in the model blob directory
+                # the logged model is a huggingface_pipeline.HuggingFacePipelineModel object
+                # the model_blob_file_or_dir_path is the pickle file that holds
+                # the huggingface_pipeline.HuggingFacePipelineModel object
+                # the snapshot of the repo is not logged
+                return cls._load_pickle_model(model_blob_file_or_dir_path)
+            else:
+                assert os.path.isdir(model_blob_file_or_dir_path)
+                # the logged model is a huggingface_pipeline.HuggingFacePipelineModel object
+                # the pickle_file holds the huggingface_pipeline.HuggingFacePipelineModel object
+                pickle_file = os.path.join(model_blob_file_or_dir_path, cls.MODEL_PICKLE_FILE)
+                m = cls._load_pickle_model(pickle_file)
+
+                # case 2: logging without the snapshot of the repo
+                if not is_repo_downloaded:
+                    # we return the huggingface_pipeline.HuggingFacePipelineModel object
+                    return m
+                # case 3: logging with the snapshot of the repo
+                else:
+                    # the model_blob_file_or_dir_path is the directory that holds
+                    # weights of the model from `huggingface_hub.snapshot_download`
+                    # the huggingface_pipeline.HuggingFacePipelineModel object is logged
+                    # with a snapshot of the repo, we create a transformers.Pipeline object
+                    # by reading the snapshot directory
+                    return _create_pipeline_from_model(model_blob_file_or_dir_path, m, **kwargs)
 
     @classmethod
     def convert_as_custom_model(
@@ -665,7 +746,7 @@ class HuggingFaceOpenAICompatibleModel:
             prompt_text,
             return_tensors="pt",
             padding=True,
-        )
+        ).to(self.model.device)
         prompt_tokens = inputs.input_ids.shape[1]
 
         from transformers import GenerationConfig
@@ -683,6 +764,7 @@ class HuggingFaceOpenAICompatibleModel:
             num_return_sequences=n,
             num_beams=max(2, n),  # must be >1
             num_beam_groups=max(2, n) if presence_penalty else 1,
+            do_sample=False,
         )
 
         # Generate text

{snowflake_ml_python-1.12.0 → snowflake_ml_python-1.13.0}/snowflake/ml/model/_packager/model_handlers/xgboost.py

@@ -229,6 +229,11 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
                     enable_categorical = False
                     for col, d_type in X.dtypes.items():
                         if pd.api.extensions.ExtensionDtype.is_dtype(d_type):
+                            if pd.CategoricalDtype.is_dtype(d_type):
+                                enable_categorical = True
+                            elif isinstance(d_type, pd.StringDtype):
+                                X[col] = X[col].astype("category")
+                                enable_categorical = True
                             continue
                         if not np.issubdtype(d_type, np.number):
                             # categorical columns are converted to numpy's str dtype