From 976cd5dd146017aaf853cc42138b158e6ca810f1 Mon Sep 17 00:00:00 2001
From: Quentin Anthony <qganthony@yahoo.com>
Date: Sun, 8 Sep 2024 20:49:24 -0700
Subject: [PATCH] precommit

---
 megatron/logging.py                           |  1 +
 megatron/model/transformer.py                 |  2 +-
 megatron/neox_arguments/arguments.py          | 22 +++++++++++--------
 .../preprocess_data_with_chat_template.py     |  2 +-
 4 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/megatron/logging.py b/megatron/logging.py
index 9e5e6448b..05945fdda 100644
--- a/megatron/logging.py
+++ b/megatron/logging.py
@@ -24,6 +24,7 @@
 from megatron import mpu, print_rank_0
 from megatron.utils import report_memory
 
+
 class Tee:
     """Duplicate output to both stdout/err and file"""
 
diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py
index 523cbe4cf..b13e6e255 100644
--- a/megatron/model/transformer.py
+++ b/megatron/model/transformer.py
@@ -230,7 +230,7 @@ def __init__(
                 #     skip_bias_add=False,
                 #     mup_rescale_parameters=is_last_layer,  # only called if neox_args.use_mup = True, despite it not being included here
                 # )
-            else: # Not using cross entropy loss for RMs
+            else:  # Not using cross entropy loss for RMs
                 self.rm_linear = mpu.RowParallelLinear(
                     neox_args=neox_args,
                     input_size=neox_args.hidden_size,
diff --git a/megatron/neox_arguments/arguments.py b/megatron/neox_arguments/arguments.py
index 98f6bb13b..52f128d9d 100644
--- a/megatron/neox_arguments/arguments.py
+++ b/megatron/neox_arguments/arguments.py
@@ -51,15 +51,15 @@
 )
 
 ### Logging colors ###
-GREEN = '\033[92m'
-RED = '\033[91m'
-YELLOW = '\033[93m'
-END = '\033[0m'
+GREEN = "\033[92m"
+RED = "\033[91m"
+YELLOW = "\033[93m"
+END = "\033[0m"
 SUCCESS = f"{GREEN} [SUCCESS] {END}"
 OKAY = f"{GREEN}[OKAY]{END}"
 WARNING = f"{YELLOW}[WARNING]{END}"
-FAIL = f'{RED}[FAIL]{END}'
-INFO = '[INFO]'
+FAIL = f"{RED}[FAIL]{END}"
+INFO = "[INFO]"
 
 # ZERO defaults by deespeed
 # These values should not be changed unless defaults in deepspeed are changed
@@ -205,11 +205,13 @@ def initialize_comet(self):
 
                 logging.info("> setting up comet ...")
             except ImportError as e:
-                logging.error(f"{FAIL} importing comet. Comet can be installed with \"pip install comet_llm\". See https://github.com/comet-ml/comet-llm for more info. Full error is:")
+                logging.error(
+                    f'{FAIL} importing comet. Comet can be installed with "pip install comet_llm". See https://github.com/comet-ml/comet-llm for more info. Full error is:'
+                )
                 raise e
             except Exception as e:
                 logging.error(
-                        f"{FAIL} Error setting up Comet. Either set \"use_comet: False\" in your configuration file, or resolve the issue with Comet. Full error is:",
+                    f'{FAIL} Error setting up Comet. Either set "use_comet: False" in your configuration file, or resolve the issue with Comet. Full error is:',
                 )
                 raise e
 
@@ -1219,7 +1221,9 @@ def validate_values(self):
 
         # learning rate
         if self.lr is None:
-            error_message = f"{FAIL} " + self.__class__.__name__ + ".validate_values() lr is None"
+            error_message = (
+                f"{FAIL} " + self.__class__.__name__ + ".validate_values() lr is None"
+            )
             logging.error(error_message)
             raise ValueError(error_message)
             return False
diff --git a/tools/datasets/preprocess_data_with_chat_template.py b/tools/datasets/preprocess_data_with_chat_template.py
index 4d058127c..3db283ca4 100644
--- a/tools/datasets/preprocess_data_with_chat_template.py
+++ b/tools/datasets/preprocess_data_with_chat_template.py
@@ -122,7 +122,7 @@ def build_chat(
         chat_tokens = tokenizer.apply_chat_template(
             chat[: i + 1], add_generation_prompt=add_gen
         )[len(tokens) :]
-  
+
         # remove previous stuff...
         tokens.extend(chat_tokens)
         if only_last_turn and (i != len(chat) - 1):