From 05c246136e6676b2f52bd5423b8bfbb021d8a2fb Mon Sep 17 00:00:00 2001
From: Vincent Chen <v.chen@databricks.com>
Date: Thu, 18 Apr 2024 21:04:53 -0700
Subject: [PATCH] code quality

---
 llmfoundry/utils/config_utils.py   | 8 ++++----
 tests/utils/test_mlflow_logging.py | 3 +--
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/llmfoundry/utils/config_utils.py b/llmfoundry/utils/config_utils.py
index 0c16c9b349..2c8b9167b9 100644
--- a/llmfoundry/utils/config_utils.py
+++ b/llmfoundry/utils/config_utils.py
@@ -184,7 +184,7 @@ def log_config(cfg: DictConfig) -> None:
 
     if 'mlflow' in cfg.get('loggers', {}):
         if not mlflow:
-            raise ImportError("MLflow is required but not installed.")
+            raise ImportError('MLflow is required but not installed.')
         if mlflow.active_run():
             mlflow.log_params(params=om.to_container(cfg, resolve=True))
             log_dataset_uri(cfg)
@@ -212,7 +212,7 @@ def parse_source_dataset(cfg: DictConfig):
                                                               {}).get('hf_name')
             backend, _, _ = parse_uri(hf_path)
             if backend:
-                hf_path = f'{hf_path.rstrip("/")}/{split}' if split else hf_path
+                hf_path = f'{hf_path.rstrip('/')}/{split}' if split else hf_path
                 data_paths.add((backend, hf_path, data_split))
             else:
                 data_paths.add(('hf', hf_path, data_split))
@@ -222,7 +222,7 @@ def parse_source_dataset(cfg: DictConfig):
             remote_path = cfg.get(f'{data_split}_loader',
                                   {}).get('dataset', {}).get('remote', None)
             backend, _, _ = parse_uri(remote_path)
-            remote_path = f'{remote_path.rstrip("/")}/{split}/' if split else remote_path
+            remote_path = f'{remote_path.rstrip('/')}/{split}/' if split else remote_path
             data_paths.add((backend, remote_path, data_split))
         # check for local path
         elif cfg.get(f'{data_split}_loader', {}).get('dataset',
@@ -239,7 +239,7 @@ def parse_source_dataset(cfg: DictConfig):
 def log_dataset_uri(cfg: DictConfig) -> mlflow.data.meta_dataset.MetaDataset:
     """Logs dataset tracking information to MLflow."""
     if mlflow is None:
-        log.warning("MLflow is not installed. Skipping dataset logging.")
+        log.warning('MLflow is not installed. Skipping dataset logging.')
         return None
     # Figure out which data source to use
     data_paths = parse_source_dataset(cfg)
diff --git a/tests/utils/test_mlflow_logging.py b/tests/utils/test_mlflow_logging.py
index 346a868292..a9ac3a4250 100644
--- a/tests/utils/test_mlflow_logging.py
+++ b/tests/utils/test_mlflow_logging.py
@@ -12,8 +12,7 @@
 
 
 def create_config(**kwargs):
-    """Helper function to create OmegaConf configurations based on input
-    kwargs."""
+    """Helper function to create OmegaConf configurations."""
     return OmegaConf.create(kwargs)