Lint

microsoft · Oct 23, 2023 · da7b36c · da7b36c
1 parent b70ca72
commit da7b36c
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 188 deletions.
diff --git a/examples/llama/generator.py b/examples/llama/generator.py
@@ -49,11 +49,7 @@ def __init__(
 
         self.runtime: ark.Runtime = None
 
-    def launch(
-        self,
-        pth_path: str,
-        tok_path: str,
-    ):
+    def launch(self, pth_path: str, tok_path: str):
         # Load a pretrained tokenizer
         self.tokenizer = Tokenizer(model_path=tok_path)
         self.args.vocab_size = self.tokenizer.n_words
@@ -163,9 +159,7 @@ def run(self, prompt: str):
 
     gen.launch(args.pth_path, args.tok_path)
 
-    prompt_list = [
-        "Where is the capital of France?",
-    ]
+    prompt_list = ["Where is the capital of France?"]
     for i, prompt in enumerate(prompt_list):
         output = gen.run(prompt)
         print(f"---\nPrompt[{i}]: {prompt}\nOutput[{i}]: {output}")
diff --git a/examples/llama/model.py b/examples/llama/model.py
@@ -204,8 +204,7 @@ def forward(self, x):
         for i in range(len(output_shape)):
             output_shape_bytes *= output_shape[i]
         output_parallel_reshape = ark.reshape(
-            output_parallel,
-            [output_shape_bytes],
+            output_parallel, [output_shape_bytes]
         )
         output_reshape = ark.all_reduce(
             output_parallel_reshape, self.local_rank, self.world_size

diff --git a/python/ark/data_type.py b/python/ark/data_type.py
@@ -7,34 +7,13 @@
 
 
 _REGISTRY_DATA_TYPE = {
-    "fp32": {
-        "np": numpy.float32,
-        "doc": """32-bit floating point.""",
-    },
-    "fp16": {
-        "np": numpy.float16,
-        "doc": """16-bit floating point.""",
-    },
-    "bf16": {
-        "np": None,
-        "doc": """bfloat16 floating point.""",
-    },
-    "int32": {
-        "np": numpy.int32,
-        "doc": """32-bit signed integer.""",
-    },
-    "uint32": {
-        "np": numpy.uint32,
-        "doc": """32-bit unsigned integer.""",
-    },
-    "int8": {
-        "np": numpy.int8,
-        "doc": """8-bit signed integer.""",
-    },
-    "uint8": {
-        "np": numpy.uint8,
-        "doc": """8-bit unsigned integer.""",
-    },
+    "fp32": {"np": numpy.float32, "doc": """32-bit floating point."""},
+    "fp16": {"np": numpy.float16, "doc": """16-bit floating point."""},
+    "bf16": {"np": None, "doc": """bfloat16 floating point."""},
+    "int32": {"np": numpy.int32, "doc": """32-bit signed integer."""},
+    "uint32": {"np": numpy.uint32, "doc": """32-bit unsigned integer."""},
+    "int8": {"np": numpy.int8, "doc": """8-bit signed integer."""},
+    "uint8": {"np": numpy.uint8, "doc": """8-bit unsigned integer."""},
     "byte": {
         "np": numpy.ubyte,
         "doc": """