update web demo

pull/624/head
ZwwWayne 2024-01-18 22:07:03 +08:00
parent e71fbcc3b5
commit 34f2f5aa96
1 changed files with 2 additions and 2 deletions

View File

@ -161,11 +161,11 @@ def on_btn_click():
@st.cache_resource
def load_model():
model = (
AutoModelForCausalLM.from_pretrained("/mnt/inspurfs/share_data/zhangwenwei/models/hf_release/internlm2-chat-7b", trust_remote_code=True)
AutoModelForCausalLM.from_pretrained("internlm/internlm2-chat-7b", trust_remote_code=True)
.to(torch.bfloat16)
.cuda()
)
tokenizer = AutoTokenizer.from_pretrained("/mnt/inspurfs/share_data/zhangwenwei/models/hf_release/internlm2-chat-7b", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("internlm/internlm2-chat-7b", trust_remote_code=True)
return model, tokenizer