github
/
ChatGLM-6B
mirror of https://github.com/THUDM/ChatGLM-6B
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Move import statement into function to avoid dependency

pull/241/head
saber 2023-03-27 22:11:57 +08:00
parent 6a5267aef7
commit c31d9c0984
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
utils.py
View File

@ -1,7 +1,6 @@
import os import os
from typing import Dict, Tuple, Union, Optional from typing import Dict, Tuple, Union, Optional
from accelerate import load_checkpoint_and_dispatch
from torch.nn import Module from torch.nn import Module
from transformers import AutoModel, AutoTokenizer from transformers import AutoModel, AutoTokenizer
from transformers.tokenization_utils import PreTrainedTokenizer from transformers.tokenization_utils import PreTrainedTokenizer
@ -40,6 +39,8 @@ def auto_configure_device_map(num_gpus: int) -> Dict[str, int]:
def load_model_on_gpus(checkpoint_path: Union[str, os.PathLike], num_gpus: int = 2, def load_model_on_gpus(checkpoint_path: Union[str, os.PathLike], num_gpus: int = 2,
multi_gpu_model_cache_dir: Union[str, os.PathLike] = "./temp_model_dir", multi_gpu_model_cache_dir: Union[str, os.PathLike] = "./temp_model_dir",
tokenizer: Optional[PreTrainedTokenizer] = None, **kwargs) -> Module: tokenizer: Optional[PreTrainedTokenizer] = None, **kwargs) -> Module:
from accelerate import load_checkpoint_and_dispatch
model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs) model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs)
model = model.eval() model = model.eval()