model_loader.py

# model_loader.py
# ==============================
# Responsible for loading models
# Base model always loads
# Core / Skill load only if enabled in config.py
# ==============================

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

import config


_model = None
_tokenizer = None


def load_model(skill: str | None = None):
    """
    Loads:
    - Base model (always)
    - Core adapter (if enabled)
    - Skill adapter (if requested & enabled)
    """

    global _model, _tokenizer

    if _model is not None and _tokenizer is not None:
        return _model, _tokenizer

    # -------- Base --------
    tokenizer = AutoTokenizer.from_pretrained(
        config.BASE_MODEL,
        trust_remote_code=True
    )

    model = AutoModelForCausalLM.from_pretrained(
        config.BASE_MODEL,
        torch_dtype=torch.float16,
        device_map="auto",
        trust_remote_code=True
    )

    # -------- Core (future: one-line enable) --------
    if config.CORE_ADAPTER:
        model = PeftModel.from_pretrained(model, config.CORE_ADAPTER)

    # -------- Skill (future: routed) --------
    if skill and skill in config.SKILL_ADAPTERS:
        model = PeftModel.from_pretrained(
            model,
            config.SKILL_ADAPTERS[skill]
        )

    model.eval()

    _model = model
    _tokenizer = tokenizer

    print("✅ Model loaded successfully")
    return model, tokenizer