model/loader.py

"""Singleton loader llama-cpp untuk Space dan lokal."""
from __future__ import annotations

import os
import threading
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from llama_cpp import Llama

NAMA_FILE_MODEL = "qwen2.5-7b-instruct-q4_k_m.gguf"
NAMA_SHARD_Q4KM_PERTAMA = "qwen2.5-7b-instruct-q4_k_m-00001-of-00002.gguf"
NAMA_SHARD_Q4KM_KEDUA = "qwen2.5-7b-instruct-q4_k_m-00002-of-00002.gguf"
NAMA_FILE_ALTERNATIF = (
    NAMA_FILE_MODEL,
    NAMA_SHARD_Q4KM_PERTAMA,
    "qwen2.5-7b-instruct-q2_k.gguf",
)

_kunci_inferensi = threading.Lock()
_model: object | None = None
_path_model_tercache: str | None = None


def kunci_inferensi():
    """Context manager untuk serialisasi panggilan model."""
    return _kunci_inferensi


def _jumlah_layer_gpu() -> int:
    nilai_env = os.getenv("N_GPU_LAYERS", "")
    if nilai_env != "":
        return int(nilai_env)
    if os.getenv("SPACE_ID"):
        return 0
    return -1


def _folder_punya_pasangan_shard(folder: Path) -> str | None:
    """llama.cpp butuh shard 00001 + 00002 di folder yang sama."""
    pertama = folder / NAMA_SHARD_Q4KM_PERTAMA
    kedua = folder / NAMA_SHARD_Q4KM_KEDUA
    if pertama.is_file() and kedua.is_file():
        return str(pertama.resolve())
    return None


def _cari_gguf_di_direktori(akar: Path, kedalaman_maks: int = 8) -> str | None:
    if not akar.is_dir():
        return None
    akar_resolved = akar.resolve()

    # Prioritas: pasangan shard Q4_K_M lengkap
    for root, dirs, berkas in os.walk(akar_resolved):
        kedalaman = len(Path(root).relative_to(akar_resolved).parts)
        if kedalaman > kedalaman_maks:
            dirs.clear()
            continue
        pasangan = _folder_punya_pasangan_shard(Path(root))
        if pasangan:
            return pasangan

    for root, dirs, berkas in os.walk(akar_resolved):
        kedalaman = len(Path(root).relative_to(akar_resolved).parts)
        if kedalaman > kedalaman_maks:
            dirs.clear()
            continue
        for nama_target in NAMA_FILE_ALTERNATIF:
            if nama_target in berkas:
                return str((Path(root) / nama_target).resolve())
    return None


def resolve_model_path() -> str:
    """Resolve path GGUF; hormati MODEL_PATH dari env Space."""
    env_path = os.getenv("MODEL_PATH")
    if env_path:
        path_env = Path(env_path)
        if path_env.is_file():
            return str(path_env.resolve())
        if path_env.is_dir():
            pasangan = _folder_punya_pasangan_shard(path_env)
            if pasangan:
                return pasangan
            ditemukan = _cari_gguf_di_direktori(path_env)
            if ditemukan:
                return ditemukan

    kandidat = [Path("./models") / n for n in NAMA_FILE_ALTERNATIF]
    kandidat += [Path(f"/data/{n}") for n in NAMA_FILE_ALTERNATIF]
    kandidat.append(Path("/data/models") / NAMA_FILE_MODEL)
    for path in kandidat:
        if path.is_file():
            return str(path.resolve())

    for folder in (Path("/data"), Path("/data/models"), Path("./models")):
        if folder.is_dir():
            pasangan = _folder_punya_pasangan_shard(folder)
            if pasangan:
                return pasangan
            ditemukan = _cari_gguf_di_direktori(folder)
            if ditemukan:
                return ditemukan

    cache_hub = Path.home() / ".cache" / "huggingface" / "hub"
    if cache_hub.is_dir():
        for repo_dir in cache_hub.glob("models--Qwen--Qwen2.5-7B-Instruct-GGUF*"):
            ditemukan = _cari_gguf_di_direktori(repo_dir)
            if ditemukan:
                return ditemukan

    return str((Path("./models") / NAMA_SHARD_Q4KM_PERTAMA).resolve())


def diagnosis_lingkungan_model() -> dict:
    """Info debug untuk log Space / skrip verifikasi."""
    path = resolve_model_path()
    folder = Path(path).parent if path else Path(".")
    ada_shard_1 = (folder / NAMA_SHARD_Q4KM_PERTAMA).is_file()
    ada_shard_2 = (folder / NAMA_SHARD_Q4KM_KEDUA).is_file()
    isi_data: list[str] = []
    if Path("/data").is_dir():
        try:
            isi_data = sorted(os.listdir("/data"))[:30]
        except OSError:
            isi_data = ["<tidak bisa baca /data>"]

    return {
        "space_id": os.getenv("SPACE_ID"),
        "model_path": path,
        "model_file_exists": Path(path).is_file() if path else False,
        "shard_00001": ada_shard_1,
        "shard_00002": ada_shard_2,
        "shard_pair_ok": ada_shard_1 and ada_shard_2,
        "n_gpu_layers": _jumlah_layer_gpu(),
        "n_ctx": os.getenv("N_CTX", "4096"),
        "skip_warmup": os.getenv("ONEIROS_SKIP_WARMUP"),
        "isi_data_root": isi_data,
    }


def get_model_path() -> str:
    global _path_model_tercache
    if _path_model_tercache is None:
        _path_model_tercache = resolve_model_path()
    return _path_model_tercache


def reset_model_path_cache() -> None:
    global _path_model_tercache, _model
    _path_model_tercache = None
    _model = None


def get_model() -> "Llama":
    global _model, _path_model_tercache
    from llama_cpp import Llama

    path = resolve_model_path()
    _path_model_tercache = path

    if _model is None:
        if not Path(path).is_file():
            diag = diagnosis_lingkungan_model()
            raise FileNotFoundError(
                f"Model GGUF tidak ditemukan di {path}. "
                f"Diagnosis: shard_pair_ok={diag['shard_pair_ok']}, "
                f"isi /data={diag['isi_data_root']}. "
                "Set MODEL_PATH ke shard 00001 atau tunggu preload."
            )
        _model = Llama(
            model_path=path,
            n_ctx=int(os.getenv("N_CTX", "4096")),
            n_gpu_layers=_jumlah_layer_gpu(),
            verbose=False,
        )
    return _model


MODEL_PATH = resolve_model_path()