From 157e059e4150e5dae7cd15e0067b77fded64bea2 Mon Sep 17 00:00:00 2001 From: wenfeiy-db Date: Tue, 14 Nov 2023 13:33:34 -0800 Subject: [PATCH] tests --- tests/test_hf_conversion_script.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/tests/test_hf_conversion_script.py b/tests/test_hf_conversion_script.py index d2c2a9e1c9..7ba2559de2 100644 --- a/tests/test_hf_conversion_script.py +++ b/tests/test_hf_conversion_script.py @@ -421,10 +421,23 @@ def test_huggingface_conversion_callback(model: str, tmp_path: pathlib.Path, trainer.fit() if dist.get_global_rank() == 0: - assert mlflow_logger_mock.save_model.call_count == (1 if log_to_mlflow - else 0) - assert mlflow_logger_mock.register_model.call_count == ( - 1 if log_to_mlflow else 0) + # assert mlflow_logger_mock.save_model.call_count == (1 if log_to_mlflow + # else 0) + # assert mlflow_logger_mock.register_model.call_count == ( + # 1 if log_to_mlflow else 0) + if log_to_mlflow: + assert mlflow_logger_mock.save_model.call_count == 1 + mlflow_logger_mock.save_model.assert_called_with( + flavor='transformers', + transformers_model=ANY, + path=ANY, + task='text-generation', + metatdata={'task': 'llm/v1/completions'} + ) + assert mlflow_logger_mock.register_model.call_count == 1 + else: + assert mlflow_logger_mock.save_model.call_count == 0 + assert mlflow_logger_mock.register_model.call_count == 0 else: assert mlflow_logger_mock.log_model.call_count == 0 assert mlflow_logger_mock.register_model.call_count == 0