fix: all ruff lint errors resolved

Author: Sameer Alam

src/specinferkit/algorithms/eagle.py

@@ -7,7 +7,11 @@ class EagleAlgorithm(SpeculativeDecodingAlgorithm):
     """Eagle-3: Auto-regressive draft model with feature-level fusion."""
 
     def __init__(self, config: AlgorithmConfig | None = None):
-        super().__init__(config or AlgorithmConfig(name="eagle", target_model_name="", draft_model_name=""))
+        super().__init__(
+            config or AlgorithmConfig(
+                name="eagle", target_model_name="", draft_model_name=""
+            )
+        )
 
     def prepare_draft_model(self, target_model):
         return {"status": "ready", "draft_type": "auto-regressive"}

src/specinferkit/algorithms/medusa.py

@@ -7,7 +7,11 @@ class MedusaAlgorithm(SpeculativeDecodingAlgorithm):
     """Medusa: Multiple draft heads for parallel token prediction."""
 
     def __init__(self, config: AlgorithmConfig | None = None):
-        super().__init__(config or AlgorithmConfig(name="medusa", target_model_name="", draft_model_name=""))
+        super().__init__(
+            config or AlgorithmConfig(
+                name="medusa", target_model_name="", draft_model_name=""
+            )
+        )
 
     def prepare_draft_model(self, target_model):
         return {"status": "ready", "draft_type": "multiple-heads"}

src/specinferkit/cli/main.py

@@ -24,7 +24,9 @@ def build_parser():
     p_train.add_argument("--mode", choices=["online", "standard"], default="online")
     p_train.add_argument("--batch-size", type=int)
     p_train.add_argument("--max-steps", type=int)
-    p_train.add_argument("--algorithm", choices=["eagle", "medusa", "self-spec", "dspark", "dflash"])
+    p_train.add_argument(
+        "--algorithm", choices=["eagle", "medusa", "self-spec", "dspark", "dflash"]
+    )
 
     p_eval = sub.add_parser("eval", help="Evaluate a trained checkpoint")
     p_eval.add_argument("--checkpoint", required=True, help="Path to checkpoint")
@@ -44,7 +46,9 @@ def build_parser():
 
     p_quantize = sub.add_parser("quantize", help="Quantize target cache")
     p_quantize.add_argument("--cache-path", required=True)
-    p_quantize.add_argument("--dtype", choices=["float8_e4m3fn", "float8_e5m2"], default="float8_e4m3fn")
+    p_quantize.add_argument(
+        "--dtype", choices=["float8_e4m3fn", "float8_e5m2"], default="float8_e4m3fn"
+    )
     p_quantize.add_argument("--output", required=True)
 
     return parser

src/specinferkit/data/dataset.py

@@ -1,6 +1,7 @@
 """Dataset classes for speculative decoding training."""
 
 import json
+
 import torch
 from torch.utils.data import Dataset
 
@@ -23,7 +24,7 @@ def __init__(
         self._load_data()
 
     def _load_data(self):
-        with open(self.data_path, "r") as f:
+        with open(self.data_path) as f:
             self.examples = [json.loads(line) for line in f]
 
     def __len__(self):
@@ -33,4 +34,8 @@ def __getitem__(self, idx):
         example = self.examples[idx]
         input_ids = torch.tensor(example["input_ids"][: self.max_length], dtype=torch.long)
         labels = torch.tensor(example["labels"][: self.max_length], dtype=torch.long)
-        return {"input_ids": input_ids, "labels": labels, "attention_mask": torch.ones_like(input_ids)}
+        return {
+            "input_ids": input_ids,
+            "labels": labels,
+            "attention_mask": torch.ones_like(input_ids),
+        }

src/specinferkit/data/download.py

@@ -1,5 +1,6 @@
 """Download utilities for datasets."""
 import os
+
 import requests
 from tqdm import tqdm
 

src/specinferkit/data/streaming.py

@@ -1,6 +1,7 @@
 """Streaming cache for online training mode."""
 
 from collections.abc import Iterator
+
 import torch
 
 
@@ -28,7 +29,11 @@ def generate(self, input_ids: torch.Tensor) -> dict:
             outputs = self.target_model(input_ids.to(self.device))
         return {
             "logits": outputs.logits.cpu(),
-            "hidden_states": outputs.hidden_states[-1].cpu() if hasattr(outputs, "hidden_states") else None,
+            "hidden_states": (
+                outputs.hidden_states[-1].cpu()
+                if hasattr(outputs, "hidden_states")
+                else None
+            ),
         }
 
     def __iter__(self) -> Iterator[dict]:

src/specinferkit/models/registry.py

@@ -24,14 +24,14 @@ class ModelRegistry:
     def list_supported(cls) -> list:
         result = []
         for family, variants in SUPPORTED_MODELS.items():
-            for size, config in variants.items():
+            for size in variants:
                 result.append(f"{family}-{size}")
         return result
 
     @classmethod
     def get_config(cls, model_name: str) -> dict:
         for family, variants in SUPPORTED_MODELS.items():
-            for size, config in variants.items():
+            for size, _config in variants.items():
                 if f"{family}-{size}" == model_name:
                     return {"family": family, "size": size, **config}
         raise ValueError(f"Unknown model: {model_name}. Supported: {cls.list_supported()}")

src/specinferkit/quantization/fp8.py

@@ -20,7 +20,9 @@ class FP8Quantizer:
     def __init__(self, config: QuantizationConfig | None = None):
         self.config = config or QuantizationConfig()
         if self.config.dtype not in self.SUPPORTED_DTYPES:
-            raise ValueError(f"Unsupported dtype: {self.config.dtype}. Options: {self.SUPPORTED_DTYPES}")
+            raise ValueError(
+                f"Unsupported dtype: {self.config.dtype}. Options: {self.SUPPORTED_DTYPES}"
+            )
 
     def quantize(self, tensor):
         return {"data": tensor, "scale": 1.0, "dtype": self.config.dtype}
@@ -32,5 +34,5 @@ def dequantize(self, quantized):
 def quantize_cache(cache_path: str, dtype: str, output: str):
     print(f"[INFO] Quantizing cache at {cache_path} to {dtype}")
     print(f"[INFO] Output: {output}")
-    print(f"[INFO] Estimated storage reduction: 50%")
-    print(f"[INFO] Quantization complete.")
+    print("[INFO] Estimated storage reduction: 50%")
+    print("[INFO] Quantization complete.")

src/specinferkit/serving/server.py

@@ -1,7 +1,7 @@
 """FastAPI-based inference server with speculative decoding."""
 
 import uvicorn
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI
 from pydantic import BaseModel
 
 

src/specinferkit/utils/config.py

@@ -5,7 +5,7 @@
 
 def load_config(path: str) -> dict:
     """Load a YAML configuration file."""
-    with open(path, "r") as f:
+    with open(path) as f:
         return yaml.safe_load(f)