Allow multiple -v for all scripts and standardize logger->LOG

Author: cg123

mergekit/architecture/__init__.py

@@ -22,7 +22,7 @@
 if TYPE_CHECKING:
     from mergekit.config import MergeConfiguration
 
-logger = logging.getLogger(__name__)
+LOG = logging.getLogger(__name__)
 
 
 def arch_info_for_config(config: PretrainedConfig) -> Optional[ModelArchitecture]:
@@ -44,11 +44,11 @@ def arch_info_for_config(config: PretrainedConfig) -> Optional[ModelArchitecture
         for c in candidates:
             if c.expected_model_type == config.model_type:
                 return c
-        logger.warning(
+        LOG.warning(
             f"Multiple architectures for {arch_name}, none match model type {config.model_type}"
         )
 
-    logger.warning(f"No JSON architecture found for {arch_name}")
+    LOG.warning(f"No JSON architecture found for {arch_name}")
     return None
 
 

mergekit/architecture/auto.py

@@ -21,7 +21,7 @@
 
 RE_LAYER_INDEX = re.compile(r"\.(\d+)\.")
 
-logger = logging.getLogger(__name__)
+LOG = logging.getLogger(__name__)
 
 
 def get_model_tensor_names(model: ModelReference, options: MergeOptions) -> List[str]:

mergekit/evo/actors.py

@@ -41,7 +41,7 @@
 from mergekit.options import MergeOptions
 from mergekit.plan import MergePlanner
 
-logger = logging.getLogger(__name__)
+LOG = logging.getLogger(__name__)
 
 
 class MergeActorBase:
@@ -91,18 +91,18 @@ def evaluate_genotype(
     ) -> dict:
         gc.collect()
         torch.cuda.empty_cache()
-        logger.info("Merging model")
+        LOG.info("Merging model")
         merged_path = merge_model(
             genotype, self.genome, self.model_storage_path, self.merge_options
         )
         if not merged_path:
-            logger.error("Model merge failed")
+            LOG.error("Model merge failed")
             return {"score": None, "results": None}
 
         model_kwargs = {}
         if self.quantization_config is not None:
             model_kwargs["quantization_config"] = self.quantization_config
-        logger.info(f"Model merged to {merged_path}")
+        LOG.info(f"Model merged to {merged_path}")
         return evaluate_model(
             merged_path,
             self.config.tasks,
@@ -167,7 +167,7 @@ def _maybe_init_model(self, config: MergeConfiguration):
                     continue
 
                 if getattr(cfg_out, key) != getattr(self.arch_info.config, key, None):
-                    logger.warning(f"Config key {key} changed, reinitializing model")
+                    LOG.warning(f"Config key {key} changed, reinitializing model")
                     different = True
                     break
 
@@ -206,7 +206,7 @@ def _maybe_init_model(self, config: MergeConfiguration):
                 del inner_model
                 tokenizer_donor = self.genome.definition.base_model
                 if tokenizer_donor is None:
-                    logger.warning(
+                    LOG.warning(
                         "Base model not set, using tokenizer from first model in genome"
                     )
                     tokenizer_donor = self.genome.definition.models[0]
@@ -224,7 +224,7 @@ def _maybe_init_model(self, config: MergeConfiguration):
                     max_model_len = min(max_model_len or 1024, window_sz)
                 if max_model_len and max_model_len > 8192:
                     max_model_len = 8192
-                    logger.warning(f"Clipping sequence length to {max_model_len}")
+                    LOG.warning(f"Clipping sequence length to {max_model_len}")
 
                 mem_util = (
                     0.7 if self.merge_options.cuda else 0.9
@@ -248,13 +248,13 @@ def _maybe_init_model(self, config: MergeConfiguration):
             if ai
             else None
         )
-        logger.info("Model initialized")
+        LOG.info("Model initialized")
 
     def evaluate(self, genotype: torch.Tensor) -> dict:
         try:
             config = self.genome.genotype_merge_config(genotype)
         except InvalidGenotypeError as e:
-            logger.error("Invalid genotype", exc_info=e)
+            LOG.error("Invalid genotype", exc_info=e)
             return {"score": None, "results": None}
 
         self._maybe_init_model(config)

mergekit/io/tensor_writer.py

@@ -10,7 +10,7 @@
 import safetensors
 import torch
 
-logger = logging.getLogger(__name__)
+LOG = logging.getLogger(__name__)
 
 
 class TensorWriter:
@@ -65,7 +65,7 @@ def _flush_current_shard(self):
         if not self.current_shard:
             return
 
-        logger.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}"
@@ -87,7 +87,7 @@ def finalize(self):
         with self.lock:
             self._flush_current_shard()
 
-            logger.info("Finalizing shard names")
+            LOG.info("Finalizing shard names")
 
             prefix, extension = self._get_name_components()
 
@@ -154,7 +154,7 @@ def _do_save():
                 and isinstance(e.args[0], str)
                 and "share memory" in e.args[0]
             ):
-                logger.warning(
+                LOG.warning(
                     "Your model has duplicated tensors but the --clone-tensors "
                     "flag is not set."
                 )

mergekit/merge.py

@@ -24,7 +24,7 @@
 from mergekit.plan import MergePlanner
 from mergekit.tokenizer import TokenizerInfo
 
-logger = logging.getLogger(__name__)
+LOG = logging.getLogger(__name__)
 
 
 def run_merge(
@@ -61,7 +61,7 @@ def run_merge(
         loader_cache.get(model)
     del pbar
 
-    logger.info("Planning operations")
+    LOG.info("Planning operations")
     targets = MergePlanner(
         merge_config,
         arch_info,
@@ -94,7 +94,7 @@ def run_merge(
             cfg_out, arch_info, tokenizer, pad_to_multiple_of=pad_to_multiple_of
         )
 
-    logger.info("Saving config")
+    LOG.info("Saving config")
     cfg_out.save_pretrained(out_path)
 
     if options.write_model_card:
@@ -115,7 +115,7 @@ def run_merge(
             fp.write(config_source)
 
     if tokenizer is not None:
-        logger.info("Saving tokenizer")
+        LOG.info("Saving tokenizer")
         _set_chat_template(tokenizer, merge_config)
         tokenizer.save_pretrained(out_path, safe_serialization=True)
     else:
@@ -125,12 +125,12 @@ def run_merge(
                     merge_config, out_path, trust_remote_code=options.trust_remote_code
                 )
             except Exception as e:
-                logger.error(
+                LOG.error(
                     "Failed to copy tokenizer. The merge was still successful, just copy it from somewhere else.",
                     exc_info=e,
                 )
         elif merge_config.chat_template:
-            logger.warning(
+            LOG.warning(
                 "Chat template specified but no tokenizer found. Chat template will not be saved."
             )
 
@@ -180,13 +180,13 @@ def _set_chat_template(
                 if template:
                     model_templates.append(template.strip())
             except Exception as e:
-                logger.warning(f"Unable to load tokenizer for {model}", exc_info=e)
+                LOG.warning(f"Unable to load tokenizer for {model}", exc_info=e)
 
         if not model_templates:
             return
 
         chat_template = Counter(model_templates).most_common(1)[0][0]
-        logger.info(f"Auto-selected chat template: {chat_template}")
+        LOG.info(f"Auto-selected chat template: {chat_template}")
 
     elif importlib.resources.is_resource(chat_templates, chat_template + ".jinja"):
         with importlib.resources.open_text(
@@ -210,7 +210,7 @@ def _copy_tagalong_files(
 
     for file_name in files:
         if os.path.exists(os.path.join(donor_model.model.path, file_name)):
-            logger.info(f"Copying {file_name} from {donor_model}")
+            LOG.info(f"Copying {file_name} from {donor_model}")
             shutil.copy(
                 os.path.join(donor_model.model.path, file_name),
                 os.path.join(out_path, file_name),
@@ -234,7 +234,7 @@ def _copy_tokenizer(
             or os.path.exists(os.path.join(donor_model.model.path, "tokenizer.model"))
         )
     ):
-        logger.info(f"Copying tokenizer from {donor_model}")
+        LOG.info(f"Copying tokenizer from {donor_model}")
 
         for file_name in [
             "tokenizer_config.json",
@@ -253,7 +253,7 @@ def _copy_tokenizer(
         return
 
     # fallback: try actually loading the tokenizer and saving it
-    logger.info(f"Reserializing tokenizer from {donor_model}")
+    LOG.info(f"Reserializing tokenizer from {donor_model}")
     tokenizer = transformers.AutoTokenizer.from_pretrained(
         donor_model.model.path,
         revision=donor_model.model.revision,
@@ -299,7 +299,7 @@ def _model_out_config(
                 cfg_key = module_info.architecture.num_layers_config_key()
                 set_config_value(res, cfg_key, module_layers[module_name])
             except Exception as e:
-                logger.warning(
+                LOG.warning(
                     f"Unable to set number of layers for module {module_name} in output config "
                     "- you may need to manually correct it.",
                     exc_info=e,
@@ -322,7 +322,7 @@ def _update_config_vocab(
             config, arch_info.vocab_size_config_key or "vocab_size", vocab_size
         )
     except Exception as e:
-        logger.warning(
+        LOG.warning(
             "Unable to set vocabulary size in output config - you may need to manually correct it.",
             exc_info=e,
         )

mergekit/multigpu_executor.py

@@ -29,7 +29,7 @@
     build_schedule,
 )
 
-logger = logging.getLogger(__name__)
+LOG = logging.getLogger(__name__)
 
 
 class MultiGPUExecutor:
@@ -61,7 +61,7 @@ def __init__(
 
         if num_gpus is None:
             num_gpus = torch.cuda.device_count()
-        logger.info(f"Using {num_gpus} GPUs for parallel execution")
+        LOG.info(f"Using {num_gpus} GPUs for parallel execution")
 
         self.universe = TaskUniverse(tasks)
         self.targets = set([self.universe.get_handle(t) for t in tasks])
@@ -82,7 +82,7 @@ def __init__(
             for t in ordered_handles
             if (t not in trailing_tasks and t not in leading_tasks)
         ]
-        logger.info(
+        LOG.info(
             f"Task breakdown: {len(self.leading_main_handles)} leading, "
             f"{len(parallel_handles)} parallel, "
             f"{len(self.trailing_main_handles)} trailing"
@@ -244,7 +244,7 @@ def _assign_islands_to_gpus(
         island_graph.add_nodes_from([t._index for t in tasks])
         island_graph.add_edges_from(edge_list)
         islands: List[Set[int]] = list(nx.weakly_connected_components(island_graph))
-        logger.info(f"Found {len(islands)} islands in parallel task graph")
+        LOG.info(f"Found {len(islands)} islands in parallel task graph")
         assignments: Dict[torch.device, List[int]] = {}
         for island in islands:
             if not island:
@@ -295,4 +295,4 @@ def _device_worker(
                         result = None
                     self.task_completion_queue.put((task_handle._index, result))
         torch.cuda.synchronize(device=device)
-        logger.debug(f"Device {device} done")
+        LOG.debug(f"Device {device} done")

mergekit/options.py

@@ -30,15 +30,23 @@ class MergeOptions(BaseModel, frozen=True):
     lazy_unpickle: bool = False
     write_model_card: bool = True
     safe_serialization: bool = True
-    verbose: bool = False
+    verbosity: int = 0
     quiet: bool = False
     read_to_gpu: bool = False
     multi_gpu: bool = False
     num_threads: Optional[int] = None
     gpu_rich: bool = False
 
     def apply_global_options(self):
-        logging.basicConfig(level=logging.INFO if self.verbose else logging.WARNING)
+        if self.verbosity > 1:
+            log_level = logging.DEBUG
+        elif self.verbosity == 1:
+            log_level = logging.INFO
+        else:
+            log_level = logging.WARNING
+        logging.basicConfig(level=log_level)
+        if self.verbosity > 5:
+            logging.debug("whoah buddy that's a lot of verbosity, two is plenty")
         if self.random_seed is not None:
             transformers.trainer_utils.set_seed(self.random_seed)
         if self.num_threads is not None:
@@ -74,7 +82,7 @@ def handle_gpu_rich(cls, value):
     "read_to_gpu": "Read model weights directly to GPU",
     "multi_gpu": "Use multi-gpu parallel graph execution engine",
     "num_threads": "Number of threads to use for parallel CPU operations",
-    "verbose": "Enable verbose logging",
+    "verbosity": "Verbose logging (repeat for more verbosity)",
     "gpu_rich": "Alias for --cuda --low-cpu-memory --read-to-gpu --multi-gpu",
 }
 
@@ -96,7 +104,7 @@ def handle_gpu_rich(cls, value):
     "trust_remote_code": "Dangerous Options",
     "allow_crimes": "Dangerous Options",
     "random_seed": "Miscellaneous",
-    "verbose": "Miscellaneous",
+    "verbosity": "Miscellaneous",
     "quiet": "Miscellaneous",
     "lora_merge_dtype": "Miscellaneous",
 }
@@ -141,8 +149,9 @@ def wrapper(*args, **kwargs):
             arg_str = f"--{arg_name}"
         param_decls = [arg_str]
         kwargs = {}
-        if field_name == "verbose":
-            param_decls = ["--verbose/--no-verbose", "-v"]
+        if field_name == "verbosity":
+            param_decls = ["-v", "verbosity"]
+            kwargs["count"] = True
         if field_name == "num_threads":
             param_decls = ["--num-threads", "-j"]
         if field_name == "gpu_rich":
@@ -155,7 +164,7 @@ def wrapper(*args, **kwargs):
             type=field_type,
             default=info.default,
             help=help_str,
-            show_default=field_name != "out_shard_size",
+            show_default=field_name not in ("out_shard_size", "verbosity"),
             **kwargs,
         )(wrapper)