From fc1058babeb3bcf342fbdb3d8ace0518481ff470 Mon Sep 17 00:00:00 2001 From: "Guo Y.K" Date: Thu, 23 Mar 2023 16:07:03 +0800 Subject: [PATCH] support windows for local model directory --- cli_demo.py | 2 +- web_demo.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cli_demo.py b/cli_demo.py index e093c24..4d6d3fb 100644 --- a/cli_demo.py +++ b/cli_demo.py @@ -2,7 +2,7 @@ import os import platform 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) model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True).half().cuda() diff --git a/web_demo.py b/web_demo.py index 969fdc4..6cb5c56 100644 --- a/web_demo.py +++ b/web_demo.py @@ -2,7 +2,7 @@ import os from transformers import AutoModel, AutoTokenizer 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) model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True).half().cuda()