File size: 3,463 Bytes
9c0a4e6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | """Sanity-check script for the ASD custom generation method.
Run from anywhere:
python scripts/example.py \
--model Qwen/Qwen2.5-1.5B-Instruct \
--assistant-model Qwen/Qwen2.5-0.5B-Instruct \
--budget 2.0 --local-ratio 0.25 --max-mismatches 2
Checks:
1. ASD runs and produces text.
2. `assistant_asd_budget=0` is token-identical to native strict assisted decoding
(the gold-standard correctness check).
3. No ASD knob set is also token-identical to native (lossless default).
"""
import argparse
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model", required=True)
parser.add_argument("--assistant-model", required=True)
parser.add_argument("--budget", type=float, default=2.0)
parser.add_argument("--local-ratio", type=float, default=0.25)
parser.add_argument("--max-mismatches", type=int, default=2)
parser.add_argument("--max-new-tokens", type=int, default=64)
parser.add_argument("--prompt", default="Give me a short introduction to large language models.")
parser.add_argument("--repo-path", default=REPO_ROOT,
help="Directory containing custom_generate/ (default: repo root)")
args = parser.parse_args()
tokenizer = AutoTokenizer.from_pretrained(args.model)
model = AutoModelForCausalLM.from_pretrained(args.model, dtype="auto", device_map="auto")
assistant_model = AutoModelForCausalLM.from_pretrained(
args.assistant_model, dtype="auto", device_map="auto"
)
inputs = tokenizer(args.prompt, return_tensors="pt").to(model.device)
gen_kwargs = dict(max_new_tokens=args.max_new_tokens, do_sample=False)
# 1. Native strict assisted decoding baseline.
strict_out = model.generate(**inputs, assistant_model=assistant_model, **gen_kwargs)
print("Strict :", tokenizer.decode(strict_out[0], skip_special_tokens=True))
# 2. ASD relaxed.
asd_out = model.generate(
**inputs,
assistant_model=assistant_model,
custom_generate=args.repo_path,
trust_remote_code=True,
assistant_asd_budget=args.budget,
assistant_asd_local_ratio=args.local_ratio,
assistant_asd_max_mismatches=args.max_mismatches,
**gen_kwargs,
)
print("ASD :", tokenizer.decode(asd_out[0], skip_special_tokens=True))
# 3. Gold check: budget=0 must be token-identical to native strict decoding.
asd_b0_out = model.generate(
**inputs,
assistant_model=assistant_model,
custom_generate=args.repo_path,
trust_remote_code=True,
assistant_asd_budget=0.0,
**gen_kwargs,
)
assert torch.equal(strict_out, asd_b0_out), "ASD(budget=0) differs from strict decoding!"
print("OK: assistant_asd_budget=0 is token-identical to strict verification.")
# 4. Lossless default: no ASD knobs must also be token-identical.
asd_default_out = model.generate(
**inputs,
assistant_model=assistant_model,
custom_generate=args.repo_path,
trust_remote_code=True,
**gen_kwargs,
)
assert torch.equal(strict_out, asd_default_out), "ASD (no knobs) differs from strict decoding!"
print("OK: no ASD knobs is token-identical to strict verification.")
if __name__ == "__main__":
main()
|