respect kwargs spacing

Author: djsaunde

.pre-commit-config.yaml

@@ -6,11 +6,12 @@ repos:
         args:
           - --fix
           - --exit-non-zero-on-fix
-      - id: ruff-format
   - repo: local
     hooks:
-      - id: enforce-kwargs-spacing
-        name: Enforce spaces around keyword equals
-        entry: scripts/enforce_kwargs_spacing.py
-        language: system
+      - id: ruff-format-with-kwargs
+        name: Ruff format with kwarg spacing
+        entry: scripts/run_ruff_format.py
+        language: python
         types: [python]
+        additional_dependencies:
+          - ruff==0.6.9

scripts/enforce_kwargs_spacing.py

@@ -13,7 +13,7 @@
 
 def enforce_spacing(text: str) -> tuple[str, bool]:
     """Return updated text with keyword '=' padded by spaces, plus change flag."""
-    lines = text.splitlines(keepends = True)
+    lines = text.splitlines(keepends=True)
     if not lines:
         return text, False
 
@@ -68,23 +68,28 @@ def process_file(path: Path) -> bool:
             original = handle.read()
             encoding = handle.encoding
     except (OSError, SyntaxError) as exc:  # SyntaxError from tokenize on invalid python
-        print(f"Failed to read {path}: {exc}", file = sys.stderr)
+        print(f"Failed to read {path}: {exc}", file=sys.stderr)
         return False
 
     updated, changed = enforce_spacing(original)
     if changed:
-        path.write_text(updated, encoding = encoding)
+        path.write_text(updated, encoding=encoding)
     return changed
 
 
 def main(argv: list[str]) -> int:
-    parser = argparse.ArgumentParser(description = __doc__)
-    parser.add_argument("files", nargs = "+", help = "Python files to fix")
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("files", nargs="+", help="Python files to fix")
     args = parser.parse_args(argv)
 
     touched: list[Path] = []
+    self_path = Path(__file__).resolve()
+
     for entry in args.files:
         path = Path(entry)
+        # Skip modifying this script to avoid self-edit loops.
+        if path.resolve() == self_path:
+            continue
         if not path.exists() or path.is_dir():
             continue
         if process_file(path):

scripts/run_ruff_format.py

@@ -0,0 +1,30 @@
+#!/usr/bin/env python3
+"""Run `ruff format` followed by kwarg spacing enforcement."""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+from pathlib import Path
+
+HERE = Path(__file__).resolve().parent
+
+
+def main(argv: list[str]) -> int:
+    files = [arg for arg in argv if Path(arg).exists()]
+    if not files:
+        return 0
+
+    ruff_cmd = [sys.executable, "-m", "ruff", "format", *files]
+    ruff_proc = subprocess.run(ruff_cmd)
+    if ruff_proc.returncode != 0:
+        return ruff_proc.returncode
+
+    spacing_script = HERE / "enforce_kwargs_spacing.py"
+    spacing_cmd = [sys.executable, str(spacing_script), *files]
+    spacing_proc = subprocess.run(spacing_cmd)
+    return spacing_proc.returncode
+
+
+if __name__ == "__main__":
+    raise SystemExit(main(sys.argv[1:]))

tests/qlora/test_hf_qlora_train_and_merge.py

@@ -54,45 +54,45 @@
     seed = 42
     batch_size = 5
     num_generations = 5
-    tokenizer = setup_tokenizer(model_name, fixup_funcs=[fix_llama3_tokenizer])
+    tokenizer = setup_tokenizer(model_name, fixup_funcs = [fix_llama3_tokenizer])
     temperature = 0.8
     max_new_tokens = 20
 
-    peft_config = get_peft_config(lora_rank=lora_rank, target_modules="all-linear")
-    model = setup_model(model_name, quantize=True, dtype=dtype, peft_config=peft_config)
+    peft_config = get_peft_config(lora_rank = lora_rank, target_modules = "all-linear")
+    model = setup_model(model_name, quantize = True, dtype = dtype, peft_config = peft_config)
 
     prompt = tokenizer.apply_chat_template(
-        [USER_MESSAGE], tokenize=False, add_generation_prompt=True
+        [USER_MESSAGE], tokenize = False, add_generation_prompt = True
     )
     with header_footer_context("Test Prompt and Answer"):
         print(f"Test Prompt:\n{prompt}\nExpected Answer:\n{ANSWER}")
 
     dataset: Dataset = create_dataset(
-        tokenizer, num_examples=num_examples, messages=DEFAULT_MESSAGES
+        tokenizer, num_examples = num_examples, messages = DEFAULT_MESSAGES
     )
     with header_footer_context("Dataset"):
         print(f"Dataset: {next(iter(dataset))}")
 
     training_args = SFTConfig(
-        output_dir=output_dir,
-        max_steps=max_steps,
-        per_device_train_batch_size=batch_size,
-        log_level="info",
-        report_to="none",
-        num_train_epochs=1,
-        logging_steps=1,
-        seed=seed,
-        bf16=dtype == torch.bfloat16,
-        fp16=dtype == torch.float16,
-        save_strategy="no",
+        output_dir = output_dir,
+        max_steps = max_steps,
+        per_device_train_batch_size = batch_size,
+        log_level = "info",
+        report_to = "none",
+        num_train_epochs = 1,
+        logging_steps = 1,
+        seed = seed,
+        bf16 = dtype == torch.bfloat16,
+        fp16 = dtype == torch.float16,
+        save_strategy = "no",
     )
 
     with header_footer_context("Train Args"):
         print(training_args)
         print(peft_config)
 
     trainer = setup_trainer(
-        model, tokenizer, dataset, training_args, peft_config=peft_config
+        model, tokenizer, dataset, training_args, peft_config = peft_config
     )
 
     with header_footer_context("Model"):
@@ -108,11 +108,11 @@
     responses = sample_responses(
         model,
         tokenizer,
-        prompt=prompt,
+        prompt = prompt,
         **generation_args,
     )
     with header_footer_context("Responses before training"):
-        check_responses(responses, answer=ANSWER, prompt=prompt)
+        check_responses(responses, answer = ANSWER, prompt = prompt)
 
     with header_footer_context("Peft Weights before training"):
         for name, stats in itertools.islice(describe_peft_weights(model), 2):
@@ -129,11 +129,11 @@
     responses = sample_responses(
         model,
         tokenizer,
-        prompt=prompt,
+        prompt = prompt,
         **generation_args,
     )
     with header_footer_context("Responses after training"):
-        check_responses(responses, answer=ANSWER, prompt=prompt)
+        check_responses(responses, answer = ANSWER, prompt = prompt)
 
     model_copy = deepcopy(model)
 
@@ -142,18 +142,18 @@
     responses = sample_responses(
         merged_model,
         tokenizer,
-        prompt=prompt,
+        prompt = prompt,
         **generation_args,
     )
     with header_footer_context("Responses after custom merging to 16bit"):
-        check_responses(responses, answer=ANSWER, prompt=prompt)
+        check_responses(responses, answer = ANSWER, prompt = prompt)
 
     merged_model_peft = model_copy.merge_and_unload()
     responses = sample_responses(
         merged_model_peft,
         tokenizer,
-        prompt=prompt,
+        prompt = prompt,
         **generation_args,
     )
     with header_footer_context("Responses after peft merge_and_unload"):
-        check_responses(responses, answer=ANSWER, prompt=prompt)
+        check_responses(responses, answer = ANSWER, prompt = prompt)

tests/qlora/test_unsloth_qlora_train_and_merge.py

@@ -50,13 +50,13 @@ def get_unsloth_model_and_tokenizer(
     dtype: torch.dtype = torch.bfloat16,
 ):
     return FastLanguageModel.from_pretrained(
-        model_name=model_name,
-        max_seq_length=max_seq_length,
-        load_in_4bit=load_in_4bit,
-        fast_inference=fast_inference,
-        max_lora_rank=max_lora_rank,
-        gpu_memory_utilization=gpu_memory_utilization,
-        dtype=dtype,
+        model_name = model_name,
+        max_seq_length = max_seq_length,
+        load_in_4bit = load_in_4bit,
+        fast_inference = fast_inference,
+        max_lora_rank = max_lora_rank,
+        gpu_memory_utilization = gpu_memory_utilization,
+        dtype = dtype,
     )
 
 
@@ -69,11 +69,11 @@ def get_unsloth_peft_model(
 ):
     return FastLanguageModel.get_peft_model(
         model,
-        r=lora_rank,
-        target_modules=target_modules,
-        lora_alpha=lora_rank,
-        use_gradient_checkpointing=use_gradient_checkpointing,
-        random_state=random_state,
+        r = lora_rank,
+        target_modules = target_modules,
+        lora_alpha = lora_rank,
+        use_gradient_checkpointing = use_gradient_checkpointing,
+        random_state = random_state,
     )
 
 
@@ -101,48 +101,48 @@ def get_unsloth_peft_model(
 
     model, tokenizer = get_unsloth_model_and_tokenizer(
         model_name,
-        max_seq_length=512,
-        load_in_4bit=True,
-        fast_inference=False,
-        max_lora_rank=lora_rank,
-        dtype=dtype,
+        max_seq_length = 512,
+        load_in_4bit = True,
+        fast_inference = False,
+        max_lora_rank = lora_rank,
+        dtype = dtype,
     )
     temperature = 0.8
     max_new_tokens = 20
 
     model = get_unsloth_peft_model(
         model,
-        lora_rank=lora_rank,
-        target_modules=target_modules,
-        use_gradient_checkpointing=gradient_checkpointing,
-        random_state=seed,
+        lora_rank = lora_rank,
+        target_modules = target_modules,
+        use_gradient_checkpointing = gradient_checkpointing,
+        random_state = seed,
     )
 
     prompt = tokenizer.apply_chat_template(
-        [USER_MESSAGE], tokenize=False, add_generation_prompt=True
+        [USER_MESSAGE], tokenize = False, add_generation_prompt = True
     )
 
     with header_footer_context("Test Prompt and Answer"):
         print(f"Test Prompt:\n{prompt}\nExpected Answer:\n{ANSWER}")
 
     dataset: Dataset = create_dataset(
-        tokenizer, num_examples=num_examples, messages=DEFAULT_MESSAGES
+        tokenizer, num_examples = num_examples, messages = DEFAULT_MESSAGES
     )
     with header_footer_context("Dataset"):
         print(f"Dataset: {next(iter(dataset))}")
 
     training_args = SFTConfig(
-        output_dir=output_dir,
-        max_steps=max_steps,
-        per_device_train_batch_size=batch_size,
-        log_level="info",
-        report_to="none",
-        num_train_epochs=1,
-        logging_steps=1,
-        seed=seed,
-        bf16=dtype == torch.bfloat16,
-        fp16=dtype == torch.float16,
-        save_strategy="no",
+        output_dir = output_dir,
+        max_steps = max_steps,
+        per_device_train_batch_size = batch_size,
+        log_level = "info",
+        report_to = "none",
+        num_train_epochs = 1,
+        logging_steps = 1,
+        seed = seed,
+        bf16 = dtype == torch.bfloat16,
+        fp16 = dtype == torch.float16,
+        save_strategy = "no",
     )
 
     with header_footer_context("Train Args"):
@@ -163,11 +163,11 @@ def get_unsloth_peft_model(
     responses = sample_responses(
         model,
         tokenizer,
-        prompt=prompt,
+        prompt = prompt,
         **generation_args,
     )
     with header_footer_context("Responses before training"):
-        check_responses(responses, answer=ANSWER, prompt=prompt)
+        check_responses(responses, answer = ANSWER, prompt = prompt)
     with header_footer_context("Peft Weights before training"):
         for name, stats in itertools.islice(describe_peft_weights(model), 2):
             print(f"{name}:\n{stats}")
@@ -183,29 +183,29 @@ def get_unsloth_peft_model(
     responses = sample_responses(
         model,
         tokenizer,
-        prompt=prompt,
+        prompt = prompt,
         **generation_args,
     )
     with header_footer_context("Responses after training"):
-        check_responses(responses, answer=ANSWER, prompt=prompt)
+        check_responses(responses, answer = ANSWER, prompt = prompt)
 
     model.save_pretrained_merged(
         unsloth_merged_path,
         tokenizer,
-        save_method="merged_16bit",
+        save_method = "merged_16bit",
     )
     merged_model_unsloth, tokenizer = get_unsloth_model_and_tokenizer(
         unsloth_merged_path,
-        max_seq_length=512,
-        load_in_4bit=False,
-        fast_inference=False,
-        dtype=dtype,
+        max_seq_length = 512,
+        load_in_4bit = False,
+        fast_inference = False,
+        dtype = dtype,
     )
     responses = sample_responses(
         merged_model_unsloth,
         tokenizer,
-        prompt=prompt,
+        prompt = prompt,
         **generation_args,
     )
     with header_footer_context("Responses after unsloth merge to 16bit"):
-        check_responses(responses, answer=ANSWER, prompt=prompt)
+        check_responses(responses, answer = ANSWER, prompt = prompt)