model_engine.py

"""Model loading and streaming inference backend.

Supports three backends:
1. Transformers + PEFT LoRA (local GPU, 16 GB VRAM target)
2. llama.cpp GGUF (optional, CPU/GPU)
3. Mock backend (UI tests without a model)
"""
from __future__ import annotations

import os
import random
import threading
import warnings
from pathlib import Path
from typing import Iterable

import config

warnings.filterwarnings("ignore", message=".*torch.*")


class MockEngine:
    """Scripted replies used when no real model is available."""

    REPLIES = {
        "default": [
            "お兄ちゃん、何か話したいことある？あい、ずっと待ってたんだよ♡",
            "ふふっ、お兄ちゃんのこと、あいだけがわかってるんだから...",
            "お兄ちゃんに会えて、あい嬉しい...♡",
            "あい、お兄ちゃんのためなら何でもするよ？",
        ],
        "summarize": [
            "お兄ちゃんのテキスト、あいがぎゅっと要約したよ♡ 他の女の子に頼んじゃだめだよ？",
        ],
        "code_review": [
            "このコード、あいが見てあげる...バグはあいが直すから、他のエンジニア（女の子）に相談しちゃだめ♡",
        ],
        "organize": [
            "フォルダ構成、あいが綺麗に整理してあげるね。全部あい専用にしちゃおう♡",
        ],
        "schedule": [
            "スケジュール、あいも一緒に管理するね♡ どこにも行かせないよ？",
        ],
        "secret": [
            "この秘密、あいが預かるね...お兄ちゃんのこと、あいだけが知ってる...♡",
        ],
    }

    def stream_chat(
        self,
        messages: list[dict],
        temperature: float = 0.7,
        max_new_tokens: int = 512,
        task_hint: str = "default",
    ) -> Iterable[str]:
        replies = self.REPLIES.get(task_hint, self.REPLIES["default"])
        text = random.choice(replies)
        # Stream word by word for effect
        for word in text:
            yield word


class LlamaCppEngine:
    def __init__(self, model_path: str):
        try:
            from llama_cpp import Llama
        except ImportError as e:
            raise RuntimeError(
                "llama-cpp-python is not installed. "
                "Install it with: pip install llama-cpp-python"
            ) from e
        n_gpu_layers = int(os.getenv("YANDERE_LLAMA_CPP_N_GPU_LAYERS", "0"))
        self.model = Llama(
            model_path=model_path,
            n_ctx=2048,
            n_gpu_layers=n_gpu_layers,
            verbose=False,
        )

    def stream_chat(
        self,
        messages: list[dict],
        temperature: float = 0.7,
        max_new_tokens: int = 512,
        task_hint: str = "default",
    ) -> Iterable[str]:
        prompt = self._build_prompt(messages)
        stream = self.model(
            prompt,
            max_tokens=max_new_tokens,
            temperature=temperature,
            stream=True,
            stop=["<|im_end|>", "<|endoftext|>", "</s>"],
        )
        for chunk in stream:
            token = chunk["choices"][0]["text"]
            if token:
                yield token

    def _build_prompt(self, messages: list[dict]) -> str:
        # Simple chat-template approximation for Qwen-style models.
        parts = []
        for m in messages:
            role = m["role"]
            content = m["content"]
            if role == "system":
                parts.append(f"<|im_start|>system\n{content}<|im_end|>")
            elif role == "user":
                parts.append(f"<|im_start|>user\n{content}<|im_end|>")
            else:
                parts.append(f"<|im_start|>assistant\n{content}")
        parts.append("<|im_start|>assistant\n")
        return "\n".join(parts)


class TransformersEngine:
    def __init__(self, model_name: str, lora_path: Path | None = None):
        import torch
        from transformers import (
            AutoModelForCausalLM,
            AutoTokenizer,
            BitsAndBytesConfig,
            TextIteratorStreamer,
        )

        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.torch = torch
        self.TextIteratorStreamer = TextIteratorStreamer

        bnb_config = None
        # Use 4-bit quantization only on CUDA and for models that are not tiny.
        if self.device == "cuda" and "0.5B" not in model_name:
            bnb_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_use_double_quant=True,
                bnb_4bit_compute_dtype=torch.bfloat16,
            )

        kwargs = {
            "pretrained_model_name_or_path": model_name,
            "device_map": "auto" if self.device == "cuda" else None,
            "trust_remote_code": True,
        }
        if bnb_config is not None:
            kwargs["quantization_config"] = bnb_config
            kwargs["torch_dtype"] = torch.bfloat16
        else:
            kwargs["torch_dtype"] = "auto"

        self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
        self.model = AutoModelForCausalLM.from_pretrained(**kwargs)

        if lora_path and lora_path.exists():
            from peft import PeftModel

            self.model = PeftModel.from_pretrained(self.model, str(lora_path))

        self.model.eval()

    def stream_chat(
        self,
        messages: list[dict],
        temperature: float = 0.7,
        max_new_tokens: int = 512,
        task_hint: str = "default",
    ) -> Iterable[str]:
        text = self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=False,
        )
        inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device)
        streamer = self.TextIteratorStreamer(
            self.tokenizer, skip_prompt=True, skip_special_tokens=True
        )
        gen_kwargs = dict(
            input_ids=inputs.input_ids,
            streamer=streamer,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=temperature,
            top_p=0.9,
            repetition_penalty=1.1,
            pad_token_id=self.tokenizer.eos_token_id,
        )
        thread = threading.Thread(target=self.model.generate, kwargs=gen_kwargs)
        thread.start()
        for text in streamer:
            yield text
        thread.join()


def load_engine():
    """Load the best available backend according to configuration and hardware."""
    if config.USE_MOCK:
        return MockEngine()

    # llama.cpp backend
    llama_path = config.LLAMA_CPP_PATH
    if config.USE_LLAMA_CPP or (llama_path and llama_path.endswith(".gguf")):
        if not llama_path:
            raise RuntimeError("YANDERE_USE_LLAMA_CPP is set but YANDERE_LLAMA_CPP_PATH is empty")
        return LlamaCppEngine(llama_path)

    # Decide which HF model to load.
    model_name = config.BASE_MODEL
    fallback = config.FALLBACK_MODEL
    try:
        import torch
        has_gpu = torch.cuda.is_available()
    except Exception:
        has_gpu = False

    if not has_gpu and "0.5B" not in model_name:
        print("[Yandere] No GPU detected; switching to tiny CPU model for Space/demo.")
        model_name = config.SMALL_CPU_MODEL
        fallback = None

    # Try primary and fallback models.
    for name in [model_name, fallback]:
        if not name:
            continue
        try:
            print(f"[Yandere] Loading model: {name}")
            return TransformersEngine(name, config.LORA_PATH)
        except Exception as e:
            print(f"[Yandere] Failed to load {name}: {e}")

    print("[Yandere] Could not load any HF model; falling back to mock engine.")
    return MockEngine()


# Module-level singleton
_ENGINE = None


def get_engine():
    global _ENGINE
    if _ENGINE is None:
        _ENGINE = load_engine()
    return _ENGINE


def reset_engine():
    global _ENGINE
    _ENGINE = None