File size: 1,456 Bytes
c30ada7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | # 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
|