| 30 | # replace_bloom_attn_with_flash_attn() |
| 31 | |
| 32 | def parse_args(): |
| 33 | parser = argparse.ArgumentParser( |
| 34 | description= |
| 35 | "Finetune a transformers model on a causal language modeling task") |
| 36 | parser.add_argument('--data_path', |
| 37 | type=str, |
| 38 | default='Dahoas/rm-static', |
| 39 | help='Path to the training dataset. A single data path.') |
| 40 | parser.add_argument( |
| 41 | '--data_output_path', |
| 42 | type=str, |
| 43 | default='./tmp/data_files/', |
| 44 | help= |
| 45 | 'Where to store the data-related files such as shuffle index. This needs to be on a local storage of a node (not on a shared storage)' |
| 46 | ) |
| 47 | parser.add_argument( |
| 48 | "--model_name_or_path", |
| 49 | type=str, |
| 50 | help= |
| 51 | "Path to pretrained model or model identifier from huggingface.co/models.", |
| 52 | required=True, |
| 53 | ) |
| 54 | parser.add_argument( |
| 55 | "--inference_model_path", |
| 56 | type=str, |
| 57 | help= |
| 58 | "Path to inference model.", |
| 59 | required=True, |
| 60 | ) |
| 61 | parser.add_argument( |
| 62 | "--max_prompt_len", |
| 63 | type=int, |
| 64 | default=512, |
| 65 | help="The maximum sequence length.", |
| 66 | ) |
| 67 | # inference params |
| 68 | parser.add_argument( |
| 69 | "--max_ans_len", |
| 70 | type=int, |
| 71 | default=256, |
| 72 | help="The maximum answer length.", |
| 73 | ) |
| 74 | parser.add_argument( |
| 75 | "--temperature", |
| 76 | type=float, |
| 77 | default=0.1, |
| 78 | help="Generate temperature params.", |
| 79 | ) |
| 80 | |
| 81 | parser.add_argument( |
| 82 | "--inference_batch", |
| 83 | type=int, |
| 84 | default=4, |
| 85 | help="Inference batch size.", |
| 86 | ) |
| 87 | # add other inference params |
| 88 | parser.add_argument( |
| 89 | "--inference_tasks", |