Parse command line arguments
()
| 39 | TRT_LOGGER = trt.Logger(trt.Logger.INFO) |
| 40 | |
| 41 | def parse_args(): |
| 42 | """ |
| 43 | Parse command line arguments |
| 44 | """ |
| 45 | parser = argparse.ArgumentParser(description=__doc__) |
| 46 | parser.add_argument('-e', '--engine', |
| 47 | help='Path to BERT TensorRT engine') |
| 48 | parser.add_argument("-b", "--batch-size", default=1, help="Batch size for inference.", type=int) |
| 49 | parser.add_argument('-p', '--passage', nargs='*', |
| 50 | help='Text for paragraph/passage for BERT QA', |
| 51 | default='') |
| 52 | parser.add_argument('-pf', '--passage-file', |
| 53 | help='File containing input passage', |
| 54 | default='') |
| 55 | parser.add_argument('-q', '--question', nargs='*', |
| 56 | help='Text for query/question for BERT QA', |
| 57 | default='') |
| 58 | parser.add_argument('-qf', '--question-file', |
| 59 | help='File containing input question', |
| 60 | default='') |
| 61 | parser.add_argument('-sq', '--squad-json', |
| 62 | help='SQuAD json file', |
| 63 | default='') |
| 64 | parser.add_argument('-o', '--output-prediction-file', |
| 65 | help='Output prediction file for SQuAD evaluation', |
| 66 | default='./predictions.json') |
| 67 | parser.add_argument('-v', '--vocab-file', |
| 68 | help='Path to file containing entire understandable vocab') |
| 69 | parser.add_argument('-s', '--sequence-length', |
| 70 | help='The sequence length to use. Defaults to 128', |
| 71 | default=128, type=int) |
| 72 | parser.add_argument('--max-query-length', |
| 73 | help='The maximum length of a query in number of tokens. Queries longer than this will be truncated', |
| 74 | default=64, type=int) |
| 75 | parser.add_argument('--max-answer-length', |
| 76 | help='The maximum length of an answer that can be generated', |
| 77 | default=30, type=int) |
| 78 | parser.add_argument('--n-best-size', |
| 79 | help='Total number of n-best predictions to generate in the nbest_predictions.json output file', |
| 80 | default=20, type=int) |
| 81 | parser.add_argument('--doc-stride', |
| 82 | help='When splitting up a long document into chunks, what stride to take between chunks', |
| 83 | default=128, type=int) |
| 84 | args, _ = parser.parse_known_args() |
| 85 | return args |
| 86 | |
| 87 | if __name__ == '__main__': |
| 88 | args = parse_args() |