浏览代码

update embedding define in local_doc_qa.py and langchain version in requirements.txt

imClumsyPanda 2 年之前
父节点
当前提交
0706a29057
共有 2 个文件被更改,包括 7 次插入5 次删除
  1. 5 3
      chains/local_doc_qa.py
  2. 2 2
      requirements.txt

+ 5 - 3
chains/local_doc_qa.py

@@ -33,6 +33,7 @@ def load_file(filepath):
 class LocalDocQA:
     llm: object = None
     embeddings: object = None
+    top_k: int = VECTOR_SEARCH_TOP_K
 
     def init_cfg(self,
                  embedding_model: str = EMBEDDING_MODEL,
@@ -49,9 +50,10 @@ class LocalDocQA:
                             use_ptuning_v2=use_ptuning_v2)
         self.llm.history_len = llm_history_len
 
-        self.embeddings = HuggingFaceEmbeddings(model_name=embedding_model_dict[embedding_model], )
-        self.embeddings.client = sentence_transformers.SentenceTransformer(self.embeddings.model_name,
-                                                                           device=embedding_device)
+        self.embeddings = HuggingFaceEmbeddings(model_name=embedding_model_dict[embedding_model],
+                                                model_kwargs={'device': embedding_device})
+        # self.embeddings.client = sentence_transformers.SentenceTransformer(self.embeddings.model_name,
+        #                                                                    device=embedding_device)
         self.top_k = top_k
 
     def init_knowledge_vector_store(self,

+ 2 - 2
requirements.txt

@@ -1,4 +1,4 @@
-langchain>=0.0.124
+langchain>=0.0.146
 transformers==4.27.1
 unstructured[local-inference]
 layoutparser[layoutmodels,tesseract]
@@ -9,4 +9,4 @@ icetk
 cpm_kernels
 faiss-cpu
 gradio>=3.25.0
-detectron2@git+https://github.com/facebookresearch/detectron2.git@v0.6#egg=detectron2
+#detectron2@git+https://github.com/facebookresearch/detectron2.git@v0.6#egg=detectron2