()
| 392 | |
| 393 | |
| 394 | def main(): |
| 395 | parser = argparse.ArgumentParser(description='Quantize model embeddings to multiple formats') |
| 396 | parser.add_argument('--input', '-i', |
| 397 | default='../models/BitNet-b1.58-2B-4T/ggml-model-f32.gguf', |
| 398 | help='Input model path (default: ../models/BitNet-b1.58-2B-4T/ggml-model-f32.gguf)') |
| 399 | parser.add_argument('--output-dir', '-o', |
| 400 | default='../models/BitNet-b1.58-2B-4T', |
| 401 | help='Output directory (default: ../models/BitNet-b1.58-2B-4T)') |
| 402 | parser.add_argument('--quantize-bin', '-q', |
| 403 | default='../build/bin/llama-quantize', |
| 404 | help='Path to llama-quantize binary (default: ../build/bin/llama-quantize)') |
| 405 | parser.add_argument('--bench-bin', '-b', |
| 406 | default='../build/bin/llama-bench', |
| 407 | help='Path to llama-bench binary (default: ../build/bin/llama-bench)') |
| 408 | parser.add_argument('--stats-dir', |
| 409 | default='../stats', |
| 410 | help='Directory to save benchmark results (default: ../stats)') |
| 411 | parser.add_argument('--csv-output', '-c', |
| 412 | help='Custom path for CSV output file (e.g., stats/my_results.csv)') |
| 413 | parser.add_argument('--types', '-t', |
| 414 | nargs='+', |
| 415 | help='Specific types to quantize (e.g., f32 q6_k q4_0)') |
| 416 | parser.add_argument('--skip-existing', '-s', |
| 417 | action='store_true', |
| 418 | help='Skip quantization if output file already exists (will still benchmark existing files)') |
| 419 | |
| 420 | args = parser.parse_args() |
| 421 | |
| 422 | # Define all supported quantization types |
| 423 | # Format: (embedding_type for command line, output_suffix for filename) |
| 424 | all_types = [ |
| 425 | ('F32', 'f32'), |
| 426 | ('F16', 'f16'), |
| 427 | ('Q8_0', 'q8_0'), |
| 428 | ('Q6_K', 'q6_k'), |
| 429 | ('Q5_0', 'q5_0'), |
| 430 | ('Q4_0', 'q4_0'), |
| 431 | ('Q3_K', 'q3_k'), |
| 432 | ('TQ2_0', 'tq2_0'), |
| 433 | ] |
| 434 | |
| 435 | # If specific types are specified, filter the list |
| 436 | if args.types: |
| 437 | types_lower = [t.lower() for t in args.types] |
| 438 | types_to_quantize = [(et, os) for et, os in all_types if os.lower() in types_lower] |
| 439 | if not types_to_quantize: |
| 440 | print(f"❌ No valid types specified. Available types: {', '.join([os for _, os in all_types])}") |
| 441 | return |
| 442 | else: |
| 443 | types_to_quantize = all_types |
| 444 | |
| 445 | # If skip existing files is enabled, no need to filter |
| 446 | # Because new logic will automatically detect and skip during quantization, but will still benchmark |
| 447 | |
| 448 | # 创建量化器并运行 |
| 449 | try: |
| 450 | quantizer = EmbeddingQuantizer( |
| 451 | args.input, |
no test coverage detected