| 157 | # ────────────────────────────────────────────────────────────────────────────── |
| 158 | |
| 159 | def parse_args() -> argparse.Namespace: |
| 160 | p = argparse.ArgumentParser(description="Run inference on a BitNet b1.58 LLaMA model.") |
| 161 | p.add_argument("--model", type=str, required=True, |
| 162 | help="Path to the saved model / checkpoint directory.") |
| 163 | p.add_argument("--prompt", type=str, default=None, |
| 164 | help="Input prompt. Omit for interactive mode.") |
| 165 | p.add_argument("--max-new-tokens", type=int, default=128) |
| 166 | p.add_argument("--temperature", type=float, default=0.8) |
| 167 | p.add_argument("--top-p", type=float, default=0.9) |
| 168 | p.add_argument("--top-k", type=int, default=50) |
| 169 | p.add_argument("--repetition-penalty", type=float, default=1.1) |
| 170 | return p.parse_args() |
| 171 | |
| 172 | |
| 173 | def main() -> None: |