llms: add cpu support

This commit is contained in:
csunny 2023-05-21 16:05:53 +08:00
parent f52c7523b5
commit 89970bd71c
3 changed files with 16 additions and 5 deletions

View File

@ -9,6 +9,8 @@ from transformers import (
AutoModel AutoModel
) )
from pilot.configs.model_config import DEVICE
class BaseLLMAdaper: class BaseLLMAdaper:
"""The Base class for multi model, in our project. """The Base class for multi model, in our project.
We will support those model, which performance resemble ChatGPT """ We will support those model, which performance resemble ChatGPT """
@ -64,6 +66,13 @@ class ChatGLMAdapater(BaseLLMAdaper):
def loader(self, model_path: str, from_pretrained_kwargs: dict): def loader(self, model_path: str, from_pretrained_kwargs: dict):
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
if DEVICE != "cuda":
model = AutoModel.from_pretrained(
model_path, trust_remote_code=True, **from_pretrained_kwargs
).float()
return model, tokenizer
else:
model = AutoModel.from_pretrained( model = AutoModel.from_pretrained(
model_path, trust_remote_code=True, **from_pretrained_kwargs model_path, trust_remote_code=True, **from_pretrained_kwargs
).half().cuda() ).half().cuda()

View File

@ -155,6 +155,7 @@ if __name__ == "__main__":
model_path = LLM_MODEL_CONFIG[CFG.LLM_MODEL] model_path = LLM_MODEL_CONFIG[CFG.LLM_MODEL]
print(model_path, DEVICE) print(model_path, DEVICE)
worker = ModelWorker( worker = ModelWorker(
model_path=model_path, model_path=model_path,
model_name=CFG.LLM_MODEL, model_name=CFG.LLM_MODEL,

View File

@ -42,6 +42,7 @@ tenacity==8.2.2
peft peft
pycocoevalcap pycocoevalcap
sentence-transformers sentence-transformers
cpm_kernels
umap-learn umap-learn
notebook notebook
gradio==3.23 gradio==3.23