mirror of https://github.com/THUDM/ChatGLM-6B
feat: add `.gitignore` and `cache_dir`
parent
c55ecd89a0
commit
d6bd5be4cd
|
@ -0,0 +1,6 @@
|
|||
# override cache dir
|
||||
.cache
|
||||
|
||||
# for ide
|
||||
.vscode
|
||||
.idea
|
|
@ -1,9 +1,13 @@
|
|||
import os
|
||||
import pathlib
|
||||
import platform
|
||||
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
|
||||
model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()
|
||||
cache_dir = os.path.join(pathlib.Path(__file__).parent.resolve(), ".cache")
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True, cache_dir=cache_dir)
|
||||
model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True, cache_dir=cache_dir).half().cuda()
|
||||
model = model.eval()
|
||||
|
||||
os_name = platform.system()
|
||||
|
|
13
web_demo.py
13
web_demo.py
|
@ -1,8 +1,13 @@
|
|||
from transformers import AutoModel, AutoTokenizer
|
||||
import gradio as gr
|
||||
import os
|
||||
import pathlib
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
|
||||
model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()
|
||||
import gradio as gr
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
cache_dir = os.path.join(pathlib.Path(__file__).parent.resolve(), ".cache")
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True, cache_dir=cache_dir)
|
||||
model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True, cache_dir=cache_dir).half().cuda()
|
||||
model = model.eval()
|
||||
|
||||
MAX_TURNS = 20
|
||||
|
|
Loading…
Reference in New Issue