|
|
|
|
|
import torch
|
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
|
|
_model = None
|
|
|
_tokenizer = None
|
|
|
|
|
|
def init_model(model_dir: str = "."):
|
|
|
"""Call this ONCE at startup to load model into memory"""
|
|
|
global _model, _tokenizer
|
|
|
|
|
|
if _model is not None:
|
|
|
print("✅ Model already loaded!")
|
|
|
return _model, _tokenizer
|
|
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
|
|
print("\n" + "=" * 70)
|
|
|
print(f"Loading model from local files on {device.upper()}...")
|
|
|
print("=" * 70)
|
|
|
|
|
|
_model = AutoModelForCausalLM.from_pretrained(
|
|
|
model_dir,
|
|
|
device_map=device,
|
|
|
torch_dtype="auto",
|
|
|
trust_remote_code=True,
|
|
|
local_files_only=True,
|
|
|
)
|
|
|
|
|
|
_tokenizer = AutoTokenizer.from_pretrained(
|
|
|
model_dir,
|
|
|
local_files_only=True,
|
|
|
)
|
|
|
|
|
|
print(f"✅ Model loaded! ({sum(p.numel() for p in _model.parameters()) / 1e9:.1f}B params)")
|
|
|
print("=" * 70 + "\n")
|
|
|
|
|
|
return _model, _tokenizer
|
|
|
|
|
|
def get_model():
|
|
|
"""Get the already-loaded model (fast)"""
|
|
|
global _model, _tokenizer
|
|
|
if _model is None:
|
|
|
raise RuntimeError("Model not initialized! Call init_model() first.")
|
|
|
return _model, _tokenizer
|
|
|
|
