Spaces:
Runtime error
Runtime error
File size: 1,457 Bytes
9d21d47 |
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 |
from transformers import AutoTokenizer, PreTrainedTokenizerFast, AutoModelForCausalLM
def build_model_signature(model_type, model_size):
if model_type == "opt":
# ["125m", "350m", "1.3b", "2.7b", "6.7b", "13b", "30b", "66b"]
return f"facebook/opt-{model_size}"
if model_type == "gpt2":
# ["sm", "medium", "large", "xl"]
if model_size == "sm":
return "gpt2"
return f"gpt2-{model_size}"
if model_type == "e-gpt":
# ["neo-125M", "neo-1.3B", "neo-2.7B", "j-6B", "neox-20b"]
return f"EleutherAI/gpt-{model_size}"
if model_type == "bloom":
# ["560m", "1b1", "1b7", "3b", "7b1"]
return f"bigscience/bloom-{model_size}"
def build_tokenizer(model_type, model_size, padding_side="left", use_fast=False):
sign = build_model_signature(model_type, model_size)
if not use_fast:
tok = AutoTokenizer.from_pretrained(sign, padding_side=padding_side)
else:
tok = PreTrainedTokenizerFast.from_pretrained(sign, padding_side=padding_side)
if model_type in ["gpt2", "e-gpt"]:
tok.pad_token_id = tok.eos_token_id
tok.pad_token = tok.eos_token
return tok
def build_model(model_type, model_size, in_8bit):
sign = build_model_signature(model_type, model_size)
model = AutoModelForCausalLM.from_pretrained(
sign,
device_map="auto",
load_in_8bit=in_8bit,
)
model.eval()
return model
|