diff --git a/scripts/eval_loop.py b/scripts/eval_loop.py index d0ac0b7..698ae9d 100644 --- a/scripts/eval_loop.py +++ b/scripts/eval_loop.py @@ -13,7 +13,8 @@ from pathlib import Path import torch -from loop_common import (DIRECT_SUFFIX, BandLooper, MergeAdapter, chat_prompt, +from loop_common import (DIRECT_SUFFIX, AdaptiveMergeAdapter, BandLooper, + MergeAdapter, chat_prompt, last_number, num_eq) import sys sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) @@ -96,6 +97,7 @@ def main(): ap.add_argument("--prompt-only", action="store_true", help="loop the prompt span only (unified regime, fast path)") ap.add_argument("--no-spider", action="store_true") + ap.add_argument("--adaptive", action="store_true") ap.add_argument("--feedforward", action="store_true", help="no-recurrence control arm: adapter(e,e) once") args = ap.parse_args() @@ -104,7 +106,8 @@ def main(): model, tok = load_model(dtype=torch.bfloat16) tok.padding_side = "left" looper = BandLooper(model) - adapter = MergeAdapter( + cls = AdaptiveMergeAdapter if args.adaptive else MergeAdapter + adapter = cls( d=model.config.get_text_config().hidden_size).cuda() if args.adapter: adapter.load_state_dict(torch.load(args.adapter, map_location="cuda"))