Parse command line arguments
()
| 38 | import infer_c |
| 39 | |
| 40 | def parse_args(): |
| 41 | """ |
| 42 | Parse command line arguments |
| 43 | """ |
| 44 | parser = argparse.ArgumentParser(description=__doc__) |
| 45 | parser.add_argument('-e', '--engine', |
| 46 | help='Path to BERT TensorRT engine') |
| 47 | parser.add_argument('-p', '--passage', nargs='*', |
| 48 | help='Text for paragraph/passage for BERT QA', |
| 49 | default='') |
| 50 | parser.add_argument('-pf', '--passage-file', |
| 51 | help='File containing input passage', |
| 52 | default='') |
| 53 | parser.add_argument('-q', '--question', nargs='*', |
| 54 | help='Text for query/question for BERT QA', |
| 55 | default='') |
| 56 | parser.add_argument('-qf', '--question-file', |
| 57 | help='File containing input question', |
| 58 | default='') |
| 59 | parser.add_argument('-sq', '--squad-json', |
| 60 | help='SQuAD json file', |
| 61 | default='') |
| 62 | parser.add_argument('-o', '--output-prediction-file', |
| 63 | help='Output prediction file for SQuAD evaluation', |
| 64 | default='./predictions.json') |
| 65 | parser.add_argument('-v', '--vocab-file', |
| 66 | help='Path to file containing entire understandable vocab') |
| 67 | parser.add_argument('-s', '--sequence-length', |
| 68 | help='The sequence length to use. Defaults to 128', |
| 69 | default=128, type=int) |
| 70 | parser.add_argument('--max-query-length', |
| 71 | help='The maximum length of a query in number of tokens. Queries longer than this will be truncated', |
| 72 | default=64, type=int) |
| 73 | parser.add_argument('--max-answer-length', |
| 74 | help='The maximum length of an answer that can be generated', |
| 75 | default=30, type=int) |
| 76 | parser.add_argument('--n-best-size', |
| 77 | help='Total number of n-best predictions to generate in the nbest_predictions.json output file', |
| 78 | default=20, type=int) |
| 79 | parser.add_argument('--enable-graph', |
| 80 | help='Enable CUDA Graph support', |
| 81 | action='store_true', |
| 82 | default=False) |
| 83 | args = parser.parse_args() |
| 84 | return args |
| 85 | |
| 86 | if __name__ == '__main__': |
| 87 | args = parse_args() |
no test coverage detected