snowflake-ml-python 1.9.2__py3-none-any.whl → 1.11.0__py3-none-any.whl

{snowflake_ml_python-1.9.2.dist-info → snowflake_ml_python-1.11.0.dist-info}/RECORD

@@ -10,7 +10,7 @@ snowflake/cortex/_sse_client.py,sha256=sLYgqAfTOPADCnaWH2RWAJi8KbU_7gSRsTUDcDD5T
 snowflake/cortex/_summarize.py,sha256=7GH8zqfIdOiHA5w4b6EvJEKEWhaTrL4YA6iDGbn7BNM,1307
 snowflake/cortex/_translate.py,sha256=9ZGjvAnJFisbzJ_bXnt4pyug5UzhHJRXW8AhGQEersM,1652
 snowflake/cortex/_util.py,sha256=krNTpbkFLXwdFqy1bd0xi7ZmOzOHRnIfHdQCPiLZJxk,3288
-snowflake/ml/version.py,sha256=h4K5WMP3n8jxN6b7-kw9Qz47L1TW7SrTH9h4orSolSY,98
+snowflake/ml/version.py,sha256=0qOy-uzl0gJVrPEVGTPIvvArEhwh9YH6ggaoY0gl0h4,99
 snowflake/ml/_internal/env.py,sha256=EY_2KVe8oR3LgKWdaeRb5rRU-NDNXJppPDsFJmMZUUY,265
 snowflake/ml/_internal/env_utils.py,sha256=x6ID94g6FYoMX3afp0zoUHzBvuvPyiE2F6RDpxx5Cq0,30967
 snowflake/ml/_internal/file_utils.py,sha256=7sA6loOeSfmGP4yx16P4usT9ZtRqG3ycnXu7_Tk7dOs,14206
@@ -45,7 +45,7 @@ snowflake/ml/_internal/utils/parallelize.py,sha256=l8Zjo-hp8zqoLgHxBlpz9Zmn2Z-MR
 snowflake/ml/_internal/utils/pkg_version_utils.py,sha256=EaY_3IsVOZ9BCH28F5VLjp-0AiEqDlL7L715vkPsgrY,5149
 snowflake/ml/_internal/utils/query_result_checker.py,sha256=1PR41Xn9BUIXvp-UmJ9FgEbj8WfgU7RUhz3PqvvVQ5E,10656
 snowflake/ml/_internal/utils/result.py,sha256=59Sz6MvhjakUNiONwg9oi2544AmORCJR3XyWTxY2vP0,2405
-snowflake/ml/_internal/utils/service_logger.py,sha256=q-Deo0nxAzqEqiqt-SUDeu02wCfraJlQWQpTsdcqL4o,5823
+snowflake/ml/_internal/utils/service_logger.py,sha256=LmADyxsSE3-TYBX1gCYtxvaEDdH_Lf6d5gRt44uue0I,6267
 snowflake/ml/_internal/utils/snowflake_env.py,sha256=k4ddzs8iJpRpVvgbbOtU8j4fUvqa77Awk65EJ5j2uxk,4253
 snowflake/ml/_internal/utils/snowpark_dataframe_utils.py,sha256=tm2leAu_oDTNUQZJ98UpKtS79k-A-c72pKxd-8AE-tg,6353
 snowflake/ml/_internal/utils/sql_identifier.py,sha256=YHIwXpb8E1U6LVUVpT8q7s9ZygONAXKPVMD4IucwXx8,4669
@@ -65,13 +65,16 @@ snowflake/ml/dataset/dataset_metadata.py,sha256=lcNvugBkP8YEkGMQqaV8SlHs5mwUKsUS
 snowflake/ml/dataset/dataset_reader.py,sha256=mZsG9HyWUGgfotrGkLrunyEsOm_659mH-Sn2OyG6A-Q,5036
 snowflake/ml/experiment/__init__.py,sha256=r7qdyPd3jwxzqvksim2ju5j_LrnYQrta0ZI6XpWUqmc,109
 snowflake/ml/experiment/_experiment_info.py,sha256=iaJ65x6nzBYJ5djleSOzBtMpZUJCUDlRpaDw0pu-dcU,2533
-snowflake/ml/experiment/callback.py,sha256=I1U-1kXLFDTdnVy-xr3H7-xuoPrPoZmAv2VSYvvGRXU,5047
 snowflake/ml/experiment/experiment_tracking.py,sha256=ljY2HA1E724qnBmvuEzQA8o3ZT0XMJkeIYJ97vXPx5A,11316
+snowflake/ml/experiment/utils.py,sha256=3bpbkilc5vvFjnti-kcyhhjAd9Ga3LqiKqJDwORiATY,628
 snowflake/ml/experiment/_client/experiment_tracking_sql_client.py,sha256=rdCBHRqTYW6I2ztCpO-Zyb9nd_0HV26QdpGMDwxZ144,4446
 snowflake/ml/experiment/_entities/__init__.py,sha256=ThrslBFuDxOUvdS8j_bVmEaEAms8nR1aY0ocYFnVPFg,155
 snowflake/ml/experiment/_entities/experiment.py,sha256=lKmQj59K8fGDWVwRqeIesxorrChb-m78vX_WUmI7PV0,225
 snowflake/ml/experiment/_entities/run.py,sha256=_bWt1YpP8iulg5jeBXMXw8zGZHr9zSE9IVIBHcCdfto,2293
 snowflake/ml/experiment/_entities/run_metadata.py,sha256=j8V2N6QBAx4TP4h7MLIPXqquYI8KyNZnkW6wzm-peuY,1589
+snowflake/ml/experiment/callback/keras.py,sha256=7oq23irYkBV7bLFBCxxKlf9pL4YuDFJDCZ8xtffVRFI,2547
+snowflake/ml/experiment/callback/lightgbm.py,sha256=5co7eR_t651cq1WTK9JCQjhSlYc2oIvxaf3aVnVOlR4,2613
+snowflake/ml/experiment/callback/xgboost.py,sha256=eZMRFAebMERwdqMFm7i6S9wkHD7_VLcwIP0OkWHDOMM,2621
 snowflake/ml/feature_store/__init__.py,sha256=MJr2Gp_EimDgDxD6DtenOEdLTzg6NYPfdNiPM-5rEtw,406
 snowflake/ml/feature_store/access_manager.py,sha256=Q5ImMXRY8WA5X5dpBMzHnIJmeyKVShjNAlbn3cQb4N8,10654
 snowflake/ml/feature_store/entity.py,sha256=ViOSlqCV17ouiO4iH-_KvkvJZqSzpf-nfsjijG6G1Uk,4047
@@ -106,45 +109,48 @@ snowflake/ml/fileset/snowfs.py,sha256=uF5QluYtiJ-HezGIhF55dONi3t0E6N7ByaVAIAlM3n
 snowflake/ml/fileset/stage_fs.py,sha256=V4pysouSKKDPLzuW3u_extxfvjkQa5OlwIRES9Srpzo,20151
 snowflake/ml/jobs/__init__.py,sha256=v-v9-SA1Vy-M98B31-NlqJgpI6uEg9jEEghJLub1RUY,468
 snowflake/ml/jobs/decorators.py,sha256=mQgdWvvCwD7q79cSFKZHKegXGh2j1u8WM64UD3lCKr4,3428
-snowflake/ml/jobs/job.py,sha256=VIRT9K5lhW3tqq1rB-H9GYq_qkraM5DEXZteHnWhdPc,21101
-snowflake/ml/jobs/manager.py,sha256=M_qhnAdMDYPWL2hQscDQqzeavzEricQ5WjztcGn5XGo,22621
-snowflake/ml/jobs/_utils/constants.py,sha256=sdidOyW2X81u0E30DU4K5aPjBTzNQmYtVyZ9D8mQaL4,4066
+snowflake/ml/jobs/job.py,sha256=jj5MtZYEPG2chAKIDy50T0noz4hbelvsLTSVQyUvO1s,21101
+snowflake/ml/jobs/manager.py,sha256=iBo0Ol21um4FM53BNaFxAc14SXCqaKhGGQ1zlUfLNHc,21945
+snowflake/ml/jobs/_utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+snowflake/ml/jobs/_utils/constants.py,sha256=PgFTpxHfqLNalfzPJqLQ7nOao8XLw1JeRqdz4nkNAjs,4168
 snowflake/ml/jobs/_utils/function_payload_utils.py,sha256=4LBaStMdhRxcqwRkwFje-WwiEKRWnBfkaOYouF3N3Kg,1308
 snowflake/ml/jobs/_utils/interop_utils.py,sha256=7mODMTjKCLXkJloACG6_9b2wvmRgjXF0Jx3wpWYyJeA,21413
-snowflake/ml/jobs/_utils/payload_utils.py,sha256=MFt6DEs5TC3uydy_wySnIFtInfReCgioYNQLWoBoVHY,28681
-snowflake/ml/jobs/_utils/query_helper.py,sha256=h5s-_MgHc_f9AmXD5C06frHdP84n9Rmevb1Yu6R1w7s,910
-snowflake/ml/jobs/_utils/spec_utils.py,sha256=VhdLXtDJXi5RJpC4exUUmx2gIb39qu-SQ5VraQi4KLc,13429
-snowflake/ml/jobs/_utils/stage_utils.py,sha256=frjXVvnzFIJCoCWeLF_5x6LsKMq20vp4q1fZvwbXONc,4734
-snowflake/ml/jobs/_utils/types.py,sha256=jXePdeg_KWVSDzs-afRTNx0m4U4MRdRF0rZxDobuNq8,2346
+snowflake/ml/jobs/_utils/payload_utils.py,sha256=WNSPw7Ph4urAw6R6qorEMU5O-e4MPhVM3WS9OM04xr4,30131
+snowflake/ml/jobs/_utils/query_helper.py,sha256=1-XK-y4iukbR1693qAELprRbHmJDM4YoEBHov8IYbHU,1115
+snowflake/ml/jobs/_utils/runtime_env_utils.py,sha256=fqa3ctf_CAOSv1zT__01Qp9T058mKgMjXuEkBZqKUqA,2247
+snowflake/ml/jobs/_utils/spec_utils.py,sha256=EhCLZURZjooHYGGL3HCd0o5N4P0Ug8bFtaL2pT9XHQY,14875
+snowflake/ml/jobs/_utils/stage_utils.py,sha256=38-LsokaGx0NzlnP8CMRioClRz-3x6xhPiZIgl2CB9g,5224
+snowflake/ml/jobs/_utils/types.py,sha256=MOyUTgzkcWqCKYgIZZwm9kcXsxqD_QpLYVTsO5L9CWs,2260
 snowflake/ml/jobs/_utils/scripts/constants.py,sha256=YyIWZqQPYOTtgCY6SfyJjk2A98I5RQVmrOuLtET5Pqg,173
-snowflake/ml/jobs/_utils/scripts/get_instance_ip.py,sha256=DmWs5cVpNmUcrqnwhrUvxE5PycDWFN88Pdut8vFDHPg,5293
-snowflake/ml/jobs/_utils/scripts/mljob_launcher.py,sha256=ZmnVYp9TGV1RKKoJDss_ExP0lkG4ueMkJ45N6YtPVg8,14955
+snowflake/ml/jobs/_utils/scripts/get_instance_ip.py,sha256=QOP2LkfaK59BzWU1-_XbVYUvuztodD-kF8vv0m13eVE,5380
+snowflake/ml/jobs/_utils/scripts/mljob_launcher.py,sha256=VUBpbwMWaDtZj5ZR5NsuTxjCW31IE6QbafoYnDU8s9I,14987
 snowflake/ml/jobs/_utils/scripts/signal_workers.py,sha256=AR1Pylkm4-FGh10WXfrCtcxaV0rI7IQ2ZiO0Li7zZ3U,7433
 snowflake/ml/jobs/_utils/scripts/worker_shutdown_listener.py,sha256=SeJ8v5XDriwHAjIGpcQkwVP-f-lO9QIdVjVD7Fkgafs,7893
 snowflake/ml/lineage/__init__.py,sha256=8p1YGynC-qOxAZ8jZX2z84Reg5bv1NoJMoJmNJCrzI4,65
 snowflake/ml/lineage/lineage_node.py,sha256=vmikk4qaZuVFhQqW-VM6DuW4tDvmQlNbACvIVZEamcU,5830
 snowflake/ml/model/__init__.py,sha256=EvPtblqPN6_T6dyVfaYUxCfo_M7D2CQ1OR5giIH4TsQ,314
 snowflake/ml/model/custom_model.py,sha256=fDhMObqlyzD_qQG1Bq6HHkBN1w3Qzg9e81JWPiqRfc4,12249
-snowflake/ml/model/event_handler.py,sha256=szMzCjo2WJcTe9bD6Prj4NrhMn6Falc9ZitArN4iGC0,4532
+snowflake/ml/model/event_handler.py,sha256=pojleQVM9TPNeDvliTvon2Sfxqbf2WWxrOebo1SaEHo,7211
+snowflake/ml/model/inference_engine.py,sha256=L0nwySY2Qwp3JzuRpPS87r0--m3HTUNUgZXYyOPJjyk,66
 snowflake/ml/model/model_signature.py,sha256=RH62vv4YmrQugTXLsh6kyuzfTs9_yz8a0TMkBR67XKY,32324
 snowflake/ml/model/target_platform.py,sha256=H5d-wtuKQyVlq9x33vPtYZAlR5ka0ytcKRYgwlKl0bQ,390
 snowflake/ml/model/task.py,sha256=Zp5JaLB-YfX5p_HSaw81P3J7UnycQq5EMa87A35VOaQ,286
-snowflake/ml/model/type_hints.py,sha256=F0EW6lbSpZyv5prXc7HBZkPga6LeeCdBpV59CfLdUI4,9309
+snowflake/ml/model/type_hints.py,sha256=G0kp85-ksnYoAUHRdXxLFQBLq3XURuqYOpu_YeKEaNA,9847
 snowflake/ml/model/_client/model/model_impl.py,sha256=Yabrbir5vPMOnsVmQJ23YN7vqhi756Jcm6pfO8Aq92o,17469
-snowflake/ml/model/_client/model/model_version_impl.py,sha256=1zZ7wLHRBQf33d6RS9ygkkLOR7yy7hqtAl2fU9nWmTQ,47482
+snowflake/ml/model/_client/model/model_version_impl.py,sha256=waHYE2GN3SHWonn9eq6vmnXUUNxuo4z3aEUY-MVWEIU,55387
 snowflake/ml/model/_client/ops/metadata_ops.py,sha256=qpK6PL3OyfuhyOmpvLCpHLy6vCxbZbp1HlEvakFGwv4,4884
 snowflake/ml/model/_client/ops/model_ops.py,sha256=z3T71w9ZNIU5eEA5G59Ous59WzEBs3YBcPO1_zeMI8M,48586
-snowflake/ml/model/_client/ops/service_ops.py,sha256=-5hhjEFHxFuiUAHR4izRWMq5fPwoXDSoSCQUfdQZijY,37518
-snowflake/ml/model/_client/service/model_deployment_spec.py,sha256=07b2vdtUEq-a6BZOqX5sqGhK8SI-L1597IgBgoX0XW0,17505
-snowflake/ml/model/_client/service/model_deployment_spec_schema.py,sha256=esuS0MsBEzEyh8ifJH4JOUkSsdX_KL1-KT8KswEHltg,1945
+snowflake/ml/model/_client/ops/service_ops.py,sha256=6aZfxY1iLe-0FpwSwgo8rgye7AoJOU3bAN17QoNWnjA,45942
+snowflake/ml/model/_client/service/model_deployment_spec.py,sha256=_lveWfsSvumzO3fJPL7gO7mwTYNeA9uORzhT5qqI3hs,20479
+snowflake/ml/model/_client/service/model_deployment_spec_schema.py,sha256=z063hgCDwkHa36SUB_PL9cjMYdxfj8Vsm3KV74D2dNM,2219
 snowflake/ml/model/_client/sql/_base.py,sha256=Qrm8M92g3MHb-QnSLUlbd8iVKCRxLhG_zr5M2qmXwJ8,1473
 snowflake/ml/model/_client/sql/model.py,sha256=nstZ8zR7MkXVEfhqLt7PWMik6dZr06nzq7VsF5NVNow,5840
 snowflake/ml/model/_client/sql/model_version.py,sha256=QwzFlDH5laTqK2qF7SJQSbt28DgspWj3R11l-yD1Da0,23496
 snowflake/ml/model/_client/sql/service.py,sha256=j-JdXeWy4oATTH_Yz6OEqhW81t1vM70b5JpAtymar5g,10360
 snowflake/ml/model/_client/sql/stage.py,sha256=2gxYNtmEXricwxeACVUr63OUDCy_iQvCi-kRT4qQtBA,887
 snowflake/ml/model/_client/sql/tag.py,sha256=9sI0VoldKmsfToWSjMQddozPPGCxYUI6n0gPBiqd6x8,4333
-snowflake/ml/model/_model_composer/model_composer.py,sha256=VIEAY9nW8t9dGKdauoa0QwHf75wEndD4cEzWSz9sENE,11842
-snowflake/ml/model/_model_composer/model_manifest/model_manifest.py,sha256=fan34qAiV13PuT_mPW3yke5KRQiLJId1emQRN2EEnhU,11220
+snowflake/ml/model/_model_composer/model_composer.py,sha256=Xqi-sxmkBoZl383LQAXhMQkq9KsAS0A3ythC5bN3EOU,8292
+snowflake/ml/model/_model_composer/model_manifest/model_manifest.py,sha256=y6lu1_4UC7pfosBKofc0dl-LWF8mpcTqhY5sKSkUH_I,9247
 snowflake/ml/model/_model_composer/model_manifest/model_manifest_schema.py,sha256=Oc_PvWG3u0E2pb35w4uMYQdDFEuHdUdOb2gnqnVLE3Q,2917
 snowflake/ml/model/_model_composer/model_method/constants.py,sha256=hoJwIopSdZiYn0fGq15_NiirC0l02d5LEs2D-4J_tPk,35
 snowflake/ml/model/_model_composer/model_method/function_generator.py,sha256=nnUJki3bJVCTF3gZ-usZW3xQ6wwlJ08EfNsPAgsnI3s,2625
@@ -195,7 +201,7 @@ snowflake/ml/model/_signatures/pytorch_handler.py,sha256=Xy-ITCCX_EgHcyIIqeYSDUI
 snowflake/ml/model/_signatures/snowpark_handler.py,sha256=YOBC_Wx-H8bQ967A47nYgqcqLjEA15FbZK69TyAEgvU,7590
 snowflake/ml/model/_signatures/tensorflow_handler.py,sha256=_yrvMg-w_jJoYuyrGXKPX4Dv7Vt8z1e6xIKiWGuZcc4,5660
 snowflake/ml/model/_signatures/utils.py,sha256=vIs12OF_UKH7qrY0JATU-yZhLTgaKt1MJoEemRULA20,17275
-snowflake/ml/model/models/huggingface_pipeline.py,sha256=as-6VMYPRiZy-HflAkmUrqu9f89U3L10BvoVz-NGNBg,17914
+snowflake/ml/model/models/huggingface_pipeline.py,sha256=VCLhlW_CBJAgU-uKSY5a6BrFjc7ANFWSzNaarR6IBq0,19658
 snowflake/ml/modeling/_internal/estimator_utils.py,sha256=oGi5qbZeV-1cM1Pl-rZLBvcr3YRoUzN_te_l-18apLI,11993
 snowflake/ml/modeling/_internal/model_specifications.py,sha256=3wFMcKPCSoiEzU7Mx6RVem89BRlBBENpX__-Rd7GwdU,4851
 snowflake/ml/modeling/_internal/model_trainer.py,sha256=5Ck1lbdyzcd-TpzAxEyovIN9fjaaVIqugyMHXt0wzH0,971
@@ -422,15 +428,16 @@ snowflake/ml/monitoring/_client/queries/rmse.ssql,sha256=OEJiSStRz9-qKoZaFvmubtY
 snowflake/ml/monitoring/_manager/model_monitor_manager.py,sha256=0jpT1-aRU2tsxSM87I-C2kfJeLevCgM-a-OwU_-VUdI,10302
 snowflake/ml/monitoring/entities/model_monitor_config.py,sha256=1W6TFTPicC6YAbjD7A0w8WMhWireyUxyuEy0RQXmqyY,1787
 snowflake/ml/registry/__init__.py,sha256=XdPQK9ejYkSJVrSQ7HD3jKQO0hKq2mC4bPCB6qrtH3U,76
-snowflake/ml/registry/registry.py,sha256=ocn2OUGs9OJzKP4YseV1XrBpuWzTzkv_-o5Eb17NkVo,33900
-snowflake/ml/registry/_manager/model_manager.py,sha256=IqeXmv-aXkMAZi8bcxyrxiOqwx38_l78RUTu17-2BrI,19480
+snowflake/ml/registry/registry.py,sha256=Ro7flVHv3FnEU9Ly3zWRnDAqWiwRSOA2uw_MSKmCBTI,32936
+snowflake/ml/registry/_manager/model_manager.py,sha256=QsEpIbg3FPEbDOQXb_oo41hBjojrdVibdrNPCyJ0Cb0,17650
+snowflake/ml/registry/_manager/model_parameter_reconciler.py,sha256=McTJSps_K2ozZhZsgPyW2k88xDOwgU2cQ36edABhYwI,14001
 snowflake/ml/utils/authentication.py,sha256=E1at4TIAQRDZDsMXSbrKvSJaT6_kSYJBkkr37vU9P2s,2606
 snowflake/ml/utils/connection_params.py,sha256=JuadbzKlgDZLZ5vJ9cnyAiSitvZT9jGSfSSNjIY9P1Q,8282
 snowflake/ml/utils/html_utils.py,sha256=L4pzpvFd20SIk4rie2kTAtcQjbxBHfjKmxonMAT2OoA,7665
 snowflake/ml/utils/sparse.py,sha256=zLBNh-ynhGpKH5TFtopk0YLkHGvv0yq1q-sV59YQKgg,3819
 snowflake/ml/utils/sql_client.py,sha256=pSe2od6Pkh-8NwG3D-xqN76_uNf-ohOtVbT55HeQg1Y,668
-snowflake_ml_python-1.9.2.dist-info/licenses/LICENSE.txt,sha256=PdEp56Av5m3_kl21iFkVTX_EbHJKFGEdmYeIO1pL_Yk,11365
-snowflake_ml_python-1.9.2.dist-info/METADATA,sha256=DwR7gSIY2qPY1l8n6FWwtEH_JCliVI-nnlGGiYVyCew,90398
-snowflake_ml_python-1.9.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-snowflake_ml_python-1.9.2.dist-info/top_level.txt,sha256=TY0gFSHKDdZy3THb0FGomyikWQasEGldIR1O0HGOHVw,10
-snowflake_ml_python-1.9.2.dist-info/RECORD,,
+snowflake_ml_python-1.11.0.dist-info/licenses/LICENSE.txt,sha256=PdEp56Av5m3_kl21iFkVTX_EbHJKFGEdmYeIO1pL_Yk,11365
+snowflake_ml_python-1.11.0.dist-info/METADATA,sha256=4LIdsNhVtufUS0AOud8IBw-7Y0duaQxLbAJWvIn_eZc,91614
+snowflake_ml_python-1.11.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+snowflake_ml_python-1.11.0.dist-info/top_level.txt,sha256=TY0gFSHKDdZy3THb0FGomyikWQasEGldIR1O0HGOHVw,10
+snowflake_ml_python-1.11.0.dist-info/RECORD,,

snowflake/ml/experiment/callback.py

@@ -1,121 +0,0 @@
-import json
-from typing import TYPE_CHECKING, Any, Optional, Union
-from warnings import warn
-
-import lightgbm as lgb
-import xgboost as xgb
-
-from snowflake.ml.model.model_signature import ModelSignature
-
-if TYPE_CHECKING:
-    from snowflake.ml.experiment.experiment_tracking import ExperimentTracking
-
-
-class SnowflakeXgboostCallback(xgb.callback.TrainingCallback):
-    def __init__(
-        self,
-        experiment_tracking: "ExperimentTracking",
-        log_model: bool = True,
-        log_metrics: bool = True,
-        log_params: bool = True,
-        model_name: Optional[str] = None,
-        model_signature: Optional[ModelSignature] = None,
-    ) -> None:
-        self._experiment_tracking = experiment_tracking
-        self.log_model = log_model
-        self.log_metrics = log_metrics
-        self.log_params = log_params
-        self.model_name = model_name
-        self.model_signature = model_signature
-
-    def before_training(self, model: xgb.Booster) -> xgb.Booster:
-        def _flatten_nested_params(params: Union[list[Any], dict[str, Any]], prefix: str = "") -> dict[str, Any]:
-            flat_params = {}
-            items = params.items() if isinstance(params, dict) else enumerate(params)
-            for key, value in items:
-                new_prefix = f"{prefix}.{key}" if prefix else str(key)
-                if isinstance(value, (dict, list)):
-                    flat_params.update(_flatten_nested_params(value, new_prefix))
-                else:
-                    flat_params[new_prefix] = value
-            return flat_params
-
-        if self.log_params:
-            params = json.loads(model.save_config())
-            self._experiment_tracking.log_params(_flatten_nested_params(params))
-
-        return model
-
-    def after_iteration(self, model: Any, epoch: int, evals_log: dict[str, dict[str, Any]]) -> bool:
-        if self.log_metrics:
-            for dataset_name, metrics in evals_log.items():
-                for metric_name, log in metrics.items():
-                    metric_key = dataset_name + ":" + metric_name
-                    self._experiment_tracking.log_metric(key=metric_key, value=log[-1], step=epoch)
-
-        return False
-
-    def after_training(self, model: xgb.Booster) -> xgb.Booster:
-        if self.log_model:
-            if not self.model_signature:
-                warn(
-                    "Model will not be logged because model signature is missing. "
-                    "To autolog the model, please specify `model_signature` when constructing SnowflakeXgboostCallback."
-                )
-                return model
-
-            model_name = self.model_name or self._experiment_tracking._get_or_set_experiment().name + "_model"
-            self._experiment_tracking.log_model(  # type: ignore[call-arg]
-                model=model,
-                model_name=model_name,
-                signatures={"predict": self.model_signature},
-            )
-
-        return model
-
-
-class SnowflakeLightgbmCallback(lgb.callback._RecordEvaluationCallback):
-    def __init__(
-        self,
-        experiment_tracking: "ExperimentTracking",
-        log_model: bool = True,
-        log_metrics: bool = True,
-        log_params: bool = True,
-        model_name: Optional[str] = None,
-        model_signature: Optional[ModelSignature] = None,
-    ) -> None:
-        self._experiment_tracking = experiment_tracking
-        self.log_model = log_model
-        self.log_metrics = log_metrics
-        self.log_params = log_params
-        self.model_name = model_name
-        self.model_signature = model_signature
-
-        super().__init__(eval_result={})
-
-    def __call__(self, env: lgb.callback.CallbackEnv) -> None:
-        if self.log_params:
-            if env.iteration == env.begin_iteration:  # Log params only at the first iteration
-                self._experiment_tracking.log_params(env.params)
-
-        if self.log_metrics:
-            super().__call__(env)
-            for dataset_name, metrics in self.eval_result.items():
-                for metric_name, log in metrics.items():
-                    metric_key = dataset_name + ":" + metric_name
-                    self._experiment_tracking.log_metric(key=metric_key, value=log[-1], step=env.iteration)
-
-        if self.log_model:
-            if env.iteration == env.end_iteration - 1:  # Log model only at the last iteration
-                if self.model_signature:
-                    model_name = self.model_name or self._experiment_tracking._get_or_set_experiment().name + "_model"
-                    self._experiment_tracking.log_model(  # type: ignore[call-arg]
-                        model=env.model,
-                        model_name=model_name,
-                        signatures={"predict": self.model_signature},
-                    )
-                else:
-                    warn(
-                        "Model will not be logged because model signature is missing. To autolog the model, "
-                        "please specify `model_signature` when constructing SnowflakeLightgbmCallback."
-                    )