tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Merge pull request #356 from tmm1/load_model-args

Browse Source 
simplify `load_model` signature
This commit is contained in:
Aman Gupta Karmani
2023-08-09 18:24:34 -07:00
committed by GitHub
parent 964312199e 718102271f
commit 11ddccb80f
2 changed files with 8 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
src/axolotl/utils/models.py
View File

@@ -78,12 +78,15 @@ def load_tokenizer(
def load_model(
base_model, base_model_config, model_type, tokenizer, cfg, adapter="lora"
):
# type: (str, str, str, PreTrainedTokenizerBase, DictDefault, Optional[str]) -> Tuple[PreTrainedModel, Optional[PeftConfig]]
cfg, tokenizer
): # type: (DictDefault, PreTrainedTokenizerBase) -> Tuple[PreTrainedModel, Optional[PeftConfig]]
"""
Load a model from a base model and a model type.
Load a model for a given configuration and tokenizer.
"""
base_model = cfg.base_model
base_model_config = cfg.base_model_config
model_type = cfg.model_type
adapter = cfg.adapter
# TODO refactor as a kwarg
load_in_8bit = cfg.load_in_8bit