infer.py

#!/usr/bin/env python3
"""
Standalone inference script for SFT checkpoints.

Example:
  python sft/infer_sft.py --interactive --mode chat
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import random
import re
import sys
from pathlib import Path
from typing import Optional, Iterator, List, Dict, Tuple

import numpy as np
import torch
import torch.nn.functional as F
from safetensors import safe_open

from config_sft import get_config

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
if hasattr(sys.stderr, "reconfigure"):
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")

# Ensure both repository root and sft module path can be imported when running this file
# directly from different working directories.
SFT_DIR = Path(__file__).resolve().parent
REPO_ROOT = SFT_DIR.parent


def _prioritize_import_path(path: Path) -> None:
    p = str(path)
    while p in sys.path:
        sys.path.remove(p)
    sys.path.insert(0, p)


# Ensure sft-local modules are resolved before same-name modules at repo root.
_prioritize_import_path(REPO_ROOT)
_prioritize_import_path(SFT_DIR)

from model import GPTModel
from toba_tokenizer import TobaTokenizer
from translation_wrapper import detect_translation_request, stream_translate_chunked, translate_chunked, translate_single

MODEL_DIR_ENV_VAR = "TOBA_SFT_MODEL_DIR"
CHECKPOINT_ENV_VAR = "TOBA_SFT_CHECKPOINT"
VOCAB_ENV_VAR = "TOBA_SFT_VOCAB"
DEFAULT_MODEL_DIR_CANDIDATES = [
    SFT_DIR / "safetensors",
]
DEFAULT_CHECKPOINT_CANDIDATES = [
    SFT_DIR / "models" / "best_model_sft.pt",
    SFT_DIR / "best_model_sft.pt",
    SFT_DIR / "checkpoint" / "best_model_sft.pt",
    SFT_DIR / "checkpoints_sft" / "best_model_sft.pt",
]
DEFAULT_VOCAB_CANDIDATES = [
    SFT_DIR / "vocab" / "toba_vocab_kbbi_2500_new.json",
    SFT_DIR / "toba_vocab_kbbi_2500_new.json",
    REPO_ROOT / "toba_vocab_kbbi_2500_new.json",
]

MODE_ALIASES = {
    "completion": "completion",
    "chat": "chat",
    "single": "chat",  # backward-compatible alias
    "multi": "chat",   # backward-compatible alias
}


def _seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def _configure_deterministic_runtime(enabled: bool) -> None:
    if not enabled:
        return

    # cuBLAS reads this env var to stabilize some CUDA kernels.
    os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    if hasattr(torch.backends.cuda.matmul, "allow_tf32"):
        torch.backends.cuda.matmul.allow_tf32 = False
    if hasattr(torch.backends.cudnn, "allow_tf32"):
        torch.backends.cudnn.allow_tf32 = False
    try:
        torch.use_deterministic_algorithms(True, warn_only=True)
    except Exception as exc:
        print(f"Warning: failed to enable deterministic algorithms fully: {exc}")


def _default_checkpoint_path() -> Path:
    return _resolve_asset_path(
        explicit_path=None,
        env_var=CHECKPOINT_ENV_VAR,
        candidates=DEFAULT_CHECKPOINT_CANDIDATES,
        label="Checkpoint",
    )


def _resolve_model_dir_path(model_dir_path: Optional[Path]) -> Optional[Path]:
    if model_dir_path is not None:
        path = Path(model_dir_path).expanduser()
        if path.exists() and path.is_dir():
            return path.resolve()
        raise FileNotFoundError(f"Model dir not found at CLI path: {path}")

    env_value = os.getenv(MODEL_DIR_ENV_VAR)
    if env_value:
        path = Path(env_value).expanduser()
        if path.exists() and path.is_dir():
            return path.resolve()
        raise FileNotFoundError(f"Model dir not found at ${MODEL_DIR_ENV_VAR}: {path}")

    for candidate in DEFAULT_MODEL_DIR_CANDIDATES:
        if candidate.exists() and candidate.is_dir():
            return candidate.resolve()

    return None


def _default_vocab_path() -> Path:
    return _resolve_asset_path(
        explicit_path=None,
        env_var=VOCAB_ENV_VAR,
        candidates=DEFAULT_VOCAB_CANDIDATES,
        label="Vocab",
    )


def _resolve_asset_path(
    explicit_path: Optional[Path],
    env_var: str,
    candidates: List[Path],
    label: str,
) -> Path:
    if explicit_path is not None:
        path = Path(explicit_path).expanduser()
        if path.exists():
            return path.resolve()
        raise FileNotFoundError(f"{label} not found at CLI path: {path}")

    env_value = os.getenv(env_var)
    if env_value:
        path = Path(env_value).expanduser()
        if path.exists():
            return path.resolve()
        raise FileNotFoundError(f"{label} not found at ${env_var}: {path}")

    for candidate in candidates:
        if candidate.exists():
            return candidate.resolve()

    return candidates[0].resolve()


def _resolve_checkpoint_path(checkpoint_path: Optional[Path]) -> Path:
    return _resolve_asset_path(
        explicit_path=checkpoint_path,
        env_var=CHECKPOINT_ENV_VAR,
        candidates=DEFAULT_CHECKPOINT_CANDIDATES,
        label="Checkpoint",
    )


def _resolve_vocab_path(vocab_path: Optional[Path]) -> Path:
    return _resolve_asset_path(
        explicit_path=vocab_path,
        env_var=VOCAB_ENV_VAR,
        candidates=DEFAULT_VOCAB_CANDIDATES,
        label="Vocab",
    )


def _file_sha256(path: Path, chunk_size: int = 1024 * 1024) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(chunk_size), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _print_file_diagnostics(label: str, path: Path) -> None:
    stat = path.stat()
    print(f"{label}: {path}")
    print(f"{label} size: {stat.st_size}")
    print(f"{label} SHA256: {_file_sha256(path)}")


def _print_model_dir_diagnostics(model_dir: Path) -> None:
    print(f"Model dir: {model_dir}")
    files = [model_dir / "model.safetensors.index.json"]
    files.extend(sorted(model_dir.glob("model*.safetensors")))
    for path in files:
        if path.exists() and path.is_file():
            _print_file_diagnostics(f"Model file {path.name}", path)


def _extract_model_state(checkpoint_obj) -> dict:
    if not isinstance(checkpoint_obj, dict):
        raise ValueError(f"Unsupported checkpoint format: {type(checkpoint_obj)}")

    for key in ("model_state", "model_state_dict", "state_dict"):
        value = checkpoint_obj.get(key)
        if isinstance(value, dict):
            return value

    raise ValueError("Checkpoint missing model state. Expected keys: model_state, model_state_dict, or state_dict.")


def _strip_module_prefix(state: dict) -> dict:
    if not state:
        return {}
    has_module = all(k.startswith("module.") for k in state.keys())
    if not has_module:
        return state
    return {k[len("module."):]: v for k, v in state.items()}


def _load_checkpoint(checkpoint_path: Path) -> dict:
    try:
        ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    except RuntimeError as exc:
        message = str(exc)
        if "failed finding central directory" in message.lower():
            raise RuntimeError(
                f"Checkpoint corrupt/terpotong: {checkpoint_path}. "
                "File zip PyTorch tidak punya central directory. "
                "Gunakan checkpoint lain atau salin ulang file checkpoint."
            ) from exc
        raise
    state = _extract_model_state(ckpt)
    return _strip_module_prefix(state)


def _load_safetensors_model_dir(model_dir: Path) -> dict:
    index_path = model_dir / "model.safetensors.index.json"
    state: dict[str, torch.Tensor] = {}

    if index_path.exists():
        with index_path.open("r", encoding="utf-8") as f:
            index = json.load(f)
        weight_map = index.get("weight_map", {})
        if not isinstance(weight_map, dict) or not weight_map:
            raise ValueError(f"Invalid safetensors index weight_map: {index_path}")

        shards: dict[str, list[str]] = {}
        for tensor_name, shard_name in weight_map.items():
            shards.setdefault(str(shard_name), []).append(str(tensor_name))

        for shard_name in sorted(shards):
            shard_path = model_dir / shard_name
            if not shard_path.exists():
                raise FileNotFoundError(f"Safetensors shard missing: {shard_path}")
            with safe_open(str(shard_path), framework="pt", device="cpu") as f:
                for tensor_name in sorted(shards[shard_name]):
                    state[tensor_name] = f.get_tensor(tensor_name).contiguous()
        return _strip_module_prefix(state)

    candidates = [
        model_dir / "model.safetensors",
        model_dir / "model-00001-of-00001.safetensors",
    ]
    model_path = next((path for path in candidates if path.exists()), None)
    if model_path is None:
        raise FileNotFoundError(f"No safetensors model found in {model_dir}")

    with safe_open(str(model_path), framework="pt", device="cpu") as f:
        for tensor_name in f.keys():
            state[tensor_name] = f.get_tensor(tensor_name).contiguous()
    return _strip_module_prefix(state)


def _infer_vocab_from_state(state: dict) -> Optional[int]:
    wte = state.get("wte.weight")
    if torch.is_tensor(wte) and wte.ndim == 2:
        return int(wte.shape[0])
    return None


def _infer_arch_signature_from_state(state: dict) -> Dict[str, Optional[int]]:
    wte = state.get("wte.weight")
    wpe = state.get("wpe.weight")
    n_emb = int(wte.shape[1]) if torch.is_tensor(wte) and wte.ndim == 2 else None
    ctx_len = int(wpe.shape[0]) if torch.is_tensor(wpe) and wpe.ndim == 2 else None

    block_indices: List[int] = []
    block_re = re.compile(r"^blocks\.(\d+)\.")
    for key in state.keys():
        match = block_re.match(key)
        if match:
            block_indices.append(int(match.group(1)))
    n_layers = (max(block_indices) + 1) if block_indices else None

    return {
        "n_layers": n_layers,
        "n_emb": n_emb,
        "ctx_len": ctx_len,
    }


def _resolve_model_size_from_state(state: dict) -> Tuple[str, Dict[str, Optional[int]]]:
    sig = _infer_arch_signature_from_state(state)
    target_layers = sig.get("n_layers")
    target_emb = sig.get("n_emb")

    candidates = []
    for size in ("small", "medium", "large"):
        cfg = get_config(size).model
        layer_delta = abs((target_layers or cfg.n_layers) - cfg.n_layers)
        emb_delta = abs((target_emb or cfg.n_emb) - cfg.n_emb)
        exact_bonus = 0
        if target_layers == cfg.n_layers:
            exact_bonus += 1
        if target_emb == cfg.n_emb:
            exact_bonus += 1
        score = (layer_delta, emb_delta, -exact_bonus)
        candidates.append((score, size))

    candidates.sort(key=lambda item: item[0])
    return candidates[0][1], sig


def _apply_repetition_penalty(logits: torch.Tensor, generated_ids: List[int], penalty: float) -> torch.Tensor:
    if penalty == 1.0:
        return logits

    used = set(generated_ids)
    for token_id in used:
        if token_id >= logits.size(-1):
            continue
        if logits[0, token_id] > 0:
            logits[0, token_id] /= penalty
        else:
            logits[0, token_id] *= penalty
    return logits


def _apply_no_repeat_ngram(logits: torch.Tensor, generated_ids: List[int], ngram_size: int) -> torch.Tensor:
    if ngram_size <= 1 or len(generated_ids) < ngram_size:
        return logits

    existing_ngrams = set()
    for i in range(len(generated_ids) - ngram_size + 1):
        existing_ngrams.add(tuple(generated_ids[i : i + ngram_size]))

    prefix = tuple(generated_ids[-(ngram_size - 1) :])
    blocked = set(
        ngram[-1]
        for ngram in existing_ngrams
        if len(ngram) == ngram_size and ngram[:-1] == prefix
    )

    for token_id in blocked:
        if token_id < logits.size(-1):
            logits[0, token_id] = -float("inf")
    return logits


def _top_k_top_p_filtering(logits: torch.Tensor, top_k: int = 50, top_p: float = 0.9) -> torch.Tensor:
    top_k = min(top_k, logits.size(-1))
    if top_k > 0:
        min_top_k = torch.topk(logits, top_k).values[..., -1, None]
        logits = torch.where(logits < min_top_k, torch.full_like(logits, float("-inf")), logits)

    if top_p < 1.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
        sorted_mask = cumulative_probs > top_p
        sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
        sorted_mask[..., 0] = 0
        remove = sorted_mask.scatter(dim=-1, index=sorted_indices, src=sorted_mask)
        logits = logits.masked_fill(remove, float("-inf"))
    return logits


def _resolve_mode(mode: str, *, announce_alias: bool = False) -> str:
    mode_raw = str(mode).strip().lower()
    resolved = MODE_ALIASES.get(mode_raw)
    if resolved is None:
        allowed = ", ".join(sorted(MODE_ALIASES.keys()))
        raise ValueError(f"Invalid mode: {mode}. Allowed: {allowed}")
    if announce_alias and mode_raw in {"single", "multi"}:
        print(f"Mode '{mode_raw}' dipetakan ke mode 'chat' (alias kompatibilitas).")
    return resolved


class SFTInferenceEngine:
    """Simple inference engine for SFT checkpoints."""

    _ROLE_STOP_PATTERN = re.compile(r"(?im)(?:^|\n)\s*(?:system|user|assistant)\s*:")
    _SPECIAL_ROLE_STOPS = ("<|system|>", "<|user|>", "<|assistant|>", "<|eos|>")

    def __init__(
        self,
        model_dir_path: Optional[Path],
        checkpoint_path: Optional[Path],
        vocab_path: Optional[Path],
        size: str = "auto",
        device: str = "auto",
        do_compile: bool = False,
        dtype: str = "auto",
        seed: Optional[int] = None,
        deterministic: bool = False,
    ) -> None:
        if model_dir_path is not None and checkpoint_path is not None:
            raise ValueError("Use either --model-dir for safetensors or --checkpoint for .pt, not both.")

        self.model_dir_path = None
        self.checkpoint_path = None
        if checkpoint_path is not None:
            self.checkpoint_path = _resolve_checkpoint_path(checkpoint_path)
        else:
            self.model_dir_path = _resolve_model_dir_path(model_dir_path)
            if self.model_dir_path is None:
                self.checkpoint_path = _resolve_checkpoint_path(None)

        self.vocab_path = _resolve_vocab_path(vocab_path)
        self.seed = seed
        self.deterministic = deterministic

        self.device = self._resolve_device(device)
        print(f"Using device: {self.device}")
        if self.model_dir_path is not None:
            _print_model_dir_diagnostics(self.model_dir_path)
        else:
            _print_file_diagnostics("Checkpoint", self.checkpoint_path)
        _print_file_diagnostics("Vocab", self.vocab_path)

        if self.model_dir_path is not None:
            state = _load_safetensors_model_dir(self.model_dir_path)
        else:
            state = _load_checkpoint(self.checkpoint_path)
        auto_size, signature = _resolve_model_size_from_state(state)
        resolved_size = auto_size if size == "auto" else size
        self.config = get_config(resolved_size)
        self.resolved_size = resolved_size
        print(
            "Resolved model size: "
            f"{self.resolved_size} "
            f"(n_layers={signature.get('n_layers')}, n_emb={signature.get('n_emb')})"
        )

        # Load tokenizer
        with open(self.vocab_path, "r", encoding="utf-8") as f:
            vocab_data = json.load(f)
        self.tokenizer = TobaTokenizer(vocab=vocab_data)

        inferred_vocab = _infer_vocab_from_state(state)
        if inferred_vocab is not None and inferred_vocab != self.config.model.vocab_size:
            print(
                f"Adjusting vocab_size from {self.config.model.vocab_size} "
                f"to {inferred_vocab} based on checkpoint."
            )
            self.config.model.vocab_size = inferred_vocab
        if signature.get("ctx_len") is not None and signature["ctx_len"] != self.config.model.ctx_len:
            print(
                f"Adjusting ctx_len from {self.config.model.ctx_len} "
                f"to {signature['ctx_len']} based on checkpoint."
            )
            self.config.model.ctx_len = int(signature["ctx_len"])

        self.model = GPTModel(self.config.model)
        missing_keys, unexpected_keys = self.model.load_state_dict(state, strict=False)
        if missing_keys:
            print(f"Missing keys: {len(missing_keys)}")
        if unexpected_keys:
            print(f"Unexpected keys: {len(unexpected_keys)}")

        self.model.to(self.device)
        resolved_dtype = self._resolve_model_dtype(dtype)
        self.model = self.model.to(dtype=resolved_dtype)
        self.model.eval()
        self.model.requires_grad_(False)

        if do_compile and hasattr(torch, "compile"):
            try:
                self.model = torch.compile(self.model)
                print("torch.compile enabled")
            except Exception as exc:
                print(f"torch.compile failed, continuing without compile: {exc}")

        self.model_dtype = resolved_dtype
        print(f"Model dtype: {self.model_dtype}")

    @staticmethod
    def _normalize_role(role: str) -> str:
        role_low = str(role).strip().lower()
        if role_low in {"assistant", "bot", "model"}:
            return "Assistant"
        if role_low == "system":
            return "System"
        return "User"

    @staticmethod
    def _clean_text(text: str) -> str:
        return str(text).replace("\r\n", "\n").replace("\r", "\n").strip()

    def _format_turn(self, role: str, text: str) -> str:
        role_norm = self._normalize_role(role)
        text_norm = self._clean_text(text)
        return f"{role_norm}: {text_norm}"

    def _compose_chat_prompt(
        self,
        user_prompt: str,
    ) -> str:
        text = self._clean_text(user_prompt)
        # Stateless chat template. The tokenizer has dedicated role tokens;
        # using literal "User:"/"Assistant:" makes those labels likely to leak.
        return f"<|user|>{text}<|assistant|>"

    def build_chat_prompt(self, user_prompt: str) -> str:
        # Stateless chat: each request only includes current user turn.
        return self._compose_chat_prompt(
            user_prompt=user_prompt,
        )

    def _resolve_device(self, device: str) -> torch.device:
        if device != "auto":
            return torch.device(device)
        if torch.cuda.is_available():
            return torch.device("cuda:0")
        return torch.device("cpu")

    def _resolve_model_dtype(self, dtype: str) -> torch.dtype:
        dtype_raw = str(dtype).strip().lower()
        if dtype_raw == "auto":
            return torch.bfloat16 if self.device.type == "cuda" else torch.float32
        if dtype_raw == "float32":
            return torch.float32
        if dtype_raw == "bfloat16":
            if self.device.type != "cuda":
                print("bfloat16 requested on non-CUDA device; falling back to float32.")
                return torch.float32
            return torch.bfloat16
        raise ValueError(f"Unsupported dtype: {dtype}. Allowed: auto, float32, bfloat16")

    def encode(self, text: str) -> torch.Tensor:
        encoded = self.tokenizer.encode(text)
        if isinstance(encoded, dict):
            token_ids = encoded.get("id", [])
        else:
            token_ids = encoded

        safe = []
        for token_id in token_ids:
            if token_id is None:
                safe.append(self.config.model.unk_token_id)
            elif token_id >= self.config.model.vocab_size:
                safe.append(self.config.model.unk_token_id)
            else:
                safe.append(int(token_id))

        safe = [self.config.model.bos_token_id] + safe
        if len(safe) > self.config.model.ctx_len:
            safe = safe[-self.config.model.ctx_len :]
        return torch.tensor([safe], device=self.device, dtype=torch.long)

    def decode(self, token_ids: List[int]) -> str:
        return self.tokenizer.decode(token_ids)

    @classmethod
    def _generation_stop_index(cls, text: str) -> Optional[int]:
        stops = []
        for marker in cls._SPECIAL_ROLE_STOPS:
            index = text.find(marker)
            if index >= 0:
                stops.append(index)

        match = cls._ROLE_STOP_PATTERN.search(text)
        if match is not None:
            stops.append(match.start())

        return min(stops) if stops else None

    @staticmethod
    def _trim_past_kv(
        past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]],
        keep_len: int,
    ) -> Optional[List[Tuple[torch.Tensor, torch.Tensor]]]:
        if past_kv is None:
            return None
        if keep_len <= 0:
            return [(k[:, :, 0:0, :], v[:, :, 0:0, :]) for (k, v) in past_kv]

        trimmed: List[Tuple[torch.Tensor, torch.Tensor]] = []
        for k, v in past_kv:
            if k.size(2) > keep_len:
                k = k[:, :, -keep_len:, :]
                v = v[:, :, -keep_len:, :]
            trimmed.append((k, v))
        return trimmed

    def generate_stream(
        self,
        prompt: str,
        max_new_tokens: int,
        temperature: float = 1.0,
        top_k: int = 40,
        top_p: float = 0.92,
        do_sample: bool = True,
        repetition_penalty: float = 1.2,
        no_repeat_ngram_size: int = 0,
        use_kv_cache: bool = True,
    ) -> Iterator[str]:
        with torch.inference_mode():
            if self.seed is not None:
                _seed_everything(self.seed)

            input_ids = self.encode(prompt)
            generated_ids: List[int] = []
            decoded_text = ""

            def emit_delta() -> Iterator[str]:
                nonlocal decoded_text
                next_text = self.decode(generated_ids) if generated_ids else ""
                stop_index = self._generation_stop_index(next_text)
                should_stop = stop_index is not None
                if should_stop:
                    next_text = next_text[:stop_index].rstrip()
                delta = next_text[len(decoded_text) :]
                decoded_text = next_text
                if delta:
                    yield delta
                return should_stop

            if not use_kv_cache:
                current_ids = input_ids.clone()
                for _ in range(max_new_tokens):
                    if current_ids.size(1) > self.config.model.ctx_len:
                        current_ids = current_ids[:, -self.config.model.ctx_len :]

                    with torch.amp.autocast(device_type="cuda", dtype=self.model_dtype, enabled=self.device.type == "cuda"):
                        logits = self.model(current_ids)[:, -1, :]

                    if repetition_penalty != 1.0:
                        logits = _apply_repetition_penalty(
                            logits,
                            generated_ids=generated_ids,
                            penalty=repetition_penalty,
                        )

                    if no_repeat_ngram_size > 0 and len(generated_ids) >= no_repeat_ngram_size - 1:
                        logits = _apply_no_repeat_ngram(
                            logits,
                            generated_ids=generated_ids,
                            ngram_size=no_repeat_ngram_size,
                        )

                    if do_sample and temperature > 0:
                        logits = logits / temperature
                        logits = _top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
                        probs = F.softmax(logits, dim=-1)
                        next_token = torch.multinomial(probs, num_samples=1)
                    else:
                        next_token = logits.argmax(dim=-1, keepdim=True)

                    next_token_id = int(next_token.item())
                    if next_token_id == self.config.model.eos_token_id or next_token_id == self.config.model.pad_token_id:
                        break

                    current_ids = torch.cat([current_ids, next_token], dim=1)
                    generated_ids.append(next_token_id)
                    should_stop = yield from emit_delta()
                    if should_stop:
                        return
                return

            input_ids = input_ids[:, -self.config.model.ctx_len :]
            with torch.amp.autocast(device_type="cuda", dtype=self.model_dtype, enabled=self.device.type == "cuda"):
                logits, past_kv, engram_state = self.model(
                    input_ids,
                    past_kv=None,
                    use_cache=True,
                    engram_state=None,
                )
            logits = logits[:, -1, :]

            for _ in range(max_new_tokens):
                if repetition_penalty != 1.0:
                    logits = _apply_repetition_penalty(
                        logits,
                        generated_ids=generated_ids,
                        penalty=repetition_penalty,
                    )

                if no_repeat_ngram_size > 0 and len(generated_ids) >= no_repeat_ngram_size - 1:
                    logits = _apply_no_repeat_ngram(
                        logits,
                        generated_ids=generated_ids,
                        ngram_size=no_repeat_ngram_size,
                    )

                if do_sample and temperature > 0:
                    logits_scaled = logits / temperature
                    logits_scaled = _top_k_top_p_filtering(logits_scaled, top_k=top_k, top_p=top_p)
                    probs = F.softmax(logits_scaled, dim=-1)
                    next_token = torch.multinomial(probs, num_samples=1)
                else:
                    next_token = logits.argmax(dim=-1, keepdim=True)

                next_token_id = int(next_token.item())
                if next_token_id == self.config.model.eos_token_id or next_token_id == self.config.model.pad_token_id:
                    break
                generated_ids.append(next_token_id)
                should_stop = yield from emit_delta()
                if should_stop:
                    return

                # Keep one slot for the current token to avoid pos-index overflow.
                past_kv = self._trim_past_kv(past_kv, keep_len=self.config.model.ctx_len - 1)

                with torch.amp.autocast(device_type="cuda", dtype=self.model_dtype, enabled=self.device.type == "cuda"):
                    step_logits, past_kv, engram_state = self.model(
                        next_token,
                        past_kv=past_kv,
                        use_cache=True,
                        engram_state=engram_state,
                    )
                logits = step_logits[:, -1, :]

    def generate(
        self,
        prompt: str,
        max_new_tokens: int,
        temperature: float = 1.0,
        top_k: int = 40,
        top_p: float = 0.92,
        do_sample: bool = True,
        repetition_penalty: float = 1.2,
        no_repeat_ngram_size: int = 0,
        use_kv_cache: bool = True,
    ) -> str:
        return "".join(
            self.generate_stream(
                prompt=prompt,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                do_sample=do_sample,
                repetition_penalty=repetition_penalty,
                no_repeat_ngram_size=no_repeat_ngram_size,
                use_kv_cache=use_kv_cache,
            )
        )

    def run_interactive(self, args) -> None:
        mode = _resolve_mode(args.mode, announce_alias=True)
        print("=" * 80)
        print("TOBA SFT Inference - Interactive")
        print("Type 'quit', 'exit', or '/q' to stop.")
        print("Commands: /temp, /topk, /topp, /tokens, /sample, /stream, /norep, /mode, /kv")
        print("Modes: completion | chat")
        print("Catatan: mode chat stateless dengan template plain.")
        print("=" * 80)

        temperature = args.temperature
        top_k = args.top_k
        top_p = args.top_p
        max_new_tokens = args.max_new_tokens
        do_sample = bool(args.do_sample)
        repetition_penalty = args.repetition_penalty
        no_repeat_ngram_size = args.no_repeat_ngram_size
        use_kv_cache = bool(args.kv_cache)
        stream = bool(args.stream)

        while True:
            try:
                prompt = input("\nAnda: ").strip()
            except (KeyboardInterrupt, EOFError):
                print("\nKeluar.")
                break

            if not prompt:
                continue
            low = prompt.lower()
            if low in {"quit", "exit", "/q"}:
                print("Selesai.")
                break

            if low.startswith("/mode "):
                value = low.split(maxsplit=1)[1].strip()
                if value in MODE_ALIASES:
                    mode = _resolve_mode(value, announce_alias=True)
                    print(f"mode={mode}")
                else:
                    print("Mode harus completion atau chat. Contoh: /mode chat")
                continue
            if low.startswith("/temp "):
                try:
                    temperature = float(low.split(maxsplit=1)[1])
                    print(f"temperature={temperature}")
                except ValueError:
                    print("Format salah. Contoh: /temp 0.8")
                continue
            if low.startswith("/topk "):
                try:
                    top_k = int(low.split(maxsplit=1)[1])
                    print(f"top_k={top_k}")
                except ValueError:
                    print("Format salah. Contoh: /topk 40")
                continue
            if low.startswith("/topp "):
                try:
                    top_p = float(low.split(maxsplit=1)[1])
                    print(f"top_p={top_p}")
                except ValueError:
                    print("Format salah. Contoh: /topp 0.92")
                continue
            if low.startswith("/tokens "):
                try:
                    max_new_tokens = int(low.split(maxsplit=1)[1])
                    print(f"max_new_tokens={max_new_tokens}")
                except ValueError:
                    print("Format salah. Contoh: /tokens 120")
                continue
            if low.startswith("/sample "):
                try:
                    do_sample = bool(int(low.split(maxsplit=1)[1]))
                    print(f"do_sample={do_sample}")
                except ValueError:
                    print("Format salah. Contoh: /sample 1")
                continue
            if low.startswith("/stream "):
                try:
                    stream = bool(int(low.split(maxsplit=1)[1]))
                    print(f"stream={stream}")
                except ValueError:
                    print("Format salah. Contoh: /stream 1")
                continue
            if low.startswith("/norep "):
                try:
                    repetition_penalty = float(low.split(maxsplit=1)[1])
                    print(f"repetition_penalty={repetition_penalty}")
                except ValueError:
                    print("Format salah. Contoh: /norep 1.2")
                continue
            if low.startswith("/kv "):
                try:
                    use_kv_cache = bool(int(low.split(maxsplit=1)[1]))
                    print(f"kv_cache={use_kv_cache}")
                except ValueError:
                    print("Format salah. Contoh: /kv 1")
                continue

            translation_request = detect_translation_request(prompt)
            if translation_request is not None:
                if stream:
                    print("TOBA: ", end="", flush=True)
                    for event in stream_translate_chunked(
                        engine=self,
                        text=translation_request.text,
                        source_lang=translation_request.source_lang,
                        target_lang=translation_request.target_lang,
                        mode=mode,
                        max_new_tokens=max_new_tokens,
                        repetition_penalty=repetition_penalty,
                        no_repeat_ngram_size=no_repeat_ngram_size,
                        use_kv_cache=use_kv_cache,
                    ):
                        if event.get("event") != "delta":
                            continue
                        print(str(event.get("delta", "")), end="", flush=True)
                    print()
                    continue

                result = translate_chunked(
                    engine=self,
                    text=translation_request.text,
                    source_lang=translation_request.source_lang,
                    target_lang=translation_request.target_lang,
                    mode=mode,
                    max_new_tokens=max_new_tokens,
                    repetition_penalty=repetition_penalty,
                    no_repeat_ngram_size=no_repeat_ngram_size,
                    use_kv_cache=use_kv_cache,
                )
                print(f"TOBA: {result.text}")
                continue

            if mode == "completion":
                model_prompt = self._clean_text(prompt)
            else:  # chat
                model_prompt = self.build_chat_prompt(
                    user_prompt=prompt,
                )

            print("TOBA: ", end="", flush=True)
            if stream:
                for delta in self.generate_stream(
                    prompt=model_prompt,
                    max_new_tokens=max_new_tokens,
                    temperature=temperature,
                    top_k=top_k,
                    top_p=top_p,
                    do_sample=do_sample,
                    repetition_penalty=repetition_penalty,
                    no_repeat_ngram_size=no_repeat_ngram_size,
                    use_kv_cache=use_kv_cache,
                ):
                    print(delta, end="", flush=True)
                print()
            else:
                response = self.generate(
                    prompt=model_prompt,
                    max_new_tokens=max_new_tokens,
                    temperature=temperature,
                    top_k=top_k,
                    top_p=top_p,
                    do_sample=do_sample,
                    repetition_penalty=repetition_penalty,
                    no_repeat_ngram_size=no_repeat_ngram_size,
                    use_kv_cache=use_kv_cache,
                )
                print(response)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser("SFT inference")
    parser.add_argument(
        "--size",
        type=str,
        choices=["auto", "small", "medium", "large"],
        default="auto",
        help="Optional override. Default auto-detect from model tensors.",
    )
    parser.add_argument(
        "--model-dir",
        type=Path,
        default=None,
        help=(
            "Path to safetensors model directory. Resolve order: "
            "--model-dir -> $TOBA_SFT_MODEL_DIR -> safetensors"
        ),
    )
    parser.add_argument(
        "--checkpoint",
        type=Path,
        default=None,
        help=(
            "Optional legacy .pt checkpoint path. Resolve order: "
            "--checkpoint -> $TOBA_SFT_CHECKPOINT -> "
            "models/best_model_sft.pt -> best_model_sft.pt -> "
            "checkpoint/best_model_sft.pt -> checkpoints_sft/best_model_sft.pt"
        ),
    )
    parser.add_argument(
        "--vocab",
        type=Path,
        default=None,
        help=(
            "Path vocab tokenizer. Resolve order: "
            "--vocab -> $TOBA_SFT_VOCAB -> "
            "vocab/toba_vocab_kbbi_2500_new.json -> "
            "toba_vocab_kbbi_2500_new.json -> ../toba_vocab_kbbi_2500_new.json"
        ),
    )
    parser.add_argument("--device", type=str, default="auto")
    parser.add_argument(
        "--dtype",
        type=str,
        choices=["auto", "float32", "bfloat16"],
        default="auto",
        help="Model weight/inference dtype. Use float32 for closer CPU-vs-GPU comparisons.",
    )
    parser.add_argument(
        "--mode",
        type=str,
        choices=sorted(MODE_ALIASES.keys()),
        default="chat",
        help="Inference mode: completion|chat (alias lama: single|multi -> chat)",
    )
    parser.add_argument(
        "--task",
        type=str,
        choices=["generate", "translate"],
        default="generate",
        help="Use translate for the chunked long-translation wrapper.",
    )
    parser.add_argument("--source-lang", type=str, default="auto", help="Source language for --task translate.")
    parser.add_argument("--target-lang", type=str, default="Indonesia", help="Target language for --task translate.")
    parser.add_argument(
        "--chunked-translation",
        type=int,
        default=1,
        help="1=split long translate requests into chunks, 0=single minimal translate prompt.",
    )
    parser.add_argument(
        "--auto-translate-detect",
        type=int,
        default=1,
        help="1=route prompts starting with terjemahkan/translate into translate wrapper.",
    )
    parser.add_argument(
        "--compare-single",
        type=int,
        default=0,
        help="1=print single-prompt translation beside chunked translation.",
    )
    parser.add_argument(
        "--consistency-pass",
        type=int,
        default=0,
        help="1=run a light terminology consistency pass for multi-chunk translation.",
    )
    parser.add_argument(
        "--chunk-source-tokens",
        type=int,
        default=None,
        help="Optional max source tokens per translation chunk.",
    )
    parser.add_argument("--interactive", action="store_true", help="Start interactive mode")
    parser.add_argument("--prompt", type=str, default=None, help="Prompt inference single-shot")
    parser.add_argument("--stream", type=int, default=1, help="1=stream output tokens/chunks, 0=print after completion")
    parser.add_argument("--max-new-tokens", type=int, default=200)
    parser.add_argument("--temperature", type=float, default=0.4)
    parser.add_argument("--top-k", type=int, default=30)
    parser.add_argument("--top-p", type=float, default=0.85)
    parser.add_argument("--do-sample", type=int, default=1, help="1=sample, 0=greedy")
    parser.add_argument("--repetition-penalty", type=float, default=1.2)
    parser.add_argument("--no-repeat-ngram-size", type=int, default=3)
    parser.add_argument("--kv-cache", type=int, default=1, help="1=enable KV cache decode, 0=disable")
    parser.add_argument("--compile", action="store_true", help="Enable torch.compile if available")
    parser.add_argument("--seed", type=int, default=42, help="Seed for reproducible generation.")
    parser.add_argument(
        "--deterministic",
        type=int,
        default=1,
        help="1=prefer deterministic runtime settings, 0=allow faster nondeterministic kernels",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    deterministic = bool(args.deterministic)
    _configure_deterministic_runtime(deterministic)
    _seed_everything(args.seed)

    if deterministic and args.compile:
        print("Deterministic mode aktif: torch.compile dimatikan untuk menjaga reproduksibilitas.")
        args.compile = False

    mode = _resolve_mode(args.mode, announce_alias=True)
    engine = SFTInferenceEngine(
        model_dir_path=args.model_dir,
        checkpoint_path=args.checkpoint,
        vocab_path=args.vocab,
        size=args.size,
        device=args.device,
        do_compile=args.compile,
        dtype=args.dtype,
        seed=args.seed,
        deterministic=deterministic,
    )

    if args.interactive or args.prompt is None:
        engine.run_interactive(args)
        return

    detected_translation = detect_translation_request(
        args.prompt,
        default_source_lang=args.source_lang,
        default_target_lang=args.target_lang,
    ) if bool(args.auto_translate_detect) else None

    if args.task == "translate" or detected_translation is not None:
        source_text = args.prompt
        source_lang = args.source_lang
        target_lang = args.target_lang
        if detected_translation is not None:
            source_text = detected_translation.text
            if args.source_lang == "auto":
                source_lang = detected_translation.source_lang
            if args.target_lang == "Indonesia":
                target_lang = detected_translation.target_lang

        translate_base = {
            "engine": engine,
            "text": source_text,
            "source_lang": source_lang,
            "target_lang": target_lang,
            "mode": mode,
        }
        translate_generation = {
            "max_new_tokens": args.max_new_tokens,
            "repetition_penalty": args.repetition_penalty,
            "no_repeat_ngram_size": args.no_repeat_ngram_size,
            "use_kv_cache": bool(args.kv_cache),
        }

        if bool(args.chunked_translation):
            if bool(args.stream) and not bool(args.compare_single):
                for event in stream_translate_chunked(
                    **translate_base,
                    **translate_generation,
                    chunk_source_tokens=args.chunk_source_tokens,
                    consistency_pass=bool(args.consistency_pass),
                    compare_single=False,
                ):
                    if event.get("event") != "delta":
                        continue
                    print(str(event.get("delta", "")), end="", flush=True)
                print()
                return

            result = translate_chunked(
                **translate_base,
                **translate_generation,
                chunk_source_tokens=args.chunk_source_tokens,
                consistency_pass=bool(args.consistency_pass),
                compare_single=bool(args.compare_single),
            )
            if bool(args.compare_single):
                print("=== SINGLE PROMPT ===")
                print(result.single_prompt_result or "")
                print("\n=== CHUNKED TRANSLATION ===")
                print(result.text)
                print("\n=== METADATA ===")
                print(json.dumps(result.metadata, ensure_ascii=False, indent=2))
            else:
                print(result.text)
            return

        response = translate_single(
            **translate_base,
            generation_params={
                "max_new_tokens": args.max_new_tokens,
                "temperature": 0.0,
                "top_k": 0,
                "top_p": 1.0,
                "do_sample": False,
                "repetition_penalty": args.repetition_penalty,
                "no_repeat_ngram_size": args.no_repeat_ngram_size,
                "use_kv_cache": bool(args.kv_cache),
            },
        )
        print(response)
        return

    if mode == "completion":
        model_prompt = engine._clean_text(args.prompt)
    else:  # chat
        model_prompt = engine.build_chat_prompt(
            user_prompt=args.prompt,
        )

    if bool(args.stream):
        for delta in engine.generate_stream(
            prompt=model_prompt,
            max_new_tokens=args.max_new_tokens,
            temperature=args.temperature,
            top_k=args.top_k,
            top_p=args.top_p,
            do_sample=bool(args.do_sample),
            repetition_penalty=args.repetition_penalty,
            no_repeat_ngram_size=args.no_repeat_ngram_size,
            use_kv_cache=bool(args.kv_cache),
        ):
            print(delta, end="", flush=True)
        print()
    else:
        response = engine.generate(
            prompt=model_prompt,
            max_new_tokens=args.max_new_tokens,
            temperature=args.temperature,
            top_k=args.top_k,
            top_p=args.top_p,
            do_sample=bool(args.do_sample),
            repetition_penalty=args.repetition_penalty,
            no_repeat_ngram_size=args.no_repeat_ngram_size,
            use_kv_cache=bool(args.kv_cache),
        )
        print(response)


if __name__ == "__main__":
    main()