wassname/alignment-handbook
1
0
Fork 0
mirror of https://github.com/wassname/alignment-handbook.git synced 2026-07-16 11:16:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add model utils tests

Browse Source
This commit is contained in:
Lewis Tunstall
2023-11-10 09:42:15 +00:00
parent 0af8011993
commit 2ed5a45d25
5 changed files with 152 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/alignment/model_utils.py
+2 -2
View File
@@ -1,4 +1,4 @@
from typing import Dict, Union
from typing import Dict
import torch
from transformers import AutoTokenizer, BitsAndBytesConfig, PreTrainedTokenizer
@@ -62,7 +62,7 @@ def get_tokenizer(model_args: ModelArguments, data_args: DataArguments) -> PreTr
return tokenizer
def get_peft_config(model_args: ModelArguments) -> Union[PeftConfig, None]:
def get_peft_config(model_args: ModelArguments) -> PeftConfig | None:
if model_args.use_peft is False:
return None