snowflake-ml-python 1.4.0__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 (203) hide show
  1. snowflake/ml/_internal/env_utils.py +11 -1
  2. snowflake/ml/_internal/utils/identifier.py +3 -1
  3. snowflake/ml/_internal/utils/sql_identifier.py +2 -6
  4. snowflake/ml/feature_store/feature_store.py +151 -78
  5. snowflake/ml/feature_store/feature_view.py +12 -24
  6. snowflake/ml/fileset/sfcfs.py +56 -50
  7. snowflake/ml/fileset/stage_fs.py +48 -13
  8. snowflake/ml/model/_client/model/model_version_impl.py +2 -50
  9. snowflake/ml/model/_client/ops/model_ops.py +78 -29
  10. snowflake/ml/model/_client/sql/model.py +23 -2
  11. snowflake/ml/model/_client/sql/model_version.py +22 -1
  12. snowflake/ml/model/_model_composer/model_manifest/model_manifest.py +19 -54
  13. snowflake/ml/model/_model_composer/model_manifest/model_manifest_schema.py +8 -1
  14. snowflake/ml/model/_model_composer/model_method/model_method.py +6 -10
  15. snowflake/ml/model/_packager/model_handlers/catboost.py +206 -0
  16. snowflake/ml/model/_packager/model_handlers/lightgbm.py +218 -0
  17. snowflake/ml/model/_packager/model_handlers/sklearn.py +3 -0
  18. snowflake/ml/model/_packager/model_meta/_core_requirements.py +1 -1
  19. snowflake/ml/model/_packager/model_meta/model_meta.py +36 -6
  20. snowflake/ml/model/_packager/model_meta/model_meta_schema.py +20 -1
  21. snowflake/ml/model/_packager/model_meta_migrator/migrator_plans.py +3 -1
  22. snowflake/ml/model/_packager/model_packager.py +2 -2
  23. snowflake/ml/model/{_model_composer/model_runtime/_runtime_requirements.py → _packager/model_runtime/_snowml_inference_alternative_requirements.py} +1 -1
  24. snowflake/ml/model/_packager/model_runtime/model_runtime.py +137 -0
  25. snowflake/ml/model/type_hints.py +21 -2
  26. snowflake/ml/modeling/_internal/estimator_utils.py +16 -11
  27. snowflake/ml/modeling/_internal/local_implementations/pandas_handlers.py +4 -1
  28. snowflake/ml/modeling/_internal/snowpark_implementations/distributed_hpo_trainer.py +13 -14
  29. snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_handlers.py +5 -5
  30. snowflake/ml/modeling/calibration/calibrated_classifier_cv.py +195 -123
  31. snowflake/ml/modeling/cluster/affinity_propagation.py +195 -123
  32. snowflake/ml/modeling/cluster/agglomerative_clustering.py +195 -123
  33. snowflake/ml/modeling/cluster/birch.py +195 -123
  34. snowflake/ml/modeling/cluster/bisecting_k_means.py +195 -123
  35. snowflake/ml/modeling/cluster/dbscan.py +195 -123
  36. snowflake/ml/modeling/cluster/feature_agglomeration.py +195 -123
  37. snowflake/ml/modeling/cluster/k_means.py +195 -123
  38. snowflake/ml/modeling/cluster/mean_shift.py +195 -123
  39. snowflake/ml/modeling/cluster/mini_batch_k_means.py +195 -123
  40. snowflake/ml/modeling/cluster/optics.py +195 -123
  41. snowflake/ml/modeling/cluster/spectral_biclustering.py +195 -123
  42. snowflake/ml/modeling/cluster/spectral_clustering.py +195 -123
  43. snowflake/ml/modeling/cluster/spectral_coclustering.py +195 -123
  44. snowflake/ml/modeling/compose/column_transformer.py +195 -123
  45. snowflake/ml/modeling/compose/transformed_target_regressor.py +195 -123
  46. snowflake/ml/modeling/covariance/elliptic_envelope.py +195 -123
  47. snowflake/ml/modeling/covariance/empirical_covariance.py +195 -123
  48. snowflake/ml/modeling/covariance/graphical_lasso.py +195 -123
  49. snowflake/ml/modeling/covariance/graphical_lasso_cv.py +195 -123
  50. snowflake/ml/modeling/covariance/ledoit_wolf.py +195 -123
  51. snowflake/ml/modeling/covariance/min_cov_det.py +195 -123
  52. snowflake/ml/modeling/covariance/oas.py +195 -123
  53. snowflake/ml/modeling/covariance/shrunk_covariance.py +195 -123
  54. snowflake/ml/modeling/decomposition/dictionary_learning.py +195 -123
  55. snowflake/ml/modeling/decomposition/factor_analysis.py +195 -123
  56. snowflake/ml/modeling/decomposition/fast_ica.py +195 -123
  57. snowflake/ml/modeling/decomposition/incremental_pca.py +195 -123
  58. snowflake/ml/modeling/decomposition/kernel_pca.py +195 -123
  59. snowflake/ml/modeling/decomposition/mini_batch_dictionary_learning.py +195 -123
  60. snowflake/ml/modeling/decomposition/mini_batch_sparse_pca.py +195 -123
  61. snowflake/ml/modeling/decomposition/pca.py +195 -123
  62. snowflake/ml/modeling/decomposition/sparse_pca.py +195 -123
  63. snowflake/ml/modeling/decomposition/truncated_svd.py +195 -123
  64. snowflake/ml/modeling/discriminant_analysis/linear_discriminant_analysis.py +195 -123
  65. snowflake/ml/modeling/discriminant_analysis/quadratic_discriminant_analysis.py +195 -123
  66. snowflake/ml/modeling/ensemble/ada_boost_classifier.py +195 -123
  67. snowflake/ml/modeling/ensemble/ada_boost_regressor.py +195 -123
  68. snowflake/ml/modeling/ensemble/bagging_classifier.py +195 -123
  69. snowflake/ml/modeling/ensemble/bagging_regressor.py +195 -123
  70. snowflake/ml/modeling/ensemble/extra_trees_classifier.py +195 -123
  71. snowflake/ml/modeling/ensemble/extra_trees_regressor.py +195 -123
  72. snowflake/ml/modeling/ensemble/gradient_boosting_classifier.py +195 -123
  73. snowflake/ml/modeling/ensemble/gradient_boosting_regressor.py +195 -123
  74. snowflake/ml/modeling/ensemble/hist_gradient_boosting_classifier.py +195 -123
  75. snowflake/ml/modeling/ensemble/hist_gradient_boosting_regressor.py +195 -123
  76. snowflake/ml/modeling/ensemble/isolation_forest.py +195 -123
  77. snowflake/ml/modeling/ensemble/random_forest_classifier.py +195 -123
  78. snowflake/ml/modeling/ensemble/random_forest_regressor.py +195 -123
  79. snowflake/ml/modeling/ensemble/stacking_regressor.py +195 -123
  80. snowflake/ml/modeling/ensemble/voting_classifier.py +195 -123
  81. snowflake/ml/modeling/ensemble/voting_regressor.py +195 -123
  82. snowflake/ml/modeling/feature_selection/generic_univariate_select.py +195 -123
  83. snowflake/ml/modeling/feature_selection/select_fdr.py +195 -123
  84. snowflake/ml/modeling/feature_selection/select_fpr.py +195 -123
  85. snowflake/ml/modeling/feature_selection/select_fwe.py +195 -123
  86. snowflake/ml/modeling/feature_selection/select_k_best.py +195 -123
  87. snowflake/ml/modeling/feature_selection/select_percentile.py +195 -123
  88. snowflake/ml/modeling/feature_selection/sequential_feature_selector.py +195 -123
  89. snowflake/ml/modeling/feature_selection/variance_threshold.py +195 -123
  90. snowflake/ml/modeling/framework/_utils.py +8 -1
  91. snowflake/ml/modeling/framework/base.py +9 -1
  92. snowflake/ml/modeling/gaussian_process/gaussian_process_classifier.py +195 -123
  93. snowflake/ml/modeling/gaussian_process/gaussian_process_regressor.py +195 -123
  94. snowflake/ml/modeling/impute/iterative_imputer.py +195 -123
  95. snowflake/ml/modeling/impute/knn_imputer.py +195 -123
  96. snowflake/ml/modeling/impute/missing_indicator.py +195 -123
  97. snowflake/ml/modeling/kernel_approximation/additive_chi2_sampler.py +195 -123
  98. snowflake/ml/modeling/kernel_approximation/nystroem.py +195 -123
  99. snowflake/ml/modeling/kernel_approximation/polynomial_count_sketch.py +195 -123
  100. snowflake/ml/modeling/kernel_approximation/rbf_sampler.py +195 -123
  101. snowflake/ml/modeling/kernel_approximation/skewed_chi2_sampler.py +195 -123
  102. snowflake/ml/modeling/kernel_ridge/kernel_ridge.py +195 -123
  103. snowflake/ml/modeling/lightgbm/lgbm_classifier.py +195 -123
  104. snowflake/ml/modeling/lightgbm/lgbm_regressor.py +195 -123
  105. snowflake/ml/modeling/linear_model/ard_regression.py +195 -123
  106. snowflake/ml/modeling/linear_model/bayesian_ridge.py +195 -123
  107. snowflake/ml/modeling/linear_model/elastic_net.py +195 -123
  108. snowflake/ml/modeling/linear_model/elastic_net_cv.py +195 -123
  109. snowflake/ml/modeling/linear_model/gamma_regressor.py +195 -123
  110. snowflake/ml/modeling/linear_model/huber_regressor.py +195 -123
  111. snowflake/ml/modeling/linear_model/lars.py +195 -123
  112. snowflake/ml/modeling/linear_model/lars_cv.py +195 -123
  113. snowflake/ml/modeling/linear_model/lasso.py +195 -123
  114. snowflake/ml/modeling/linear_model/lasso_cv.py +195 -123
  115. snowflake/ml/modeling/linear_model/lasso_lars.py +195 -123
  116. snowflake/ml/modeling/linear_model/lasso_lars_cv.py +195 -123
  117. snowflake/ml/modeling/linear_model/lasso_lars_ic.py +195 -123
  118. snowflake/ml/modeling/linear_model/linear_regression.py +195 -123
  119. snowflake/ml/modeling/linear_model/logistic_regression.py +195 -123
  120. snowflake/ml/modeling/linear_model/logistic_regression_cv.py +195 -123
  121. snowflake/ml/modeling/linear_model/multi_task_elastic_net.py +195 -123
  122. snowflake/ml/modeling/linear_model/multi_task_elastic_net_cv.py +195 -123
  123. snowflake/ml/modeling/linear_model/multi_task_lasso.py +195 -123
  124. snowflake/ml/modeling/linear_model/multi_task_lasso_cv.py +195 -123
  125. snowflake/ml/modeling/linear_model/orthogonal_matching_pursuit.py +195 -123
  126. snowflake/ml/modeling/linear_model/passive_aggressive_classifier.py +195 -123
  127. snowflake/ml/modeling/linear_model/passive_aggressive_regressor.py +195 -123
  128. snowflake/ml/modeling/linear_model/perceptron.py +195 -123
  129. snowflake/ml/modeling/linear_model/poisson_regressor.py +195 -123
  130. snowflake/ml/modeling/linear_model/ransac_regressor.py +195 -123
  131. snowflake/ml/modeling/linear_model/ridge.py +195 -123
  132. snowflake/ml/modeling/linear_model/ridge_classifier.py +195 -123
  133. snowflake/ml/modeling/linear_model/ridge_classifier_cv.py +195 -123
  134. snowflake/ml/modeling/linear_model/ridge_cv.py +195 -123
  135. snowflake/ml/modeling/linear_model/sgd_classifier.py +195 -123
  136. snowflake/ml/modeling/linear_model/sgd_one_class_svm.py +195 -123
  137. snowflake/ml/modeling/linear_model/sgd_regressor.py +195 -123
  138. snowflake/ml/modeling/linear_model/theil_sen_regressor.py +195 -123
  139. snowflake/ml/modeling/linear_model/tweedie_regressor.py +195 -123
  140. snowflake/ml/modeling/manifold/isomap.py +195 -123
  141. snowflake/ml/modeling/manifold/mds.py +195 -123
  142. snowflake/ml/modeling/manifold/spectral_embedding.py +195 -123
  143. snowflake/ml/modeling/manifold/tsne.py +195 -123
  144. snowflake/ml/modeling/mixture/bayesian_gaussian_mixture.py +195 -123
  145. snowflake/ml/modeling/mixture/gaussian_mixture.py +195 -123
  146. snowflake/ml/modeling/model_selection/grid_search_cv.py +42 -18
  147. snowflake/ml/modeling/model_selection/randomized_search_cv.py +42 -18
  148. snowflake/ml/modeling/multiclass/one_vs_one_classifier.py +195 -123
  149. snowflake/ml/modeling/multiclass/one_vs_rest_classifier.py +195 -123
  150. snowflake/ml/modeling/multiclass/output_code_classifier.py +195 -123
  151. snowflake/ml/modeling/naive_bayes/bernoulli_nb.py +195 -123
  152. snowflake/ml/modeling/naive_bayes/categorical_nb.py +195 -123
  153. snowflake/ml/modeling/naive_bayes/complement_nb.py +195 -123
  154. snowflake/ml/modeling/naive_bayes/gaussian_nb.py +195 -123
  155. snowflake/ml/modeling/naive_bayes/multinomial_nb.py +195 -123
  156. snowflake/ml/modeling/neighbors/k_neighbors_classifier.py +195 -123
  157. snowflake/ml/modeling/neighbors/k_neighbors_regressor.py +195 -123
  158. snowflake/ml/modeling/neighbors/kernel_density.py +195 -123
  159. snowflake/ml/modeling/neighbors/local_outlier_factor.py +195 -123
  160. snowflake/ml/modeling/neighbors/nearest_centroid.py +195 -123
  161. snowflake/ml/modeling/neighbors/nearest_neighbors.py +195 -123
  162. snowflake/ml/modeling/neighbors/neighborhood_components_analysis.py +195 -123
  163. snowflake/ml/modeling/neighbors/radius_neighbors_classifier.py +195 -123
  164. snowflake/ml/modeling/neighbors/radius_neighbors_regressor.py +195 -123
  165. snowflake/ml/modeling/neural_network/bernoulli_rbm.py +195 -123
  166. snowflake/ml/modeling/neural_network/mlp_classifier.py +195 -123
  167. snowflake/ml/modeling/neural_network/mlp_regressor.py +195 -123
  168. snowflake/ml/modeling/pipeline/pipeline.py +4 -4
  169. snowflake/ml/modeling/preprocessing/binarizer.py +1 -5
  170. snowflake/ml/modeling/preprocessing/k_bins_discretizer.py +1 -5
  171. snowflake/ml/modeling/preprocessing/label_encoder.py +1 -5
  172. snowflake/ml/modeling/preprocessing/max_abs_scaler.py +1 -5
  173. snowflake/ml/modeling/preprocessing/min_max_scaler.py +10 -12
  174. snowflake/ml/modeling/preprocessing/normalizer.py +1 -5
  175. snowflake/ml/modeling/preprocessing/one_hot_encoder.py +1 -5
  176. snowflake/ml/modeling/preprocessing/ordinal_encoder.py +1 -5
  177. snowflake/ml/modeling/preprocessing/polynomial_features.py +195 -123
  178. snowflake/ml/modeling/preprocessing/robust_scaler.py +1 -5
  179. snowflake/ml/modeling/preprocessing/standard_scaler.py +11 -11
  180. snowflake/ml/modeling/semi_supervised/label_propagation.py +195 -123
  181. snowflake/ml/modeling/semi_supervised/label_spreading.py +195 -123
  182. snowflake/ml/modeling/svm/linear_svc.py +195 -123
  183. snowflake/ml/modeling/svm/linear_svr.py +195 -123
  184. snowflake/ml/modeling/svm/nu_svc.py +195 -123
  185. snowflake/ml/modeling/svm/nu_svr.py +195 -123
  186. snowflake/ml/modeling/svm/svc.py +195 -123
  187. snowflake/ml/modeling/svm/svr.py +195 -123
  188. snowflake/ml/modeling/tree/decision_tree_classifier.py +195 -123
  189. snowflake/ml/modeling/tree/decision_tree_regressor.py +195 -123
  190. snowflake/ml/modeling/tree/extra_tree_classifier.py +195 -123
  191. snowflake/ml/modeling/tree/extra_tree_regressor.py +195 -123
  192. snowflake/ml/modeling/xgboost/xgb_classifier.py +195 -123
  193. snowflake/ml/modeling/xgboost/xgb_regressor.py +195 -123
  194. snowflake/ml/modeling/xgboost/xgbrf_classifier.py +195 -123
  195. snowflake/ml/modeling/xgboost/xgbrf_regressor.py +195 -123
  196. snowflake/ml/registry/registry.py +1 -1
  197. snowflake/ml/version.py +1 -1
  198. {snowflake_ml_python-1.4.0.dist-info → snowflake_ml_python-1.4.1.dist-info}/METADATA +68 -57
  199. {snowflake_ml_python-1.4.0.dist-info → snowflake_ml_python-1.4.1.dist-info}/RECORD +202 -200
  200. snowflake/ml/model/_model_composer/model_runtime/model_runtime.py +0 -97
  201. {snowflake_ml_python-1.4.0.dist-info → snowflake_ml_python-1.4.1.dist-info}/LICENSE.txt +0 -0
  202. {snowflake_ml_python-1.4.0.dist-info → snowflake_ml_python-1.4.1.dist-info}/WHEEL +0 -0
  203. {snowflake_ml_python-1.4.0.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
@@ -305,12 +304,7 @@ class RANSACRegressor(BaseTransformer):
305
304
  )
306
305
  return selected_cols
307
306
 
308
- @telemetry.send_api_usage_telemetry(
309
- project=_PROJECT,
310
- subproject=_SUBPROJECT,
311
- custom_tags=dict([("autogen", True)]),
312
- )
313
- def fit(self, dataset: Union[DataFrame, pd.DataFrame]) -> "RANSACRegressor":
307
+ def _fit(self, dataset: Union[DataFrame, pd.DataFrame]) -> "RANSACRegressor":
314
308
  """Fit estimator using RANSAC algorithm
315
309
  For more details on this function, see [sklearn.linear_model.RANSACRegressor.fit]
316
310
  (https://scikit-learn.org/stable/modules/generated/sklearn.linear_model.RANSACRegressor.html#sklearn.linear_model.RANSACRegressor.fit)
@@ -337,12 +331,14 @@ class RANSACRegressor(BaseTransformer):
337
331
 
338
332
  self._snowpark_cols = dataset.select(self.input_cols).columns
339
333
 
340
- # If we are already in a stored procedure, no need to kick off another one.
334
+ # If we are already in a stored procedure, no need to kick off another one.
341
335
  if SNOWML_SPROC_ENV in os.environ:
342
336
  statement_params = telemetry.get_function_usage_statement_params(
343
337
  project=_PROJECT,
344
338
  subproject=_SUBPROJECT,
345
- function_name=telemetry.get_statement_params_full_func_name(inspect.currentframe(), RANSACRegressor.__class__.__name__),
339
+ function_name=telemetry.get_statement_params_full_func_name(
340
+ inspect.currentframe(), RANSACRegressor.__class__.__name__
341
+ ),
346
342
  api_calls=[Session.call],
347
343
  custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
348
344
  )
@@ -363,7 +359,7 @@ class RANSACRegressor(BaseTransformer):
363
359
  )
364
360
  self._sklearn_object = model_trainer.train()
365
361
  self._is_fitted = True
366
- self._get_model_signatures(dataset)
362
+ self._generate_model_signatures(dataset)
367
363
  return self
368
364
 
369
365
  def _batch_inference_validate_snowpark(
@@ -439,7 +435,9 @@ class RANSACRegressor(BaseTransformer):
439
435
  # when it is classifier, infer the datatype from label columns
440
436
  if expected_type_inferred == "" and 'predict' in self.model_signatures:
441
437
  # Batch inference takes a single expected output column type. Use the first columns type for now.
442
- label_cols_signatures = [row for row in self.model_signatures['predict'].outputs if row.name in self.output_cols]
438
+ label_cols_signatures = [
439
+ row for row in self.model_signatures['predict'].outputs if row.name in self.output_cols
440
+ ]
443
441
  if len(label_cols_signatures) == 0:
444
442
  error_str = f"Output columns {self.output_cols} do not match model signatures {self.model_signatures['predict'].outputs}."
445
443
  raise exceptions.SnowflakeMLException(
@@ -447,25 +445,22 @@ class RANSACRegressor(BaseTransformer):
447
445
  original_exception=ValueError(error_str),
448
446
  )
449
447
 
450
- expected_type_inferred = convert_sp_to_sf_type(
451
- label_cols_signatures[0].as_snowpark_type()
452
- )
448
+ expected_type_inferred = convert_sp_to_sf_type(label_cols_signatures[0].as_snowpark_type())
453
449
 
454
450
  self._deps = self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
455
- assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
451
+ assert isinstance(
452
+ dataset._session, Session
453
+ ) # mypy does not recognize the check in _batch_inference_validate_snowpark()
456
454
 
457
455
  transform_kwargs = dict(
458
- session = dataset._session,
459
- dependencies = self._deps,
460
- drop_input_cols = self._drop_input_cols,
461
- expected_output_cols_type = expected_type_inferred,
456
+ session=dataset._session,
457
+ dependencies=self._deps,
458
+ drop_input_cols=self._drop_input_cols,
459
+ expected_output_cols_type=expected_type_inferred,
462
460
  )
463
461
 
464
462
  elif isinstance(dataset, pd.DataFrame):
465
- transform_kwargs = dict(
466
- snowpark_input_cols = self._snowpark_cols,
467
- drop_input_cols = self._drop_input_cols
468
- )
463
+ transform_kwargs = dict(snowpark_input_cols=self._snowpark_cols, drop_input_cols=self._drop_input_cols)
469
464
 
470
465
  transform_handlers = ModelTransformerBuilder.build(
471
466
  dataset=dataset,
@@ -505,7 +500,7 @@ class RANSACRegressor(BaseTransformer):
505
500
  Transformed dataset.
506
501
  """
507
502
  super()._check_dataset_type(dataset)
508
- inference_method="transform"
503
+ inference_method = "transform"
509
504
 
510
505
  # This dictionary contains optional kwargs for batch inference. These kwargs
511
506
  # are specific to the type of dataset used.
@@ -542,17 +537,14 @@ class RANSACRegressor(BaseTransformer):
542
537
  assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
543
538
 
544
539
  transform_kwargs = dict(
545
- session = dataset._session,
546
- dependencies = self._deps,
547
- drop_input_cols = self._drop_input_cols,
548
- expected_output_cols_type = expected_dtype,
540
+ session=dataset._session,
541
+ dependencies=self._deps,
542
+ drop_input_cols=self._drop_input_cols,
543
+ expected_output_cols_type=expected_dtype,
549
544
  )
550
545
 
551
546
  elif isinstance(dataset, pd.DataFrame):
552
- transform_kwargs = dict(
553
- snowpark_input_cols = self._snowpark_cols,
554
- drop_input_cols = self._drop_input_cols
555
- )
547
+ transform_kwargs = dict(snowpark_input_cols=self._snowpark_cols, drop_input_cols=self._drop_input_cols)
556
548
 
557
549
  transform_handlers = ModelTransformerBuilder.build(
558
550
  dataset=dataset,
@@ -571,7 +563,11 @@ class RANSACRegressor(BaseTransformer):
571
563
  return output_df
572
564
 
573
565
  @available_if(original_estimator_has_callable("fit_predict")) # type: ignore[misc]
574
- def fit_predict(self, dataset: Union[DataFrame, pd.DataFrame], output_cols_prefix: str = "fit_predict_",) -> Union[DataFrame, pd.DataFrame]:
566
+ def fit_predict(
567
+ self,
568
+ dataset: Union[DataFrame, pd.DataFrame],
569
+ output_cols_prefix: str = "fit_predict_",
570
+ ) -> Union[DataFrame, pd.DataFrame]:
575
571
  """ Method not supported for this class.
576
572
 
577
573
 
@@ -596,7 +592,9 @@ class RANSACRegressor(BaseTransformer):
596
592
  )
597
593
  output_result, fitted_estimator = model_trainer.train_fit_predict(
598
594
  drop_input_cols=self._drop_input_cols,
599
- expected_output_cols_list=self.output_cols if self.output_cols else self._get_output_column_names(output_cols_prefix),
595
+ expected_output_cols_list=(
596
+ self.output_cols if self.output_cols else self._get_output_column_names(output_cols_prefix)
597
+ ),
600
598
  )
601
599
  self._sklearn_object = fitted_estimator
602
600
  self._is_fitted = True
@@ -613,6 +611,62 @@ class RANSACRegressor(BaseTransformer):
613
611
  assert self._sklearn_object is not None
614
612
  return self._sklearn_object.embedding_
615
613
 
614
+
615
+ def _get_output_column_names(self, output_cols_prefix: str, output_cols: Optional[List[str]] = None) -> List[str]:
616
+ """ Returns the list of output columns for predict_proba(), decision_function(), etc.. functions.
617
+ Returns a list with output_cols_prefix as the only element if the estimator is not a classifier.
618
+ """
619
+ output_cols_prefix = identifier.resolve_identifier(output_cols_prefix)
620
+ # The following condition is introduced for kneighbors methods, and not used in other methods
621
+ if output_cols:
622
+ output_cols = [
623
+ identifier.concat_names([output_cols_prefix, identifier.resolve_identifier(c)])
624
+ for c in output_cols
625
+ ]
626
+ elif getattr(self._sklearn_object, "classes_", None) is None:
627
+ output_cols = [output_cols_prefix]
628
+ elif self._sklearn_object is not None:
629
+ classes = self._sklearn_object.classes_
630
+ if isinstance(classes, numpy.ndarray):
631
+ output_cols = [f'{output_cols_prefix}{str(c)}' for c in classes.tolist()]
632
+ elif isinstance(classes, list) and len(classes) > 0 and isinstance(classes[0], numpy.ndarray):
633
+ # If the estimator is a multioutput estimator, classes_ will be a list of ndarrays.
634
+ output_cols = []
635
+ for i, cl in enumerate(classes):
636
+ # For binary classification, there is only one output column for each class
637
+ # ndarray as the two classes are complementary.
638
+ if len(cl) == 2:
639
+ output_cols.append(f'{output_cols_prefix}{i}_{cl[0]}')
640
+ else:
641
+ output_cols.extend([
642
+ f'{output_cols_prefix}{i}_{c}' for c in cl.tolist()
643
+ ])
644
+ else:
645
+ output_cols = []
646
+
647
+ # Make sure column names are valid snowflake identifiers.
648
+ assert output_cols is not None # Make MyPy happy
649
+ rv = [identifier.rename_to_valid_snowflake_identifier(c) for c in output_cols]
650
+
651
+ return rv
652
+
653
+ def _align_expected_output_names(
654
+ self, method: str, dataset: DataFrame, expected_output_cols_list: List[str], output_cols_prefix: str
655
+ ) -> List[str]:
656
+ # in case the inferred output column names dimension is different
657
+ # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
658
+ output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
659
+ output_df_columns = list(output_df_pd.columns)
660
+ output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
661
+ if self.sample_weight_col:
662
+ output_df_columns_set -= set(self.sample_weight_col)
663
+ # if the dimension of inferred output column names is correct; use it
664
+ if len(expected_output_cols_list) == len(output_df_columns_set):
665
+ return expected_output_cols_list
666
+ # otherwise, use the sklearn estimator's output
667
+ else:
668
+ return sorted(list(output_df_columns_set), key=lambda x: output_df_columns.index(x))
669
+
616
670
  @available_if(original_estimator_has_callable("predict_proba")) # type: ignore[misc]
617
671
  @telemetry.send_api_usage_telemetry(
618
672
  project=_PROJECT,
@@ -643,24 +697,28 @@ class RANSACRegressor(BaseTransformer):
643
697
  # are specific to the type of dataset used.
644
698
  transform_kwargs: BatchInferenceKwargsTypedDict = dict()
645
699
 
700
+ expected_output_cols = self._get_output_column_names(output_cols_prefix)
701
+
646
702
  if isinstance(dataset, DataFrame):
647
703
  self._deps = self._batch_inference_validate_snowpark(
648
704
  dataset=dataset,
649
705
  inference_method=inference_method,
650
706
  )
651
- assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
707
+ assert isinstance(
708
+ dataset._session, Session
709
+ ) # mypy does not recognize the check in _batch_inference_validate_snowpark()
652
710
  transform_kwargs = dict(
653
711
  session=dataset._session,
654
712
  dependencies=self._deps,
655
- drop_input_cols = self._drop_input_cols,
713
+ drop_input_cols=self._drop_input_cols,
656
714
  expected_output_cols_type="float",
657
715
  )
716
+ expected_output_cols = self._align_expected_output_names(
717
+ inference_method, dataset, expected_output_cols, output_cols_prefix
718
+ )
658
719
 
659
720
  elif isinstance(dataset, pd.DataFrame):
660
- transform_kwargs = dict(
661
- snowpark_input_cols = self._snowpark_cols,
662
- drop_input_cols = self._drop_input_cols
663
- )
721
+ transform_kwargs = dict(snowpark_input_cols=self._snowpark_cols, drop_input_cols=self._drop_input_cols)
664
722
 
665
723
  transform_handlers = ModelTransformerBuilder.build(
666
724
  dataset=dataset,
@@ -672,7 +730,7 @@ class RANSACRegressor(BaseTransformer):
672
730
  output_df: DATAFRAME_TYPE = transform_handlers.batch_inference(
673
731
  inference_method=inference_method,
674
732
  input_cols=self.input_cols,
675
- expected_output_cols=self._get_output_column_names(output_cols_prefix),
733
+ expected_output_cols=expected_output_cols,
676
734
  **transform_kwargs
677
735
  )
678
736
  return output_df
@@ -702,7 +760,8 @@ class RANSACRegressor(BaseTransformer):
702
760
  Output dataset with log probability of the sample for each class in the model.
703
761
  """
704
762
  super()._check_dataset_type(dataset)
705
- inference_method="predict_log_proba"
763
+ inference_method = "predict_log_proba"
764
+ expected_output_cols = self._get_output_column_names(output_cols_prefix)
706
765
 
707
766
  # This dictionary contains optional kwargs for batch inference. These kwargs
708
767
  # are specific to the type of dataset used.
@@ -713,18 +772,20 @@ class RANSACRegressor(BaseTransformer):
713
772
  dataset=dataset,
714
773
  inference_method=inference_method,
715
774
  )
716
- assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
775
+ assert isinstance(
776
+ dataset._session, Session
777
+ ) # mypy does not recognize the check in _batch_inference_validate_snowpark()
717
778
  transform_kwargs = dict(
718
779
  session=dataset._session,
719
780
  dependencies=self._deps,
720
- drop_input_cols = self._drop_input_cols,
781
+ drop_input_cols=self._drop_input_cols,
721
782
  expected_output_cols_type="float",
722
783
  )
784
+ expected_output_cols = self._align_expected_output_names(
785
+ inference_method, dataset, expected_output_cols, output_cols_prefix
786
+ )
723
787
  elif isinstance(dataset, pd.DataFrame):
724
- transform_kwargs = dict(
725
- snowpark_input_cols = self._snowpark_cols,
726
- drop_input_cols = self._drop_input_cols
727
- )
788
+ transform_kwargs = dict(snowpark_input_cols=self._snowpark_cols, drop_input_cols=self._drop_input_cols)
728
789
 
729
790
  transform_handlers = ModelTransformerBuilder.build(
730
791
  dataset=dataset,
@@ -737,7 +798,7 @@ class RANSACRegressor(BaseTransformer):
737
798
  output_df: DATAFRAME_TYPE = transform_handlers.batch_inference(
738
799
  inference_method=inference_method,
739
800
  input_cols=self.input_cols,
740
- expected_output_cols=self._get_output_column_names(output_cols_prefix),
801
+ expected_output_cols=expected_output_cols,
741
802
  **transform_kwargs
742
803
  )
743
804
  return output_df
@@ -763,30 +824,34 @@ class RANSACRegressor(BaseTransformer):
763
824
  Output dataset with results of the decision function for the samples in input dataset.
764
825
  """
765
826
  super()._check_dataset_type(dataset)
766
- inference_method="decision_function"
827
+ inference_method = "decision_function"
767
828
 
768
829
  # This dictionary contains optional kwargs for batch inference. These kwargs
769
830
  # are specific to the type of dataset used.
770
831
  transform_kwargs: BatchInferenceKwargsTypedDict = dict()
771
832
 
833
+ expected_output_cols = self._get_output_column_names(output_cols_prefix)
834
+
772
835
  if isinstance(dataset, DataFrame):
773
836
  self._deps = self._batch_inference_validate_snowpark(
774
837
  dataset=dataset,
775
838
  inference_method=inference_method,
776
839
  )
777
- assert isinstance(dataset._session, Session) # mypy does not recognize the check in _batch_inference_validate_snowpark()
840
+ assert isinstance(
841
+ dataset._session, Session
842
+ ) # mypy does not recognize the check in _batch_inference_validate_snowpark()
778
843
  transform_kwargs = dict(
779
844
  session=dataset._session,
780
845
  dependencies=self._deps,
781
- drop_input_cols = self._drop_input_cols,
846
+ drop_input_cols=self._drop_input_cols,
782
847
  expected_output_cols_type="float",
783
848
  )
849
+ expected_output_cols = self._align_expected_output_names(
850
+ inference_method, dataset, expected_output_cols, output_cols_prefix
851
+ )
784
852
 
785
853
  elif isinstance(dataset, pd.DataFrame):
786
- transform_kwargs = dict(
787
- snowpark_input_cols = self._snowpark_cols,
788
- drop_input_cols = self._drop_input_cols
789
- )
854
+ transform_kwargs = dict(snowpark_input_cols=self._snowpark_cols, drop_input_cols=self._drop_input_cols)
790
855
 
791
856
  transform_handlers = ModelTransformerBuilder.build(
792
857
  dataset=dataset,
@@ -799,7 +864,7 @@ class RANSACRegressor(BaseTransformer):
799
864
  output_df: DATAFRAME_TYPE = transform_handlers.batch_inference(
800
865
  inference_method=inference_method,
801
866
  input_cols=self.input_cols,
802
- expected_output_cols=self._get_output_column_names(output_cols_prefix),
867
+ expected_output_cols=expected_output_cols,
803
868
  **transform_kwargs
804
869
  )
805
870
  return output_df
@@ -828,12 +893,14 @@ class RANSACRegressor(BaseTransformer):
828
893
  Output dataset with probability of the sample for each class in the model.
829
894
  """
830
895
  super()._check_dataset_type(dataset)
831
- inference_method="score_samples"
896
+ inference_method = "score_samples"
832
897
 
833
898
  # This dictionary contains optional kwargs for batch inference. These kwargs
834
899
  # are specific to the type of dataset used.
835
900
  transform_kwargs: BatchInferenceKwargsTypedDict = dict()
836
901
 
902
+ expected_output_cols = self._get_output_column_names(output_cols_prefix)
903
+
837
904
  if isinstance(dataset, DataFrame):
838
905
  self._deps = self._batch_inference_validate_snowpark(
839
906
  dataset=dataset,
@@ -846,6 +913,9 @@ class RANSACRegressor(BaseTransformer):
846
913
  drop_input_cols = self._drop_input_cols,
847
914
  expected_output_cols_type="float",
848
915
  )
916
+ expected_output_cols = self._align_expected_output_names(
917
+ inference_method, dataset, expected_output_cols, output_cols_prefix
918
+ )
849
919
 
850
920
  elif isinstance(dataset, pd.DataFrame):
851
921
  transform_kwargs = dict(
@@ -864,7 +934,7 @@ class RANSACRegressor(BaseTransformer):
864
934
  output_df: DATAFRAME_TYPE = transform_handlers.batch_inference(
865
935
  inference_method=inference_method,
866
936
  input_cols=self.input_cols,
867
- expected_output_cols=self._get_output_column_names(output_cols_prefix),
937
+ expected_output_cols=expected_output_cols,
868
938
  **transform_kwargs
869
939
  )
870
940
  return output_df
@@ -1011,50 +1081,84 @@ class RANSACRegressor(BaseTransformer):
1011
1081
  )
1012
1082
  return output_df
1013
1083
 
1084
+
1085
+
1086
+ def to_sklearn(self) -> Any:
1087
+ """Get sklearn.linear_model.RANSACRegressor object.
1088
+ """
1089
+ if self._sklearn_object is None:
1090
+ self._sklearn_object = self._create_sklearn_object()
1091
+ return self._sklearn_object
1092
+
1093
+ def to_xgboost(self) -> Any:
1094
+ raise exceptions.SnowflakeMLException(
1095
+ error_code=error_codes.METHOD_NOT_ALLOWED,
1096
+ original_exception=AttributeError(
1097
+ modeling_error_messages.UNSUPPORTED_MODEL_CONVERSION.format(
1098
+ "to_xgboost()",
1099
+ "to_sklearn()"
1100
+ )
1101
+ ),
1102
+ )
1103
+
1104
+ def to_lightgbm(self) -> Any:
1105
+ raise exceptions.SnowflakeMLException(
1106
+ error_code=error_codes.METHOD_NOT_ALLOWED,
1107
+ original_exception=AttributeError(
1108
+ modeling_error_messages.UNSUPPORTED_MODEL_CONVERSION.format(
1109
+ "to_lightgbm()",
1110
+ "to_sklearn()"
1111
+ )
1112
+ ),
1113
+ )
1014
1114
 
1015
- def _get_model_signatures(self, dataset: Union[DataFrame, pd.DataFrame]) -> None:
1115
+ def _get_dependencies(self) -> List[str]:
1116
+ return self._deps
1117
+
1118
+
1119
+ def _generate_model_signatures(self, dataset: Union[DataFrame, pd.DataFrame]) -> None:
1016
1120
  self._model_signature_dict = dict()
1017
1121
 
1018
1122
  PROB_FUNCTIONS = ["predict_log_proba", "predict_proba", "decision_function"]
1019
1123
 
1020
- inputs = list(_infer_signature(dataset[self.input_cols], "input"))
1124
+ inputs = list(_infer_signature(dataset[self.input_cols], "input", use_snowflake_identifiers=True))
1021
1125
  outputs: List[BaseFeatureSpec] = []
1022
1126
  if hasattr(self, "predict"):
1023
1127
  # keep mypy happy
1024
- assert self._sklearn_object is not None and hasattr(self._sklearn_object, "_estimator_type")
1128
+ assert self._sklearn_object is not None and hasattr(self._sklearn_object, "_estimator_type")
1025
1129
  # For classifier, the type of predict is the same as the type of label
1026
- if self._sklearn_object._estimator_type == 'classifier':
1027
- # label columns is the desired type for output
1130
+ if self._sklearn_object._estimator_type == "classifier":
1131
+ # label columns is the desired type for output
1028
1132
  outputs = list(_infer_signature(dataset[self.label_cols], "output", use_snowflake_identifiers=True))
1029
1133
  # rename the output columns
1030
1134
  outputs = list(model_signature_utils.rename_features(outputs, self.output_cols))
1031
- self._model_signature_dict["predict"] = ModelSignature(inputs,
1032
- ([] if self._drop_input_cols else inputs)
1033
- + outputs)
1135
+ self._model_signature_dict["predict"] = ModelSignature(
1136
+ inputs, ([] if self._drop_input_cols else inputs) + outputs
1137
+ )
1034
1138
  # For mixture models that use the density mixin, `predict` returns the argmax of the log prob.
1035
1139
  # For outlier models, returns -1 for outliers and 1 for inliers.
1036
- # Clusterer returns int64 cluster labels.
1140
+ # Clusterer returns int64 cluster labels.
1037
1141
  elif self._sklearn_object._estimator_type in ["DensityEstimator", "clusterer", "outlier_detector"]:
1038
1142
  outputs = [FeatureSpec(dtype=DataType.INT64, name=c) for c in self.output_cols]
1039
- self._model_signature_dict["predict"] = ModelSignature(inputs,
1040
- ([] if self._drop_input_cols else inputs)
1041
- + outputs)
1042
-
1143
+ self._model_signature_dict["predict"] = ModelSignature(
1144
+ inputs, ([] if self._drop_input_cols else inputs) + outputs
1145
+ )
1146
+
1043
1147
  # For regressor, the type of predict is float64
1044
- elif self._sklearn_object._estimator_type == 'regressor':
1148
+ elif self._sklearn_object._estimator_type == "regressor":
1045
1149
  outputs = [FeatureSpec(dtype=DataType.DOUBLE, name=c) for c in self.output_cols]
1046
- self._model_signature_dict["predict"] = ModelSignature(inputs,
1047
- ([] if self._drop_input_cols else inputs)
1048
- + outputs)
1049
-
1150
+ self._model_signature_dict["predict"] = ModelSignature(
1151
+ inputs, ([] if self._drop_input_cols else inputs) + outputs
1152
+ )
1153
+
1050
1154
  for prob_func in PROB_FUNCTIONS:
1051
1155
  if hasattr(self, prob_func):
1052
1156
  output_cols_prefix: str = f"{prob_func}_"
1053
1157
  output_column_names = self._get_output_column_names(output_cols_prefix)
1054
1158
  outputs = [FeatureSpec(dtype=DataType.DOUBLE, name=c) for c in output_column_names]
1055
- self._model_signature_dict[prob_func] = ModelSignature(inputs,
1056
- ([] if self._drop_input_cols else inputs)
1057
- + outputs)
1159
+ self._model_signature_dict[prob_func] = ModelSignature(
1160
+ inputs, ([] if self._drop_input_cols else inputs) + outputs
1161
+ )
1058
1162
 
1059
1163
  # Output signature names may still need to be renamed, since they were not created with `_infer_signature`.
1060
1164
  items = list(self._model_signature_dict.items())
@@ -1067,10 +1171,10 @@ class RANSACRegressor(BaseTransformer):
1067
1171
  """Returns model signature of current class.
1068
1172
 
1069
1173
  Raises:
1070
- exceptions.SnowflakeMLException: If estimator is not fitted, then model signature cannot be inferred
1174
+ SnowflakeMLException: If estimator is not fitted, then model signature cannot be inferred
1071
1175
 
1072
1176
  Returns:
1073
- Dict[str, ModelSignature]: each method and its input output signature
1177
+ Dict with each method and its input output signature
1074
1178
  """
1075
1179
  if self._model_signature_dict is None:
1076
1180
  raise exceptions.SnowflakeMLException(
@@ -1078,35 +1182,3 @@ class RANSACRegressor(BaseTransformer):
1078
1182
  original_exception=RuntimeError("Estimator not fitted before accessing property model_signatures!"),
1079
1183
  )
1080
1184
  return self._model_signature_dict
1081
-
1082
- def to_sklearn(self) -> Any:
1083
- """Get sklearn.linear_model.RANSACRegressor object.
1084
- """
1085
- if self._sklearn_object is None:
1086
- self._sklearn_object = self._create_sklearn_object()
1087
- return self._sklearn_object
1088
-
1089
- def to_xgboost(self) -> Any:
1090
- raise exceptions.SnowflakeMLException(
1091
- error_code=error_codes.METHOD_NOT_ALLOWED,
1092
- original_exception=AttributeError(
1093
- modeling_error_messages.UNSUPPORTED_MODEL_CONVERSION.format(
1094
- "to_xgboost()",
1095
- "to_sklearn()"
1096
- )
1097
- ),
1098
- )
1099
-
1100
- def to_lightgbm(self) -> Any:
1101
- raise exceptions.SnowflakeMLException(
1102
- error_code=error_codes.METHOD_NOT_ALLOWED,
1103
- original_exception=AttributeError(
1104
- modeling_error_messages.UNSUPPORTED_MODEL_CONVERSION.format(
1105
- "to_lightgbm()",
1106
- "to_sklearn()"
1107
- )
1108
- ),
1109
- )
1110
-
1111
- def _get_dependencies(self) -> List[str]:
1112
- return self._deps