Apply isort and black reformatting

Signed-off-by: Laplasjan107 <[email protected]>
NVIDIA · Nov 21, 2024 · a66e20a · a66e20a
1 parent df99faa
commit a66e20a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/nemo/collections/llm/quantization/quantizer.py b/nemo/collections/llm/quantization/quantizer.py
@@ -147,9 +147,8 @@ def _generate_sample(model: MegatronParallel):
 
         mcore_tokenizer = MCoreTokenizerWrappper(model.tokenizer)
         mcore_inference = model.get_inference_wrapper(
-            params_dtype=torch.bfloat16,
-            inference_batch_times_seqlen_threshold=30
-            )
+            params_dtype=torch.bfloat16, inference_batch_times_seqlen_threshold=30
+        )
 
         generated = [r.generated_text for r in generate(mcore_inference, mcore_tokenizer, prompts)]
         outputs = [prompt + generation for prompt, generation in zip(prompts, generated)]

diff --git a/nemo/collections/llm/quantization/utils.py b/nemo/collections/llm/quantization/utils.py
@@ -23,7 +23,6 @@
 from nemo.utils import logging
 
 
-
 def get_modelopt_decoder_type(model: llm.GPTModel) -> str:
     """Infers the modelopt decoder type from GPTModel subclass."""
     mapping = [
@@ -36,7 +35,7 @@ def get_modelopt_decoder_type(model: llm.GPTModel) -> str:
         (llm.MixtralModel, "llama"),
         (llm.NemotronModel, "gptnext"),
         (llm.Qwen2Model, "qwen"),
-        (llm.StarcoderModel,   "gptnext"),
+        (llm.StarcoderModel, "gptnext"),
         (llm.Starcoder2Model, "gptnext"),
         (llm.Phi3Model, "phi3"),
     ]