adaptive flag for GSM eval

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Nils
2026-07-14 12:33:03 +02:00
co-authored by Claude Fable 5
parent 95be35d20f
commit 12ec01becc
+5 -2
View File
@@ -13,7 +13,8 @@ from pathlib import Path
import torch 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) last_number, num_eq)
import sys import sys
sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
@@ -96,6 +97,7 @@ def main():
ap.add_argument("--prompt-only", action="store_true", ap.add_argument("--prompt-only", action="store_true",
help="loop the prompt span only (unified regime, fast path)") help="loop the prompt span only (unified regime, fast path)")
ap.add_argument("--no-spider", action="store_true") ap.add_argument("--no-spider", action="store_true")
ap.add_argument("--adaptive", action="store_true")
ap.add_argument("--feedforward", action="store_true", ap.add_argument("--feedforward", action="store_true",
help="no-recurrence control arm: adapter(e,e) once") help="no-recurrence control arm: adapter(e,e) once")
args = ap.parse_args() args = ap.parse_args()
@@ -104,7 +106,8 @@ def main():
model, tok = load_model(dtype=torch.bfloat16) model, tok = load_model(dtype=torch.bfloat16)
tok.padding_side = "left" tok.padding_side = "left"
looper = BandLooper(model) looper = BandLooper(model)
adapter = MergeAdapter( cls = AdaptiveMergeAdapter if args.adaptive else MergeAdapter
adapter = cls(
d=model.config.get_text_config().hidden_size).cuda() d=model.config.get_text_config().hidden_size).cuda()
if args.adapter: if args.adapter:
adapter.load_state_dict(torch.load(args.adapter, map_location="cuda")) adapter.load_state_dict(torch.load(args.adapter, map_location="cuda"))