weiweiw
/
Langchain-Chatchat
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update chatglm_llm.py

This commit is contained in:
littlepanda0716 2023-04-07 09:28:45 +08:00
parent 3ab6dbee00
commit 51c44e3e0a
1 changed files with 13 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
chatglm_llm.py
View File

@ -3,7 +3,17 @@ from typing import Optional, List
from langchain.llms.utils import enforce_stop_tokens from langchain.llms.utils import enforce_stop_tokens
from transformers import AutoTokenizer, AutoModel from transformers import AutoTokenizer, AutoModel
"""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):
@ -11,18 +21,6 @@ class ChatGLM(LLM):
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,13 +32,12 @@ 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,
temperature=self.temperature, temperature=self.temperature,
) )
print("history: ", self.history) print("history: ", self.history)
if stop is not None: if stop is not None: