From 1b1107315c65b3da856902f16914a9cfb231a330 Mon Sep 17 00:00:00 2001
From: zhangwei <18907196110@189.cn>
Date: Fri, 7 Apr 2023 12:01:15 +0800
Subject: [PATCH] =?UTF-8?q?=E6=A0=B9=E6=8D=AE4.7=E6=97=A5=E6=9B=B4?=
 =?UTF-8?q?=E6=96=B0=E4=BB=A3=E7=A0=81=EF=BC=8C=E8=A7=A3=E5=86=B3=E5=86=B2?=
 =?UTF-8?q?=E7=AA=81=20=E6=96=B0=E5=A2=9Eernie=E7=B3=BB=E5=88=97=E6=A8=A1?=
 =?UTF-8?q?=E5=9E=8B=E4=BD=9C=E4=B8=BAembedding=20model,=E9=BB=98=E8=AE=A4?=
 =?UTF-8?q?=E4=BE=9D=E6=97=A7=E4=BD=BF=E7=94=A8GanymedeNil/text2vec-large-?=
 =?UTF-8?q?chinese=20=E6=98=BE=E5=AD=98=E4=B8=8D=E8=B6=B3=E6=97=B6?=
 =?UTF-8?q?=E5=8F=AF=E5=B0=9D=E8=AF=95=E4=BD=BF=E7=94=A8ernie-tiny?=
 =?UTF-8?q?=E4=BD=9C=E4=B8=BAembedding=E6=A8=A1=E5=9E=8B=20=E8=B0=83?=
 =?UTF-8?q?=E7=94=A8=E6=96=B9=E5=BC=8F=E4=B8=BAembeddings=20=3D=20HuggingF?=
 =?UTF-8?q?aceEmbeddings(model=5Fname=3Dembedding=5Fmodel=5Fdict["ernie-ti?=
 =?UTF-8?q?ny"],=20)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 knowledge_based_chatglm.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/knowledge_based_chatglm.py b/knowledge_based_chatglm.py
index 8f36403..d1f49c7 100644
--- a/knowledge_based_chatglm.py
+++ b/knowledge_based_chatglm.py
@@ -10,11 +10,16 @@ from langchain.vectorstores import FAISS
 from langchain.document_loaders import UnstructuredFileLoader
 from chatglm_llm import ChatGLM
 
+embedding_model_dict = {
+    "ernie-tiny": "nghuyong/ernie-3.0-nano-zh",
+    "ernie-base": "nghuyong/ernie-3.0-base-zh",
+    "text2vec": "GanymedeNil/text2vec-large-chinese"
+}
 chatglm = ChatGLM()
 
 
 def init_knowledge_vector_store(filepath):
-    embeddings = HuggingFaceEmbeddings(model_name="GanymedeNil/text2vec-large-chinese", )
+    embeddings = HuggingFaceEmbeddings(model_name=embedding_model_dict["text2vec"], )
     loader = UnstructuredFileLoader(filepath, mode="elements")
     docs = loader.load()