mirror of https://github.com/THUDM/ChatGLM-6B
support windows for local model directory
parent
da09ca4dff
commit
fc1058babe
|
@ -2,7 +2,7 @@ import os
|
||||||
import platform
|
import platform
|
||||||
from transformers import AutoTokenizer, AutoModel
|
from transformers import AutoTokenizer, AutoModel
|
||||||
|
|
||||||
MODEL_ID = "./model" if os.path.exists('./model') else "THUDM/chatglm-6b"
|
MODEL_ID = "model" if os.path.isdir('model') else "THUDM/chatglm-6b"
|
||||||
|
|
||||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
|
||||||
model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True).half().cuda()
|
model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True).half().cuda()
|
||||||
|
|
|
@ -2,7 +2,7 @@ import os
|
||||||
from transformers import AutoModel, AutoTokenizer
|
from transformers import AutoModel, AutoTokenizer
|
||||||
import gradio as gr
|
import gradio as gr
|
||||||
|
|
||||||
MODEL_ID = "./model" if os.path.exists('./model') else "THUDM/chatglm-6b"
|
MODEL_ID = "model" if os.path.isdir('model') else "THUDM/chatglm-6b"
|
||||||
|
|
||||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
|
||||||
model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True).half().cuda()
|
model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True).half().cuda()
|
||||||
|
|
Loading…
Reference in New Issue