| 28 | ] |
| 29 | |
| 30 | def parse_args(): |
| 31 | parser = argparse.ArgumentParser() |
| 32 | parser.add_argument("--target_name_or_path", type=str, required=True) |
| 33 | parser.add_argument("--draft_name_or_path",type=str,required=True) |
| 34 | parser.add_argument("--max-new-tokens", type=int, default=2048) |
| 35 | parser.add_argument("--temperature", type=float, default=1.0) |
| 36 | parser.add_argument( |
| 37 | "--confidence-threshold", |
| 38 | type=float, |
| 39 | default=0.0, |
| 40 | help=("Confidence-head early-stop threshold. Confidence calibration metrics are collected only when this is 0.0."), |
| 41 | ) |
| 42 | parser.add_argument("--tensorboard-dir", type=str, default=None) |
| 43 | parser.add_argument("--step", type=int, default=None,help=("step for tensorboard logging"),) |
| 44 | parser.add_argument("--seed", type=int, default=980406) |
| 45 | args = parser.parse_args() |
| 46 | args.tasks = list(TASKS) |
| 47 | return args |
| 48 | |
| 49 | |
| 50 | def main(local_rank: int, args): |