| 60 | |
| 61 | |
| 62 | def get_argument_parser(): |
| 63 | parser = argparse.ArgumentParser( |
| 64 | description="Run Export to ET for suppoorted LLM models.", allow_abbrev=False |
| 65 | ) |
| 66 | parser.add_argument( |
| 67 | "config", |
| 68 | type=str, |
| 69 | help="[Required] Model config json file. " |
| 70 | "Model config must be in same directory as all model weight bins and tokenizer files.", |
| 71 | ) |
| 72 | parser.add_argument( |
| 73 | "-p", |
| 74 | "--precision", |
| 75 | type=str, |
| 76 | default="A16W8", |
| 77 | choices=["A16W4", "A16W8", "A16W16", "A8W4", "A8W8"], |
| 78 | help="Precision to quantize entire model to.", |
| 79 | ) |
| 80 | parser.add_argument( |
| 81 | "--platform", |
| 82 | type=str, |
| 83 | default="DX4", |
| 84 | choices=["DX3", "DX4"], |
| 85 | help="Chip model of the inference device. " |
| 86 | "DX3 for Dimensity 9300, DX4 for Dimensity 9400.", |
| 87 | ) |
| 88 | parser.add_argument( |
| 89 | "-d", |
| 90 | "--dataset", |
| 91 | type=str, |
| 92 | default=None, |
| 93 | help="Calibration dataset name or path to dataset. Defaults to None to use random inputs", |
| 94 | ) |
| 95 | parser.add_argument( |
| 96 | "-n", |
| 97 | "--num_chunks", |
| 98 | type=int, |
| 99 | default=4, |
| 100 | help="Number of chunks to cut the model into. Defaults to 4.", |
| 101 | ) |
| 102 | parser.add_argument( |
| 103 | "-r", |
| 104 | "--response_cap", |
| 105 | type=int, |
| 106 | default=9, |
| 107 | help="Max Number of Response Tokens to save during calibration. Defaults to 9.", |
| 108 | ) |
| 109 | parser.add_argument( |
| 110 | "--preformatter", |
| 111 | type=str, |
| 112 | default=None, |
| 113 | help="Preformatter Template to use to wrap input with. Defaults to None.", |
| 114 | ) |
| 115 | parser.add_argument( |
| 116 | "-shapes", |
| 117 | nargs="+", |
| 118 | help="[Required] Expected input shapes to reconfigure TFLites to. Space separated list of " |
| 119 | "shapes in the format: xtyc (e.g. 32t512c)", |