Apply isort and black reformatting

Signed-off-by: cuichenx <[email protected]>
NVIDIA · Jul 1, 2024 · 66646b8 · 66646b8
1 parent 53ce64a
commit 66646b8
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/scripts/checkpoint_converters/convert_gpt_nemo_to_mcore.py b/scripts/checkpoint_converters/convert_gpt_nemo_to_mcore.py
@@ -128,19 +128,19 @@ def build_key_mapping(nemo_cfg):
         f"{model_str}.decoder.final_layernorm.weight": "model.language_model.encoder.final_layernorm.weight",
     }
     if has_layernorm_bias:
-        mcore_to_nemo_mapping[
-            f"{model_str}.decoder.final_layernorm.bias"
-        ] = "model.language_model.encoder.final_layernorm.bias"
+        mcore_to_nemo_mapping[f"{model_str}.decoder.final_layernorm.bias"] = (
+            "model.language_model.encoder.final_layernorm.bias"
+        )
 
     if not nemo_cfg.get("share_embeddings_and_output_weights", True):
         mcore_to_nemo_mapping[f"{model_str}.output_layer.weight"] = "model.language_model.output_layer.weight"
 
     if nemo_cfg.get("position_embedding_type", 'learned_absolute') == 'rope':
         mcore_to_nemo_mapping[f"{model_str}.rotary_pos_emb.inv_freq"] = "model.language_model.rotary_pos_emb.inv_freq"
     else:
-        mcore_to_nemo_mapping[
-            f"{model_str}.embedding.position_embeddings.weight"
-        ] = "model.language_model.embedding.position_embeddings.weight"
+        mcore_to_nemo_mapping[f"{model_str}.embedding.position_embeddings.weight"] = (
+            "model.language_model.embedding.position_embeddings.weight"
+        )
 
     nemo_prefix = "model.language_model.encoder.layers"
     mcore_prefix = f"{model_str}.decoder.layers"
@@ -338,5 +338,7 @@ def run_sanity_checks(nemo_file, mcore_file, cpu_only=False, ignore_if_missing=t
     try:
         run_sanity_checks(input_nemo_file, output_nemo_file, cpu_only=cpu_only, ignore_if_missing=ignore_if_missing)
     except torch.cuda.OutOfMemoryError:
-        logging.info("✅ Conversion was successful, but could not run sanity check due to torch.cuda.OutOfMemoryError.")
+        logging.info(
+            "✅ Conversion was successful, but could not run sanity check due to torch.cuda.OutOfMemoryError."
+        )
         logging.info("Please run the script with the same command again to run sanity check.")