format style

Former-commit-id: 53b683531b83cd1d19de97c6565f16c1eca6f5e1

tests/cal_flops.py

@@ -3,11 +3,12 @@
 # Usage: python cal_flops.py --model_name_or_path path_to_model --batch_size 1 --seq_length 512
 # Inspired by: https://www.deepspeed.ai/tutorials/flops-profiler/
 
+from typing import Optional
+
 import fire
 import torch
-from typing import Optional
-from deepspeed.accelerator import get_accelerator # type: ignore
-from deepspeed.profiling.flops_profiler import get_model_profile # type: ignore
+from deepspeed.accelerator import get_accelerator  # type: ignore
+from deepspeed.profiling.flops_profiler import get_model_profile  # type: ignore
 
 from llmtuner import ChatModel
 
@@ -16,25 +17,13 @@ def calculate_flops(
     model_name_or_path: str,
     batch_size: Optional[int] = 1,
     seq_length: Optional[int] = 256,
-    flash_attn: Optional[bool] = False
+    flash_attn: Optional[bool] = False,
 ):
     with get_accelerator().device(0):
-        chat_model = ChatModel(dict(
-            model_name_or_path=model_name_or_path,
-            template="vanilla",
-            flash_attn=flash_attn
-        ))
+        chat_model = ChatModel(dict(model_name_or_path=model_name_or_path, template="vanilla", flash_attn=flash_attn))
         fake_input = torch.ones((batch_size, seq_length), dtype=torch.long, device=chat_model.model.device)
-        input_dict = {
-            "input_ids": fake_input,
-            "labels": fake_input.clone()
-        }
-        flops, macs, params = get_model_profile(
-            chat_model.model,
-            kwargs=input_dict,
-            print_profile=True,
-            detailed=True
-        )
+        input_dict = {"input_ids": fake_input, "labels": fake_input.clone()}
+        flops, macs, params = get_model_profile(chat_model.model, kwargs=input_dict, print_profile=True, detailed=True)
         print("FLOPs:", flops)
         print("MACs:", macs)
         print("Params:", params)