From 7d837078e7dcc6d3dc8d5b680b96226c2162f6e0 Mon Sep 17 00:00:00 2001
From: imClumsyPanda <littlepanda0716@gmail.com>
Date: Fri, 19 May 2023 23:12:39 +0800
Subject: [PATCH] update model_config.py

---
 configs/model_config.py | 28 ++++++++++++++--------------
 models/shared.py        |  4 ++--
 2 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/configs/model_config.py b/configs/model_config.py
index c2b0db3..5dcf2d3 100644
--- a/configs/model_config.py
+++ b/configs/model_config.py
@@ -30,44 +30,44 @@ llm_model_dict 处理了loader的一些预设行为，如加载位置，模型
 llm_model_dict = {
     "chatglm-6b-int4-qe": {
         "name": "chatglm-6b-int4-qe",
-        "remote-checkpoint": "THUDM/chatglm-6b-int4-qe",
-        "path": None,
+        "pretrained_model_name": "THUDM/chatglm-6b-int4-qe",
+        "local_model_path": None,
         "provides": "ChatGLM"
     },
     "chatglm-6b-int4": {
         "name": "chatglm-6b-int4",
-        "remote-checkpoint": "THUDM/chatglm-6b-int4",
-        "path": None,
+        "pretrained_model_name": "THUDM/chatglm-6b-int4",
+        "local_model_path": None,
         "provides": "ChatGLM"
     },
     "chatglm-6b": {
         "name": "chatglm-6b",
-        "remote-checkpoint": "THUDM/chatglm-6b-int4",
-        "path": None,
+        "pretrained_model_name": "THUDM/chatglm-6b-int4",
+        "local_model_path": None,
         "provides": "ChatGLM"
     },
     "llama-7b-hf": {
         "name": "llama-7b-hf",
-        "remote-checkpoint": "llama-7b-hf",
-        "path": None,
+        "pretrained_model_name": "llama-7b-hf",
+        "local_model_path": None,
         "provides": "LLamaLLM"
     },
     "vicuna-13b-hf": {
         "name": "vicuna-13b-hf",
-        "remote-checkpoint": "vicuna-13b-hf",
-        "path": None,
+        "pretrained_model_name": "vicuna-13b-hf",
+        "local_model_path": None,
         "provides": "LLamaLLM"
     },
     "chatyuan": {
         "name": "chatyuan",
-        "remote-checkpoint": "ClueAI/ChatYuan-large-v2",
-        "path": None,
+        "pretrained_model_name": "ClueAI/ChatYuan-large-v2",
+        "local_model_path": None,
         "provides": None
     },
     "chatglm-6b-int8":{
         "name": "chatglm-6b-int8",
-        "remote-checkpoint":  "THUDM/chatglm-6b-int8",
-        "path": None,
+        "pretrained_model_name":  "THUDM/chatglm-6b-int8",
+        "local_model_path": None,
         "provides": "ChatGLM"
     },
 }
diff --git a/models/shared.py b/models/shared.py
index 68325b0..1f5f6fb 100644
--- a/models/shared.py
+++ b/models/shared.py
@@ -32,9 +32,9 @@ def loaderLLM(llm_model: str = None, no_remote_model: bool = False, use_ptuning_
     if loaderCheckPoint.no_remote_model:
         loaderCheckPoint.model_name = llm_model_info['name']
     else:
-        loaderCheckPoint.model_name = llm_model_info['remote-checkpoint']
+        loaderCheckPoint.model_name = llm_model_info['pretrained_model_name']
 
-    loaderCheckPoint.model_path = llm_model_info['path']
+    loaderCheckPoint.model_path = llm_model_info["local_model_path"]
 
     loaderCheckPoint.reload_model()