snowflake-ml-python 1.3.1__py3-none-any.whl → 1.4.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (219) hide show
  1. snowflake/ml/_internal/env_utils.py +11 -1
  2. snowflake/ml/_internal/human_readable_id/adjectives.txt +128 -0
  3. snowflake/ml/_internal/human_readable_id/animals.txt +128 -0
  4. snowflake/ml/_internal/human_readable_id/hrid_generator.py +40 -0
  5. snowflake/ml/_internal/human_readable_id/hrid_generator_base.py +135 -0
  6. snowflake/ml/_internal/utils/formatting.py +1 -1
  7. snowflake/ml/_internal/utils/identifier.py +3 -1
  8. snowflake/ml/_internal/utils/sql_identifier.py +2 -6
  9. snowflake/ml/feature_store/feature_store.py +166 -184
  10. snowflake/ml/feature_store/feature_view.py +12 -24
  11. snowflake/ml/fileset/sfcfs.py +56 -50
  12. snowflake/ml/fileset/stage_fs.py +48 -13
  13. snowflake/ml/model/_client/model/model_version_impl.py +6 -49
  14. snowflake/ml/model/_client/ops/model_ops.py +78 -29
  15. snowflake/ml/model/_client/sql/model.py +23 -2
  16. snowflake/ml/model/_client/sql/model_version.py +22 -1
  17. snowflake/ml/model/_deploy_client/image_builds/server_image_builder.py +1 -3
  18. snowflake/ml/model/_deploy_client/snowservice/deploy.py +5 -2
  19. snowflake/ml/model/_model_composer/model_composer.py +7 -5
  20. snowflake/ml/model/_model_composer/model_manifest/model_manifest.py +19 -54
  21. snowflake/ml/model/_model_composer/model_manifest/model_manifest_schema.py +8 -1
  22. snowflake/ml/model/_model_composer/model_method/infer_table_function.py_template +1 -1
  23. snowflake/ml/model/_model_composer/model_method/model_method.py +6 -10
  24. snowflake/ml/model/_packager/model_handlers/catboost.py +206 -0
  25. snowflake/ml/model/_packager/model_handlers/lightgbm.py +218 -0
  26. snowflake/ml/model/_packager/model_handlers/sklearn.py +3 -0
  27. snowflake/ml/model/_packager/model_handlers/snowmlmodel.py +13 -1
  28. snowflake/ml/model/_packager/model_handlers/xgboost.py +1 -1
  29. snowflake/ml/model/_packager/model_meta/_core_requirements.py +1 -1
  30. snowflake/ml/model/_packager/model_meta/model_meta.py +36 -6
  31. snowflake/ml/model/_packager/model_meta/model_meta_schema.py +20 -1
  32. snowflake/ml/model/_packager/model_meta_migrator/migrator_plans.py +3 -1
  33. snowflake/ml/model/_packager/model_packager.py +2 -2
  34. snowflake/ml/model/{_model_composer/model_runtime/_runtime_requirements.py → _packager/model_runtime/_snowml_inference_alternative_requirements.py} +1 -1
  35. snowflake/ml/model/_packager/model_runtime/model_runtime.py +137 -0
  36. snowflake/ml/model/custom_model.py +3 -1
  37. snowflake/ml/model/type_hints.py +21 -2
  38. snowflake/ml/modeling/_internal/estimator_utils.py +16 -11
  39. snowflake/ml/modeling/_internal/local_implementations/pandas_handlers.py +4 -1
  40. snowflake/ml/modeling/_internal/model_specifications.py +3 -1
  41. snowflake/ml/modeling/_internal/snowpark_implementations/distributed_hpo_trainer.py +545 -0
  42. snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_handlers.py +8 -5
  43. snowflake/ml/modeling/calibration/calibrated_classifier_cv.py +195 -123
  44. snowflake/ml/modeling/cluster/affinity_propagation.py +195 -123
  45. snowflake/ml/modeling/cluster/agglomerative_clustering.py +195 -123
  46. snowflake/ml/modeling/cluster/birch.py +195 -123
  47. snowflake/ml/modeling/cluster/bisecting_k_means.py +195 -123
  48. snowflake/ml/modeling/cluster/dbscan.py +195 -123
  49. snowflake/ml/modeling/cluster/feature_agglomeration.py +195 -123
  50. snowflake/ml/modeling/cluster/k_means.py +195 -123
  51. snowflake/ml/modeling/cluster/mean_shift.py +195 -123
  52. snowflake/ml/modeling/cluster/mini_batch_k_means.py +195 -123
  53. snowflake/ml/modeling/cluster/optics.py +195 -123
  54. snowflake/ml/modeling/cluster/spectral_biclustering.py +195 -123
  55. snowflake/ml/modeling/cluster/spectral_clustering.py +195 -123
  56. snowflake/ml/modeling/cluster/spectral_coclustering.py +195 -123
  57. snowflake/ml/modeling/compose/column_transformer.py +195 -123
  58. snowflake/ml/modeling/compose/transformed_target_regressor.py +195 -123
  59. snowflake/ml/modeling/covariance/elliptic_envelope.py +195 -123
  60. snowflake/ml/modeling/covariance/empirical_covariance.py +195 -123
  61. snowflake/ml/modeling/covariance/graphical_lasso.py +195 -123
  62. snowflake/ml/modeling/covariance/graphical_lasso_cv.py +195 -123
  63. snowflake/ml/modeling/covariance/ledoit_wolf.py +195 -123
  64. snowflake/ml/modeling/covariance/min_cov_det.py +195 -123
  65. snowflake/ml/modeling/covariance/oas.py +195 -123
  66. snowflake/ml/modeling/covariance/shrunk_covariance.py +195 -123
  67. snowflake/ml/modeling/decomposition/dictionary_learning.py +195 -123
  68. snowflake/ml/modeling/decomposition/factor_analysis.py +195 -123
  69. snowflake/ml/modeling/decomposition/fast_ica.py +195 -123
  70. snowflake/ml/modeling/decomposition/incremental_pca.py +195 -123
  71. snowflake/ml/modeling/decomposition/kernel_pca.py +195 -123
  72. snowflake/ml/modeling/decomposition/mini_batch_dictionary_learning.py +195 -123
  73. snowflake/ml/modeling/decomposition/mini_batch_sparse_pca.py +195 -123
  74. snowflake/ml/modeling/decomposition/pca.py +195 -123
  75. snowflake/ml/modeling/decomposition/sparse_pca.py +195 -123
  76. snowflake/ml/modeling/decomposition/truncated_svd.py +195 -123
  77. snowflake/ml/modeling/discriminant_analysis/linear_discriminant_analysis.py +195 -123
  78. snowflake/ml/modeling/discriminant_analysis/quadratic_discriminant_analysis.py +195 -123
  79. snowflake/ml/modeling/ensemble/ada_boost_classifier.py +195 -123
  80. snowflake/ml/modeling/ensemble/ada_boost_regressor.py +195 -123
  81. snowflake/ml/modeling/ensemble/bagging_classifier.py +195 -123
  82. snowflake/ml/modeling/ensemble/bagging_regressor.py +195 -123
  83. snowflake/ml/modeling/ensemble/extra_trees_classifier.py +195 -123
  84. snowflake/ml/modeling/ensemble/extra_trees_regressor.py +195 -123
  85. snowflake/ml/modeling/ensemble/gradient_boosting_classifier.py +195 -123
  86. snowflake/ml/modeling/ensemble/gradient_boosting_regressor.py +195 -123
  87. snowflake/ml/modeling/ensemble/hist_gradient_boosting_classifier.py +195 -123
  88. snowflake/ml/modeling/ensemble/hist_gradient_boosting_regressor.py +195 -123
  89. snowflake/ml/modeling/ensemble/isolation_forest.py +195 -123
  90. snowflake/ml/modeling/ensemble/random_forest_classifier.py +195 -123
  91. snowflake/ml/modeling/ensemble/random_forest_regressor.py +195 -123
  92. snowflake/ml/modeling/ensemble/stacking_regressor.py +195 -123
  93. snowflake/ml/modeling/ensemble/voting_classifier.py +195 -123
  94. snowflake/ml/modeling/ensemble/voting_regressor.py +195 -123
  95. snowflake/ml/modeling/feature_selection/generic_univariate_select.py +195 -123
  96. snowflake/ml/modeling/feature_selection/select_fdr.py +195 -123
  97. snowflake/ml/modeling/feature_selection/select_fpr.py +195 -123
  98. snowflake/ml/modeling/feature_selection/select_fwe.py +195 -123
  99. snowflake/ml/modeling/feature_selection/select_k_best.py +195 -123
  100. snowflake/ml/modeling/feature_selection/select_percentile.py +195 -123
  101. snowflake/ml/modeling/feature_selection/sequential_feature_selector.py +195 -123
  102. snowflake/ml/modeling/feature_selection/variance_threshold.py +195 -123
  103. snowflake/ml/modeling/framework/_utils.py +8 -1
  104. snowflake/ml/modeling/framework/base.py +24 -6
  105. snowflake/ml/modeling/gaussian_process/gaussian_process_classifier.py +195 -123
  106. snowflake/ml/modeling/gaussian_process/gaussian_process_regressor.py +195 -123
  107. snowflake/ml/modeling/impute/iterative_imputer.py +195 -123
  108. snowflake/ml/modeling/impute/knn_imputer.py +195 -123
  109. snowflake/ml/modeling/impute/missing_indicator.py +195 -123
  110. snowflake/ml/modeling/impute/simple_imputer.py +4 -15
  111. snowflake/ml/modeling/kernel_approximation/additive_chi2_sampler.py +195 -123
  112. snowflake/ml/modeling/kernel_approximation/nystroem.py +195 -123
  113. snowflake/ml/modeling/kernel_approximation/polynomial_count_sketch.py +195 -123
  114. snowflake/ml/modeling/kernel_approximation/rbf_sampler.py +195 -123
  115. snowflake/ml/modeling/kernel_approximation/skewed_chi2_sampler.py +195 -123
  116. snowflake/ml/modeling/kernel_ridge/kernel_ridge.py +195 -123
  117. snowflake/ml/modeling/lightgbm/lgbm_classifier.py +198 -125
  118. snowflake/ml/modeling/lightgbm/lgbm_regressor.py +198 -125
  119. snowflake/ml/modeling/linear_model/ard_regression.py +195 -123
  120. snowflake/ml/modeling/linear_model/bayesian_ridge.py +195 -123
  121. snowflake/ml/modeling/linear_model/elastic_net.py +195 -123
  122. snowflake/ml/modeling/linear_model/elastic_net_cv.py +195 -123
  123. snowflake/ml/modeling/linear_model/gamma_regressor.py +195 -123
  124. snowflake/ml/modeling/linear_model/huber_regressor.py +195 -123
  125. snowflake/ml/modeling/linear_model/lars.py +195 -123
  126. snowflake/ml/modeling/linear_model/lars_cv.py +195 -123
  127. snowflake/ml/modeling/linear_model/lasso.py +195 -123
  128. snowflake/ml/modeling/linear_model/lasso_cv.py +195 -123
  129. snowflake/ml/modeling/linear_model/lasso_lars.py +195 -123
  130. snowflake/ml/modeling/linear_model/lasso_lars_cv.py +195 -123
  131. snowflake/ml/modeling/linear_model/lasso_lars_ic.py +195 -123
  132. snowflake/ml/modeling/linear_model/linear_regression.py +195 -123
  133. snowflake/ml/modeling/linear_model/logistic_regression.py +195 -123
  134. snowflake/ml/modeling/linear_model/logistic_regression_cv.py +195 -123
  135. snowflake/ml/modeling/linear_model/multi_task_elastic_net.py +195 -123
  136. snowflake/ml/modeling/linear_model/multi_task_elastic_net_cv.py +195 -123
  137. snowflake/ml/modeling/linear_model/multi_task_lasso.py +195 -123
  138. snowflake/ml/modeling/linear_model/multi_task_lasso_cv.py +195 -123
  139. snowflake/ml/modeling/linear_model/orthogonal_matching_pursuit.py +195 -123
  140. snowflake/ml/modeling/linear_model/passive_aggressive_classifier.py +195 -123
  141. snowflake/ml/modeling/linear_model/passive_aggressive_regressor.py +195 -123
  142. snowflake/ml/modeling/linear_model/perceptron.py +195 -123
  143. snowflake/ml/modeling/linear_model/poisson_regressor.py +195 -123
  144. snowflake/ml/modeling/linear_model/ransac_regressor.py +195 -123
  145. snowflake/ml/modeling/linear_model/ridge.py +195 -123
  146. snowflake/ml/modeling/linear_model/ridge_classifier.py +195 -123
  147. snowflake/ml/modeling/linear_model/ridge_classifier_cv.py +195 -123
  148. snowflake/ml/modeling/linear_model/ridge_cv.py +195 -123
  149. snowflake/ml/modeling/linear_model/sgd_classifier.py +195 -123
  150. snowflake/ml/modeling/linear_model/sgd_one_class_svm.py +195 -123
  151. snowflake/ml/modeling/linear_model/sgd_regressor.py +195 -123
  152. snowflake/ml/modeling/linear_model/theil_sen_regressor.py +195 -123
  153. snowflake/ml/modeling/linear_model/tweedie_regressor.py +195 -123
  154. snowflake/ml/modeling/manifold/isomap.py +195 -123
  155. snowflake/ml/modeling/manifold/mds.py +195 -123
  156. snowflake/ml/modeling/manifold/spectral_embedding.py +195 -123
  157. snowflake/ml/modeling/manifold/tsne.py +195 -123
  158. snowflake/ml/modeling/mixture/bayesian_gaussian_mixture.py +195 -123
  159. snowflake/ml/modeling/mixture/gaussian_mixture.py +195 -123
  160. snowflake/ml/modeling/model_selection/grid_search_cv.py +42 -18
  161. snowflake/ml/modeling/model_selection/randomized_search_cv.py +42 -18
  162. snowflake/ml/modeling/multiclass/one_vs_one_classifier.py +195 -123
  163. snowflake/ml/modeling/multiclass/one_vs_rest_classifier.py +195 -123
  164. snowflake/ml/modeling/multiclass/output_code_classifier.py +195 -123
  165. snowflake/ml/modeling/naive_bayes/bernoulli_nb.py +195 -123
  166. snowflake/ml/modeling/naive_bayes/categorical_nb.py +195 -123
  167. snowflake/ml/modeling/naive_bayes/complement_nb.py +195 -123
  168. snowflake/ml/modeling/naive_bayes/gaussian_nb.py +195 -123
  169. snowflake/ml/modeling/naive_bayes/multinomial_nb.py +195 -123
  170. snowflake/ml/modeling/neighbors/k_neighbors_classifier.py +195 -123
  171. snowflake/ml/modeling/neighbors/k_neighbors_regressor.py +195 -123
  172. snowflake/ml/modeling/neighbors/kernel_density.py +195 -123
  173. snowflake/ml/modeling/neighbors/local_outlier_factor.py +195 -123
  174. snowflake/ml/modeling/neighbors/nearest_centroid.py +195 -123
  175. snowflake/ml/modeling/neighbors/nearest_neighbors.py +195 -123
  176. snowflake/ml/modeling/neighbors/neighborhood_components_analysis.py +195 -123
  177. snowflake/ml/modeling/neighbors/radius_neighbors_classifier.py +195 -123
  178. snowflake/ml/modeling/neighbors/radius_neighbors_regressor.py +195 -123
  179. snowflake/ml/modeling/neural_network/bernoulli_rbm.py +195 -123
  180. snowflake/ml/modeling/neural_network/mlp_classifier.py +195 -123
  181. snowflake/ml/modeling/neural_network/mlp_regressor.py +195 -123
  182. snowflake/ml/modeling/pipeline/pipeline.py +4 -4
  183. snowflake/ml/modeling/preprocessing/binarizer.py +1 -5
  184. snowflake/ml/modeling/preprocessing/k_bins_discretizer.py +1 -5
  185. snowflake/ml/modeling/preprocessing/label_encoder.py +1 -5
  186. snowflake/ml/modeling/preprocessing/max_abs_scaler.py +1 -5
  187. snowflake/ml/modeling/preprocessing/min_max_scaler.py +10 -12
  188. snowflake/ml/modeling/preprocessing/normalizer.py +1 -5
  189. snowflake/ml/modeling/preprocessing/one_hot_encoder.py +1 -5
  190. snowflake/ml/modeling/preprocessing/ordinal_encoder.py +1 -5
  191. snowflake/ml/modeling/preprocessing/polynomial_features.py +195 -123
  192. snowflake/ml/modeling/preprocessing/robust_scaler.py +1 -5
  193. snowflake/ml/modeling/preprocessing/standard_scaler.py +11 -11
  194. snowflake/ml/modeling/semi_supervised/label_propagation.py +195 -123
  195. snowflake/ml/modeling/semi_supervised/label_spreading.py +195 -123
  196. snowflake/ml/modeling/svm/linear_svc.py +195 -123
  197. snowflake/ml/modeling/svm/linear_svr.py +195 -123
  198. snowflake/ml/modeling/svm/nu_svc.py +195 -123
  199. snowflake/ml/modeling/svm/nu_svr.py +195 -123
  200. snowflake/ml/modeling/svm/svc.py +195 -123
  201. snowflake/ml/modeling/svm/svr.py +195 -123
  202. snowflake/ml/modeling/tree/decision_tree_classifier.py +195 -123
  203. snowflake/ml/modeling/tree/decision_tree_regressor.py +195 -123
  204. snowflake/ml/modeling/tree/extra_tree_classifier.py +195 -123
  205. snowflake/ml/modeling/tree/extra_tree_regressor.py +195 -123
  206. snowflake/ml/modeling/xgboost/xgb_classifier.py +195 -123
  207. snowflake/ml/modeling/xgboost/xgb_regressor.py +195 -123
  208. snowflake/ml/modeling/xgboost/xgbrf_classifier.py +195 -123
  209. snowflake/ml/modeling/xgboost/xgbrf_regressor.py +195 -123
  210. snowflake/ml/registry/_manager/model_manager.py +5 -1
  211. snowflake/ml/registry/model_registry.py +99 -26
  212. snowflake/ml/registry/registry.py +3 -2
  213. snowflake/ml/version.py +1 -1
  214. {snowflake_ml_python-1.3.1.dist-info → snowflake_ml_python-1.4.1.dist-info}/METADATA +94 -55
  215. {snowflake_ml_python-1.3.1.dist-info → snowflake_ml_python-1.4.1.dist-info}/RECORD +218 -212
  216. snowflake/ml/model/_model_composer/model_runtime/model_runtime.py +0 -97
  217. {snowflake_ml_python-1.3.1.dist-info → snowflake_ml_python-1.4.1.dist-info}/LICENSE.txt +0 -0
  218. {snowflake_ml_python-1.3.1.dist-info → snowflake_ml_python-1.4.1.dist-info}/WHEEL +0 -0
  219. {snowflake_ml_python-1.3.1.dist-info → snowflake_ml_python-1.4.1.dist-info}/top_level.txt +0 -0
@@ -33,6 +33,15 @@ from snowflake.ml.modeling._internal.transformer_protocols import (
33
33
  BatchInferenceKwargsTypedDict,
34
34
  ScoreKwargsTypedDict
35
35
  )
36
+ from snowflake.ml.model._signatures import utils as model_signature_utils
37
+ from snowflake.ml.model.model_signature import (
38
+ BaseFeatureSpec,
39
+ DataType,
40
+ FeatureSpec,
41
+ ModelSignature,
42
+ _infer_signature,
43
+ _rename_signature_with_snowflake_identifiers,
44
+ )
36
45
 
37
46
  from snowflake.ml.modeling._internal.model_transformer_builder import ModelTransformerBuilder
38
47
 
@@ -43,16 +52,6 @@ from snowflake.ml.modeling._internal.estimator_utils import (
43
52
  validate_sklearn_args,
44
53
  )
45
54
 
46
- from snowflake.ml.model.model_signature import (
47
- DataType,
48
- FeatureSpec,
49
- ModelSignature,
50
- _infer_signature,
51
- _rename_signature_with_snowflake_identifiers,
52
- BaseFeatureSpec,
53
- )
54
- from snowflake.ml.model._signatures import utils as model_signature_utils
55
-
56
55
  _PROJECT = "ModelDevelopment"
57
56
  # Derive subproject from module name by removing "sklearn"
58
57
  # and converting module name from underscore to CamelCase
@@ -283,12 +282,7 @@ class LassoLarsCV(BaseTransformer):
283
282
  )
284
283
  return selected_cols
285
284
 
286
- @telemetry.send_api_usage_telemetry(
287
- project=_PROJECT,
288
- subproject=_SUBPROJECT,
289
- custom_tags=dict([("autogen", True)]),
290
- )
291
- def fit(self, dataset: Union[DataFrame, pd.DataFrame]) -> "LassoLarsCV":
285
+ def _fit(self, dataset: Union[DataFrame, pd.DataFrame]) -> "LassoLarsCV":
292
286
  """Fit the model using X, y as training data
293
287
  For more details on this function, see [sklearn.linear_model.LassoLarsCV.fit]
294
288
  (https://scikit-learn.org/stable/modules/generated/sklearn.linear_model.LassoLarsCV.html#sklearn.linear_model.LassoLarsCV.fit)
@@ -315,12 +309,14 @@ class LassoLarsCV(BaseTransformer):
315
309
 
316
310
  self._snowpark_cols = dataset.select(self.input_cols).columns
317
311
 
318
- # If we are already in a stored procedure, no need to kick off another one.
312
+ # If we are already in a stored procedure, no need to kick off another one.
319
313
  if SNOWML_SPROC_ENV in os.environ:
320
314
  statement_params = telemetry.get_function_usage_statement_params(
321
315
  project=_PROJECT,
322
316
  subproject=_SUBPROJECT,
323
- function_name=telemetry.get_statement_params_full_func_name(inspect.currentframe(), LassoLarsCV.__class__.__name__),
317
+ function_name=telemetry.get_statement_params_full_func_name(
318
+ inspect.currentframe(), LassoLarsCV.__class__.__name__
319
+ ),
324
320
  api_calls=[Session.call],
325
321
  custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
326
322
  )
@@ -341,7 +337,7 @@ class LassoLarsCV(BaseTransformer):
341
337
  )
342
338
  self._sklearn_object = model_trainer.train()
343
339
  self._is_fitted = True
344
- self._get_model_signatures(dataset)
340
+ self._generate_model_signatures(dataset)
345
341
  return self
346
342
 
347
343
  def _batch_inference_validate_snowpark(
@@ -417,7 +413,9 @@ class LassoLarsCV(BaseTransformer):
417
413
  # when it is classifier, infer the datatype from label columns
418
414
  if expected_type_inferred == "" and 'predict' in self.model_signatures:
419
415
  # Batch inference takes a single expected output column type. Use the first columns type for now.
420
- label_cols_signatures = [row for row in self.model_signatures['predict'].outputs if row.name in self.output_cols]
416
+ label_cols_signatures = [
417
+ row for row in self.model_signatures['predict'].outputs if row.name in self.output_cols
418
+ ]
421
419
  if len(label_cols_signatures) == 0:
422
420
  error_str = f"Output columns {self.output_cols} do not match model signatures {self.model_signatures['predict'].outputs}."
423
421
  raise exceptions.SnowflakeMLException(
@@ -425,25 +423,22 @@ class LassoLarsCV(BaseTransformer):
425
423
  original_exception=ValueError(error_str),
426
424
  )
427
425
 
428
- expected_type_inferred = convert_sp_to_sf_type(
429
- label_cols_signatures[0].as_snowpark_type()
430
- )
426
+ expected_type_inferred = convert_sp_to_sf_type(label_cols_signatures[0].as_snowpark_type())
431
427
 
432
428
  self._deps = self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
433
- assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
429
+ assert isinstance(
430
+ dataset._session, Session
431
+ ) # mypy does not recognize the check in _batch_inference_validate_snowpark()
434
432
 
435
433
  transform_kwargs = dict(
436
- session = dataset._session,
437
- dependencies = self._deps,
438
- drop_input_cols = self._drop_input_cols,
439
- expected_output_cols_type = expected_type_inferred,
434
+ session=dataset._session,
435
+ dependencies=self._deps,
436
+ drop_input_cols=self._drop_input_cols,
437
+ expected_output_cols_type=expected_type_inferred,
440
438
  )
441
439
 
442
440
  elif isinstance(dataset, pd.DataFrame):
443
- transform_kwargs = dict(
444
- snowpark_input_cols = self._snowpark_cols,
445
- drop_input_cols = self._drop_input_cols
446
- )
441
+ transform_kwargs = dict(snowpark_input_cols=self._snowpark_cols, drop_input_cols=self._drop_input_cols)
447
442
 
448
443
  transform_handlers = ModelTransformerBuilder.build(
449
444
  dataset=dataset,
@@ -483,7 +478,7 @@ class LassoLarsCV(BaseTransformer):
483
478
  Transformed dataset.
484
479
  """
485
480
  super()._check_dataset_type(dataset)
486
- inference_method="transform"
481
+ inference_method = "transform"
487
482
 
488
483
  # This dictionary contains optional kwargs for batch inference. These kwargs
489
484
  # are specific to the type of dataset used.
@@ -520,17 +515,14 @@ class LassoLarsCV(BaseTransformer):
520
515
  assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
521
516
 
522
517
  transform_kwargs = dict(
523
- session = dataset._session,
524
- dependencies = self._deps,
525
- drop_input_cols = self._drop_input_cols,
526
- expected_output_cols_type = expected_dtype,
518
+ session=dataset._session,
519
+ dependencies=self._deps,
520
+ drop_input_cols=self._drop_input_cols,
521
+ expected_output_cols_type=expected_dtype,
527
522
  )
528
523
 
529
524
  elif isinstance(dataset, pd.DataFrame):
530
- transform_kwargs = dict(
531
- snowpark_input_cols = self._snowpark_cols,
532
- drop_input_cols = self._drop_input_cols
533
- )
525
+ transform_kwargs = dict(snowpark_input_cols=self._snowpark_cols, drop_input_cols=self._drop_input_cols)
534
526
 
535
527
  transform_handlers = ModelTransformerBuilder.build(
536
528
  dataset=dataset,
@@ -549,7 +541,11 @@ class LassoLarsCV(BaseTransformer):
549
541
  return output_df
550
542
 
551
543
  @available_if(original_estimator_has_callable("fit_predict")) # type: ignore[misc]
552
- def fit_predict(self, dataset: Union[DataFrame, pd.DataFrame], output_cols_prefix: str = "fit_predict_",) -> Union[DataFrame, pd.DataFrame]:
544
+ def fit_predict(
545
+ self,
546
+ dataset: Union[DataFrame, pd.DataFrame],
547
+ output_cols_prefix: str = "fit_predict_",
548
+ ) -> Union[DataFrame, pd.DataFrame]:
553
549
  """ Method not supported for this class.
554
550
 
555
551
 
@@ -574,7 +570,9 @@ class LassoLarsCV(BaseTransformer):
574
570
  )
575
571
  output_result, fitted_estimator = model_trainer.train_fit_predict(
576
572
  drop_input_cols=self._drop_input_cols,
577
- expected_output_cols_list=self.output_cols if self.output_cols else self._get_output_column_names(output_cols_prefix),
573
+ expected_output_cols_list=(
574
+ self.output_cols if self.output_cols else self._get_output_column_names(output_cols_prefix)
575
+ ),
578
576
  )
579
577
  self._sklearn_object = fitted_estimator
580
578
  self._is_fitted = True
@@ -591,6 +589,62 @@ class LassoLarsCV(BaseTransformer):
591
589
  assert self._sklearn_object is not None
592
590
  return self._sklearn_object.embedding_
593
591
 
592
+
593
+ def _get_output_column_names(self, output_cols_prefix: str, output_cols: Optional[List[str]] = None) -> List[str]:
594
+ """ Returns the list of output columns for predict_proba(), decision_function(), etc.. functions.
595
+ Returns a list with output_cols_prefix as the only element if the estimator is not a classifier.
596
+ """
597
+ output_cols_prefix = identifier.resolve_identifier(output_cols_prefix)
598
+ # The following condition is introduced for kneighbors methods, and not used in other methods
599
+ if output_cols:
600
+ output_cols = [
601
+ identifier.concat_names([output_cols_prefix, identifier.resolve_identifier(c)])
602
+ for c in output_cols
603
+ ]
604
+ elif getattr(self._sklearn_object, "classes_", None) is None:
605
+ output_cols = [output_cols_prefix]
606
+ elif self._sklearn_object is not None:
607
+ classes = self._sklearn_object.classes_
608
+ if isinstance(classes, numpy.ndarray):
609
+ output_cols = [f'{output_cols_prefix}{str(c)}' for c in classes.tolist()]
610
+ elif isinstance(classes, list) and len(classes) > 0 and isinstance(classes[0], numpy.ndarray):
611
+ # If the estimator is a multioutput estimator, classes_ will be a list of ndarrays.
612
+ output_cols = []
613
+ for i, cl in enumerate(classes):
614
+ # For binary classification, there is only one output column for each class
615
+ # ndarray as the two classes are complementary.
616
+ if len(cl) == 2:
617
+ output_cols.append(f'{output_cols_prefix}{i}_{cl[0]}')
618
+ else:
619
+ output_cols.extend([
620
+ f'{output_cols_prefix}{i}_{c}' for c in cl.tolist()
621
+ ])
622
+ else:
623
+ output_cols = []
624
+
625
+ # Make sure column names are valid snowflake identifiers.
626
+ assert output_cols is not None # Make MyPy happy
627
+ rv = [identifier.rename_to_valid_snowflake_identifier(c) for c in output_cols]
628
+
629
+ return rv
630
+
631
+ def _align_expected_output_names(
632
+ self, method: str, dataset: DataFrame, expected_output_cols_list: List[str], output_cols_prefix: str
633
+ ) -> List[str]:
634
+ # in case the inferred output column names dimension is different
635
+ # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
636
+ output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
637
+ output_df_columns = list(output_df_pd.columns)
638
+ output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
639
+ if self.sample_weight_col:
640
+ output_df_columns_set -= set(self.sample_weight_col)
641
+ # if the dimension of inferred output column names is correct; use it
642
+ if len(expected_output_cols_list) == len(output_df_columns_set):
643
+ return expected_output_cols_list
644
+ # otherwise, use the sklearn estimator's output
645
+ else:
646
+ return sorted(list(output_df_columns_set), key=lambda x: output_df_columns.index(x))
647
+
594
648
  @available_if(original_estimator_has_callable("predict_proba")) # type: ignore[misc]
595
649
  @telemetry.send_api_usage_telemetry(
596
650
  project=_PROJECT,
@@ -621,24 +675,28 @@ class LassoLarsCV(BaseTransformer):
621
675
  # are specific to the type of dataset used.
622
676
  transform_kwargs: BatchInferenceKwargsTypedDict = dict()
623
677
 
678
+ expected_output_cols = self._get_output_column_names(output_cols_prefix)
679
+
624
680
  if isinstance(dataset, DataFrame):
625
681
  self._deps = self._batch_inference_validate_snowpark(
626
682
  dataset=dataset,
627
683
  inference_method=inference_method,
628
684
  )
629
- assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
685
+ assert isinstance(
686
+ dataset._session, Session
687
+ ) # mypy does not recognize the check in _batch_inference_validate_snowpark()
630
688
  transform_kwargs = dict(
631
689
  session=dataset._session,
632
690
  dependencies=self._deps,
633
- drop_input_cols = self._drop_input_cols,
691
+ drop_input_cols=self._drop_input_cols,
634
692
  expected_output_cols_type="float",
635
693
  )
694
+ expected_output_cols = self._align_expected_output_names(
695
+ inference_method, dataset, expected_output_cols, output_cols_prefix
696
+ )
636
697
 
637
698
  elif isinstance(dataset, pd.DataFrame):
638
- transform_kwargs = dict(
639
- snowpark_input_cols = self._snowpark_cols,
640
- drop_input_cols = self._drop_input_cols
641
- )
699
+ transform_kwargs = dict(snowpark_input_cols=self._snowpark_cols, drop_input_cols=self._drop_input_cols)
642
700
 
643
701
  transform_handlers = ModelTransformerBuilder.build(
644
702
  dataset=dataset,
@@ -650,7 +708,7 @@ class LassoLarsCV(BaseTransformer):
650
708
  output_df: DATAFRAME_TYPE = transform_handlers.batch_inference(
651
709
  inference_method=inference_method,
652
710
  input_cols=self.input_cols,
653
- expected_output_cols=self._get_output_column_names(output_cols_prefix),
711
+ expected_output_cols=expected_output_cols,
654
712
  **transform_kwargs
655
713
  )
656
714
  return output_df
@@ -680,7 +738,8 @@ class LassoLarsCV(BaseTransformer):
680
738
  Output dataset with log probability of the sample for each class in the model.
681
739
  """
682
740
  super()._check_dataset_type(dataset)
683
- inference_method="predict_log_proba"
741
+ inference_method = "predict_log_proba"
742
+ expected_output_cols = self._get_output_column_names(output_cols_prefix)
684
743
 
685
744
  # This dictionary contains optional kwargs for batch inference. These kwargs
686
745
  # are specific to the type of dataset used.
@@ -691,18 +750,20 @@ class LassoLarsCV(BaseTransformer):
691
750
  dataset=dataset,
692
751
  inference_method=inference_method,
693
752
  )
694
- assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
753
+ assert isinstance(
754
+ dataset._session, Session
755
+ ) # mypy does not recognize the check in _batch_inference_validate_snowpark()
695
756
  transform_kwargs = dict(
696
757
  session=dataset._session,
697
758
  dependencies=self._deps,
698
- drop_input_cols = self._drop_input_cols,
759
+ drop_input_cols=self._drop_input_cols,
699
760
  expected_output_cols_type="float",
700
761
  )
762
+ expected_output_cols = self._align_expected_output_names(
763
+ inference_method, dataset, expected_output_cols, output_cols_prefix
764
+ )
701
765
  elif isinstance(dataset, pd.DataFrame):
702
- transform_kwargs = dict(
703
- snowpark_input_cols = self._snowpark_cols,
704
- drop_input_cols = self._drop_input_cols
705
- )
766
+ transform_kwargs = dict(snowpark_input_cols=self._snowpark_cols, drop_input_cols=self._drop_input_cols)
706
767
 
707
768
  transform_handlers = ModelTransformerBuilder.build(
708
769
  dataset=dataset,
@@ -715,7 +776,7 @@ class LassoLarsCV(BaseTransformer):
715
776
  output_df: DATAFRAME_TYPE = transform_handlers.batch_inference(
716
777
  inference_method=inference_method,
717
778
  input_cols=self.input_cols,
718
- expected_output_cols=self._get_output_column_names(output_cols_prefix),
779
+ expected_output_cols=expected_output_cols,
719
780
  **transform_kwargs
720
781
  )
721
782
  return output_df
@@ -741,30 +802,34 @@ class LassoLarsCV(BaseTransformer):
741
802
  Output dataset with results of the decision function for the samples in input dataset.
742
803
  """
743
804
  super()._check_dataset_type(dataset)
744
- inference_method="decision_function"
805
+ inference_method = "decision_function"
745
806
 
746
807
  # This dictionary contains optional kwargs for batch inference. These kwargs
747
808
  # are specific to the type of dataset used.
748
809
  transform_kwargs: BatchInferenceKwargsTypedDict = dict()
749
810
 
811
+ expected_output_cols = self._get_output_column_names(output_cols_prefix)
812
+
750
813
  if isinstance(dataset, DataFrame):
751
814
  self._deps = self._batch_inference_validate_snowpark(
752
815
  dataset=dataset,
753
816
  inference_method=inference_method,
754
817
  )
755
- assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
818
+ assert isinstance(
819
+ dataset._session, Session
820
+ ) # mypy does not recognize the check in _batch_inference_validate_snowpark()
756
821
  transform_kwargs = dict(
757
822
  session=dataset._session,
758
823
  dependencies=self._deps,
759
- drop_input_cols = self._drop_input_cols,
824
+ drop_input_cols=self._drop_input_cols,
760
825
  expected_output_cols_type="float",
761
826
  )
827
+ expected_output_cols = self._align_expected_output_names(
828
+ inference_method, dataset, expected_output_cols, output_cols_prefix
829
+ )
762
830
 
763
831
  elif isinstance(dataset, pd.DataFrame):
764
- transform_kwargs = dict(
765
- snowpark_input_cols = self._snowpark_cols,
766
- drop_input_cols = self._drop_input_cols
767
- )
832
+ transform_kwargs = dict(snowpark_input_cols=self._snowpark_cols, drop_input_cols=self._drop_input_cols)
768
833
 
769
834
  transform_handlers = ModelTransformerBuilder.build(
770
835
  dataset=dataset,
@@ -777,7 +842,7 @@ class LassoLarsCV(BaseTransformer):
777
842
  output_df: DATAFRAME_TYPE = transform_handlers.batch_inference(
778
843
  inference_method=inference_method,
779
844
  input_cols=self.input_cols,
780
- expected_output_cols=self._get_output_column_names(output_cols_prefix),
845
+ expected_output_cols=expected_output_cols,
781
846
  **transform_kwargs
782
847
  )
783
848
  return output_df
@@ -806,12 +871,14 @@ class LassoLarsCV(BaseTransformer):
806
871
  Output dataset with probability of the sample for each class in the model.
807
872
  """
808
873
  super()._check_dataset_type(dataset)
809
- inference_method="score_samples"
874
+ inference_method = "score_samples"
810
875
 
811
876
  # This dictionary contains optional kwargs for batch inference. These kwargs
812
877
  # are specific to the type of dataset used.
813
878
  transform_kwargs: BatchInferenceKwargsTypedDict = dict()
814
879
 
880
+ expected_output_cols = self._get_output_column_names(output_cols_prefix)
881
+
815
882
  if isinstance(dataset, DataFrame):
816
883
  self._deps = self._batch_inference_validate_snowpark(
817
884
  dataset=dataset,
@@ -824,6 +891,9 @@ class LassoLarsCV(BaseTransformer):
824
891
  drop_input_cols = self._drop_input_cols,
825
892
  expected_output_cols_type="float",
826
893
  )
894
+ expected_output_cols = self._align_expected_output_names(
895
+ inference_method, dataset, expected_output_cols, output_cols_prefix
896
+ )
827
897
 
828
898
  elif isinstance(dataset, pd.DataFrame):
829
899
  transform_kwargs = dict(
@@ -842,7 +912,7 @@ class LassoLarsCV(BaseTransformer):
842
912
  output_df: DATAFRAME_TYPE = transform_handlers.batch_inference(
843
913
  inference_method=inference_method,
844
914
  input_cols=self.input_cols,
845
- expected_output_cols=self._get_output_column_names(output_cols_prefix),
915
+ expected_output_cols=expected_output_cols,
846
916
  **transform_kwargs
847
917
  )
848
918
  return output_df
@@ -989,50 +1059,84 @@ class LassoLarsCV(BaseTransformer):
989
1059
  )
990
1060
  return output_df
991
1061
 
1062
+
1063
+
1064
+ def to_sklearn(self) -> Any:
1065
+ """Get sklearn.linear_model.LassoLarsCV object.
1066
+ """
1067
+ if self._sklearn_object is None:
1068
+ self._sklearn_object = self._create_sklearn_object()
1069
+ return self._sklearn_object
1070
+
1071
+ def to_xgboost(self) -> Any:
1072
+ raise exceptions.SnowflakeMLException(
1073
+ error_code=error_codes.METHOD_NOT_ALLOWED,
1074
+ original_exception=AttributeError(
1075
+ modeling_error_messages.UNSUPPORTED_MODEL_CONVERSION.format(
1076
+ "to_xgboost()",
1077
+ "to_sklearn()"
1078
+ )
1079
+ ),
1080
+ )
1081
+
1082
+ def to_lightgbm(self) -> Any:
1083
+ raise exceptions.SnowflakeMLException(
1084
+ error_code=error_codes.METHOD_NOT_ALLOWED,
1085
+ original_exception=AttributeError(
1086
+ modeling_error_messages.UNSUPPORTED_MODEL_CONVERSION.format(
1087
+ "to_lightgbm()",
1088
+ "to_sklearn()"
1089
+ )
1090
+ ),
1091
+ )
992
1092
 
993
- def _get_model_signatures(self, dataset: Union[DataFrame, pd.DataFrame]) -> None:
1093
+ def _get_dependencies(self) -> List[str]:
1094
+ return self._deps
1095
+
1096
+
1097
+ def _generate_model_signatures(self, dataset: Union[DataFrame, pd.DataFrame]) -> None:
994
1098
  self._model_signature_dict = dict()
995
1099
 
996
1100
  PROB_FUNCTIONS = ["predict_log_proba", "predict_proba", "decision_function"]
997
1101
 
998
- inputs = list(_infer_signature(dataset[self.input_cols], "input"))
1102
+ inputs = list(_infer_signature(dataset[self.input_cols], "input", use_snowflake_identifiers=True))
999
1103
  outputs: List[BaseFeatureSpec] = []
1000
1104
  if hasattr(self, "predict"):
1001
1105
  # keep mypy happy
1002
- assert self._sklearn_object is not None and hasattr(self._sklearn_object, "_estimator_type")
1106
+ assert self._sklearn_object is not None and hasattr(self._sklearn_object, "_estimator_type")
1003
1107
  # For classifier, the type of predict is the same as the type of label
1004
- if self._sklearn_object._estimator_type == 'classifier':
1005
- # label columns is the desired type for output
1108
+ if self._sklearn_object._estimator_type == "classifier":
1109
+ # label columns is the desired type for output
1006
1110
  outputs = list(_infer_signature(dataset[self.label_cols], "output", use_snowflake_identifiers=True))
1007
1111
  # rename the output columns
1008
1112
  outputs = list(model_signature_utils.rename_features(outputs, self.output_cols))
1009
- self._model_signature_dict["predict"] = ModelSignature(inputs,
1010
- ([] if self._drop_input_cols else inputs)
1011
- + outputs)
1113
+ self._model_signature_dict["predict"] = ModelSignature(
1114
+ inputs, ([] if self._drop_input_cols else inputs) + outputs
1115
+ )
1012
1116
  # For mixture models that use the density mixin, `predict` returns the argmax of the log prob.
1013
1117
  # For outlier models, returns -1 for outliers and 1 for inliers.
1014
- # Clusterer returns int64 cluster labels.
1118
+ # Clusterer returns int64 cluster labels.
1015
1119
  elif self._sklearn_object._estimator_type in ["DensityEstimator", "clusterer", "outlier_detector"]:
1016
1120
  outputs = [FeatureSpec(dtype=DataType.INT64, name=c) for c in self.output_cols]
1017
- self._model_signature_dict["predict"] = ModelSignature(inputs,
1018
- ([] if self._drop_input_cols else inputs)
1019
- + outputs)
1020
-
1121
+ self._model_signature_dict["predict"] = ModelSignature(
1122
+ inputs, ([] if self._drop_input_cols else inputs) + outputs
1123
+ )
1124
+
1021
1125
  # For regressor, the type of predict is float64
1022
- elif self._sklearn_object._estimator_type == 'regressor':
1126
+ elif self._sklearn_object._estimator_type == "regressor":
1023
1127
  outputs = [FeatureSpec(dtype=DataType.DOUBLE, name=c) for c in self.output_cols]
1024
- self._model_signature_dict["predict"] = ModelSignature(inputs,
1025
- ([] if self._drop_input_cols else inputs)
1026
- + outputs)
1027
-
1128
+ self._model_signature_dict["predict"] = ModelSignature(
1129
+ inputs, ([] if self._drop_input_cols else inputs) + outputs
1130
+ )
1131
+
1028
1132
  for prob_func in PROB_FUNCTIONS:
1029
1133
  if hasattr(self, prob_func):
1030
1134
  output_cols_prefix: str = f"{prob_func}_"
1031
1135
  output_column_names = self._get_output_column_names(output_cols_prefix)
1032
1136
  outputs = [FeatureSpec(dtype=DataType.DOUBLE, name=c) for c in output_column_names]
1033
- self._model_signature_dict[prob_func] = ModelSignature(inputs,
1034
- ([] if self._drop_input_cols else inputs)
1035
- + outputs)
1137
+ self._model_signature_dict[prob_func] = ModelSignature(
1138
+ inputs, ([] if self._drop_input_cols else inputs) + outputs
1139
+ )
1036
1140
 
1037
1141
  # Output signature names may still need to be renamed, since they were not created with `_infer_signature`.
1038
1142
  items = list(self._model_signature_dict.items())
@@ -1045,10 +1149,10 @@ class LassoLarsCV(BaseTransformer):
1045
1149
  """Returns model signature of current class.
1046
1150
 
1047
1151
  Raises:
1048
- exceptions.SnowflakeMLException: If estimator is not fitted, then model signature cannot be inferred
1152
+ SnowflakeMLException: If estimator is not fitted, then model signature cannot be inferred
1049
1153
 
1050
1154
  Returns:
1051
- Dict[str, ModelSignature]: each method and its input output signature
1155
+ Dict with each method and its input output signature
1052
1156
  """
1053
1157
  if self._model_signature_dict is None:
1054
1158
  raise exceptions.SnowflakeMLException(
@@ -1056,35 +1160,3 @@ class LassoLarsCV(BaseTransformer):
1056
1160
  original_exception=RuntimeError("Estimator not fitted before accessing property model_signatures!"),
1057
1161
  )
1058
1162
  return self._model_signature_dict
1059
-
1060
- def to_sklearn(self) -> Any:
1061
- """Get sklearn.linear_model.LassoLarsCV object.
1062
- """
1063
- if self._sklearn_object is None:
1064
- self._sklearn_object = self._create_sklearn_object()
1065
- return self._sklearn_object
1066
-
1067
- def to_xgboost(self) -> Any:
1068
- raise exceptions.SnowflakeMLException(
1069
- error_code=error_codes.METHOD_NOT_ALLOWED,
1070
- original_exception=AttributeError(
1071
- modeling_error_messages.UNSUPPORTED_MODEL_CONVERSION.format(
1072
- "to_xgboost()",
1073
- "to_sklearn()"
1074
- )
1075
- ),
1076
- )
1077
-
1078
- def to_lightgbm(self) -> Any:
1079
- raise exceptions.SnowflakeMLException(
1080
- error_code=error_codes.METHOD_NOT_ALLOWED,
1081
- original_exception=AttributeError(
1082
- modeling_error_messages.UNSUPPORTED_MODEL_CONVERSION.format(
1083
- "to_lightgbm()",
1084
- "to_sklearn()"
1085
- )
1086
- ),
1087
- )
1088
-
1089
- def _get_dependencies(self) -> List[str]:
1090
- return self._deps