Run the linter on most of the repo (#895)

* Initial linter run

* undoing changes to unsupported files

Author: Jack-Khuu

.ci/scripts/extract-sequence.py

@@ -1,5 +1,6 @@
 import sys
 
+
 def print_until_equals(filename):
     output = False
     past_output = False

build/builder.py

@@ -10,18 +10,18 @@
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Dict, Optional, Union
-from utils.measure_time import measure_time
 
 import torch
 import torch._dynamo.config
 import torch._inductor.config
 
 from config.model_config import resolve_model_config
+from distributed import init_distributed, ParallelDims, parallelize_llama
 from quantization.quantize import quantize_model
+from utils.measure_time import measure_time
 
 from build.model import Transformer
 from build.utils import device_sync, is_cpu_device, is_cuda_or_cpu_device, name_to_dtype
-from distributed import parallelize_llama, ParallelDims, init_distributed
 
 
 @dataclass
@@ -287,6 +287,7 @@ def _init_model_on_meta_device(builder_args):
         else:
             return Transformer.from_name(builder_args.checkpoint_path.parent.name)
 
+
 def _load_model_gguf(builder_args, only_config=False):
     assert builder_args.gguf_path
     if builder_args.gguf_kwargs is None:

build/utils.py

@@ -136,7 +136,7 @@ def name_to_dtype(name, device):
         import platform
 
         if platform.processor() == "arm":
-            device=get_device_str(device)
+            device = get_device_str(device)
             # ARM CPU is faster with float16, MPS with bf16 if supported
             if device == "cpu" or int(platform.mac_ver()[0].split(".")[0]) < 14:
                 return torch.float16

cli.py

@@ -15,7 +15,7 @@
 from build.utils import allowable_dtype_names, allowable_params_table, get_device_str
 from download import download_and_convert, is_model_downloaded
 
-logging.basicConfig(level=logging.INFO,format="%(message)s")
+logging.basicConfig(level=logging.INFO, format="%(message)s")
 logger = logging.getLogger(__name__)
 
 default_device = os.getenv("TORCHCHAT_DEVICE", "fast")
@@ -24,18 +24,19 @@
 ).expanduser()
 
 
-# Subcommands related to downloading and managing model artifacts 
+# Subcommands related to downloading and managing model artifacts
 INVENTORY_VERBS = ["download", "list", "remove", "where"]
 
 # List of all supported subcommands in torchchat
 KNOWN_VERBS = ["chat", "browser", "generate", "eval", "export"] + INVENTORY_VERBS
 
+
 # Handle CLI arguments that are common to a majority of subcommands.
 def check_args(args, verb: str) -> None:
     # Handle model download. Skip this for download, since it has slightly
     # different semantics.
     if (
-        verb not in INVENTORY_VERBS         
+        verb not in INVENTORY_VERBS
         and args.model
         and not is_model_downloaded(args.model, args.model_directory)
     ):
@@ -47,11 +48,11 @@ def add_arguments_for_verb(parser, verb: str) -> None:
     # A model can be specified using a positional model name or HuggingFace
     # path. Alternatively, the model can be specified via --gguf-path or via
     # an explicit --checkpoint-dir, --checkpoint-path, or --tokenizer-path.
-    
+
     if verb in INVENTORY_VERBS:
         _configure_artifact_inventory_args(parser, verb)
         _add_cli_metadata_args(parser)
-        return 
+        return
 
     parser.add_argument(
         "model",
@@ -164,10 +165,10 @@ def add_arguments_for_verb(parser, verb: str) -> None:
         choices=["fast", "cpu", "cuda", "mps"],
         help="Hardware device to use. Options: cpu, cuda, mps",
     )
-    
+
     if verb == "eval":
         _add_evaluation_args(parser)
-    
+
     parser.add_argument(
         "--hf-token",
         type=str,
@@ -350,7 +351,7 @@ def arg_init(args):
         )
 
     if sys.version_info.major != 3 or sys.version_info.minor < 10:
-       raise RuntimeError("Please use Python 3.10 or later.")
+        raise RuntimeError("Please use Python 3.10 or later.")
 
     if hasattr(args, "quantize") and Path(args.quantize).is_file():
         with open(args.quantize, "r") as f:

download.py

@@ -26,10 +26,7 @@ def _download_hf_snapshot(
     from requests.exceptions import HTTPError
 
     # Download and store the HF model artifacts.
-    print(
-        f"Downloading {model_config.name} from HuggingFace...",
-        file=sys.stderr
-    )
+    print(f"Downloading {model_config.name} from HuggingFace...", file=sys.stderr)
     try:
         snapshot_download(
             model_config.distribution_path,
@@ -54,10 +51,7 @@ def _download_hf_snapshot(
             raise e
 
     # Convert the model to the torchchat format.
-    print(
-        f"Converting {model_config.name} to torchchat format...",
-        file=sys.stderr
-    )
+    print(f"Converting {model_config.name} to torchchat format...", file=sys.stderr)
     convert_hf_checkpoint(
         model_dir=artifact_dir, model_name=model_config.name, remove_bin_files=True
     )
@@ -177,6 +171,7 @@ def remove_main(args) -> None:
     shutil.rmtree(model_dir)
     print("Done.")
 
+
 # Subcommand to print downloaded model artifacts directory.
 # Asking for location will/should trigger download of model if not available.
 def where_main(args) -> None:
@@ -196,6 +191,7 @@ def where_main(args) -> None:
     print(str(os.path.abspath(model_dir)))
     exit(0)
 
+
 # Subcommand to download model artifacts.
 def download_main(args) -> None:
     download_and_convert(args.model, args.model_directory, args.hf_token)

eval.py

@@ -10,7 +10,6 @@
 import torch
 import torch._dynamo.config
 import torch._inductor.config
-from utils.measure_time import measure_time
 from build.builder import (
     _initialize_model,
     _initialize_tokenizer,
@@ -22,6 +21,7 @@
 from build.utils import set_precision
 from cli import add_arguments_for_verb, arg_init
 from generate import encode_tokens, model_forward
+from utils.measure_time import measure_time
 
 torch._dynamo.config.automatic_dynamic_shapes = True
 torch._inductor.config.triton.unique_kernel_names = True

generate.py

@@ -6,9 +6,9 @@
 import argparse
 import itertools
 import logging
+import os
 import sys
 import time
-import os
 from dataclasses import dataclass
 from pathlib import Path
 from typing import List, Optional, Tuple

runner/run.cpp

@@ -196,16 +196,16 @@ float* forward(Transformer* transformer, int token, int pos) {
   torch::Tensor token_tensor =
       torch::from_blob(token_buffer, {1, 1}, torch::kLong);
   torch::Tensor pos_tensor = torch::from_blob(pos_buffer, {1}, torch::kLong);
-  std::vector<torch::Tensor> inputs{token_tensor.to(aoti_device), pos_tensor.to(aoti_device)};
+  std::vector<torch::Tensor> inputs{
+      token_tensor.to(aoti_device), pos_tensor.to(aoti_device)};
 
   torch::Tensor result = transformer->runner->run(inputs)[0]
                              .to(torch::dtype(torch::kFloat32))
                              .to(torch::kCPU);
   auto logits = result[0].data_ptr();
 #else // __ET_MODEL__
   ManagedTensor pos_managed(pos_buffer, {1}, ScalarType::Long);
-  ManagedTensor tokens_managed(
-      token_buffer, {1, 1}, ScalarType::Long);
+  ManagedTensor tokens_managed(token_buffer, {1, 1}, ScalarType::Long);
   std::vector<EValue> inputs;
   auto tmp1 = EValue(tokens_managed.get_aliasing_tensor());
   auto tmp2 = EValue(pos_managed.get_aliasing_tensor());
@@ -811,7 +811,9 @@ void error_usage() {
       "  -v <int>    (optional) vocab size, default is model-specific.\n");
   fprintf(
       stderr, "  -l <int>    (optional) llama version (2 or 3), default 2.\n");
-  fprintf(stderr, "  -d <string> (optional) device(CUDA or CPU)  model was exported for\n");
+  fprintf(
+      stderr,
+      "  -d <string> (optional) device(CUDA or CPU)  model was exported for\n");
   exit(EXIT_FAILURE);
 }
 
@@ -884,16 +886,16 @@ int main(int argc, char* argv[]) {
 #ifdef __AOTI_MODEL__
     } else if (argv[i][1] == 'd') {
 #ifdef USE_CUDA
-       if (strcasecmp(argv[i + 1], "CUDA") == 0) {
-          aoti_device = torch::Device(torch::kCUDA);
-       } else
+      if (strcasecmp(argv[i + 1], "CUDA") == 0) {
+        aoti_device = torch::Device(torch::kCUDA);
+      } else
 #endif
-       if (strcasecmp(argv[i + 1], "CPU") == 0) {
-          aoti_device = torch::Device(torch::kCPU);
-       } else {
-         fprintf(stderr, "Unknown device %s", argv[i + 1]);
-         exit(1);
-       }
+          if (strcasecmp(argv[i + 1], "CPU") == 0) {
+        aoti_device = torch::Device(torch::kCPU);
+      } else {
+        fprintf(stderr, "Unknown device %s", argv[i + 1]);
+        exit(1);
+      }
 #endif
     } else {
       error_usage();

scripts/patch_triton.py

@@ -3,20 +3,21 @@
 
 
 from pathlib import Path
+
 import triton
 
 
 def patch_def_search_in_jit_py(jit_py: Path) -> None:
-  with jit_py.open() as f:
-    lines = f.readlines()
-  old_match = 'self.src = self.src[self.src.find("def"):]'
-  new_match = 'self.src = self.src[re.search(r"^def\s+\w+\s*\(", self.src, re.MULTILINE).start():]'
-  lines.insert(4, "import re\n")
-  for idx, line in enumerate(lines):
-    if old_match in line:
-      lines[idx] = line.replace(old_match, new_match) 
-  jit_py.write_text("".join(lines))
+    with jit_py.open() as f:
+        lines = f.readlines()
+    old_match = 'self.src = self.src[self.src.find("def"):]'
+    new_match = 'self.src = self.src[re.search(r"^def\s+\w+\s*\(", self.src, re.MULTILINE).start():]'
+    lines.insert(4, "import re\n")
+    for idx, line in enumerate(lines):
+        if old_match in line:
+            lines[idx] = line.replace(old_match, new_match)
+    jit_py.write_text("".join(lines))
+
 
 jit_py = Path(triton.__file__).parent / "runtime" / "jit.py"
 patch_def_search_in_jit_py(jit_py)
-

scripts/updown.py

@@ -86,12 +86,12 @@ def updown_process_line(
         # [ x1 | c2 | x3 ] means "pick one", so we may have to check that and pick one
         # of the options.  Probably pick the last option because testing has more likely
         # been performed with the first option!
-        last=True
+        last = True
         if last:
-            line=select_last_option_between_brackets(line)
+            line = select_last_option_between_brackets(line)
         else:
-            line=select_first_option_between_brackets(line)
-            
+            line = select_first_option_between_brackets(line)
+
         output(
             remove_text_between_brackets(line),
             replace_list=replace_list,

torchchat.py

@@ -11,10 +11,10 @@
 
 from cli import (
     add_arguments_for_verb,
-    KNOWN_VERBS,
-    INVENTORY_VERBS,
     arg_init,
     check_args,
+    INVENTORY_VERBS,
+    KNOWN_VERBS,
 )
 
 default_device = "cpu"
@@ -34,19 +34,19 @@
     subparsers.required = True
 
     VERB_HELP = {
-       "chat": "Chat interactively with a model",
-       "browser": "Chat interactively in a browser",
-       "download": "Download a model from Hugging Face or others",
-       "generate": "Generate responses from a model given a prompt",
-       "eval": "Evaluate a model given a prompt",
-       "export": "Export a model for AOT Inductor or ExecuTorch",
-       "list": "List supported models",
-       "remove": "Remove downloaded model artifacts",
-       "where": "Return directory containing downloaded model artifacts",
+        "chat": "Chat interactively with a model",
+        "browser": "Chat interactively in a browser",
+        "download": "Download a model from Hugging Face or others",
+        "generate": "Generate responses from a model given a prompt",
+        "eval": "Evaluate a model given a prompt",
+        "export": "Export a model for AOT Inductor or ExecuTorch",
+        "list": "List supported models",
+        "remove": "Remove downloaded model artifacts",
+        "where": "Return directory containing downloaded model artifacts",
     }
     for verb in KNOWN_VERBS:
-       subparser = subparsers.add_parser(verb, help=VERB_HELP[verb])
-       add_arguments_for_verb(subparser, verb)
+        subparser = subparsers.add_parser(verb, help=VERB_HELP[verb])
+        add_arguments_for_verb(subparser, verb)
 
     # Now parse the arguments
     args = parser.parse_args()
@@ -72,7 +72,7 @@
         args.gui = True
         check_args(args, "browser")
 
-        from browser.browser import main as browser_main 
+        from browser.browser import main as browser_main
 
         browser_main(args)
     elif args.command == "generate":

utils/measure_time.py

@@ -1,15 +1,13 @@
 from time import perf_counter
 from typing import Optional
 
+
 class measure_time:
-    def __init__(
-            self,
-            message: Optional[str] = 'Time: {time:.3f} seconds'
-    ):
+    def __init__(self, message: Optional[str] = "Time: {time:.3f} seconds"):
         self.message = message
 
     def __enter__(
-            self,
+        self,
     ):
         self.start = perf_counter()
         self.message