Эх сурвалжийг харах

cli_demo.py: 兼容loader加载器

glide-the 2 жил өмнө
parent
commit
372024a92e

+ 5 - 6
cli_demo.py

@@ -4,7 +4,7 @@ import os
 import nltk
 from models.loader.args import parser
 import models.shared as shared
-from models.loader import LoaderLLM
+from models.loader import LoaderCheckPoint
 from models.chatglm_llm import ChatGLM
 
 nltk.data.path = [os.path.join(os.path.dirname(__file__), "nltk_data")] + nltk.data.path
@@ -22,13 +22,12 @@ if __name__ == "__main__":
     args = None
     args = parser.parse_args()
     args_dict = vars(args)
-
-    shared.loaderLLM = LoaderLLM(args_dict)
-    chatGLMLLM = ChatGLM(shared.loaderLLM)
-    chatGLMLLM.history_len = LLM_HISTORY_LEN
+    shared.loaderCheckPoint = LoaderCheckPoint(args_dict)
+    llm_model_ins = shared.loaderLLM()
+    llm_model_ins.history_len = LLM_HISTORY_LEN
 
     local_doc_qa = LocalDocQA()
-    local_doc_qa.init_cfg(llm_model=chatGLMLLM,
+    local_doc_qa.init_cfg(llm_model=llm_model_ins,
                           embedding_model=EMBEDDING_MODEL,
                           embedding_device=EMBEDDING_DEVICE,
                           top_k=VECTOR_SEARCH_TOP_K)

+ 1 - 1
configs/model_config.py

@@ -6,7 +6,7 @@ embedding_model_dict = {
     "ernie-tiny": "nghuyong/ernie-3.0-nano-zh",
     "ernie-base": "nghuyong/ernie-3.0-base-zh",
     "text2vec-base": "shibing624/text2vec-base-chinese",
-    "text2vec": "GanymedeNil/text2vec-large-chinese",
+    "text2vec": "/media/mnt/text2vec-large-chinese/",
 }
 
 # Embedding model name