tests

mosaicml · Nov 14, 2023 · 157e059 · 157e059
1 parent 3adaf11
commit 157e059
Showing 1 changed file with 17 additions and 4 deletions.
diff --git 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