| 15 | |
| 16 | |
| 17 | class EmbeddingQuantizer: |
| 18 | def __init__(self, input_model, output_dir, quantize_bin="../build/bin/llama-quantize", |
| 19 | bench_bin="../build/bin/llama-bench", stats_dir="../stats", csv_output=None): |
| 20 | self.input_model = Path(input_model) |
| 21 | self.output_dir = Path(output_dir) |
| 22 | self.quantize_bin = Path(quantize_bin) |
| 23 | self.bench_bin = Path(bench_bin) |
| 24 | self.stats_dir = Path(stats_dir) |
| 25 | self.csv_output = Path(csv_output) if csv_output else None |
| 26 | |
| 27 | # Verify input file exists |
| 28 | if not self.input_model.exists(): |
| 29 | raise FileNotFoundError(f"Input model not found: {self.input_model}") |
| 30 | |
| 31 | # Verify quantize tool exists |
| 32 | if not self.quantize_bin.exists(): |
| 33 | raise FileNotFoundError(f"Quantize binary not found: {self.quantize_bin}") |
| 34 | |
| 35 | # Verify bench tool exists |
| 36 | if not self.bench_bin.exists(): |
| 37 | raise FileNotFoundError(f"Benchmark binary not found: {self.bench_bin}") |
| 38 | |
| 39 | # Create output directories |
| 40 | self.output_dir.mkdir(parents=True, exist_ok=True) |
| 41 | self.stats_dir.mkdir(parents=True, exist_ok=True) |
| 42 | |
| 43 | self.results = [] |
| 44 | self.newly_created_files = set() # Track newly created files |
| 45 | |
| 46 | def quantize(self, embedding_type, output_suffix): |
| 47 | """ |
| 48 | Perform single quantization |
| 49 | |
| 50 | Args: |
| 51 | embedding_type: Token embedding type (uppercase format, e.g., Q6_K) |
| 52 | output_suffix: Output file suffix (lowercase format, e.g., q6_k) |
| 53 | |
| 54 | Returns: |
| 55 | bool: Whether successful |
| 56 | """ |
| 57 | output_file = self.output_dir / f"ggml-model-i2_s-embed-{output_suffix}.gguf" |
| 58 | |
| 59 | # Check if file already exists |
| 60 | file_already_existed = output_file.exists() |
| 61 | |
| 62 | if file_already_existed: |
| 63 | print(f"ℹ️ File already exists: {output_file}") |
| 64 | print(f" Skipping quantization, will use existing file for benchmark") |
| 65 | return True |
| 66 | |
| 67 | cmd = [ |
| 68 | str(self.quantize_bin), |
| 69 | "--token-embedding-type", embedding_type, |
| 70 | str(self.input_model), |
| 71 | str(output_file), |
| 72 | "I2_S", |
| 73 | "1", |
| 74 | "1" |