TQ_engine_lib/quantizer.py

import torch
import numpy as np
import math
import os
import gc
import json
import subprocess
from typing import Union, List, Optional, Tuple, Sequence
from dataclasses import dataclass
from .codebook import ScalarQuantizer
from .rotation import get_orthogonal_matrix, rotate_forward, rotate_backward
from .tq_bridge import tq_native

@dataclass
class ProdQuantized:
    sq_codes: np.ndarray
    qjl_signs: np.ndarray
    norms: np.ndarray
    centroids: np.ndarray
    dim: int
    sq_bits: int
    total_bits: int
    qjl_scale: float
    rot_op: np.ndarray
    res_norms: np.ndarray

@dataclass
class IVFData:
    coarse_centroids: torch.Tensor
    pq_data: ProdQuantized
    vector_ids: np.ndarray
    list_offsets: np.ndarray
    n_list: int
    n_probe: int

class TQEngine:
    def __init__(self, dim: int = 768, bits: int = 4, device: str = None, use_ivf: bool = False, ivf_nlist: int = 1024, ivf_nprobe: int = 32):
        if bits not in [2, 4]:
            raise ValueError(f"TurboQuant currently only supports 2-bit (1+1) and 4-bit (3+1) configurations. Received: {bits}")
            
        self.dim = dim
        self.bits = bits
        self.sq_bits = bits - 1
        self.device = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
        self.use_ivf = use_ivf
        self.ivf_nlist = ivf_nlist
        self.ivf_nprobe = int(ivf_nprobe)
        self.ivf_target_candidates = 20000
        self.ivf_nprobe_min = 1
        self.ivf_nprobe_max = None
        self.deleted_ids = set()
        self.dynamic_shards = {}
        self.current_ivf_data = None
        self.raw_vectors = None
        self.sq_quantizer = ScalarQuantizer(dim=dim, bits=self.sq_bits, device=self.device)
        self.rot_op_t = get_orthogonal_matrix(dim, device=self.device)
        self.rot_op_np = self.rot_op_t.cpu().numpy().astype(np.float32)
        self.qjl_scale = 1.0 
        self.hnsw_navigator = None
        self.use_hnsw = False 

    def _auto_nprobe(self) -> int:
        ivf = self.current_ivf_data
        if ivf is None or not isinstance(ivf, IVFData):
            return max(1, self.ivf_nprobe)
        offsets = np.asarray(ivf.list_offsets, dtype=np.int64)
        if offsets.ndim != 1 or offsets.size < 2:
            return max(1, self.ivf_nprobe)
        counts = offsets[1:] - offsets[:-1]
        if counts.size == 0:
            return max(1, self.ivf_nprobe)
        med = float(np.median(counts))
        if med <= 0:
            med = float(np.mean(counts)) if float(np.mean(counts)) > 0 else 1.0
        target = max(1, int(self.ivf_target_candidates))
        nprobe = int(np.ceil(target / med))
        nlist = int(ivf.n_list) if hasattr(ivf, "n_list") else int(counts.size)
        nprobe = max(self.ivf_nprobe_min, nprobe)
        if self.ivf_nprobe_max is not None:
            nprobe = min(int(self.ivf_nprobe_max), nprobe)
        nprobe = min(nlist, nprobe)
        return max(1, nprobe)

    def _train_kmeans(self, x_sample: torch.Tensor, n_list: int, iters: int = 10):
        N, D = x_sample.shape
        indices = torch.randperm(N)[:n_list]
        centroids = x_sample[indices].clone()
        for i in range(iters):
            assignments = torch.zeros(N, dtype=torch.long, device=self.device)
            chunk_size = 8192
            for s in range(0, N, chunk_size):
                e = min(s + chunk_size, N)
                batch = x_sample[s:e]
                scores = torch.mm(batch, centroids.t())
                assignments[s:e] = scores.argmax(dim=1)
                del scores, batch
            new_centroids = torch.zeros_like(centroids)
            counts = torch.zeros(n_list, 1, device=self.device)
            ones = torch.ones(N, 1, device=self.device)
            new_centroids.index_add_(0, assignments, x_sample)
            counts.index_add_(0, assignments, ones)
            counts = torch.clamp(counts, min=1.0)
            centroids = new_centroids / counts
            centroids = centroids / (torch.norm(centroids, dim=1, keepdim=True) + 1e-12)
            gc.collect()
        return centroids

    def bind_raw_data(self, npy_path: str):
        if os.path.exists(npy_path):
            self.raw_vectors = np.load(npy_path, mmap_mode='r')
        else:
            print(f"⚠️ Warning: Raw data file not found: {npy_path}")

    def index(self, x: Union[torch.Tensor, np.ndarray], online_clustering: bool = False, 
              save_path: str = None, n_train: int = None, train_iters: int = 20,
              init_centroids: np.ndarray = None):
        if isinstance(x, torch.Tensor):
            x_np = x.detach().cpu().numpy().astype(np.float32, copy=False)
        else:
            x_np = np.asarray(x, dtype=np.float32)
        N, D = x_np.shape
        self.dim = D
        if self.use_ivf:
            if init_centroids is not None:
                print(f"  Using pre-trained coarse centroids ({init_centroids.shape[0]} clusters)...")
                centroids_np = np.asarray(init_centroids, dtype=np.float32)
                n_sample = n_train if n_train is not None else 180000
            else:
                if n_train is None:
                    n_train = min(N, 256 * self.ivf_nlist)
                
                n_sample = min(N, n_train)
                indices = np.random.choice(N, n_sample, replace=False)
                x_sample_np = x_np[indices].copy()
                
                print(f"  Training {self.ivf_nlist} coarse centroids on {n_sample:,} samples (iters={train_iters})...")
                centroids_np = tq_native.tq_kmeans_train(x_sample_np, self.ivf_nlist, iters=train_iters)
                del x_sample_np

            coarse_centroids = torch.from_numpy(centroids_np).to(self.device)
            
            print(f"  Assigning {N:,} vectors to clusters...")
            assignments = tq_native.tq_assign_clusters(x_np, centroids_np)
            
            with torch.no_grad():
                n_scale = n_sample
                scale_idx = np.random.choice(N, n_scale, replace=False)
                x_scale = torch.from_numpy(np.asarray(x_np[scale_idx], dtype=np.float32)).to(self.device)
                c_idx = torch.from_numpy(assignments[scale_idx]).to(self.device).long()
                c = coarse_centroids.index_select(0, c_idx)
                residual = x_scale - c
                residual_rot = rotate_forward(residual, self.rot_op_t)
                
                print(f"  Calibrating Lloyd-Max Codebook for {self.sq_bits}-bit SQ...")
                self.sq_quantizer.fit(residual_rot, iterations=30)
                
                self.qjl_scale = float(torch.mean(torch.abs(residual_rot)).item())
                del x_scale, c_idx, c, residual, residual_rot
            index_dir = save_path if save_path else "tq_index_temp"
            os.makedirs(index_dir, exist_ok=True)
            counts = np.bincount(assignments, minlength=self.ivf_nlist)
            offsets = np.zeros(self.ivf_nlist + 1, dtype=np.int64)
            offsets[1:] = np.cumsum(counts)
            sq_packed_dim = self.dim
            if self.sq_bits == 1: sq_packed_dim = self.dim // 8
            elif self.sq_bits == 3: sq_packed_dim = self.dim // 2
            qjl_packed_dim = self.dim // 8
            from numpy.lib.format import open_memmap
            f_sq = open_memmap(os.path.join(index_dir, "sq_codes.npy"), mode='w+', dtype=np.uint8, shape=(N, sq_packed_dim))
            f_signs = open_memmap(os.path.join(index_dir, "qjl_signs.npy"), mode='w+', dtype=np.uint8, shape=(N, qjl_packed_dim))
            f_norms = open_memmap(os.path.join(index_dir, "norms.npy"), mode='w+', dtype=np.float32, shape=(N,))
            f_res_norms = open_memmap(os.path.join(index_dir, "res_norms.npy"), mode='w+', dtype=np.float32, shape=(N,))
            f_ids = open_memmap(os.path.join(index_dir, "vector_ids.npy"), mode='w+', dtype=np.int64, shape=(N,))
            dummy_pq = self._quantize_flat(torch.zeros((1, self.dim), device=self.device))
            
            print(f"  Streaming Indexing: Rotating & Quantizing {N:,} vectors in batches...")
            cluster_counters = np.zeros(self.ivf_nlist, dtype=np.int64)
            batch_size = 100000
            with torch.no_grad():
                for s in range(0, N, batch_size):
                    e = min(s + batch_size, N)
                    batch_x = torch.from_numpy(x_np[s:e]).to(self.device)
                    batch_c = torch.from_numpy(centroids_np[assignments[s:e]]).to(self.device)
                    batch_res = batch_x - batch_c
                    batch_rot_res = rotate_forward(batch_res, self.rot_op_t)
                    batch_q_res = self._quantize_flat(batch_x, online_clustering=False, x_rot=batch_rot_res)
                    target_positions = np.zeros(e - s, dtype=np.int64)
                    batch_ass = assignments[s:e]
                    for i, c_id in enumerate(batch_ass):
                        target_positions[i] = offsets[c_id] + cluster_counters[c_id]
                        cluster_counters[c_id] += 1
                    f_sq[target_positions] = batch_q_res.sq_codes
                    f_signs[target_positions] = batch_q_res.qjl_signs
                    f_norms[target_positions] = batch_q_res.norms
                    f_res_norms[target_positions] = batch_q_res.res_norms
                    f_ids[target_positions] = np.arange(s, e, dtype=np.int64)
                    if s % 1000000 == 0:
                        f_sq.flush(); f_signs.flush(); f_norms.flush(); f_res_norms.flush(); f_ids.flush()
                        gc.collect()
                    del batch_x, batch_c, batch_res, batch_rot_res, batch_q_res
            
            import faiss
            self.hnsw_navigator = faiss.IndexHNSWFlat(self.dim, 32, faiss.METRIC_INNER_PRODUCT)
            self.hnsw_navigator.add(centroids_np)
            
            print(f"  Finalizing index files...")
            f_sq.flush(); f_signs.flush(); f_norms.flush(); f_res_norms.flush(); f_ids.flush()
            
            # Lưu các file metadata quan trọng
            np.save(os.path.join(index_dir, "list_offsets.npy"), offsets.astype(np.int32))
            np.save(os.path.join(index_dir, "coarse_centroids.npy"), centroids_np)
            np.save(os.path.join(index_dir, "sq_centroids.npy"), dummy_pq.centroids)
            np.save(os.path.join(index_dir, "rot_op.npy"), self.rot_op_np)
            
            if hasattr(self, "hnsw_navigator") and self.hnsw_navigator is not None:
                faiss.write_index(self.hnsw_navigator, os.path.join(index_dir, "centroids.hnsw"))
                
            import json
            with open(os.path.join(index_dir, "metadata.json"), "w", encoding='utf-8') as f:
                json.dump({
                    "dim": int(self.dim), "bits": int(self.bits), "qjl_scale": float(self.qjl_scale),
                    "n_list": int(self.ivf_nlist), "n_probe": int(self.ivf_nprobe), "deleted_ids": []
                }, f, indent=2)
            flat_pq = ProdQuantized(
                sq_codes=f_sq, qjl_signs=f_signs, norms=f_norms,
                centroids=dummy_pq.centroids, dim=self.dim, sq_bits=self.sq_bits, total_bits=self.bits,
                qjl_scale=self.qjl_scale, rot_op=self.rot_op_np, res_norms=f_res_norms
            )
            self.current_ivf_data = IVFData(
                coarse_centroids=coarse_centroids, pq_data=flat_pq, vector_ids=f_ids,
                list_offsets=offsets.astype(np.int32), n_list=self.ivf_nlist, n_probe=self.ivf_nprobe
            )
            return self.current_ivf_data
        else:
            return self._quantize_flat(torch.from_numpy(np.array(x, dtype=np.float32)).to(self.device), online_clustering)

    def add(self, vector: torch.Tensor, vector_id: int):
        if self.current_ivf_data is None:
            raise ValueError("Cần gọi index() hoặc load_index() trước khi add().")
        if vector.device.type != self.device:
            vector = vector.to(self.device)
        if vector.dim() == 1:
            vector = vector.unsqueeze(0)
        scores = torch.mm(vector, self.current_ivf_data.coarse_centroids.t())
        c_idx = scores.argmax(dim=1).item()
        centroid = self.current_ivf_data.coarse_centroids[c_idx].unsqueeze(0)
        pq_single = self._quantize_flat(vector, online_clustering=False, centroid=centroid)
        if c_idx not in self.dynamic_shards:
            self.dynamic_shards[c_idx] = []
        self.dynamic_shards[c_idx].append((vector_id, pq_single))
        if vector_id in self.deleted_ids:
            self.deleted_ids.remove(vector_id)

    def merge_dynamic_shards(self):
        ivf = self.current_ivf_data
        if not isinstance(ivf, IVFData) or not self.dynamic_shards:
            return
        new_total_size = len(ivf.vector_ids) + sum(len(v) for v in self.dynamic_shards.values())
        updated_sq_codes = torch.zeros((new_total_size, ivf.pq_data.sq_codes.shape[1]), dtype=torch.uint8)
        updated_qjl_signs = torch.zeros((new_total_size, ivf.pq_data.qjl_signs.shape[1]), dtype=torch.int8)
        updated_norms = torch.zeros(new_total_size)
        updated_res_norms = torch.zeros(new_total_size)
        updated_vector_ids = np.zeros(new_total_size, dtype=np.int64)
        updated_offsets = [0]
        curr_pos = 0
        for c_idx in range(ivf.n_list):
            old_start = ivf.list_offsets[c_idx]
            old_end = ivf.list_offsets[c_idx+1]
            old_size = old_end - old_start
            if old_size > 0:
                updated_sq_codes[curr_pos:curr_pos+old_size] = torch.from_numpy(ivf.pq_data.sq_codes[old_start:old_end])
                updated_qjl_signs[curr_pos:curr_pos+old_size] = torch.from_numpy(ivf.pq_data.qjl_signs[old_start:old_end])
                updated_norms[curr_pos:curr_pos+old_size] = torch.from_numpy(ivf.pq_data.norms[old_start:old_end])
                updated_res_norms[curr_pos:curr_pos+old_size] = torch.from_numpy(ivf.pq_data.res_norms[old_start:old_end])
                updated_vector_ids[curr_pos:curr_pos+old_size] = ivf.vector_ids[old_start:old_end]
                curr_pos += old_size
            if c_idx in self.dynamic_shards:
                for vid, dpq in self.dynamic_shards[c_idx]:
                    updated_sq_codes[curr_pos] = torch.from_numpy(dpq.sq_codes) if isinstance(dpq.sq_codes, np.ndarray) else dpq.sq_codes
                    updated_qjl_signs[curr_pos] = torch.from_numpy(dpq.qjl_signs) if isinstance(dpq.qjl_signs, np.ndarray) else dpq.qjl_signs
                    updated_norms[curr_pos] = float(dpq.norms[0]) if isinstance(dpq.norms, np.ndarray) else float(dpq.norms)
                    updated_res_norms[curr_pos] = float(dpq.res_norms[0]) if isinstance(dpq.res_norms, np.ndarray) else float(dpq.res_norms)
                    updated_vector_ids[curr_pos] = vid
                    curr_pos += 1
            updated_offsets.append(curr_pos)
        ivf.pq_data.sq_codes = updated_sq_codes
        ivf.pq_data.qjl_signs = updated_qjl_signs
        ivf.pq_data.norms = updated_norms
        ivf.pq_data.res_norms = updated_res_norms
        ivf.vector_ids = updated_vector_ids
        ivf.list_offsets = torch.tensor(updated_offsets, dtype=torch.long)
        self.dynamic_shards.clear()

    def delete(self, vector_id: int):
        self.deleted_ids.add(vector_id)

    def _quantize_flat(self, x: torch.Tensor, online_clustering: bool = False, x_rot: torch.Tensor = None, centroid: torch.Tensor = None) -> ProdQuantized:
        if x.device.type != self.device:
            x = x.to(self.device)
        if centroid is not None:
            if centroid.device.type != self.device:
                centroid = centroid.to(self.device)
            x_target = x - centroid
        else:
            x_target = x
        if x_rot is None:
            x_rot = rotate_forward(x_target, self.rot_op_t)
        norms = torch.norm(x, dim=-1)
        res_norms = torch.norm(x_target, dim=-1)
        if online_clustering:
            self.sq_quantizer.fit(x_rot)
        x_rot_np = np.ascontiguousarray(x_rot.detach().cpu().numpy(), dtype=np.float32)
        sq_centroids_np = np.ascontiguousarray(self.sq_quantizer.centroids.detach().cpu().numpy(), dtype=np.float32)
        try:
            sq_codes_packed, qjl_signs_packed, res_norms_np = tq_native.tq_quantize_rotated(x_rot_np, sq_centroids_np, int(self.sq_bits))
            sq_codes_np = np.asarray(sq_codes_packed, dtype=np.uint8)
            qjl_signs = np.asarray(qjl_signs_packed, dtype=np.uint8)
            res_norms_np = np.asarray(res_norms_np, dtype=np.float32)
        except Exception:
            sq_q = self.sq_quantizer.quantize(x_rot)
            x_hat_1 = self.sq_quantizer.reconstruct(sq_q.indices)
            residual = x_rot - x_hat_1
            res_norms_np = torch.norm(residual, dim=-1).detach().cpu().numpy().astype(np.float32)
            signs = (residual > 0).to(torch.uint8).cpu().numpy()
            qjl_signs = np.packbits(signs, axis=-1, bitorder='little').astype(np.uint8)
            sq_codes_np = sq_q.indices.cpu().numpy().astype(np.uint8)
        return ProdQuantized(
            sq_codes=sq_codes_np, qjl_signs=qjl_signs.astype(np.uint8),
            norms=norms.cpu().numpy().astype(np.float32),
            centroids=self.sq_quantizer.centroids.cpu().numpy().astype(np.float32),
            dim=self.dim, sq_bits=self.sq_bits, total_bits=self.bits,
            qjl_scale=self.qjl_scale, rot_op=self.rot_op_np, res_norms=res_norms_np
        )

    def search_batch(self, queries: torch.Tensor, top_k: int = 100, n_probe: int = None, 
                     allowed_ids: Optional[List[int]] = None, 
                     raw_corpus: Optional[np.ndarray] = None, 
                     rerank_factor: Optional[int] = None) -> List[Tuple[torch.Tensor, torch.Tensor]]:
        ivf = self.current_ivf_data
        if ivf is None:
            raise ValueError("No data indexed. Call index() first.")
        if queries.device.type != self.device:
            queries = queries.to(self.device)
        if queries.dim() == 1:
            queries = queries.unsqueeze(0)
        
        nprobe = int(n_probe) if n_probe is not None else (self._auto_nprobe() if self.ivf_nprobe <= 0 else int(self.ivf_nprobe))
        num_queries = queries.shape[0]
        queries_np = np.ascontiguousarray(queries.detach().cpu().numpy(), dtype=np.float32)
        pq = ivf.pq_data
        
        # Fast Path: If HNSW is not active, run 100% in Rust using tq_unified_search
        if not self.use_hnsw or self.hnsw_navigator is None:
            allowed_arr = np.array(allowed_ids, dtype=np.int64) if allowed_ids is not None else None
            raw_arr = raw_corpus if raw_corpus is not None else (self.raw_vectors if hasattr(self, "raw_vectors") else None)
            
            scores, indices = tq_native.tq_unified_search(
                queries_np,
                self.rot_op_np,
                ivf.coarse_centroids.cpu().numpy(),
                ivf.list_offsets,
                ivf.vector_ids,
                pq.sq_codes,
                pq.centroids,
                pq.norms,
                pq.qjl_signs,
                pq.res_norms,
                float(pq.qjl_scale),
                int(self.dim),
                int(self.sq_bits),
                int(nprobe),
                int(top_k),
                allowed_arr,
                raw_arr,
                int(rerank_factor) if rerank_factor is not None else None
            )
            
            results = []
            for i in range(num_queries):
                final_ids = torch.from_numpy(indices[i]).to(self.device)
                final_scores = torch.from_numpy(scores[i]).to(self.device)
                results.append((final_ids, final_scores))
            return results
        
        # Slow/Fallback Path: Standard HNSW scan (uses tq_ivf_scan_with_clusters)
        if self.use_hnsw and self.hnsw_navigator is not None:
            cluster_scores_np, cluster_ids_np = self.hnsw_navigator.search(queries_np, nprobe)
            cluster_ids = torch.from_numpy(cluster_ids_np).to(self.device).long()
            cluster_scores = torch.from_numpy(cluster_scores_np).to(self.device).float()
        else:
            scores_c = torch.mm(queries, ivf.coarse_centroids.t())
            cluster_scores, cluster_ids = torch.topk(scores_c, nprobe, dim=1)
            
        q_rot = rotate_forward(queries, self.rot_op_t)
        q_rot_np = np.ascontiguousarray(q_rot.cpu().numpy(), dtype=np.float32)
        
        scores, indices = tq_native.tq_ivf_scan_with_clusters(
            queries_np, pq.sq_codes,
            pq.centroids,
            pq.norms,
            pq.qjl_signs,
            pq.res_norms,
            q_rot_np, ivf.list_offsets,
            np.ascontiguousarray(cluster_ids.cpu().numpy(), dtype=np.int32),
            np.ascontiguousarray(cluster_scores.cpu().numpy(), dtype=np.float32),
            float(pq.qjl_scale), int(self.dim), int(self.sq_bits),
            int(top_k if allowed_ids is None else top_k * 10)
        )
        
        results = []
        global_ids = ivf.vector_ids
        allowed_set = set(allowed_ids) if allowed_ids is not None else None
        for i in range(num_queries):
            valid_mask = indices[i] != -1
            q_indices = indices[i][valid_mask]
            q_scores = scores[i][valid_mask]
            q_global_ids = global_ids[q_indices]
            
            # Apply rerank if requested in slow path
            if raw_corpus is not None and rerank_factor is not None:
                candidates_raw = torch.from_numpy(raw_corpus[q_global_ids]).to(self.device)
                exact_scores = torch.mm(queries[i:i+1], candidates_raw.t()).view(-1)
                _, final_idx = torch.topk(exact_scores, min(top_k, len(exact_scores)))
                q_global_ids = q_global_ids[final_idx.cpu().numpy()]
                q_scores = exact_scores[final_idx].cpu().numpy()
                
            if allowed_set is not None:
                mask = np.isin(q_global_ids, list(allowed_set))
                q_global_ids = q_global_ids[mask]
                q_scores = q_scores[mask]
                q_global_ids = q_global_ids[:top_k]
                q_scores = q_scores[:top_k]
                
            final_ids = torch.from_numpy(q_global_ids.copy()).to(self.device)
            final_scores = torch.from_numpy(q_scores.copy()).to(self.device)
            results.append((final_ids, final_scores))
        return results

    def search(self, query: torch.Tensor, top_k: int = 100, n_probe: int = None, allowed_ids: Optional[List[int]] = None) -> tuple[torch.Tensor, torch.Tensor]:
        ivf = self.current_ivf_data
        if ivf is None:
            raise ValueError("No data indexed. Call index() first.")
        if query.device.type != self.device:
            query = query.to(self.device)
        if isinstance(ivf, IVFData):
            if query.dim() == 1:
                query = query.unsqueeze(0)
            results = self.search_batch(query, top_k=top_k, n_probe=n_probe, allowed_ids=allowed_ids)
            return results[0]
        else:
            return self._native_cosine_search_flat(query, ivf, top_k)

    def _native_cosine_search_flat(self, query: torch.Tensor, pq: ProdQuantized, top_k: int = 100) -> tuple[torch.Tensor, torch.Tensor]:
        if query.device.type != self.device:
            query = query.to(self.device)
        if query.dim() == 1:
            query = query.unsqueeze(0)
        q_rot = rotate_forward(query, self.rot_op_t).squeeze(0)
        q_np = q_rot.cpu().numpy().astype(np.float32)
        query_1d = np.array(q_np, dtype=np.float32, order='C')
        total_vectors = pq.sq_codes.shape[0]
        ram_gb = 4.0
        h = 10**(len(str(total_vectors)) - 1)
        raw_batch_size = int((0.3 * total_vectors * (h * 100) / total_vectors) / (ram_gb * (ram_gb / 0.4))) + 1
        compression_ratio = 32 // self.bits
        tq_batch_size = raw_batch_size * compression_ratio
        all_scores = []
        for start_idx in range(0, total_vectors, tq_batch_size):
            end_idx = min(start_idx + tq_batch_size, total_vectors)
            sq_batch = np.ascontiguousarray(pq.sq_codes[start_idx:end_idx], dtype=np.uint8)
            qjl_batch = np.ascontiguousarray(pq.qjl_signs[start_idx:end_idx], dtype=np.uint8)
            norms_batch = np.ascontiguousarray(pq.norms[start_idx:end_idx], dtype=np.float32)
            res_norms_batch = np.ascontiguousarray(pq.res_norms[start_idx:end_idx], dtype=np.float32)
            centroids_1d = np.ascontiguousarray(pq.centroids, dtype=np.float32)
            batch_scores = tq_native.tq_scan(query_1d, sq_batch, centroids_1d, norms_batch, qjl_batch, res_norms_batch, query_1d, float(pq.qjl_scale), int(self.dim), int(self.sq_bits))
            all_scores.append(batch_scores)
        final_scores = np.concatenate(all_scores)
        scores_t = torch.from_numpy(final_scores).view(-1)
        top_scores, top_indices = torch.topk(scores_t, min(top_k, len(scores_t)))
        return top_indices, top_scores

    def save_index(self, path: str):
        from pathlib import Path
        import faiss
        save_path = Path(path).resolve()
        if self.current_ivf_data is None:
            raise ValueError("Không có dữ liệu để lưu. Hãy gọi index() trước.")
        gc.collect()
        if not save_path.exists():
            save_path.mkdir(parents=True, exist_ok=True)
        ivf = self.current_ivf_data
        pq = ivf.pq_data
        def to_np(obj):
            if isinstance(obj, torch.Tensor):
                return obj.detach().cpu().numpy()
            return obj
        files_to_save = {}
        if not isinstance(pq.sq_codes, np.memmap):
            files_to_save["sq_codes.npy"] = to_np(pq.sq_codes)
            files_to_save["qjl_signs.npy"] = to_np(pq.qjl_signs)
            files_to_save["norms.npy"] = to_np(pq.norms)
            files_to_save["res_norms.npy"] = to_np(pq.res_norms)
            files_to_save["vector_ids.npy"] = to_np(ivf.vector_ids)
            
        # Các file này phải luôn được lưu
        files_to_save["list_offsets.npy"] = to_np(ivf.list_offsets)
        files_to_save["coarse_centroids.npy"] = to_np(ivf.coarse_centroids)
        files_to_save["rot_op.npy"] = to_np(pq.rot_op)
        files_to_save["sq_centroids.npy"] = to_np(pq.centroids)
        for filename, data in files_to_save.items():
            np.save(str(save_path / filename), data)
        if self.hnsw_navigator is not None:
            faiss.write_index(self.hnsw_navigator, str(save_path / "centroids.hnsw"))
        meta = {
            "dim": int(self.dim), "bits": int(self.bits), "qjl_scale": float(self.qjl_scale),
            "n_list": int(ivf.n_list), "n_probe": int(ivf.n_probe), "deleted_ids": list(self.deleted_ids)
        }
        with open(str(save_path / "metadata.json"), "w", encoding='utf-8') as f:
            json.dump(meta, f, indent=2)

    def load_index(self, path: str, use_mmap: bool = True):
        from pathlib import Path
        import platform
        import faiss
        load_path = Path(path).resolve()
        if not load_path.exists():
            raise FileNotFoundError(f"Thư mục index không tồn tại: {load_path}")
        with open(str(load_path / "metadata.json"), "r", encoding='utf-8') as f:
            meta = json.load(f)
        self.dim = meta["dim"]
        self.bits = meta["bits"]
        self.sq_bits = self.bits - 1
        self.qjl_scale = meta["qjl_scale"]
        self.ivf_nlist = meta["n_list"]
        self.ivf_nprobe = meta["n_probe"]
        self.deleted_ids = set(meta["deleted_ids"])
        
        mmap_val = 'r'
        if platform.system() == "Windows" and not use_mmap:
            mmap_val = None
        elif platform.system() != "Windows":
            mmap_val = 'r'
            
        coarse_centroids = torch.from_numpy(np.load(os.path.join(path, "coarse_centroids.npy"))).to(self.device)
        sq_codes = np.load(os.path.join(path, "sq_codes.npy"), mmap_mode=mmap_val)
        qjl_signs = np.load(os.path.join(path, "qjl_signs.npy"), mmap_mode=mmap_val)
        norms = np.load(os.path.join(path, "norms.npy"), mmap_mode=mmap_val)
        res_norms = np.load(os.path.join(path, "res_norms.npy"), mmap_mode=mmap_val)
        vector_ids = np.load(os.path.join(path, "vector_ids.npy"), mmap_mode=mmap_val)
        list_offsets = np.load(os.path.join(path, "list_offsets.npy"), mmap_mode=mmap_val)
        rot_op = np.load(os.path.join(path, "rot_op.npy"), mmap_mode=mmap_val)
        sq_centroids = np.load(os.path.join(path, "sq_centroids.npy"), mmap_mode=mmap_val)
        self.rot_op_np = rot_op
        self.rot_op_t = torch.from_numpy(rot_op).to(self.device)
        pq_data = ProdQuantized(
            sq_codes=sq_codes, qjl_signs=qjl_signs, norms=norms, centroids=sq_centroids,
            dim=self.dim, sq_bits=self.sq_bits, total_bits=self.bits,
            qjl_scale=self.qjl_scale, rot_op=rot_op, res_norms=res_norms
        )
        self.current_ivf_data = IVFData(
            coarse_centroids=coarse_centroids, pq_data=pq_data, vector_ids=vector_ids,
            list_offsets=list_offsets, n_list=self.ivf_nlist, n_probe=self.ivf_nprobe
        )
        # Load HNSW Navigator nếu có
        hnsw_path = os.path.join(path, "centroids.hnsw")
        if os.path.exists(hnsw_path):
            import faiss
            idx = faiss.read_index(str(hnsw_path))
            # Kiểm tra metric, nếu là L2 (0) thì phải tạo lại vì TQ dùng IP (1)
            if idx.metric_type != faiss.METRIC_INNER_PRODUCT:
                print("  Warning: HNSW index uses L2 metric. Rebuilding with Inner Product...")
                self.hnsw_navigator = faiss.IndexHNSWFlat(self.dim, 32, faiss.METRIC_INNER_PRODUCT)
                self.hnsw_navigator.add(coarse_centroids.cpu().numpy())
                faiss.write_index(self.hnsw_navigator, str(hnsw_path))
            else:
                self.hnsw_navigator = idx
        elif self.use_hnsw:
             print("  Warning: HNSW index missing. Building now...")
             import faiss
             self.hnsw_navigator = faiss.IndexHNSWFlat(self.dim, 32, faiss.METRIC_INNER_PRODUCT)
             self.hnsw_navigator.add(coarse_centroids.cpu().numpy())
             faiss.write_index(self.hnsw_navigator, str(hnsw_path))

        print(f"Loaded index from: {path} (mmap={mmap_val})")