| 26 | sys.exit(1) |
| 27 | |
| 28 | def get_args(): |
| 29 | parser = argparse.ArgumentParser(description="FireRedASR Inference") |
| 30 | parser.add_argument( |
| 31 | "--engine_dir", |
| 32 | type=str, |
| 33 | default="/workspace_yuekai/asr/FireRedASR/examples/FireRedASR-AED-L-TensorRT", |
| 34 | help="TensorRT engine directory" |
| 35 | ) |
| 36 | parser.add_argument( |
| 37 | "--checkpoint_dir", |
| 38 | type=str, |
| 39 | default="/workspace_yuekai/asr/FireRedASR/examples/pretrained_models/FireRedASR-AED-L", |
| 40 | help="Checkpoint directory containing cmvn, dict, spm" |
| 41 | ) |
| 42 | parser.add_argument( |
| 43 | "--huggingface_dataset", |
| 44 | type=str, |
| 45 | required=True, |
| 46 | help="Dataset name" |
| 47 | ) |
| 48 | parser.add_argument( |
| 49 | "--subset_name", |
| 50 | type=str, |
| 51 | default=None, |
| 52 | help="Dataset subset name" |
| 53 | ) |
| 54 | parser.add_argument( |
| 55 | "--split_name", |
| 56 | type=str, |
| 57 | default="test", |
| 58 | help="Dataset split name" |
| 59 | ) |
| 60 | parser.add_argument( |
| 61 | "--output_file", |
| 62 | type=str, |
| 63 | default="results.txt", |
| 64 | help="Output file path" |
| 65 | ) |
| 66 | parser.add_argument( |
| 67 | "--batch_size", |
| 68 | type=int, |
| 69 | default=1, |
| 70 | help="Batch size (per-device)" |
| 71 | ) |
| 72 | parser.add_argument( |
| 73 | "--num_workers", |
| 74 | type=int, |
| 75 | default=4, |
| 76 | help="Num workers for dataloader" |
| 77 | ) |
| 78 | parser.add_argument( |
| 79 | "--prefetch", |
| 80 | type=int, |
| 81 | default=5, |
| 82 | help="Prefetch factor for dataloader" |
| 83 | ) |
| 84 | return parser.parse_args() |
| 85 | |