From 5bc4f10ee97abf38730d33fa88ca9d25066670d3 Mon Sep 17 00:00:00 2001 From: Brian Date: Fri, 10 May 2024 01:17:56 +1000 Subject: [PATCH] Update convert_grok.py to use logging module --- convert_grok.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/convert_grok.py b/convert_grok.py index f3aa6d369..aba0793b2 100644 --- a/convert_grok.py +++ b/convert_grok.py @@ -34,6 +34,8 @@ if "NO_LOCAL_GGUF" not in os.environ: import gguf +logger = logging.getLogger("convert_grok") + GGML_QK8_0 = 32 GGML_QK4_0 = 32 GGML_QK4_1 = 32 @@ -214,7 +216,7 @@ def dump_state_dict(f, ggml_type, input_dir, config): tensor_ggml_type, ) weights[name] = weight, scales - print("Loaded", len(weight_names), "files") + logger.info("Loaded", len(weight_names), "files") f.write_header_to_file() f.write_kv_data_to_file() @@ -230,7 +232,7 @@ def dump_state_dict(f, ggml_type, input_dir, config): _, tensor_ggml_type = get_dtype_and_ggml_type(tensor, ggml_type) array = maybe_quantize_tensor(tensor, tensor_ggml_type).numpy() - print( + logger.debug( f"dumping {name}:", f"{tensor_ggml_type.name}/{array.dtype}, {list(tensor.shape)}, {array.nbytes} bytes", ) @@ -239,12 +241,12 @@ def dump_state_dict(f, ggml_type, input_dir, config): tensor_info.append((name, list(tensor.shape), tensor_ggml_type.name)) try: - print(tabulate(tensor_info, headers=["name", "shape", "dtype"], tablefmt="psql")) + print(tabulate(tensor_info, headers=["name", "shape", "dtype"], tablefmt="psql")) # noqa: NP100 except NameError: pass if len(tensor_info) != len(weight_names): - print("Warning: not all tensors are converted") + logger.warning("Not all tensors are converted") def from_numpy(array): @@ -377,7 +379,7 @@ def convert_grok(args, vocab, ggml_type): config.num_experts = len(config.experts) assert config.num_experts >= 2, "need at least 2 experts" - print("experts to export:", config.experts) + logger.info("experts to export:", config.experts) f = gguf.GGUFWriter(args.save_path, "grok", endianess=gguf.GGUFEndian.LITTLE) @@ -409,12 +411,12 @@ def convert_grok(args, vocab, ggml_type): delta = time.time() - start - print(f"grok GGUF model saved to {args.save_path}. Total time {delta:.2f} sec") + logger.info(f"grok GGUF model saved to {args.save_path}. Total time {delta:.2f} sec") def load_vocab(path): def load_spm(p): - print(f"Loading vocab file {p}") + logger.info(f"Loading vocab file {p}") return SentencePieceVocab(p) # Be extra-friendly and accept either a file or a directory. Also, if it's @@ -445,8 +447,12 @@ def main(): ) parser.add_argument("--vocab_dir", type=str, default="") parser.add_argument("--experts", type=str, default="") + parser.add_argument("--verbose", action="store_true", help="increase output verbosity") + args = parser.parse_args() + logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO) + vocab = load_vocab( pathlib.Path(args.vocab_dir) if args.vocab_dir else pathlib.Path(args.input_dir) )