From 16054e72d9b24aeb6953f0f8108aab1f5eb8af08 Mon Sep 17 00:00:00 2001 From: EgrorBs <44816767+EgrorBs@users.noreply.github.com> Date: Sun, 19 Mar 2023 13:01:08 +0300 Subject: [PATCH 1/2] Replace dots on loading with numbers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit I love counting dots, it's cool, but reading numbers is even cooler! 💯 :) --- rwkv_pip_package/src/rwkv/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rwkv_pip_package/src/rwkv/model.py b/rwkv_pip_package/src/rwkv/model.py index b74b259e..ab99581b 100644 --- a/rwkv_pip_package/src/rwkv/model.py +++ b/rwkv_pip_package/src/rwkv/model.py @@ -277,7 +277,7 @@ def __init__(self, model, strategy): print(x.ljust(32), dt.rjust(4), str(w[x].device).rjust(8), shape, ' (pinned)' if w[x].is_pinned() else '') else: print_need_newline = True - print('.', end = '', flush = True) + print(f'\r[{i}/{len(keys)-1}] {x.ljust(32)} ', end = '', flush = True) assert len(keys) == 4 + (4+9+5) * args.n_layer, 'Error: not a RWKV-4 model (4a and 4b models are not supported as of now)' gc.collect() if 'cuda' in args.strategy_string: From 4c7bd893ff4fff66a1b5152a1569c55d73df6e0e Mon Sep 17 00:00:00 2001 From: EgrorBs <44816767+EgrorBs@users.noreply.github.com> Date: Sun, 19 Mar 2023 13:12:59 +0300 Subject: [PATCH 2/2] quick fix --- rwkv_pip_package/src/rwkv/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rwkv_pip_package/src/rwkv/model.py b/rwkv_pip_package/src/rwkv/model.py index ab99581b..82acb4b3 100644 --- a/rwkv_pip_package/src/rwkv/model.py +++ b/rwkv_pip_package/src/rwkv/model.py @@ -180,7 +180,7 @@ def __init__(self, model, strategy): # Load weights print_need_newline = False - for x in keys: + for i, x in enumerate(keys): w[x].requires_grad = False layer_id = int(x.split('.')[1]) if ('blocks.' in x) else 0 if ('ln_out.' in x) or ('head.' in x):