57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
import os
|
|
from typing import Union
|
|
|
|
from transformers import (
|
|
AutoConfig,
|
|
AutoModel,
|
|
AutoModelForCausalLM,
|
|
AutoTokenizer,
|
|
PreTrainedModel,
|
|
PreTrainedTokenizer,
|
|
)
|
|
|
|
import custom_models
|
|
|
|
|
|
def init_model(model_name: Union[str, os.PathLike]) -> PreTrainedModel:
|
|
config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
|
|
|
|
if model_name in custom_models.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
|
|
model = custom_models.AutoModelForCausalLM.from_config(config)
|
|
elif model_name in custom_models.MODEL_MAPPING_NAMES:
|
|
model = custom_models.AutoModel.from_config(config)
|
|
else:
|
|
try:
|
|
model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
|
|
except ValueError:
|
|
model = AutoModel.from_config(config, trust_remote_code=True)
|
|
return model
|
|
|
|
|
|
def load_model(model_name_or_path: Union[str, os.PathLike]) -> PreTrainedModel:
|
|
if model_name_or_path in custom_models.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
|
|
model = custom_models.AutoModelForCausalLM.from_pretrained(model_name_or_path)
|
|
elif model_name_or_path in custom_models.MODEL_MAPPING_NAMES:
|
|
model = custom_models.AutoModel.from_pretrained(model_name_or_path)
|
|
else:
|
|
try:
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_name_or_path, trust_remote_code=True
|
|
)
|
|
except ValueError:
|
|
model = AutoModel.from_pretrained(
|
|
model_name_or_path, trust_remote_code=True
|
|
)
|
|
return model
|
|
|
|
|
|
def load_tokenizer(
|
|
tokenizer_name_or_path: Union[str, os.PathLike]
|
|
) -> PreTrainedTokenizer:
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
tokenizer_name_or_path, padding_side='left', trust_remote_code=True
|
|
)
|
|
if tokenizer.pad_token_id is None:
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
return tokenizer
|