提交 bed03a6f 作者: myml

fix: 修复chatglm模型被复制,显存占用过多

model作为类成员会在类实例化时进行一次复制
这导致每询问一个问题显存占用就会翻倍
通过将model改成全局变量修复这个问题
上级 1851b4f0
...@@ -5,24 +5,14 @@ from transformers import AutoTokenizer, AutoModel ...@@ -5,24 +5,14 @@ from transformers import AutoTokenizer, AutoModel
"""ChatGLM_G is a wrapper around the ChatGLM model to fit LangChain framework. May not be an optimal implementation""" """ChatGLM_G is a wrapper around the ChatGLM model to fit LangChain framework. May not be an optimal implementation"""
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()
class ChatGLM(LLM): class ChatGLM(LLM):
max_token: int = 10000 max_token: int = 10000
temperature: float = 0.1 temperature: float = 0.1
top_p = 0.9 top_p = 0.9
history = [] history = []
tokenizer = AutoTokenizer.from_pretrained(
"THUDM/chatglm-6b",
trust_remote_code=True
)
model = (
AutoModel.from_pretrained(
"THUDM/chatglm-6b",
trust_remote_code=True)
.half()
.cuda()
)
def __init__(self): def __init__(self):
super().__init__() super().__init__()
...@@ -34,8 +24,8 @@ class ChatGLM(LLM): ...@@ -34,8 +24,8 @@ class ChatGLM(LLM):
def _call(self, def _call(self,
prompt: str, prompt: str,
stop: Optional[List[str]] = None) -> str: stop: Optional[List[str]] = None) -> str:
response, updated_history = self.model.chat( response, updated_history = model.chat(
self.tokenizer, tokenizer,
prompt, prompt,
history=self.history, history=self.history,
max_length=self.max_token, max_length=self.max_token,
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论