提交 3dc5860c 作者: imClumsyPanda

Add llm_model_dict to choose llm and add chatglm-6b-int4 as an option

上级 6d5b1438
...@@ -17,7 +17,6 @@ def torch_gc(): ...@@ -17,7 +17,6 @@ def torch_gc():
class ChatGLM(LLM): class ChatGLM(LLM):
model_name: str
max_token: int = 10000 max_token: int = 10000
temperature: float = 0.1 temperature: float = 0.1
top_p = 0.9 top_p = 0.9
...@@ -28,20 +27,6 @@ class ChatGLM(LLM): ...@@ -28,20 +27,6 @@ class ChatGLM(LLM):
def __init__(self): def __init__(self):
super().__init__() super().__init__()
def load_model(self,
model_name_or_path: str = "THUDM/chatglm-6b"):
self.tokenizer = AutoTokenizer.from_pretrained(
model_name_or_path,
trust_remote_code=True
)
self.model = (
AutoModel.from_pretrained(
model_name_or_path,
trust_remote_code=True)
.half()
.cuda()
)
@property @property
def _llm_type(self) -> str: def _llm_type(self) -> str:
return "ChatGLM" return "ChatGLM"
...@@ -62,3 +47,17 @@ class ChatGLM(LLM): ...@@ -62,3 +47,17 @@ class ChatGLM(LLM):
response = enforce_stop_tokens(response, stop) response = enforce_stop_tokens(response, stop)
self.history = updated_history self.history = updated_history
return response return response
def load_model(self,
model_name_or_path: str = "THUDM/chatglm-6b"):
self.tokenizer = AutoTokenizer.from_pretrained(
model_name_or_path,
trust_remote_code=True
)
self.model = (
AutoModel.from_pretrained(
model_name_or_path,
trust_remote_code=True)
.half()
.cuda()
)
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论