test

mosaicml · Nov 7, 2023 · 3e17fd0 · 3e17fd0
1 parent 5872f3f
commit 3e17fd0
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 9 deletions.
diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
@@ -74,11 +74,12 @@ def __init__(
         if self.mlflow_registered_model_name is not None:
             # Both the metadata and the task are needed in order for mlflow
             # and databricks optimized model serving to work
-            default_metadata = {
-                'task': 'llm/v1/completions'
-            }
+            default_metadata = {'task': 'llm/v1/completions'}
             passed_metadata = mlflow_logging_config.get('metadata', {})
-            mlflow_logging_config['metadata'] = {**default_metadata, **passed_metadata}
+            mlflow_logging_config['metadata'] = {
+                **default_metadata,
+                **passed_metadata
+            }
             mlflow_logging_config.setdefault('task', 'text-generation')
         self.mlflow_logging_config = mlflow_logging_config
 

diff --git a/tests/test_hf_conversion_script.py b/tests/test_hf_conversion_script.py
@@ -5,7 +5,7 @@
 import os
 import pathlib
 import sys
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, patch, ANY
 
 from composer import Trainer
 from composer.loggers import MLFlowLogger
@@ -421,10 +421,24 @@ 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)
+        if log_to_mlflow:
+            # assert mlflow_logger_mock.save_model.call_count == 1
+            mlflow_logger_mock.save_model.assert_called_once_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
+            # mlflow_logger.save_model(
+            #     flavor='transformers',
+            #     transformers_model=components,
+            #     path=local_save_path,
+            #     **self.mlflow_logging_config,
+            # )
+        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