(args=None)
| 55 | |
| 56 | |
| 57 | def parse_arguments(args=None): |
| 58 | parser = argparse.ArgumentParser() |
| 59 | parser.add_argument('--batch_size', type=int, default=1) |
| 60 | parser.add_argument('--max_input_length', type=int, default=923) |
| 61 | parser.add_argument('--output_log_probs_npy', |
| 62 | type=str, |
| 63 | help='Numpy file where the log_probs are stored', |
| 64 | default=None) |
| 65 | |
| 66 | parser.add_argument('--output_cum_log_probs_npy', |
| 67 | type=str, |
| 68 | help='Numpy file where the cum_log_probs are stored', |
| 69 | default=None) |
| 70 | |
| 71 | parser.add_argument( |
| 72 | "--task", |
| 73 | type=str, |
| 74 | choices=['passkey', 'kv_retrieval'], |
| 75 | required=True, |
| 76 | help= |
| 77 | "Which task to use. Note that \"all\" can only be used in `compute_scores.py`.", # noqa |
| 78 | ) |
| 79 | parser.add_argument('--data_dir', |
| 80 | type=str, |
| 81 | default='./', |
| 82 | help="The directory of data.") |
| 83 | parser.add_argument("--output_dir", |
| 84 | type=str, |
| 85 | default=None, |
| 86 | help="Where to dump the prediction results.") # noqa |
| 87 | parser.add_argument( |
| 88 | "--start_idx", |
| 89 | type=int, |
| 90 | default=0, |
| 91 | help= |
| 92 | "The index of the first example to infer on. This is used if you want to evaluate on a (contiguous) subset of the data." |
| 93 | ) # noqa |
| 94 | parser.add_argument( |
| 95 | "--stop_idx", |
| 96 | type=int, |
| 97 | help= |
| 98 | "The index of the last example to infer on. This is used if you want to evaluate on a (contiguous) subset of the data. Defaults to the length of dataset." |
| 99 | ) # noqa |
| 100 | parser.add_argument('--tensorrt_llm_accuracy_threshold', |
| 101 | type=float, |
| 102 | default=99) |
| 103 | parser = add_common_args(parser) |
| 104 | |
| 105 | return parser.parse_args(args=args) |
| 106 | |
| 107 | |
| 108 | def parse_input(tokenizer, |
no test coverage detected