Main benchmark runner.
()
| 789 | |
| 790 | |
| 791 | def main(): |
| 792 | """Main benchmark runner.""" |
| 793 | import argparse |
| 794 | |
| 795 | parser = argparse.ArgumentParser(description="TokenDagger vs TikToken Multithreaded Throughput Benchmark") |
| 796 | parser.add_argument("--tokenizer", choices=["llama", "mistral"], default="llama", |
| 797 | help="Tokenizer configuration to use (default: llama)") |
| 798 | parser.add_argument("--threads", type=str, default="1,2,4,8,16,32", |
| 799 | help="Comma-separated list of thread counts to test (default: 1,2,4,8,16,32)") |
| 800 | parser.add_argument("--text-size", type=int, default=1024, |
| 801 | help="Size of test text in MB (default: 1024 = 1GB)") |
| 802 | parser.add_argument("--iterations", type=int, default=10, |
| 803 | help="Iterations per thread (default: 10)") |
| 804 | parser.add_argument("--quick", action="store_true", |
| 805 | help="Run quick benchmark (smaller text size and fewer iterations)") |
| 806 | parser.add_argument("--output", type=str, default=None, |
| 807 | help="Output SVG filename (default: auto-generated)") |
| 808 | |
| 809 | args = parser.parse_args() |
| 810 | |
| 811 | # Parse thread counts |
| 812 | thread_counts = [int(x.strip()) for x in args.threads.split(',')] |
| 813 | |
| 814 | # Quick mode adjustments |
| 815 | if args.quick: |
| 816 | args.text_size = 64 # 64 MB instead of 1 GB |
| 817 | args.iterations = 3 # 3 iterations instead of 10 |
| 818 | print("Quick mode enabled: reduced text size and iterations") |
| 819 | |
| 820 | benchmark = ThroughputBenchmark( |
| 821 | tokenizer_type=args.tokenizer, |
| 822 | thread_counts=thread_counts, |
| 823 | text_size_mb=args.text_size, |
| 824 | iterations_per_thread=args.iterations |
| 825 | ) |
| 826 | |
| 827 | success = benchmark.run_full_benchmark() |
| 828 | sys.exit(0 if success else 1) |
| 829 | |
| 830 | |
| 831 | if __name__ == "__main__": |
no test coverage detected