TQ_engine_lib/__init__.py

import os
import sys
import subprocess
import shutil
import time
import torch

def should_rebuild(core_dir, binary_path):
    """Kiểm tra xem code Rust có mới hơn file binary hiện tại không."""
    if not os.path.exists(binary_path):
        return True
    
    binary_time = os.path.getmtime(binary_path)
    src_dir = os.path.join(core_dir, "src")
    
    if not os.path.exists(src_dir):
        return False
        
    for root, _, files in os.walk(src_dir):
        for f in files:
            if f.endswith(".rs"):
                if os.path.getmtime(os.path.join(root, f)) > binary_time:
                    return True
    return False

def build_native_core(force=False):
    """Tự động biên dịch lõi Rust SIMD nếu thiếu hoặc code đã thay đổi."""
    base_dir = os.path.dirname(os.path.abspath(__file__))
    core_dir = os.path.join(base_dir, "core")
    binary_name = "tq_native_lib.pyd" if os.name == "nt" else "tq_native_lib.so"
    binary_path = os.path.join(base_dir, binary_name)
    
    if not os.path.exists(core_dir):
        return False

    if not force and not should_rebuild(core_dir, binary_path):
        return True

    print(f"--- TurboQuant: {'Code changed! ' if os.path.exists(binary_path) else ''}Compiling Rust SIMD core ---")
    
    # 1. Check Cargo
    try:
        subprocess.run(["cargo", "--version"], check=True, capture_output=True)
    except (subprocess.CalledProcessError, FileNotFoundError):
        print("TurboQuant Error: 'cargo' not found. Please install Rust.")
        return False

    # 2. Build
    try:
        subprocess.run(["cargo", "build", "--release"], cwd=core_dir, check=True)
    except subprocess.CalledProcessError as e:
        print(f"TurboQuant Error: Build failed: {e}")
        return False

    # 3. Copy & Cleanup
    target_dir = os.path.join(core_dir, "target", "release")
    search_ext = ".dll" if os.name == "nt" else ".so"
    search_prefix = "tq_native_lib" if os.name == "nt" else "libtq_native_lib"
    
    found_lib = None
    if os.path.exists(target_dir):
        for f in os.listdir(target_dir):
            if (f.startswith("tq_native_lib") or f.startswith(search_prefix)) and f.endswith(search_ext):
                found_lib = f
                break
    
    if found_lib:
        try:
            shutil.copy2(os.path.join(target_dir, found_lib), binary_path)
            print(f"--- TurboQuant: Native core updated! ---")
        except PermissionError:
            print(f"--- TurboQuant: Native core locked but exists, skipping copy ---")
        except Exception as e:
            print(f"--- TurboQuant: Copy failed: {e} ---")
        
        # Cleanup
        try:
            shutil.rmtree(os.path.join(core_dir, "target"), ignore_errors=True)
            if os.path.exists(os.path.join(core_dir, "Cargo.lock")):
                os.remove(os.path.join(core_dir, "Cargo.lock"))
        except: pass
        return True
    
    return False

# --- AUTO-INIT PHASE ---
# Nếu tq_native_lib đã được cài sẵn (pip install từ maturin wheel trong Docker),
# bỏ qua hoàn toàn bước build Rust để tránh lỗi 'cargo: Permission denied'.
_native_available = False
try:
    import tq_native_lib as _tq_check  # Thử import global (pip-installed)
    _native_available = True
except ImportError:
    pass

if not _native_available:
    # Chỉ build khi module chưa được cài sẵn (dev mode, local)
    build_native_core()

try:
    from . import tq_native_lib  # Thử local .pyd/.so trước
except ImportError:
    try:
        import tq_native_lib  # Fallback: pip-installed (HF Space / Docker)
        import sys
        # Gắn vào namespace của package để các submodule dùng được
        sys.modules[__name__ + '.tq_native_lib'] = tq_native_lib
    except ImportError:
        print("TurboQuant Warning: Native SIMD mode disabled (Build failed or file missing).")

from .quantizer import TQEngine, ProdQuantized
from .codebook import ScalarQuantizer

__version__ = "0.4.1"

class TurboQuant:
    """
    High-level API for TurboQuant Vector Search.
    
    Hỗ trợ tự động biên dịch và nén SQ+QJL với IVF.
    """
    def __init__(self, dim: int, bits: int = 4, device: str = None, 
                 use_ivf: bool = False, ivf_nlist: int = 1024, ivf_nprobe: int = 32):
        self.engine = TQEngine(dim=dim, bits=bits, device=device, 
                               use_ivf=use_ivf, ivf_nlist=ivf_nlist, ivf_nprobe=ivf_nprobe)
        self.pq_data = None

    def index(self, vectors: torch.Tensor, online_clustering: bool = False):
        """Lập chỉ mục dữ liệu."""
        self.pq_data = self.engine.quantize(vectors, online_clustering=online_clustering)
        print(f"TurboQuant: Đã lập chỉ mục {vectors.shape[0]} vectors.")

    def search(self, query: torch.Tensor, top_k: int = 10):
        """Tìm kiếm Top-K."""
        if self.pq_data is None:
            raise ValueError("Index trống. Vui lòng gọi .index() trước.")
        return self.engine.native_cosine_search(query, self.pq_data, top_k=top_k)

    def save_index(self, directory: str, prefix: str):
        """Lưu chỉ mục xuống đĩa."""
        import os
        import numpy as np
        from .quantizer import IVFData, ProdQuantized
        
        os.makedirs(directory, exist_ok=True)
        if self.pq_data is None: raise ValueError("Index trống.")
            
        is_ivf = isinstance(self.pq_data, IVFData)
        config = {"dim": self.engine.dim, "bits": self.engine.bits, "use_ivf": is_ivf}
        
        if is_ivf:
            config["ivf_nlist"] = self.pq_data.n_list
            config["ivf_nprobe"] = self.pq_data.n_probe
            np.savez(os.path.join(directory, f"{prefix}_ivf_meta.npz"), 
                     coarse_centroids=self.pq_data.coarse_centroids.cpu().numpy(),
                     list_offsets=self.pq_data.list_offsets,
                     vector_ids=self.pq_data.vector_ids, **config)
            pq = self.pq_data.pq_data
        else:
            np.savez(os.path.join(directory, f"{prefix}_meta.npz"), **config)
            pq = self.pq_data
            
        np.save(os.path.join(directory, f"{prefix}_sq_codes.npy"), pq.sq_codes)
        np.save(os.path.join(directory, f"{prefix}_qjl_signs.npy"), pq.qjl_signs)
        np.save(os.path.join(directory, f"{prefix}_norms.npy"), pq.norms)
        np.save(os.path.join(directory, f"{prefix}_res_norms.npy"), pq.res_norms)
        
        np.savez(os.path.join(directory, f"{prefix}_pq_meta.npz"),
                 centroids=pq.centroids, dim=pq.dim, sq_bits=pq.sq_bits,
                 total_bits=pq.total_bits, qjl_scale=pq.qjl_scale, rot_op=pq.rot_op)
        print(f"TurboQuant: Đã lưu tại {directory} với prefix '{prefix}'.")

    def load_index(self, directory: str, prefix: str, use_mmap: bool = True):
        """Tải chỉ mục từ đĩa."""
        import os
        import numpy as np
        from .quantizer import IVFData, ProdQuantized
        
        meta_ivf_path = os.path.join(directory, f"{prefix}_ivf_meta.npz")
        meta_flat_path = os.path.join(directory, f"{prefix}_meta.npz")
        is_ivf = os.path.exists(meta_ivf_path)
        
        if is_ivf:
            meta = np.load(meta_ivf_path)
            self.engine.use_ivf = True
            self.engine.ivf_nlist = int(meta["ivf_nlist"])
            self.engine.ivf_nprobe = int(meta["ivf_nprobe"])
        else:
            meta = np.load(meta_flat_path)
            self.engine.use_ivf = False
            
        self.engine.dim = int(meta["dim"])
        self.engine.bits = int(meta["bits"])
        pq_meta = np.load(os.path.join(directory, f"{prefix}_pq_meta.npz"))
        
        import platform
        mmap_val = 'r'
        if platform.system() == "Windows" and not use_mmap:
            mmap_val = None
        elif platform.system() != "Windows":
            mmap_val = 'r'

        pq = ProdQuantized(
            sq_codes=np.load(os.path.join(directory, f"{prefix}_sq_codes.npy"), mmap_mode=mmap_val),
            qjl_signs=np.load(os.path.join(directory, f"{prefix}_qjl_signs.npy"), mmap_mode=mmap_val),
            norms=np.load(os.path.join(directory, f"{prefix}_norms.npy"), mmap_mode=mmap_val),
            res_norms=np.load(os.path.join(directory, f"{prefix}_res_norms.npy"), mmap_mode=mmap_val),
            centroids=pq_meta["centroids"], dim=int(pq_meta["dim"]),
            sq_bits=int(pq_meta["sq_bits"]), total_bits=int(pq_meta["total_bits"]),
            qjl_scale=float(pq_meta["qjl_scale"]), rot_op=pq_meta["rot_op"]
        )
        
        if is_ivf:
            self.pq_data = IVFData(
                coarse_centroids=torch.from_numpy(meta["coarse_centroids"]).to(self.engine.device),
                pq_data=pq, vector_ids=meta["vector_ids"], list_offsets=meta["list_offsets"],
                n_list=self.engine.ivf_nlist, n_probe=self.engine.ivf_nprobe)
        else:
            self.pq_data = pq
        print(f"TurboQuant: Đã tải chỉ mục từ {directory}.")