snowflake-ml-python 1.5.1__py3-none-any.whl → 1.5.3__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 (207) hide show
  1. snowflake/cortex/_complete.py +26 -5
  2. snowflake/cortex/_sentiment.py +7 -4
  3. snowflake/cortex/_sse_client.py +81 -0
  4. snowflake/cortex/_util.py +105 -8
  5. snowflake/ml/_internal/lineage/lineage_utils.py +34 -25
  6. snowflake/ml/_internal/utils/temp_file_utils.py +5 -2
  7. snowflake/ml/dataset/dataset.py +15 -12
  8. snowflake/ml/dataset/dataset_factory.py +3 -4
  9. snowflake/ml/feature_store/access_manager.py +34 -30
  10. snowflake/ml/feature_store/feature_store.py +3 -3
  11. snowflake/ml/feature_store/feature_view.py +12 -11
  12. snowflake/ml/fileset/snowfs.py +2 -31
  13. snowflake/ml/model/_client/ops/model_ops.py +43 -0
  14. snowflake/ml/model/_client/sql/model_version.py +55 -3
  15. snowflake/ml/model/_model_composer/model_composer.py +7 -3
  16. snowflake/ml/model/_model_composer/model_method/infer_table_function.py_template +3 -1
  17. snowflake/ml/model/_packager/model_meta/_core_requirements.py +1 -1
  18. snowflake/ml/model/_packager/model_meta/model_meta.py +1 -3
  19. snowflake/ml/model/_packager/model_runtime/_snowml_inference_alternative_requirements.py +1 -1
  20. snowflake/ml/model/_packager/model_runtime/model_runtime.py +3 -27
  21. snowflake/ml/model/_signatures/builtins_handler.py +2 -1
  22. snowflake/ml/model/_signatures/core.py +13 -1
  23. snowflake/ml/model/_signatures/pandas_handler.py +2 -0
  24. snowflake/ml/model/_signatures/snowpark_handler.py +3 -3
  25. snowflake/ml/model/model_signature.py +2 -0
  26. snowflake/ml/model/type_hints.py +1 -0
  27. snowflake/ml/modeling/_internal/estimator_utils.py +58 -1
  28. snowflake/ml/modeling/_internal/snowpark_implementations/distributed_hpo_trainer.py +196 -242
  29. snowflake/ml/modeling/_internal/snowpark_implementations/distributed_search_udf_file.py +161 -0
  30. snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_handlers.py +38 -18
  31. snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_trainer.py +82 -134
  32. snowflake/ml/modeling/_internal/snowpark_implementations/xgboost_external_memory_trainer.py +21 -17
  33. snowflake/ml/modeling/calibration/calibrated_classifier_cv.py +9 -2
  34. snowflake/ml/modeling/cluster/affinity_propagation.py +9 -2
  35. snowflake/ml/modeling/cluster/agglomerative_clustering.py +9 -2
  36. snowflake/ml/modeling/cluster/birch.py +9 -2
  37. snowflake/ml/modeling/cluster/bisecting_k_means.py +9 -2
  38. snowflake/ml/modeling/cluster/dbscan.py +9 -2
  39. snowflake/ml/modeling/cluster/feature_agglomeration.py +9 -2
  40. snowflake/ml/modeling/cluster/k_means.py +9 -2
  41. snowflake/ml/modeling/cluster/mean_shift.py +9 -2
  42. snowflake/ml/modeling/cluster/mini_batch_k_means.py +9 -2
  43. snowflake/ml/modeling/cluster/optics.py +9 -2
  44. snowflake/ml/modeling/cluster/spectral_biclustering.py +9 -2
  45. snowflake/ml/modeling/cluster/spectral_clustering.py +9 -2
  46. snowflake/ml/modeling/cluster/spectral_coclustering.py +9 -2
  47. snowflake/ml/modeling/compose/column_transformer.py +9 -2
  48. snowflake/ml/modeling/compose/transformed_target_regressor.py +9 -2
  49. snowflake/ml/modeling/covariance/elliptic_envelope.py +9 -2
  50. snowflake/ml/modeling/covariance/empirical_covariance.py +9 -2
  51. snowflake/ml/modeling/covariance/graphical_lasso.py +9 -2
  52. snowflake/ml/modeling/covariance/graphical_lasso_cv.py +9 -2
  53. snowflake/ml/modeling/covariance/ledoit_wolf.py +9 -2
  54. snowflake/ml/modeling/covariance/min_cov_det.py +9 -2
  55. snowflake/ml/modeling/covariance/oas.py +9 -2
  56. snowflake/ml/modeling/covariance/shrunk_covariance.py +9 -2
  57. snowflake/ml/modeling/decomposition/dictionary_learning.py +9 -2
  58. snowflake/ml/modeling/decomposition/factor_analysis.py +9 -2
  59. snowflake/ml/modeling/decomposition/fast_ica.py +9 -2
  60. snowflake/ml/modeling/decomposition/incremental_pca.py +9 -2
  61. snowflake/ml/modeling/decomposition/kernel_pca.py +9 -2
  62. snowflake/ml/modeling/decomposition/mini_batch_dictionary_learning.py +9 -2
  63. snowflake/ml/modeling/decomposition/mini_batch_sparse_pca.py +9 -2
  64. snowflake/ml/modeling/decomposition/pca.py +9 -2
  65. snowflake/ml/modeling/decomposition/sparse_pca.py +9 -2
  66. snowflake/ml/modeling/decomposition/truncated_svd.py +9 -2
  67. snowflake/ml/modeling/discriminant_analysis/linear_discriminant_analysis.py +9 -2
  68. snowflake/ml/modeling/discriminant_analysis/quadratic_discriminant_analysis.py +9 -2
  69. snowflake/ml/modeling/ensemble/ada_boost_classifier.py +9 -2
  70. snowflake/ml/modeling/ensemble/ada_boost_regressor.py +9 -2
  71. snowflake/ml/modeling/ensemble/bagging_classifier.py +9 -2
  72. snowflake/ml/modeling/ensemble/bagging_regressor.py +9 -2
  73. snowflake/ml/modeling/ensemble/extra_trees_classifier.py +9 -2
  74. snowflake/ml/modeling/ensemble/extra_trees_regressor.py +9 -2
  75. snowflake/ml/modeling/ensemble/gradient_boosting_classifier.py +9 -2
  76. snowflake/ml/modeling/ensemble/gradient_boosting_regressor.py +9 -2
  77. snowflake/ml/modeling/ensemble/hist_gradient_boosting_classifier.py +9 -2
  78. snowflake/ml/modeling/ensemble/hist_gradient_boosting_regressor.py +9 -2
  79. snowflake/ml/modeling/ensemble/isolation_forest.py +9 -2
  80. snowflake/ml/modeling/ensemble/random_forest_classifier.py +9 -2
  81. snowflake/ml/modeling/ensemble/random_forest_regressor.py +9 -2
  82. snowflake/ml/modeling/ensemble/stacking_regressor.py +9 -2
  83. snowflake/ml/modeling/ensemble/voting_classifier.py +9 -2
  84. snowflake/ml/modeling/ensemble/voting_regressor.py +9 -2
  85. snowflake/ml/modeling/feature_selection/generic_univariate_select.py +9 -2
  86. snowflake/ml/modeling/feature_selection/select_fdr.py +9 -2
  87. snowflake/ml/modeling/feature_selection/select_fpr.py +9 -2
  88. snowflake/ml/modeling/feature_selection/select_fwe.py +9 -2
  89. snowflake/ml/modeling/feature_selection/select_k_best.py +9 -2
  90. snowflake/ml/modeling/feature_selection/select_percentile.py +9 -2
  91. snowflake/ml/modeling/feature_selection/sequential_feature_selector.py +9 -2
  92. snowflake/ml/modeling/feature_selection/variance_threshold.py +9 -2
  93. snowflake/ml/modeling/framework/base.py +3 -8
  94. snowflake/ml/modeling/gaussian_process/gaussian_process_classifier.py +9 -2
  95. snowflake/ml/modeling/gaussian_process/gaussian_process_regressor.py +9 -2
  96. snowflake/ml/modeling/impute/iterative_imputer.py +9 -2
  97. snowflake/ml/modeling/impute/knn_imputer.py +9 -2
  98. snowflake/ml/modeling/impute/missing_indicator.py +9 -2
  99. snowflake/ml/modeling/impute/simple_imputer.py +28 -5
  100. snowflake/ml/modeling/kernel_approximation/additive_chi2_sampler.py +9 -2
  101. snowflake/ml/modeling/kernel_approximation/nystroem.py +9 -2
  102. snowflake/ml/modeling/kernel_approximation/polynomial_count_sketch.py +9 -2
  103. snowflake/ml/modeling/kernel_approximation/rbf_sampler.py +9 -2
  104. snowflake/ml/modeling/kernel_approximation/skewed_chi2_sampler.py +9 -2
  105. snowflake/ml/modeling/kernel_ridge/kernel_ridge.py +9 -2
  106. snowflake/ml/modeling/lightgbm/lgbm_classifier.py +9 -2
  107. snowflake/ml/modeling/lightgbm/lgbm_regressor.py +9 -2
  108. snowflake/ml/modeling/linear_model/ard_regression.py +9 -2
  109. snowflake/ml/modeling/linear_model/bayesian_ridge.py +9 -2
  110. snowflake/ml/modeling/linear_model/elastic_net.py +9 -2
  111. snowflake/ml/modeling/linear_model/elastic_net_cv.py +9 -2
  112. snowflake/ml/modeling/linear_model/gamma_regressor.py +9 -2
  113. snowflake/ml/modeling/linear_model/huber_regressor.py +9 -2
  114. snowflake/ml/modeling/linear_model/lars.py +9 -2
  115. snowflake/ml/modeling/linear_model/lars_cv.py +9 -2
  116. snowflake/ml/modeling/linear_model/lasso.py +9 -2
  117. snowflake/ml/modeling/linear_model/lasso_cv.py +9 -2
  118. snowflake/ml/modeling/linear_model/lasso_lars.py +9 -2
  119. snowflake/ml/modeling/linear_model/lasso_lars_cv.py +9 -2
  120. snowflake/ml/modeling/linear_model/lasso_lars_ic.py +9 -2
  121. snowflake/ml/modeling/linear_model/linear_regression.py +9 -2
  122. snowflake/ml/modeling/linear_model/logistic_regression.py +9 -2
  123. snowflake/ml/modeling/linear_model/logistic_regression_cv.py +9 -2
  124. snowflake/ml/modeling/linear_model/multi_task_elastic_net.py +9 -2
  125. snowflake/ml/modeling/linear_model/multi_task_elastic_net_cv.py +9 -2
  126. snowflake/ml/modeling/linear_model/multi_task_lasso.py +9 -2
  127. snowflake/ml/modeling/linear_model/multi_task_lasso_cv.py +9 -2
  128. snowflake/ml/modeling/linear_model/orthogonal_matching_pursuit.py +9 -2
  129. snowflake/ml/modeling/linear_model/passive_aggressive_classifier.py +9 -2
  130. snowflake/ml/modeling/linear_model/passive_aggressive_regressor.py +9 -2
  131. snowflake/ml/modeling/linear_model/perceptron.py +9 -2
  132. snowflake/ml/modeling/linear_model/poisson_regressor.py +9 -2
  133. snowflake/ml/modeling/linear_model/ransac_regressor.py +9 -2
  134. snowflake/ml/modeling/linear_model/ridge.py +9 -2
  135. snowflake/ml/modeling/linear_model/ridge_classifier.py +9 -2
  136. snowflake/ml/modeling/linear_model/ridge_classifier_cv.py +9 -2
  137. snowflake/ml/modeling/linear_model/ridge_cv.py +9 -2
  138. snowflake/ml/modeling/linear_model/sgd_classifier.py +9 -2
  139. snowflake/ml/modeling/linear_model/sgd_one_class_svm.py +9 -2
  140. snowflake/ml/modeling/linear_model/sgd_regressor.py +9 -2
  141. snowflake/ml/modeling/linear_model/theil_sen_regressor.py +9 -2
  142. snowflake/ml/modeling/linear_model/tweedie_regressor.py +9 -2
  143. snowflake/ml/modeling/manifold/isomap.py +9 -2
  144. snowflake/ml/modeling/manifold/mds.py +9 -2
  145. snowflake/ml/modeling/manifold/spectral_embedding.py +9 -2
  146. snowflake/ml/modeling/manifold/tsne.py +9 -2
  147. snowflake/ml/modeling/mixture/bayesian_gaussian_mixture.py +9 -2
  148. snowflake/ml/modeling/mixture/gaussian_mixture.py +9 -2
  149. snowflake/ml/modeling/model_selection/grid_search_cv.py +1 -5
  150. snowflake/ml/modeling/model_selection/randomized_search_cv.py +1 -5
  151. snowflake/ml/modeling/multiclass/one_vs_one_classifier.py +9 -2
  152. snowflake/ml/modeling/multiclass/one_vs_rest_classifier.py +9 -2
  153. snowflake/ml/modeling/multiclass/output_code_classifier.py +9 -2
  154. snowflake/ml/modeling/naive_bayes/bernoulli_nb.py +9 -2
  155. snowflake/ml/modeling/naive_bayes/categorical_nb.py +9 -2
  156. snowflake/ml/modeling/naive_bayes/complement_nb.py +9 -2
  157. snowflake/ml/modeling/naive_bayes/gaussian_nb.py +9 -2
  158. snowflake/ml/modeling/naive_bayes/multinomial_nb.py +9 -2
  159. snowflake/ml/modeling/neighbors/k_neighbors_classifier.py +9 -2
  160. snowflake/ml/modeling/neighbors/k_neighbors_regressor.py +9 -2
  161. snowflake/ml/modeling/neighbors/kernel_density.py +9 -2
  162. snowflake/ml/modeling/neighbors/local_outlier_factor.py +9 -2
  163. snowflake/ml/modeling/neighbors/nearest_centroid.py +9 -2
  164. snowflake/ml/modeling/neighbors/nearest_neighbors.py +9 -2
  165. snowflake/ml/modeling/neighbors/neighborhood_components_analysis.py +9 -2
  166. snowflake/ml/modeling/neighbors/radius_neighbors_classifier.py +9 -2
  167. snowflake/ml/modeling/neighbors/radius_neighbors_regressor.py +9 -2
  168. snowflake/ml/modeling/neural_network/bernoulli_rbm.py +9 -2
  169. snowflake/ml/modeling/neural_network/mlp_classifier.py +9 -2
  170. snowflake/ml/modeling/neural_network/mlp_regressor.py +9 -2
  171. snowflake/ml/modeling/parameters/enable_anonymous_sproc.py +5 -0
  172. snowflake/ml/modeling/pipeline/pipeline.py +5 -0
  173. snowflake/ml/modeling/preprocessing/binarizer.py +7 -3
  174. snowflake/ml/modeling/preprocessing/k_bins_discretizer.py +7 -2
  175. snowflake/ml/modeling/preprocessing/label_encoder.py +8 -7
  176. snowflake/ml/modeling/preprocessing/max_abs_scaler.py +7 -3
  177. snowflake/ml/modeling/preprocessing/min_max_scaler.py +7 -4
  178. snowflake/ml/modeling/preprocessing/normalizer.py +7 -3
  179. snowflake/ml/modeling/preprocessing/one_hot_encoder.py +10 -2
  180. snowflake/ml/modeling/preprocessing/ordinal_encoder.py +8 -5
  181. snowflake/ml/modeling/preprocessing/polynomial_features.py +9 -2
  182. snowflake/ml/modeling/preprocessing/robust_scaler.py +7 -4
  183. snowflake/ml/modeling/preprocessing/standard_scaler.py +7 -3
  184. snowflake/ml/modeling/semi_supervised/label_propagation.py +9 -2
  185. snowflake/ml/modeling/semi_supervised/label_spreading.py +9 -2
  186. snowflake/ml/modeling/svm/linear_svc.py +9 -2
  187. snowflake/ml/modeling/svm/linear_svr.py +9 -2
  188. snowflake/ml/modeling/svm/nu_svc.py +9 -2
  189. snowflake/ml/modeling/svm/nu_svr.py +9 -2
  190. snowflake/ml/modeling/svm/svc.py +9 -2
  191. snowflake/ml/modeling/svm/svr.py +9 -2
  192. snowflake/ml/modeling/tree/decision_tree_classifier.py +9 -2
  193. snowflake/ml/modeling/tree/decision_tree_regressor.py +9 -2
  194. snowflake/ml/modeling/tree/extra_tree_classifier.py +9 -2
  195. snowflake/ml/modeling/tree/extra_tree_regressor.py +9 -2
  196. snowflake/ml/modeling/xgboost/xgb_classifier.py +9 -2
  197. snowflake/ml/modeling/xgboost/xgb_regressor.py +9 -2
  198. snowflake/ml/modeling/xgboost/xgbrf_classifier.py +9 -2
  199. snowflake/ml/modeling/xgboost/xgbrf_regressor.py +9 -2
  200. snowflake/ml/registry/_manager/model_manager.py +59 -1
  201. snowflake/ml/registry/registry.py +10 -1
  202. snowflake/ml/version.py +1 -1
  203. {snowflake_ml_python-1.5.1.dist-info → snowflake_ml_python-1.5.3.dist-info}/METADATA +32 -4
  204. {snowflake_ml_python-1.5.1.dist-info → snowflake_ml_python-1.5.3.dist-info}/RECORD +207 -204
  205. {snowflake_ml_python-1.5.1.dist-info → snowflake_ml_python-1.5.3.dist-info}/LICENSE.txt +0 -0
  206. {snowflake_ml_python-1.5.1.dist-info → snowflake_ml_python-1.5.3.dist-info}/WHEEL +0 -0
  207. {snowflake_ml_python-1.5.1.dist-info → snowflake_ml_python-1.5.3.dist-info}/top_level.txt +0 -0
@@ -67,11 +67,14 @@ class OrdinalEncoder(base.BaseTransformer):
67
67
  The value to be used to encode unknown categories.
68
68
 
69
69
  input_cols: Optional[Union[str, List[str]]], default=None
70
- The name(s) of one or more columns in a DataFrame containing a feature to be encoded.
70
+ The name(s) of one or more columns in the input DataFrame containing feature(s) to be encoded. Input
71
+ columns must be specified before fit with this argument or after initialization with the
72
+ `set_input_cols` method. This argument is optional for API consistency.
71
73
 
72
74
  output_cols: Optional[Union[str, List[str]]], default=None
73
- The name(s) of one or more columns in a DataFrame in which results will be stored. The number of
74
- columns specified must match the number of input columns.
75
+ The prefix to be used for encoded output for each input column. The number of
76
+ output column prefixes specified must equal the number of input columns. Output column prefixes must be
77
+ specified before transform with this argument or after initialization with the `set_output_cols` method.
75
78
 
76
79
  passthrough_cols: Optional[Union[str, List[str]]], default=None
77
80
  A string or a list of strings indicating column names to be excluded from any
@@ -247,7 +250,7 @@ class OrdinalEncoder(base.BaseTransformer):
247
250
  # columns: COLUMN_NAME, CATEGORY, INDEX
248
251
  state_df = self._get_category_index_state_df(dataset)
249
252
  # save the dataframe on server side so that transform doesn't need to upload
250
- state_df.write.save_as_table( # type: ignore[call-overload]
253
+ state_df.write.save_as_table(
251
254
  self._vocab_table_name,
252
255
  mode="overwrite",
253
256
  table_type="temporary",
@@ -520,7 +523,7 @@ class OrdinalEncoder(base.BaseTransformer):
520
523
  )
521
524
 
522
525
  batch_table_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.TABLE)
523
- transformed_dataset.write.save_as_table( # type: ignore[call-overload]
526
+ transformed_dataset.write.save_as_table(
524
527
  batch_table_name,
525
528
  mode="overwrite",
526
529
  table_type="temporary",
@@ -251,7 +251,7 @@ class PolynomialFeatures(BaseTransformer):
251
251
  inspect.currentframe(), PolynomialFeatures.__class__.__name__
252
252
  ),
253
253
  api_calls=[Session.call],
254
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
254
+ custom_tags={"autogen": True} if self._autogenerated else None,
255
255
  )
256
256
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
257
257
  pd_df.columns = dataset.columns
@@ -586,7 +586,14 @@ class PolynomialFeatures(BaseTransformer):
586
586
  ) -> List[str]:
587
587
  # in case the inferred output column names dimension is different
588
588
  # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
589
- output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
589
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
590
+
591
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
592
+ # seen during the fit.
593
+ snowpark_column_names = dataset.select(self.input_cols).columns
594
+ sample_pd_df.columns = snowpark_column_names
595
+
596
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
590
597
  output_df_columns = list(output_df_pd.columns)
591
598
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
592
599
  if self.sample_weight_col:
@@ -37,12 +37,15 @@ class RobustScaler(base.BaseTransformer):
37
37
  the dataset is scaled down. If less than 1, the dataset is scaled up.
38
38
 
39
39
  input_cols: Optional[Union[str, List[str]]], default=None
40
- The name(s) of one or more columns in a DataFrame containing a feature to be scaled.
40
+ The name(s) of one or more columns in the input DataFrame containing feature(s) to be scaled. Input
41
+ columns must be specified before fit with this argument or after initialization with the
42
+ `set_input_cols` method. This argument is optional for API consistency.
41
43
 
42
44
  output_cols: Optional[Union[str, List[str]]], default=None
43
- The name(s) of one or more columns in a DataFrame in which results will be stored. The number of
44
- columns specified must match the number of input columns. For dense output, the column names specified are
45
- used as base names for the columns created for each category.
45
+ The name(s) to assign output columns in the output DataFrame. The number of
46
+ columns specified must equal the number of input columns. Output columns must be specified before transform
47
+ with this argument or after initialization with the `set_output_cols` method. This argument is optional for
48
+ API consistency.
46
49
 
47
50
  passthrough_cols: Optional[Union[str, List[str]]], default=None
48
51
  A string or a list of strings indicating column names to be excluded from any
@@ -26,11 +26,15 @@ class StandardScaler(base.BaseTransformer):
26
26
  If True, scale the data unit variance (i.e. unit standard deviation).
27
27
 
28
28
  input_cols: Optional[Union[str, List[str]]], default=None
29
- The name(s) of one or more columns in a DataFrame containing a feature to be scaled.
29
+ The name(s) of one or more columns in the input DataFrame containing feature(s) to be scaled. Input
30
+ columns must be specified before fit with this argument or after initialization with the
31
+ `set_input_cols` method. This argument is optional for API consistency.
30
32
 
31
33
  output_cols: Optional[Union[str, List[str]]], default=None
32
- The name(s) of one or more columns in a DataFrame in which results will be stored. The number of
33
- columns specified must match the number of input columns.
34
+ The name(s) to assign output columns in the output DataFrame. The number of
35
+ columns specified must equal the number of input columns. Output columns must be specified before transform
36
+ with this argument or after initialization with the `set_output_cols` method. This argument is optional for
37
+ API consistency.
34
38
 
35
39
  passthrough_cols: Optional[Union[str, List[str]]], default=None
36
40
  A string or a list of strings indicating column names to be excluded from any
@@ -257,7 +257,7 @@ class LabelPropagation(BaseTransformer):
257
257
  inspect.currentframe(), LabelPropagation.__class__.__name__
258
258
  ),
259
259
  api_calls=[Session.call],
260
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
260
+ custom_tags={"autogen": True} if self._autogenerated else None,
261
261
  )
262
262
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
263
263
  pd_df.columns = dataset.columns
@@ -590,7 +590,14 @@ class LabelPropagation(BaseTransformer):
590
590
  ) -> List[str]:
591
591
  # in case the inferred output column names dimension is different
592
592
  # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
593
- output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
593
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
594
+
595
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
596
+ # seen during the fit.
597
+ snowpark_column_names = dataset.select(self.input_cols).columns
598
+ sample_pd_df.columns = snowpark_column_names
599
+
600
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
594
601
  output_df_columns = list(output_df_pd.columns)
595
602
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
596
603
  if self.sample_weight_col:
@@ -266,7 +266,7 @@ class LabelSpreading(BaseTransformer):
266
266
  inspect.currentframe(), LabelSpreading.__class__.__name__
267
267
  ),
268
268
  api_calls=[Session.call],
269
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
269
+ custom_tags={"autogen": True} if self._autogenerated else None,
270
270
  )
271
271
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
272
272
  pd_df.columns = dataset.columns
@@ -599,7 +599,14 @@ class LabelSpreading(BaseTransformer):
599
599
  ) -> List[str]:
600
600
  # in case the inferred output column names dimension is different
601
601
  # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
602
- output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
602
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
603
+
604
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
605
+ # seen during the fit.
606
+ snowpark_column_names = dataset.select(self.input_cols).columns
607
+ sample_pd_df.columns = snowpark_column_names
608
+
609
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
603
610
  output_df_columns = list(output_df_pd.columns)
604
611
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
605
612
  if self.sample_weight_col:
@@ -322,7 +322,7 @@ class LinearSVC(BaseTransformer):
322
322
  inspect.currentframe(), LinearSVC.__class__.__name__
323
323
  ),
324
324
  api_calls=[Session.call],
325
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
325
+ custom_tags={"autogen": True} if self._autogenerated else None,
326
326
  )
327
327
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
328
328
  pd_df.columns = dataset.columns
@@ -655,7 +655,14 @@ class LinearSVC(BaseTransformer):
655
655
  ) -> List[str]:
656
656
  # in case the inferred output column names dimension is different
657
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)
658
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
659
+
660
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
661
+ # seen during the fit.
662
+ snowpark_column_names = dataset.select(self.input_cols).columns
663
+ sample_pd_df.columns = snowpark_column_names
664
+
665
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
659
666
  output_df_columns = list(output_df_pd.columns)
660
667
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
661
668
  if self.sample_weight_col:
@@ -294,7 +294,7 @@ class LinearSVR(BaseTransformer):
294
294
  inspect.currentframe(), LinearSVR.__class__.__name__
295
295
  ),
296
296
  api_calls=[Session.call],
297
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
297
+ custom_tags={"autogen": True} if self._autogenerated else None,
298
298
  )
299
299
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
300
300
  pd_df.columns = dataset.columns
@@ -627,7 +627,14 @@ class LinearSVR(BaseTransformer):
627
627
  ) -> List[str]:
628
628
  # in case the inferred output column names dimension is different
629
629
  # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
630
- output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
630
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
631
+
632
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
633
+ # seen during the fit.
634
+ snowpark_column_names = dataset.select(self.input_cols).columns
635
+ sample_pd_df.columns = snowpark_column_names
636
+
637
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
631
638
  output_df_columns = list(output_df_pd.columns)
632
639
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
633
640
  if self.sample_weight_col:
@@ -328,7 +328,7 @@ class NuSVC(BaseTransformer):
328
328
  inspect.currentframe(), NuSVC.__class__.__name__
329
329
  ),
330
330
  api_calls=[Session.call],
331
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
331
+ custom_tags={"autogen": True} if self._autogenerated else None,
332
332
  )
333
333
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
334
334
  pd_df.columns = dataset.columns
@@ -661,7 +661,14 @@ class NuSVC(BaseTransformer):
661
661
  ) -> List[str]:
662
662
  # in case the inferred output column names dimension is different
663
663
  # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
664
- output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
664
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
665
+
666
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
667
+ # seen during the fit.
668
+ snowpark_column_names = dataset.select(self.input_cols).columns
669
+ sample_pd_df.columns = snowpark_column_names
670
+
671
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
665
672
  output_df_columns = list(output_df_pd.columns)
666
673
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
667
674
  if self.sample_weight_col:
@@ -289,7 +289,7 @@ class NuSVR(BaseTransformer):
289
289
  inspect.currentframe(), NuSVR.__class__.__name__
290
290
  ),
291
291
  api_calls=[Session.call],
292
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
292
+ custom_tags={"autogen": True} if self._autogenerated else None,
293
293
  )
294
294
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
295
295
  pd_df.columns = dataset.columns
@@ -622,7 +622,14 @@ class NuSVR(BaseTransformer):
622
622
  ) -> List[str]:
623
623
  # in case the inferred output column names dimension is different
624
624
  # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
625
- output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
625
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
626
+
627
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
628
+ # seen during the fit.
629
+ snowpark_column_names = dataset.select(self.input_cols).columns
630
+ sample_pd_df.columns = snowpark_column_names
631
+
632
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
626
633
  output_df_columns = list(output_df_pd.columns)
627
634
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
628
635
  if self.sample_weight_col:
@@ -331,7 +331,7 @@ class SVC(BaseTransformer):
331
331
  inspect.currentframe(), SVC.__class__.__name__
332
332
  ),
333
333
  api_calls=[Session.call],
334
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
334
+ custom_tags={"autogen": True} if self._autogenerated else None,
335
335
  )
336
336
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
337
337
  pd_df.columns = dataset.columns
@@ -664,7 +664,14 @@ class SVC(BaseTransformer):
664
664
  ) -> List[str]:
665
665
  # in case the inferred output column names dimension is different
666
666
  # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
667
- output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
667
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
668
+
669
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
670
+ # seen during the fit.
671
+ snowpark_column_names = dataset.select(self.input_cols).columns
672
+ sample_pd_df.columns = snowpark_column_names
673
+
674
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
668
675
  output_df_columns = list(output_df_pd.columns)
669
676
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
670
677
  if self.sample_weight_col:
@@ -292,7 +292,7 @@ class SVR(BaseTransformer):
292
292
  inspect.currentframe(), SVR.__class__.__name__
293
293
  ),
294
294
  api_calls=[Session.call],
295
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
295
+ custom_tags={"autogen": True} if self._autogenerated else None,
296
296
  )
297
297
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
298
298
  pd_df.columns = dataset.columns
@@ -625,7 +625,14 @@ class SVR(BaseTransformer):
625
625
  ) -> List[str]:
626
626
  # in case the inferred output column names dimension is different
627
627
  # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
628
- output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
628
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
629
+
630
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
631
+ # seen during the fit.
632
+ snowpark_column_names = dataset.select(self.input_cols).columns
633
+ sample_pd_df.columns = snowpark_column_names
634
+
635
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
629
636
  output_df_columns = list(output_df_pd.columns)
630
637
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
631
638
  if self.sample_weight_col:
@@ -359,7 +359,7 @@ class DecisionTreeClassifier(BaseTransformer):
359
359
  inspect.currentframe(), DecisionTreeClassifier.__class__.__name__
360
360
  ),
361
361
  api_calls=[Session.call],
362
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
362
+ custom_tags={"autogen": True} if self._autogenerated else None,
363
363
  )
364
364
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
365
365
  pd_df.columns = dataset.columns
@@ -692,7 +692,14 @@ class DecisionTreeClassifier(BaseTransformer):
692
692
  ) -> List[str]:
693
693
  # in case the inferred output column names dimension is different
694
694
  # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
695
- output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
695
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
696
+
697
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
698
+ # seen during the fit.
699
+ snowpark_column_names = dataset.select(self.input_cols).columns
700
+ sample_pd_df.columns = snowpark_column_names
701
+
702
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
696
703
  output_df_columns = list(output_df_pd.columns)
697
704
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
698
705
  if self.sample_weight_col:
@@ -341,7 +341,7 @@ class DecisionTreeRegressor(BaseTransformer):
341
341
  inspect.currentframe(), DecisionTreeRegressor.__class__.__name__
342
342
  ),
343
343
  api_calls=[Session.call],
344
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
344
+ custom_tags={"autogen": True} if self._autogenerated else None,
345
345
  )
346
346
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
347
347
  pd_df.columns = dataset.columns
@@ -674,7 +674,14 @@ class DecisionTreeRegressor(BaseTransformer):
674
674
  ) -> List[str]:
675
675
  # in case the inferred output column names dimension is different
676
676
  # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
677
- output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
677
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
678
+
679
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
680
+ # seen during the fit.
681
+ snowpark_column_names = dataset.select(self.input_cols).columns
682
+ sample_pd_df.columns = snowpark_column_names
683
+
684
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
678
685
  output_df_columns = list(output_df_pd.columns)
679
686
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
680
687
  if self.sample_weight_col:
@@ -351,7 +351,7 @@ class ExtraTreeClassifier(BaseTransformer):
351
351
  inspect.currentframe(), ExtraTreeClassifier.__class__.__name__
352
352
  ),
353
353
  api_calls=[Session.call],
354
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
354
+ custom_tags={"autogen": True} if self._autogenerated else None,
355
355
  )
356
356
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
357
357
  pd_df.columns = dataset.columns
@@ -684,7 +684,14 @@ class ExtraTreeClassifier(BaseTransformer):
684
684
  ) -> List[str]:
685
685
  # in case the inferred output column names dimension is different
686
686
  # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
687
- output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
687
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
688
+
689
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
690
+ # seen during the fit.
691
+ snowpark_column_names = dataset.select(self.input_cols).columns
692
+ sample_pd_df.columns = snowpark_column_names
693
+
694
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
688
695
  output_df_columns = list(output_df_pd.columns)
689
696
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
690
697
  if self.sample_weight_col:
@@ -333,7 +333,7 @@ class ExtraTreeRegressor(BaseTransformer):
333
333
  inspect.currentframe(), ExtraTreeRegressor.__class__.__name__
334
334
  ),
335
335
  api_calls=[Session.call],
336
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
336
+ custom_tags={"autogen": True} if self._autogenerated else None,
337
337
  )
338
338
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
339
339
  pd_df.columns = dataset.columns
@@ -666,7 +666,14 @@ class ExtraTreeRegressor(BaseTransformer):
666
666
  ) -> List[str]:
667
667
  # in case the inferred output column names dimension is different
668
668
  # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
669
- output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
669
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
670
+
671
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
672
+ # seen during the fit.
673
+ snowpark_column_names = dataset.select(self.input_cols).columns
674
+ sample_pd_df.columns = snowpark_column_names
675
+
676
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
670
677
  output_df_columns = list(output_df_pd.columns)
671
678
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
672
679
  if self.sample_weight_col:
@@ -451,7 +451,7 @@ class XGBClassifier(BaseTransformer):
451
451
  inspect.currentframe(), XGBClassifier.__class__.__name__
452
452
  ),
453
453
  api_calls=[Session.call],
454
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
454
+ custom_tags={"autogen": True} if self._autogenerated else None,
455
455
  )
456
456
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
457
457
  pd_df.columns = dataset.columns
@@ -784,7 +784,14 @@ class XGBClassifier(BaseTransformer):
784
784
  ) -> List[str]:
785
785
  # in case the inferred output column names dimension is different
786
786
  # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
787
- output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
787
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
788
+
789
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
790
+ # seen during the fit.
791
+ snowpark_column_names = dataset.select(self.input_cols).columns
792
+ sample_pd_df.columns = snowpark_column_names
793
+
794
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
788
795
  output_df_columns = list(output_df_pd.columns)
789
796
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
790
797
  if self.sample_weight_col:
@@ -450,7 +450,7 @@ class XGBRegressor(BaseTransformer):
450
450
  inspect.currentframe(), XGBRegressor.__class__.__name__
451
451
  ),
452
452
  api_calls=[Session.call],
453
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
453
+ custom_tags={"autogen": True} if self._autogenerated else None,
454
454
  )
455
455
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
456
456
  pd_df.columns = dataset.columns
@@ -783,7 +783,14 @@ class XGBRegressor(BaseTransformer):
783
783
  ) -> List[str]:
784
784
  # in case the inferred output column names dimension is different
785
785
  # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
786
- output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
786
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
787
+
788
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
789
+ # seen during the fit.
790
+ snowpark_column_names = dataset.select(self.input_cols).columns
791
+ sample_pd_df.columns = snowpark_column_names
792
+
793
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
787
794
  output_df_columns = list(output_df_pd.columns)
788
795
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
789
796
  if self.sample_weight_col:
@@ -455,7 +455,7 @@ class XGBRFClassifier(BaseTransformer):
455
455
  inspect.currentframe(), XGBRFClassifier.__class__.__name__
456
456
  ),
457
457
  api_calls=[Session.call],
458
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
458
+ custom_tags={"autogen": True} if self._autogenerated else None,
459
459
  )
460
460
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
461
461
  pd_df.columns = dataset.columns
@@ -788,7 +788,14 @@ class XGBRFClassifier(BaseTransformer):
788
788
  ) -> List[str]:
789
789
  # in case the inferred output column names dimension is different
790
790
  # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
791
- output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
791
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
792
+
793
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
794
+ # seen during the fit.
795
+ snowpark_column_names = dataset.select(self.input_cols).columns
796
+ sample_pd_df.columns = snowpark_column_names
797
+
798
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
792
799
  output_df_columns = list(output_df_pd.columns)
793
800
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
794
801
  if self.sample_weight_col:
@@ -455,7 +455,7 @@ class XGBRFRegressor(BaseTransformer):
455
455
  inspect.currentframe(), XGBRFRegressor.__class__.__name__
456
456
  ),
457
457
  api_calls=[Session.call],
458
- custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
458
+ custom_tags={"autogen": True} if self._autogenerated else None,
459
459
  )
460
460
  pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
461
461
  pd_df.columns = dataset.columns
@@ -788,7 +788,14 @@ class XGBRFRegressor(BaseTransformer):
788
788
  ) -> List[str]:
789
789
  # in case the inferred output column names dimension is different
790
790
  # we use one line of snowpark dataframe and put it into sklearn estimator using pandas
791
- output_df_pd = getattr(self, method)(dataset.limit(1).to_pandas(), output_cols_prefix)
791
+ sample_pd_df = dataset.select(self.input_cols).limit(1).to_pandas()
792
+
793
+ # Rename the pandas df column names to snowflake identifiers and reorder columns to match the order
794
+ # seen during the fit.
795
+ snowpark_column_names = dataset.select(self.input_cols).columns
796
+ sample_pd_df.columns = snowpark_column_names
797
+
798
+ output_df_pd = getattr(self, method)(sample_pd_df, output_cols_prefix)
792
799
  output_df_columns = list(output_df_pd.columns)
793
800
  output_df_columns_set: Set[str] = set(output_df_columns) - set(dataset.columns)
794
801
  if self.sample_weight_col:
@@ -1,5 +1,5 @@
1
1
  from types import ModuleType
2
- from typing import Any, Dict, List, Optional
2
+ from typing import Any, Dict, List, Optional, Union
3
3
 
4
4
  import pandas as pd
5
5
  from absl.logging import logging
@@ -31,6 +31,64 @@ class ModelManager:
31
31
  self._hrid_generator = hrid_generator.HRID16()
32
32
 
33
33
  def log_model(
34
+ self,
35
+ *,
36
+ model: Union[model_types.SupportedModelType, model_version_impl.ModelVersion],
37
+ model_name: str,
38
+ version_name: Optional[str] = None,
39
+ comment: Optional[str] = None,
40
+ metrics: Optional[Dict[str, Any]] = None,
41
+ conda_dependencies: Optional[List[str]] = None,
42
+ pip_requirements: Optional[List[str]] = None,
43
+ python_version: Optional[str] = None,
44
+ signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
45
+ sample_input_data: Optional[model_types.SupportedDataType] = None,
46
+ code_paths: Optional[List[str]] = None,
47
+ ext_modules: Optional[List[ModuleType]] = None,
48
+ options: Optional[model_types.ModelSaveOption] = None,
49
+ statement_params: Optional[Dict[str, Any]] = None,
50
+ ) -> model_version_impl.ModelVersion:
51
+ if not version_name:
52
+ version_name = self._hrid_generator.generate()[1]
53
+
54
+ if isinstance(model, model_version_impl.ModelVersion):
55
+ (
56
+ source_database_name_id,
57
+ source_schema_name_id,
58
+ source_model_name_id,
59
+ ) = sql_identifier.parse_fully_qualified_name(model.fully_qualified_model_name)
60
+
61
+ self._model_ops.create_from_model_version(
62
+ source_database_name=source_database_name_id,
63
+ source_schema_name=source_schema_name_id,
64
+ source_model_name=source_model_name_id,
65
+ source_version_name=sql_identifier.SqlIdentifier(model.version_name),
66
+ database_name=None,
67
+ schema_name=None,
68
+ model_name=sql_identifier.SqlIdentifier(model_name),
69
+ version_name=sql_identifier.SqlIdentifier(version_name),
70
+ statement_params=statement_params,
71
+ )
72
+ return self.get_model(model_name=model_name, statement_params=statement_params).version(version_name)
73
+
74
+ return self._log_model(
75
+ model=model,
76
+ model_name=model_name,
77
+ version_name=version_name,
78
+ comment=comment,
79
+ metrics=metrics,
80
+ conda_dependencies=conda_dependencies,
81
+ pip_requirements=pip_requirements,
82
+ python_version=python_version,
83
+ signatures=signatures,
84
+ sample_input_data=sample_input_data,
85
+ code_paths=code_paths,
86
+ ext_modules=ext_modules,
87
+ options=options,
88
+ statement_params=statement_params,
89
+ )
90
+
91
+ def _log_model(
34
92
  self,
35
93
  model: model_types.SupportedModelType,
36
94
  *,
@@ -71,6 +71,16 @@ class Registry:
71
71
  @telemetry.send_api_usage_telemetry(
72
72
  project=_TELEMETRY_PROJECT,
73
73
  subproject=_MODEL_TELEMETRY_SUBPROJECT,
74
+ func_params_to_log=[
75
+ "model_name",
76
+ "version_name",
77
+ "comment",
78
+ "metrics",
79
+ "conda_dependencies",
80
+ "pip_requirements",
81
+ "python_version",
82
+ "signatures",
83
+ ],
74
84
  )
75
85
  def log_model(
76
86
  self,
@@ -142,7 +152,6 @@ class Registry:
142
152
  Returns:
143
153
  ModelVersion: ModelVersion object corresponding to the model just logged.
144
154
  """
145
-
146
155
  statement_params = telemetry.get_statement_params(
147
156
  project=_TELEMETRY_PROJECT,
148
157
  subproject=_MODEL_TELEMETRY_SUBPROJECT,
snowflake/ml/version.py CHANGED
@@ -1 +1 @@
1
- VERSION="1.5.1"
1
+ VERSION="1.5.3"