| 21 | |
| 22 | |
| 23 | def build_parser(): |
| 24 | parser = argparse.ArgumentParser(description="LoRA or QLoRA finetuning.") |
| 25 | parser.add_argument( |
| 26 | "--model", |
| 27 | default="mlx_model", |
| 28 | help="The path to the local model directory or Hugging Face repo.", |
| 29 | ) |
| 30 | # Generation args |
| 31 | parser.add_argument( |
| 32 | "--max-tokens", |
| 33 | "-m", |
| 34 | type=int, |
| 35 | default=100, |
| 36 | help="The maximum number of tokens to generate", |
| 37 | ) |
| 38 | parser.add_argument( |
| 39 | "--temp", type=float, default=0.8, help="The sampling temperature" |
| 40 | ) |
| 41 | parser.add_argument( |
| 42 | "--prompt", |
| 43 | "-p", |
| 44 | type=str, |
| 45 | help="The prompt for generation", |
| 46 | default=None, |
| 47 | ) |
| 48 | |
| 49 | # Training args |
| 50 | parser.add_argument( |
| 51 | "--train", |
| 52 | action="store_true", |
| 53 | help="Do training", |
| 54 | ) |
| 55 | parser.add_argument( |
| 56 | "--add-eos-token", |
| 57 | type=int, |
| 58 | default=1, |
| 59 | help="Enable add_eos_token for tokenizer", |
| 60 | ) |
| 61 | parser.add_argument( |
| 62 | "--data", |
| 63 | type=str, |
| 64 | default="data/", |
| 65 | help="Directory with {train, valid, test}.jsonl files", |
| 66 | ) |
| 67 | parser.add_argument( |
| 68 | "--lora-layers", |
| 69 | type=int, |
| 70 | default=16, |
| 71 | help="Number of layers to fine-tune", |
| 72 | ) |
| 73 | parser.add_argument("--batch-size", type=int, default=4, help="Minibatch size.") |
| 74 | parser.add_argument( |
| 75 | "--iters", type=int, default=1000, help="Iterations to train for." |
| 76 | ) |
| 77 | parser.add_argument( |
| 78 | "--val-batches", |
| 79 | type=int, |
| 80 | default=25, |