Spaces:
Sleeping
Sleeping
File size: 1,147 Bytes
15fa2e5 0885208 15fa2e5 8c9ab96 15fa2e5 8c9ab96 15fa2e5 a6c748c 15fa2e5 |
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 |
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
import torch
def load_local_model(model_path: str, device: int = -1, token: str = None):
"""
Load a Hugging Face model (CPU by default) with optional token for private repos.
Args:
model_path (str): Hugging Face repo ID or local path.
device (int): -1 for CPU, >=0 for GPU index.
token (str): HF token for private models.
Returns:
model, tokenizer
"""
try:
tokenizer = AutoTokenizer.from_pretrained(model_path, use_auth_token=token)
except Exception as e:
raise RuntimeError(f"Failed to load tokenizer: {e}")
try:
config = AutoConfig.from_pretrained(model_path, use_auth_token=token)
model = AutoModelForCausalLM.from_pretrained(
model_path, config=config, use_auth_token=token
)
# Device mapping
if device >= 0 and torch.cuda.is_available():
model.to(f"cuda:{device}")
else:
model.to("cpu")
except Exception as e:
raise RuntimeError(f"Failed to load model: {e}")
return model, tokenizer
|