Skip to content

Commit b16d543

Browse files
authored
Update convert_grok.py to use logging module
1 parent c9c8952 commit b16d543

File tree

1 file changed

+13
-7
lines changed

1 file changed

+13
-7
lines changed

convert_grok.py

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@
3434

3535
import gguf
3636

37+
logger = logging.getLogger("convert_grok")
38+
3739
GGML_QK8_0 = 32
3840
GGML_QK4_0 = 32
3941
GGML_QK4_1 = 32
@@ -214,7 +216,7 @@ def dump_state_dict(f, ggml_type, input_dir, config):
214216
tensor_ggml_type,
215217
)
216218
weights[name] = weight, scales
217-
print("Loaded", len(weight_names), "files")
219+
logger.info("Loaded", len(weight_names), "files")
218220

219221
f.write_header_to_file()
220222
f.write_kv_data_to_file()
@@ -230,7 +232,7 @@ def dump_state_dict(f, ggml_type, input_dir, config):
230232
_, tensor_ggml_type = get_dtype_and_ggml_type(tensor, ggml_type)
231233
array = maybe_quantize_tensor(tensor, tensor_ggml_type).numpy()
232234

233-
print(
235+
logger.debug(
234236
f"dumping {name}:",
235237
f"{tensor_ggml_type.name}/{array.dtype}, {list(tensor.shape)}, {array.nbytes} bytes",
236238
)
@@ -239,12 +241,12 @@ def dump_state_dict(f, ggml_type, input_dir, config):
239241
tensor_info.append((name, list(tensor.shape), tensor_ggml_type.name))
240242

241243
try:
242-
print(tabulate(tensor_info, headers=["name", "shape", "dtype"], tablefmt="psql"))
244+
print(tabulate(tensor_info, headers=["name", "shape", "dtype"], tablefmt="psql")) # noqa: NP100
243245
except NameError:
244246
pass
245247

246248
if len(tensor_info) != len(weight_names):
247-
print("Warning: not all tensors are converted")
249+
logger.warning("Not all tensors are converted")
248250

249251

250252
def from_numpy(array):
@@ -377,7 +379,7 @@ def ffn_size(emb_size, widening_factor):
377379
config.num_experts = len(config.experts)
378380

379381
assert config.num_experts >= 2, "need at least 2 experts"
380-
print("experts to export:", config.experts)
382+
logger.info("experts to export:", config.experts)
381383

382384
f = gguf.GGUFWriter(args.save_path, "grok", endianess=gguf.GGUFEndian.LITTLE)
383385

@@ -409,12 +411,12 @@ def ffn_size(emb_size, widening_factor):
409411

410412
delta = time.time() - start
411413

412-
print(f"grok GGUF model saved to {args.save_path}. Total time {delta:.2f} sec")
414+
logger.info(f"grok GGUF model saved to {args.save_path}. Total time {delta:.2f} sec")
413415

414416

415417
def load_vocab(path):
416418
def load_spm(p):
417-
print(f"Loading vocab file {p}")
419+
logger.info(f"Loading vocab file {p}")
418420
return SentencePieceVocab(p)
419421

420422
# Be extra-friendly and accept either a file or a directory. Also, if it's
@@ -445,8 +447,12 @@ def main():
445447
)
446448
parser.add_argument("--vocab_dir", type=str, default="")
447449
parser.add_argument("--experts", type=str, default="")
450+
parser.add_argument("--verbose", action="store_true", help="increase output verbosity")
451+
448452
args = parser.parse_args()
449453

454+
logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)
455+
450456
vocab = load_vocab(
451457
pathlib.Path(args.vocab_dir) if args.vocab_dir else pathlib.Path(args.input_dir)
452458
)

0 commit comments

Comments
 (0)