Spaces:
Running on Zero
Running on Zero
File size: 1,065 Bytes
c0a2ac4 | 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 | from transformers import AutoModel, AutoTokenizer
import torch
def load_model(model_name: str, device: str = "auto"):
print(f"Loading model: {model_name}")
tokenizer = AutoTokenizer.from_pretrained(model_name)
if device == "auto":
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cuda":
try:
model = AutoModel.from_pretrained(
model_name
)
model = model.cuda()
print(
"Model loaded on CUDA:",
torch.cuda.get_device_name(0)
)
except Exception as e:
print("Standard CUDA loading failed.")
print(e)
print("Trying automatic device mapping...")
model = AutoModel.from_pretrained(
model_name,
device_map="auto"
)
else:
model = AutoModel.from_pretrained(
model_name
)
model = model.cpu()
print("Model loaded on CPU.")
return model, tokenizer
|