Spaces:
Runtime error
Runtime error
File size: 1,345 Bytes
88f55d9 |
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 |
import torch
from transformers import AutoModelForSeq2SeqLM, T5Tokenizer
from optimum.bettertransformer import BetterTransformer
def load_model(
base,
finetuned,
mode_cpu,
mode_mps,
mode_full_gpu,
mode_8bit,
mode_4bit,
force_download_ckpt
):
tokenizer = T5Tokenizer.from_pretrained(base, use_fast=False)
tokenizer.padding_side = "left"
if mode_cpu:
print("cpu mode")
model = AutoModelForSeq2SeqLM.from_pretrained(
base,
device_map={"": "cpu"},
use_safetensors=False,
)
elif mode_mps:
print("mps mode")
model = AutoModelForSeq2SeqLM.from_pretrained(
base,
device_map={"": "mps"},
torch_dtype=torch.float16,
use_safetensors=False,
)
else:
print("gpu mode")
print(f"8bit = {mode_8bit}, 4bit = {mode_4bit}")
model = AutoModelForSeq2SeqLM.from_pretrained(
base,
load_in_8bit=mode_8bit,
load_in_4bit=mode_4bit,
device_map="auto",
torch_dtype=torch.float16,
use_safetensors=False,
)
if not mode_8bit and not mode_4bit:
model.half()
model = BetterTransformer.transform(model)
return model, tokenizer |