From 2d7e0df7a8fa658169b38c5bcdd159e6ffb2e26a Mon Sep 17 00:00:00 2001 From: vyom-hai Date: Thu, 30 Apr 2026 16:51:39 +0000 Subject: [PATCH] Fix --do-sample / --enable-bucketing argparse type=bool bug --- main.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/main.py b/main.py index c8283fb..9109db8 100644 --- a/main.py +++ b/main.py @@ -110,7 +110,14 @@ def parse_args(): parser.add_argument("--top-p", type=float, default=0.95) parser.add_argument("--temperature", type=float, default=0.6) parser.add_argument("--global-topk", type=int) - parser.add_argument("--do-sample", type=bool, default=True) + parser.add_argument( + "--do-sample", + action=argparse.BooleanOptionalAction, + default=True, + help="Enable multinomial sampling on the on-device sampler. " + "Pass --no-do-sample to force greedy/argmax (required for " + "deterministic logit_validation accuracy checks).", + ) parser.add_argument("--dynamic", action="store_true") parser.add_argument("--pad-token-id", type=int, default=2) parser.add_argument("--top-k-kernel-enabled", action="store_true", default=False) @@ -154,7 +161,12 @@ def parse_args(): parser.add_argument("--on-device-sampling", action="store_true") # Bucketing - parser.add_argument("--enable-bucketing", type=bool, default=True) + parser.add_argument( + "--enable-bucketing", + action=argparse.BooleanOptionalAction, + default=True, + help="Enable bucketed compilation. Pass --no-enable-bucketing to disable.", + ) parser.add_argument("--bucket-n-active-tokens", action="store_true") parser.add_argument("--context-encoding-buckets", nargs="+", type=int) parser.add_argument("--prefix-buckets", nargs="+", type=int)