local_model.py

"""CPU-optimized local LoRA backend for DesignGym.

Loads Qwen2.5-0.5B-Instruct + a PEFT LoRA adapter on CPU and exposes
an OpenAI-shaped .chat.completions.create() interface so inference.py
needs minimal changes.
"""
from __future__ import annotations

import os
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

BASE_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
ADAPTERS: Dict[str, Optional[str]] = {
    "sft": "yashvyasop/designgym2-sft-qwen05-lora",
    "grpo": "yashvyasop/designgym2-grpo-qwen05-lora",
    "smoke": "yashvyasop/designgym2-grpo-qwen05-lora-smoke",
    "base": None,
}


@dataclass
class _Msg:
    role: str = "assistant"
    content: str = ""


@dataclass
class _Choice:
    index: int = 0
    message: _Msg = field(default_factory=_Msg)
    finish_reason: str = "stop"


@dataclass
class _Completion:
    choices: List[_Choice] = field(default_factory=list)
    model: str = ""
    backend: str = ""


class _Chat:
    def __init__(self, parent: "LocalLoRAClient"):
        self._parent = parent

    @property
    def completions(self):
        return self

    def create(self, *, model: str = "", messages: List[Dict[str, str]] = None,
               temperature: float = 0.0, max_tokens: int = 24, **_kw) -> _Completion:
        text = self._parent._generate(
            messages=messages or [],
            temperature=temperature,
            max_new_tokens=max_tokens,
        )
        comp = _Completion(
            choices=[_Choice(message=_Msg(content=text))],
            model=model or self._parent.model_id,
            backend=self._parent._backend_label(),
        )
        return comp


class LocalLoRAClient:
    _instance_lock = threading.Lock()

    def __init__(self, adapter_key: str = "sft"):
        if adapter_key not in ADAPTERS:
            raise ValueError(f"Unknown adapter key {adapter_key!r}; valid: {list(ADAPTERS)}")
        self.adapter_key: str = adapter_key
        self.adapter_id: Optional[str] = ADAPTERS[adapter_key]
        self.model_id: str = f"{BASE_MODEL}+{adapter_key}" if self.adapter_id else BASE_MODEL
        self._model = None
        self._tok = None
        self._device: str = "cpu"
        self._dtype = None
        self._ready: bool = False
        self._loading: bool = False
        self._load_error: Optional[str] = None
        self._load_seconds: Optional[float] = None

    def _backend_label(self) -> str:
        return "local-lora" if self.adapter_id else "local-base"

    @staticmethod
    def _pick_device_and_dtype():
        import torch
        if torch.cuda.is_available():
            return "cuda", torch.float16
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return "mps", torch.float16
        return "cpu", torch.float32

    def _ensure_loaded(self) -> None:
        if self._ready:
            return
        with self._instance_lock:
            if self._ready:
                return
            if self._loading:
                return
            self._loading = True

        try:
            import torch
            from transformers import AutoModelForCausalLM, AutoTokenizer

            if self._device == "cpu":
                try:
                    torch.set_num_threads(int(os.getenv("TORCH_NUM_THREADS", "2")))
                except Exception:
                    pass

            device, dtype = self._pick_device_and_dtype()
            self._device = device
            self._dtype = dtype

            t0 = time.time()
            print(f"[local_model] loading tokenizer {BASE_MODEL} ...", flush=True)
            self._tok = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)

            print(f"[local_model] loading base model {BASE_MODEL} ({dtype}, {device}) ...", flush=True)
            self._model = AutoModelForCausalLM.from_pretrained(
                BASE_MODEL,
                torch_dtype=dtype,
                low_cpu_mem_usage=True,
                trust_remote_code=True,
            )

            if self.adapter_id:
                from peft import PeftModel
                print(f"[local_model] applying LoRA adapter {self.adapter_id} ...", flush=True)
                self._model = PeftModel.from_pretrained(self._model, self.adapter_id)

            self._model = self._model.to(device)
            self._model.eval()
            self._load_seconds = round(time.time() - t0, 2)
            self._ready = True
            self._loading = False
            print(f"[local_model] ready in {self._load_seconds}s  device={device}  backend={self._backend_label()}", flush=True)

        except Exception as exc:
            self._load_error = f"{type(exc).__name__}: {exc}"
            self._loading = False
            print(f"[local_model] LOAD FAILED: {self._load_error}", flush=True)
            raise

    def _generate(self, messages: List[Dict[str, str]], temperature: float = 0.0,
                  max_new_tokens: int = 24) -> str:
        import torch

        self._ensure_loaded()
        assert self._tok is not None and self._model is not None

        prompt_text = self._tok.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        inputs = self._tok(prompt_text, return_tensors="pt")
        inputs = {k: v.to(self._device) for k, v in inputs.items()}
        input_len = inputs["input_ids"].shape[-1]

        gen_kwargs: Dict[str, Any] = dict(
            max_new_tokens=max_new_tokens,
            pad_token_id=self._tok.eos_token_id,
        )
        if temperature > 0:
            gen_kwargs["do_sample"] = True
            gen_kwargs["temperature"] = temperature
        else:
            gen_kwargs["do_sample"] = False
            gen_kwargs["temperature"] = 1.0
            gen_kwargs["top_p"] = 1.0
            gen_kwargs["top_k"] = 50

        with torch.no_grad():
            output_ids = self._model.generate(**inputs, **gen_kwargs)

        new_ids = output_ids[0][input_len:]
        return self._tok.decode(new_ids, skip_special_tokens=True).strip()

    @property
    def chat(self) -> _Chat:
        return _Chat(self)

    def describe(self) -> Dict[str, Any]:
        return {
            "backend": self._backend_label(),
            "base_model": BASE_MODEL,
            "adapter_key": self.adapter_key,
            "adapter_id": self.adapter_id,
            "device": self._device,
            "dtype": str(self._dtype) if self._dtype else None,
            "ready": self._ready,
            "loading": self._loading,
            "load_seconds": self._load_seconds,
            "load_error": self._load_error,
        }


_GLOBAL_CLIENT: Optional[LocalLoRAClient] = None
_GLOBAL_LOCK = threading.Lock()


def get_client(adapter_key: Optional[str] = None) -> LocalLoRAClient:
    global _GLOBAL_CLIENT
    with _GLOBAL_LOCK:
        if adapter_key is None:
            if _GLOBAL_CLIENT is not None:
                return _GLOBAL_CLIENT
            adapter_key = os.getenv("DESIGNGYM_ADAPTER", "sft")
        if adapter_key not in ADAPTERS:
            raise ValueError(f"Unknown adapter key {adapter_key!r}; valid: {list(ADAPTERS)}")
        if _GLOBAL_CLIENT is None or _GLOBAL_CLIENT.adapter_key != adapter_key:
            _GLOBAL_CLIENT = LocalLoRAClient(adapter_key=adapter_key)
        return _GLOBAL_CLIENT


def warm_up_async(adapter_key: Optional[str] = None) -> None:
    client = get_client(adapter_key)
    if client._ready or client._loading:
        return
    t = threading.Thread(target=client._ensure_loaded, daemon=True)
    t.start()


def describe_client(client) -> Dict[str, Any]:
    if client is None:
        return {"backend": "none", "ready": False}
    if isinstance(client, LocalLoRAClient):
        return client.describe()
    return {
        "backend": "router",
        "base_model": BASE_MODEL,
        "adapter_key": None,
        "adapter_id": None,
        "device": "remote",
        "ready": True,
        "loading": False,
        "load_seconds": None,
        "load_error": None,
    }