github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-08-09 12:18:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

add chatglm support

Browse Source
This commit is contained in:
csunny 2023-05-21 16:30:03 +08:00
parent 5ec1f413b6
commit a3fae0bdf2
2 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
pilot/conversation.py
View File

@ -15,6 +15,9 @@ DB_SETTINGS = {
"port": CFG.LOCAL_DB_PORT "port": CFG.LOCAL_DB_PORT
} }
ROLE_USER = "USER"
ROLE_ASSISTANT = "Assistant"
class SeparatorStyle(Enum): class SeparatorStyle(Enum):
SINGLE = auto() SINGLE = auto()
TWO = auto() TWO = auto()

6
pilot/model/chatglm_llm.py
View File

@ -3,6 +3,8 @@
import torch import torch
from pilot.conversation import ROLE_USER, ROLE_ASSISTANT
@torch.inference_mode() @torch.inference_mode()
def chatglm_generate_stream(model, tokenizer, params, device, context_len=2048, stream_interval=2): def chatglm_generate_stream(model, tokenizer, params, device, context_len=2048, stream_interval=2):
@ -30,9 +32,9 @@ def chatglm_generate_stream(model, tokenizer, params, device, context_len=2048,
# Add history chat to hist for model. # Add history chat to hist for model.
for i in range(1, len(messages) - 2, 2): for i in range(1, len(messages) - 2, 2):
hist.append((messages[i].split(":")[1], messages[i+1].split(":")[1])) hist.append((messages[i].split(ROLE_USER + ":")[1], messages[i+1].split(ROLE_ASSISTANT + ":")[1]))
query = messages[-2].split(":")[1] query = messages[-2].split(ROLE_USER + ":")[1]
print("Query Message: ", query) print("Query Message: ", query)
output = "" output = ""
i = 0 i = 0