Formatting

Author: cg123

mergekit/common.py

@@ -39,7 +39,7 @@ def set_config_value(config: PretrainedConfig, key: str, value: Any):
     for idx, part in enumerate(parts[:-1]):
         if not hasattr(obj, part):
             raise RuntimeError(
-                f"Config {config} has no attribute {'.'.join(parts[:idx+1])}"
+                f"Config {config} has no attribute {'.'.join(parts[: idx + 1])}"
             )
         obj = getattr(obj, part)
     setattr(obj, parts[-1], value)
@@ -52,7 +52,7 @@ def get_config_value(config: PretrainedConfig, key: str) -> Any:
     for idx, part in enumerate(parts):
         if not hasattr(obj, part):
             raise RuntimeError(
-                f"Config {config} has no attribute {'.'.join(parts[:idx+1])}"
+                f"Config {config} has no attribute {'.'.join(parts[: idx + 1])}"
             )
         obj = getattr(obj, part)
     return obj

mergekit/io/tasks.py

@@ -64,7 +64,7 @@ def _normalized_shard_name(path: str) -> int:
     name = name.lower().replace("pytorch_model", "model")
     if m := shard_name_re.search(name):
         frac = int(m.group(1)) / int(m.group(2))
-        name = f"model-{int(frac*100):03d}pct"
+        name = f"model-{int(frac * 100):03d}pct"
     return name
 
 

mergekit/io/tensor_writer.py

@@ -65,10 +65,10 @@ def _flush_current_shard(self):
         if not self.current_shard:
             return
 
-        LOG.info(f"Writing shard #{self.shards_written+1} to disk")
+        LOG.info(f"Writing shard #{self.shards_written + 1} to disk")
 
         prefix, extension = self._get_name_components()
-        shard_name = f"{prefix}-{self.shards_written+1}.{extension}"
+        shard_name = f"{prefix}-{self.shards_written + 1}.{extension}"
 
         for key in self.current_shard:
             self.weight_map[key] = shard_name
@@ -95,8 +95,8 @@ def finalize(self):
             total_shards = self.shards_written
             name_remap = {}
             for idx in range(total_shards):
-                name_remap[f"{prefix}-{idx+1}.{extension}"] = (
-                    f"{prefix}-{idx+1:05d}-of-{total_shards:05d}.{extension}"
+                name_remap[f"{prefix}-{idx + 1}.{extension}"] = (
+                    f"{prefix}-{idx + 1:05d}-of-{total_shards:05d}.{extension}"
                 )
 
             if total_shards < 2:

mergekit/merge_methods/easy_define.py

@@ -167,7 +167,7 @@ def _execute(self, tensors: Dict[ModelReference, torch.Tensor], **_kwargs):
 
     tt_fields["execute"] = _execute
 
-    tt_name = f"{name.title().replace(' ','')}MergeTask"
+    tt_name = f"{name.title().replace(' ', '')}MergeTask"
     tt_cls = pydantic.create_model(tt_name, __base__=Task[torch.Tensor], **tt_fields)
 
     mm_fields = {}
@@ -220,7 +220,7 @@ def _parameters(self) -> List[ConfigParameterDef]:
 
     mm_fields["parameters"] = _parameters
 
-    mm_name = f"{name.title().replace(' ','')}MergeMethod"
+    mm_name = f"{name.title().replace(' ', '')}MergeMethod"
     mm_cls = type(mm_name, (MergeMethod,), mm_fields)
     REGISTERED_MERGE_METHODS[name] = mm_cls()
     return func

mergekit/multigpu_executor.py

@@ -21,8 +21,6 @@
 import torch
 import tqdm
 
-from mergekit.io.tasks import TensorWriterTask
-
 from .graph import (
     Executor,
     Task,
@@ -101,7 +99,7 @@ def __init__(
             offending = [
                 t.task() for t in parallel_handles if t.task().main_thread_only()
             ]
-            logging.error(f"Main-thread-only tasks in parallel section:")
+            logging.error("Main-thread-only tasks in parallel section:")
             for task in offending:
                 logging.error(f"  {type(task).__name__}")
             raise RuntimeError(

mergekit/scripts/tokensurgeon.py

@@ -532,7 +532,7 @@ def build_embedding_matrix(
     )
     shared_numeric_tokens = set(orig_vocab.keys()) & set(donor_numeric_tokens)
     LOG.debug(
-        f"{len(shared_numeric_tokens)} shared numeric tokens ({100.0*len(shared_numeric_tokens)/len(donor_numeric_tokens):.2f}%)"
+        f"{len(shared_numeric_tokens)} shared numeric tokens ({100.0 * len(shared_numeric_tokens) / len(donor_numeric_tokens):.2f}%)"
     )
     LOG.debug(
         [donor_tokenizer.decode([donor_vocab[tok]]) for tok in shared_numeric_tokens]

mergekit/tokenizer/build.py

@@ -257,7 +257,7 @@ def build_tokenizer(
             orig_idx = model_vocab[tok]
             if orig_idx >= vocab_size:
                 LOG.warning(
-                    f"{model} token {repr(tok)} has index {orig_idx}>{vocab_size-1} (padding?)"
+                    f"{model} token {repr(tok)} has index {orig_idx}>{vocab_size - 1} (padding?)"
                 )
                 continue
 

mergekit/tokensurgeon/magikarp.py

@@ -70,7 +70,7 @@ def well_trained_tokens(
         ).float()
         threshold = torch.quantile(cos_sim, 1 - quantile, dim=0)
         LOG.debug(
-            f"Unused token threshold in embed_tokens: {threshold.item():.4f} ({int((1-quantile) * 100)}th percentile)"
+            f"Unused token threshold in embed_tokens: {threshold.item():.4f} ({int((1 - quantile) * 100)}th percentile)"
         )
         if threshold < 0.5:
             threshold = 0.5
@@ -89,7 +89,7 @@ def well_trained_tokens(
         ).float()
         threshold = torch.quantile(cos_sim, 1 - quantile, dim=0)
         LOG.debug(
-            f"Unused token threshold in lm_head: {threshold.item():.4f} ({int((1-quantile) * 100)}th percentile)"
+            f"Unused token threshold in lm_head: {threshold.item():.4f} ({int((1 - quantile) * 100)}th percentile)"
         )
         if threshold < 0.5:
             threshold = 0.5