clean up ruff check

Signed-off-by: kcirred <16872435+kcirred@users.noreply.github.com>

Author: kcirred

aiu_fms_testing_utils/testing/validation.py

@@ -288,8 +288,8 @@ def extract_validation_information(
 
     if hasattr(post_iteration_hook, "extracted_logits"):
         validation_info = [
-            {"tokens": t.to("cpu"), "logits": l.to("cpu")}
-            for t, l in zip(
+            {"tokens": t.to("cpu"), "logits": logits.to("cpu")}
+            for t, logits in zip(
                 torch.unbind(result), torch.unbind(post_iteration_hook.extracted_logits)
             )
         ]

aiu_fms_testing_utils/utils/aiu_setup.py

@@ -49,9 +49,9 @@ def aiu_setup(rank=0, world_size=1, local_rank=0, local_size=1, verbose=False):
     #     )  # directory needs to exist
 
     if os.getenv("FLEX_COMPUTE") == "SENTIENT":
-        dprint(f"Sentient AIU: Enabled")
+        dprint("Sentient AIU: Enabled")
     else:
-        dprint(f"Sentient AIU: Disabled (Senulator)")
+        dprint("Sentient AIU: Disabled (Senulator)")
 
 
 # ==============================================================
@@ -67,6 +67,6 @@ def aiu_dist_setup(rank, world_size, local_rank=-0, local_size=-1, verbose=False
         os.environ["MASTER_ADDR"] = "localhost"
         os.environ["MASTER_PORT"] = "12355"
     elif rank == 0 or verbose:
-        dprint(f"Detected running via torchrun")
+        dprint("Detected running via torchrun")
 
     aiu_setup(rank, world_size)

aiu_fms_testing_utils/utils/paged.py

@@ -3,7 +3,7 @@
 import time
 from typing import Any, Callable, List, MutableMapping, Optional, Tuple, Union
 import torch
-import fms.utils.spyre.paged
+import torch.nn.functional as F
 
 
 def adjust_inputs_to_batch(input_ids: torch.Tensor, **padding_kwargs):

scripts/generate_metrics.py

@@ -1,9 +1,6 @@
 import argparse
 import ast
-import json
 import os
-import random
-from typing import List, Optional, Tuple
 
 import torch
 from torch import distributed as dist
@@ -178,11 +175,11 @@ def find_eos_index(reference_tokens, eos_token_id):
     return result
 
 
-def filter_before_eos(l, filter_indexes):
+def filter_before_eos(metrics, filter_indexes):
     from itertools import groupby
 
     filtered_results = [
-        list(g)[: filter_indexes[k]] for k, g in groupby(l, key=lambda x: x[0])
+        list(g)[: filter_indexes[k]] for k, g in groupby(metrics, key=lambda x: x[0])
     ]
     return [item for sublist in filtered_results for item in sublist]
 
@@ -199,10 +196,10 @@ def __prepare_inputs(batch_size, seq_length, tokenizer, seed=0):
     return input_ids, padding_kwargs
 
 
-def write_csv(l, path, metric):
+def write_csv(metrics, path, metric_name):
     with open(path, "w") as f:
-        f.write(f"{metric}\n")
-        for t in l:
+        f.write(f"{metric_name}\n")
+        for t in metrics:
             f.write(f"{t[2].item()}\n")
         f.close()
 
@@ -279,20 +276,20 @@ def write_csv(l, path, metric):
 if num_test_tokens_per_sequence is None:
     num_test_tokens_per_sequence = args.max_new_tokens
 
-cross_entropy = lambda r, t: torch.nn.CrossEntropyLoss()(
+cross_entropy = lambda r, t: torch.nn.CrossEntropyLoss()(  # noqa: E731
     r, t.softmax(dim=1).to(dtype=torch.float32)
 )
-prob_mean = lambda r, t: torch.mean(
+prob_mean = lambda r, t: torch.mean(  # noqa: E731
     (
         r.softmax(dim=1).to(dtype=torch.float32)
         / t.softmax(dim=1).to(dtype=torch.float32)
     )
     - 1.0
 )
-prob_std = lambda r, t: torch.std(
+prob_std = lambda r, t: torch.std(  # noqa: E731
     r.softmax(dim=1).to(dtype=torch.float32) / t.softmax(dim=1).to(dtype=torch.float32)
 )
-diff_mean = lambda r, t: torch.mean(
+diff_mean = lambda r, t: torch.mean(  # noqa: E731
     torch.abs(
         r.softmax(dim=1).to(dtype=torch.float32)
         - t.softmax(dim=1).to(dtype=torch.float32)

scripts/inference.py

@@ -7,7 +7,6 @@
 from pathlib import Path
 import random
 import time
-import contextlib
 
 # Third Party
 from aiu_fms_testing_utils.utils import aiu_setup, warmup_model
@@ -236,17 +235,13 @@
 if args.quantization == "gptq":
     if "aiu" in args.device_type:
         try:
-            from fms_mo.aiu_addons.gptq import gptq_aiu_adapter, gptq_aiu_linear
-
             print("Loaded `aiu_addons` functionalities")
-        except:
+        except ImportError:
             raise ImportError("Failed to import GPTQ addons from fms-mo.")
 elif args.quantization == "int8":
     try:
-        from fms_mo.aiu_addons.i8i8 import i8i8_aiu_adapter, i8i8_aiu_linear
-
         print("Loaded `aiu_addons` functionalities")
-    except:
+    except ImportError:
         raise ImportError("Failed to import INT8 addons from fms-mo.")
 
 # this is a test model config
@@ -284,8 +279,6 @@
     device = torch.device(args.device_type, local_rank)
     torch.cuda.set_device(device)
 elif is_aiu_backend:
-    from torch_sendnn import torch_sendnn
-
     if not args.distributed:
         aiu_setup.aiu_setup(rank, world_size)
 
@@ -617,7 +610,7 @@ def truncate_prompts_to_max_length(prompts, max_len, max_allowed_length):
 
 if args.fixed_prompt_length != 0 and args.fixed_prompt_length < max_len:
     dprint(
-        f"One or more prompts require truncation. Truncation has been disabled as fixed_prompt_length has been set."
+        "One or more prompts require truncation. Truncation has been disabled as fixed_prompt_length has been set."
     )
     exit(1)
 prompts = truncate_prompts_to_max_length(prompts, max_len, max_allowed_length)
@@ -734,7 +727,7 @@ def infer(use_cache, do_sample, warmup):
 ]  # True/False are identical with greedy iff `torch.use_deterministic_algorithms(True)`
 
 if args.compile:
-    dprint(f"compilation warmup")
+    dprint("compilation warmup")
     pt_compile_model_time = time.time()
     if args.device_type == "aiu":  # only run warmup for AIU, no need for senulator
         warmup_model(
@@ -756,7 +749,7 @@ def infer(use_cache, do_sample, warmup):
     pt_compile_model_time = time.time() - pt_compile_model_time
     dprint(f"PT compile complete, took {pt_compile_model_time:.3f}s")
 
-dprint(f"generating output")
+dprint("generating output")
 
 for sample, cache in itertools.product(do_sample, use_cache):
     for _ in range(args.iters):

scripts/roberta.py

@@ -29,7 +29,6 @@
 from fms.models.hf import to_hf_api
 
 # Import AIU Libraries
-from torch_sendnn import torch_sendnn
 
 # ==============================================================
 # Main
@@ -99,7 +98,7 @@
     # Create the model
     # -------------
     if 0 == world_rank:
-        dprint(f"Creating the model...")
+        dprint("Creating the model...")
     # model_name = "roberta-base"
     # model_name = "deepset/roberta-base-squad2-distilled"
     model_name = "FacebookAI/roberta-base"
@@ -128,7 +127,7 @@
     # Compile the model
     # -------------
     if 0 == world_rank:
-        dprint(f"Compiling the model...")
+        dprint("Compiling the model...")
     the_compiled_model = torch.compile(hf_model_fms, backend=dynamo_backend)
     the_compiled_model.eval()  # inference only mode
     torch.set_grad_enabled(False)
@@ -149,15 +148,15 @@
 
     # First run will create compiled artifacts
     if 0 == world_rank:
-        dprint(f"Running model: First Time...")
+        dprint("Running model: First Time...")
     unmasker = pipeline("fill-mask", model=the_compiled_model, tokenizer=tokenizer)
     the_output = unmasker(prompt)
     if 0 == world_rank:
         dprint(f"Answer: ({the_output[0]['score']:6.5f}) {the_output[0]['sequence']}")
 
     # Second run will be faster as it uses the cached artifacts
     if 0 == world_rank:
-        dprint(f"Running model: Second Time...")
+        dprint("Running model: Second Time...")
     unmasker = pipeline("fill-mask", model=the_compiled_model, tokenizer=tokenizer)
     the_output = unmasker(prompt)
     if 0 == world_rank:
@@ -167,7 +166,7 @@
     # Cleanup
     # -------------
     if 0 == world_rank:
-        dprint(f"Done")
+        dprint("Done")
     if is_distributed:
         torch.distributed.barrier()
         torch.distributed.destroy_process_group()

scripts/small-toy.py

@@ -16,7 +16,6 @@
 from fms.utils.tp_wrapping import apply_tp
 
 # Import AIU Libraries
-from torch_sendnn import torch_sendnn
 
 
 # ==============================================================
@@ -136,7 +135,7 @@ def forward(self, x):
     # Create the model
     # -------------
     if 0 == world_rank:
-        dprint(f"Creating the model...")
+        dprint("Creating the model...")
     the_model = ToyModelFM()
     if is_distributed:
         # Create a Tensor Parallel version of the model
@@ -146,7 +145,7 @@ def forward(self, x):
     # Compile the model
     # -------------
     if 0 == world_rank:
-        dprint(f"Compiling the model...")
+        dprint("Compiling the model...")
     the_compiled_model = torch.compile(the_model, backend=dynamo_backend)
     the_compiled_model.eval()  # inference only mode
     torch.set_grad_enabled(False)
@@ -164,19 +163,19 @@ def forward(self, x):
 
     # First run will create compiled artifacts
     if 0 == world_rank:
-        dprint(f"Running model: First Time...")
+        dprint("Running model: First Time...")
     the_outputs = the_compiled_model(the_inputs)
 
     # Second run will be faster as it uses the cached artifacts
     if 0 == world_rank:
-        dprint(f"Running model: Second Time...")
+        dprint("Running model: Second Time...")
     the_outputs = the_compiled_model(the_inputs)
 
     # -------------
     # Cleanup
     # -------------
     if 0 == world_rank:
-        dprint(f"Done")
+        dprint("Done")
     if is_distributed:
         torch.distributed.barrier()
         torch.distributed.destroy_process_group()

scripts/validation.py

@@ -19,7 +19,6 @@
     LogitsExtractorHook,
     capture_level_1_metrics,
     extract_validation_information,
-    StaticTokenInjectorHook,
     GoldenTokenHook,
     filter_failed_level_1_cases,
     validate_level_0,
@@ -264,7 +263,7 @@
 if args.quantization == "gptq":
     try:
         # validation script always loads AIU addon
-        from fms_mo.aiu_addons.gptq import gptq_aiu_adapter, gptq_aiu_linear
+        from fms_mo.aiu_addons.gptq import gptq_aiu_adapter, gptq_aiu_linear  # noqa: F401
 
         print("Loaded `aiu_addons` functionalities")
 
@@ -295,7 +294,6 @@
     aiu_setup.aiu_dist_setup(dist.get_rank(), dist.get_world_size())
 
 # Always initialize AIU in this script
-from torch_sendnn import torch_sendnn
 
 if not args.distributed:
     aiu_setup.aiu_setup(rank, world_size)
@@ -595,7 +593,7 @@ def truncate_prompts_to_max_length(prompts, max_len, max_allowed_length):
 
 if args.fixed_prompt_length != 0 and args.fixed_prompt_length < max_len:
     dprint(
-        f"One or more prompts require truncation. Truncation has been disabled as fixed_prompt_length has been set."
+        "One or more prompts require truncation. Truncation has been disabled as fixed_prompt_length has been set."
     )
     exit(1)
 prompts = truncate_prompts_to_max_length(prompts, max_len, max_allowed_length)
@@ -645,7 +643,7 @@ def print_result(result, result_idx: int = 0, file_prefix: str = ""):
         tokenizer,
     )
 
-    val_tokens = [torch.tensor(l) for l in validation_info.get_info("tokens")]
+    val_tokens = [torch.tensor(_) for _ in validation_info.get_info("tokens")]
     max_val_len = max([prompt.size(0) for prompt in val_tokens])
     val_num_gen_tokens = int(args.max_new_tokens)
     if max_allowed_length is not None:

tests/models/conftest.py

@@ -2,7 +2,6 @@
 
 from aiu_fms_testing_utils.utils.aiu_setup import aiu_setup, rank, world_size
 import os
-import pytest
 
 
 def pytest_sessionstart(session):

tests/models/test_decoders.py

@@ -28,7 +28,7 @@
 import os
 
 try:
-    from fms_mo.aiu_addons.gptq import gptq_aiu_adapter, gptq_aiu_linear
+    from fms_mo.aiu_addons.gptq import gptq_aiu_adapter, gptq_aiu_linear  # noqa
 
     GPTQ_ENABLED = True
 except ImportError:
@@ -278,11 +278,11 @@ def __find_eos_index(reference_tokens, eos_token_id, seq_length, max_new_tokens)
     return result
 
 
-def __filter_before_eos(l, filter_indexes):
+def __filter_before_eos(metrics, filter_indexes):
     from itertools import groupby
 
     filtered_results = [
-        list(g)[: filter_indexes[k]] for k, g in groupby(l, key=lambda x: x[0])
+        list(g)[: filter_indexes[k]] for k, g in groupby(metrics, key=lambda x: x[0])
     ]
     return [item for sublist in filtered_results for item in sublist]