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 os
|
||||||
|
import pathlib
|
||||||
import platform
|
import platform
|
||||||
|
|
||||||
from transformers import AutoTokenizer, AutoModel
|
from transformers import AutoTokenizer, AutoModel
|
||||||
|
|
||||||
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
|
cache_dir = os.path.join(pathlib.Path(__file__).parent.resolve(), ".cache")
|
||||||
model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()
|
|
||||||
|
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()
|
model = model.eval()
|
||||||
|
|
||||||
os_name = platform.system()
|
os_name = platform.system()
|
||||||
|
|
13
web_demo.py
13
web_demo.py
|
@ -1,8 +1,13 @@
|
||||||
from transformers import AutoModel, AutoTokenizer
|
import os
|
||||||
import gradio as gr
|
import pathlib
|
||||||
|
|
||||||
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
|
import gradio as gr
|
||||||
model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()
|
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()
|
model = model.eval()
|
||||||
|
|
||||||
MAX_TURNS = 20
|
MAX_TURNS = 20
|
||||||
|
|
Loading…
Reference in New Issue