pre-commit

mosaicml · Nov 9, 2023 · e416b96 · e416b96
1 parent 622f141
commit e416b96
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
@@ -69,7 +69,9 @@ def __init__(
 
         # mlflow config setup
         self.mlflow_registered_model_name = mlflow_registered_model_name
-        print(f"__init__::mlflow_logging_config={mlflow_logging_config}\n[{type(mlflow_logging_config)}]")
+        print(
+            f'__init__::mlflow_logging_config={mlflow_logging_config}\n[{type(mlflow_logging_config)}]'
+        )
         if mlflow_logging_config is None:
             mlflow_logging_config = {}
         if self.mlflow_registered_model_name is not None:
@@ -260,8 +262,12 @@ def _save_checkpoint(self, state: State, logger: Logger):
                         # TODO: Remove after mlflow fixes the bug that makes this necessary
                         import mlflow
                         mlflow.store._unity_catalog.registry.rest_store.get_feature_dependencies = lambda *args, **kwargs: ''
-                        print(f"_save_checkpoint::self.mlflow_logging_config={self.mlflow_logging_config}\n[{type(self.mlflow_logging_config)}]")
-                        print(f"_save_checkpoint::self.mlflow_logging_config['metadata']={self.mlflow_logging_config['metadata']}\n{type(self.mlflow_logging_config['metadata'])}]")
+                        print(
+                            f'_save_checkpoint::self.mlflow_logging_config={self.mlflow_logging_config}\n[{type(self.mlflow_logging_config)}]'
+                        )
+                        print(
+                            f"_save_checkpoint::self.mlflow_logging_config['metadata']={self.mlflow_logging_config['metadata']}\n{type(self.mlflow_logging_config['metadata'])}]"
+                        )
                         mlflow_logger.save_model(
                             flavor='transformers',
                             transformers_model=components,

diff --git a/llmfoundry/utils/builders.py b/llmfoundry/utils/builders.py
@@ -117,10 +117,10 @@ def build_callback(name: str, kwargs: Dict[str, Any]) -> Callback:
     elif name == 'early_stopper':
         return EarlyStopper(**kwargs)
     elif name == 'hf_checkpointer':
-        mlflow_logging_config = kwargs.pop("mlflow_logging_config", None)
+        mlflow_logging_config = kwargs.pop('mlflow_logging_config', None)
         if isinstance(mlflow_logging_config, omegaconf.dictconfig.DictConfig):
             mlflow_logging_config = om.to_object(mlflow_logging_config)
-        kwargs["mlflow_logging_config"] = mlflow_logging_config
+        kwargs['mlflow_logging_config'] = mlflow_logging_config
         return HuggingFaceCheckpointer(**kwargs)
     else:
         raise ValueError(f'Not sure how to build callback: {name}')