()
| 132 | |
| 133 | |
| 134 | def main(): |
| 135 | parser = argparse.ArgumentParser( |
| 136 | formatter_class=argparse.ArgumentDefaultsHelpFormatter |
| 137 | ) |
| 138 | parser.add_argument( |
| 139 | "--mode", |
| 140 | choices=["sequence", "parallel"], |
| 141 | default="sequence", |
| 142 | help="benchmark in parallel or sequence mode", |
| 143 | ) |
| 144 | parser.add_argument("--model_path", type=str, help="model path") |
| 145 | parser.add_argument("--src", type=str, help="source file") |
| 146 | parser.add_argument("--target", type=str, help="target file") |
| 147 | parser.add_argument("--batch_size", type=int, help="batch size") |
| 148 | args = parser.parse_args() |
| 149 | |
| 150 | print("Loading the model...") |
| 151 | generator = ctranslate2.Generator(args.model_path, device="cuda", tensor_parallel=True, |
| 152 | flash_attention=False, inter_threads=2) |
| 153 | sp = spm.SentencePieceProcessor(os.path.join(args.model_path, "tokenizer.model")) |
| 154 | |
| 155 | if not os.path.exists(args.src): |
| 156 | raise Exception("No source file found: " + args.src) |
| 157 | # Open the file in read mode |
| 158 | with open(args.src, 'r') as file: |
| 159 | # Read all lines from the file and create a list |
| 160 | inputs = file.readlines() |
| 161 | |
| 162 | prompt_tokens = build_prompt(sp, inputs) |
| 163 | result = benchmark_generation(generator, sp, prompt_tokens, args.target, args.mode, args.batch_size) |
| 164 | if ctranslate2.MpiInfo.getCurRank() == 0: |
| 165 | print("Benchmark result (%d sample(s)):" % len(prompt_tokens)) |
| 166 | print("- Generation time: %.2f s" % result.generation_time) |
| 167 | print("- Number of tokens: %d" % result.num_tokens) |
| 168 | print("- Throughput: %.1f" % (result.num_tokens / result.generation_time)) |
| 169 | print("- max. GPU memory usage: %dMB" % int(result.max_gpu_mem)) |
| 170 | |
| 171 | |
| 172 | if __name__ == "__main__": |
no test coverage detected