From 7e68150f9e0fd6141951dc558d22409d2c0baac4 Mon Sep 17 00:00:00 2001 From: Degao Chu Date: Wed, 7 Jun 2023 11:25:00 +0800 Subject: [PATCH] Update dataset.py --- RWKV-v4neo/src/dataset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/RWKV-v4neo/src/dataset.py b/RWKV-v4neo/src/dataset.py index 934091321..8f9e803b4 100644 --- a/RWKV-v4neo/src/dataset.py +++ b/RWKV-v4neo/src/dataset.py @@ -57,13 +57,13 @@ def __init__(self, args): elif args.data_type == "numpy": self.data = np.load(args.data_file).astype("int") self.vocab_size = args.vocab_size - rank_zero_info("Current vocab size =", self.vocab_size, "(make sure it's correct)") + rank_zero_info(f"Current vocab size = {self.vocab_size} (make sure it's correct)") self.data_size = len(self.data) rank_zero_info(f"Data has {self.data_size} tokens.") elif args.data_type == "uint16": self.data = np.fromfile(args.data_file, dtype=np.uint16).astype("int32").reshape(-1, args.my_sample_len) self.vocab_size = args.vocab_size - rank_zero_info("Current vocab size =", self.vocab_size, "(make sure it's correct)") + rrank_zero_info(f"Current vocab size = {self.vocab_size} (make sure it's correct)") self.data_size = self.data.shape[0] rank_zero_info(f"Data has {self.data_size} samples.") elif args.data_type == "wds_img":