支持本地embedding

modules/base_model.py

@@ -132,8 +132,8 @@ class BaseLLMModel:
             status_text = self.token_message()
             yield get_return_value()
             if self.interrupted:
-                    self.recover()
-                    break
+                self.recover()
+                break
         self.history.append(construct_assistant(partial_text))
 
     def next_chatbot_at_once(self, inputs, chatbot, fake_input=None, display_append=""):
@@ -170,7 +170,14 @@ class BaseLLMModel:
     ):  # repetition_penalty, top_k
         from llama_index.indices.vector_store.base_query import GPTVectorStoreIndexQuery
         from llama_index.indices.query.schema import QueryBundle
-        from langchain.llms import OpenAIChat
+        from langchain.embeddings.huggingface import HuggingFaceEmbeddings
+        from langchain.chat_models import ChatOpenAI
+        from llama_index import (
+            GPTSimpleVectorIndex,
+            ServiceContext,
+            LangchainEmbedding,
+            OpenAIEmbedding,
+        )
 
         logging.info(
             "输入为：" + colorama.Fore.BLUE + f"{inputs}" + colorama.Style.RESET_ALL
@@ -182,20 +189,22 @@ class BaseLLMModel:
         old_inputs = None
         display_reference = []
         limited_context = False
-        if files and self.api_key:
+        if files:
             limited_context = True
             old_inputs = inputs
             msg = "加载索引中……（这可能需要几分钟）"
             logging.info(msg)
             yield chatbot + [(inputs, "")], msg
             index = construct_index(self.api_key, file_src=files)
+            assert index is not None, "索引构建失败"
             msg = "索引构建完成，获取回答中……"
+            if local_embedding:
+                embed_model = LangchainEmbedding(HuggingFaceEmbeddings())
+            else:
+                embed_model = OpenAIEmbedding()
             logging.info(msg)
             yield chatbot + [(inputs, "")], msg
             with retrieve_proxy():
-                llm_predictor = LLMPredictor(
-                    llm=OpenAIChat(temperature=0, model_name=self.model_name)
-                )
                 prompt_helper = PromptHelper(
                     max_input_size=4096,
                     num_output=5,
@@ -205,7 +214,7 @@ class BaseLLMModel:
                 from llama_index import ServiceContext
 
                 service_context = ServiceContext.from_defaults(
-                    llm_predictor=llm_predictor, prompt_helper=prompt_helper
+                    prompt_helper=prompt_helper, embed_model=embed_model
                 )
                 query_object = GPTVectorStoreIndexQuery(
                     index.index_struct,
@@ -249,7 +258,11 @@ class BaseLLMModel:
         else:
             display_reference = ""
 
-        if self.api_key is not None and len(self.api_key) == 0 and not shared.state.multi_api_key:
+        if (
+            self.api_key is not None
+            and len(self.api_key) == 0
+            and not shared.state.multi_api_key
+        ):
             status_text = STANDARD_ERROR_MSG + NO_APIKEY_MSG
             logging.info(status_text)
             chatbot.append((inputs, ""))

modules/config.py

@@ -117,6 +117,8 @@ https_proxy = os.environ.get("HTTPS_PROXY", https_proxy)
 os.environ["HTTP_PROXY"] = ""
 os.environ["HTTPS_PROXY"] = ""
 
+local_embedding = config.get("local_embedding", False) # 是否使用本地embedding
+
 @contextmanager
 def retrieve_proxy(proxy=None):
     """

modules/llama_func.py

@@ -15,6 +15,8 @@ from tqdm import tqdm
 
 from modules.presets import *
 from modules.utils import *
+from modules.config import local_embedding
+
 
 def get_index_name(file_src):
     file_paths = [x.name for x in file_src]
@@ -28,6 +30,7 @@ def get_index_name(file_src):
 
     return md5_hash.hexdigest()
 
+
 def block_split(text):
     blocks = []
     while len(text) > 0:
@@ -35,6 +38,7 @@ def block_split(text):
         text = text[1000:]
     return blocks
 
+
 def get_documents(file_src):
     documents = []
     logging.debug("Loading documents...")
@@ -50,11 +54,12 @@ def get_documents(file_src):
                 try:
                     from modules.pdf_func import parse_pdf
                     from modules.config import advance_docs
+
                     two_column = advance_docs["pdf"].get("two_column", False)
                     pdftext = parse_pdf(filepath, two_column).text
                 except:
                     pdftext = ""
-                    with open(filepath, 'rb') as pdfFileObj:
+                    with open(filepath, "rb") as pdfFileObj:
                         pdfReader = PyPDF2.PdfReader(pdfFileObj)
                         for page in tqdm(pdfReader.pages):
                             pdftext += page.extract_text()
@@ -91,19 +96,21 @@ def get_documents(file_src):
 
 
 def construct_index(
-        api_key,
-        file_src,
-        max_input_size=4096,
-        num_outputs=5,
-        max_chunk_overlap=20,
-        chunk_size_limit=600,
-        embedding_limit=None,
-        separator=" "
+    api_key,
+    file_src,
+    max_input_size=4096,
+    num_outputs=5,
+    max_chunk_overlap=20,
+    chunk_size_limit=600,
+    embedding_limit=None,
+    separator=" ",
 ):
     from langchain.chat_models import ChatOpenAI
-    from llama_index import GPTSimpleVectorIndex, ServiceContext
+    from langchain.embeddings.huggingface import HuggingFaceEmbeddings
+    from llama_index import GPTSimpleVectorIndex, ServiceContext, LangchainEmbedding, OpenAIEmbedding
 
-    os.environ["OPENAI_API_KEY"] = api_key
+    if api_key:
+        os.environ["OPENAI_API_KEY"] = api_key
     chunk_size_limit = None if chunk_size_limit == 0 else chunk_size_limit
     embedding_limit = None if embedding_limit == 0 else embedding_limit
     separator = " " if separator == "" else separator
@@ -111,7 +118,14 @@ def construct_index(
     llm_predictor = LLMPredictor(
         llm=ChatOpenAI(model_name="gpt-3.5-turbo-0301", openai_api_key=api_key)
     )
-    prompt_helper = PromptHelper(max_input_size = max_input_size, num_output = num_outputs, max_chunk_overlap = max_chunk_overlap, embedding_limit=embedding_limit, chunk_size_limit=600, separator=separator)
+    prompt_helper = PromptHelper(
+        max_input_size=max_input_size,
+        num_output=num_outputs,
+        max_chunk_overlap=max_chunk_overlap,
+        embedding_limit=embedding_limit,
+        chunk_size_limit=600,
+        separator=separator,
+    )
     index_name = get_index_name(file_src)
     if os.path.exists(f"./index/{index_name}.json"):
         logging.info("找到了缓存的索引文件，加载中……")
@@ -119,11 +133,20 @@ def construct_index(
     else:
         try:
             documents = get_documents(file_src)
+            if local_embedding:
+                embed_model = LangchainEmbedding(HuggingFaceEmbeddings())
+            else:
+                embed_model = OpenAIEmbedding()
             logging.info("构建索引中……")
             with retrieve_proxy():
-                service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper, chunk_size_limit=chunk_size_limit)
+                service_context = ServiceContext.from_defaults(
+                    llm_predictor=llm_predictor,
+                    prompt_helper=prompt_helper,
+                    chunk_size_limit=chunk_size_limit,
+                    embed_model=LangchainEmbedding(HuggingFaceEmbeddings()),
+                )
                 index = GPTSimpleVectorIndex.from_documents(
-                    documents,  service_context=service_context
+                    documents, service_context=service_context
                 )
             logging.debug("索引构建完成！")
             os.makedirs("./index", exist_ok=True)

modules/models.py

@@ -30,6 +30,7 @@ from .llama_func import *
 from .utils import *
 from . import shared
 from .config import retrieve_proxy
+from modules import config
 from .base_model import BaseLLMModel, ModelType
 
 
@@ -379,6 +380,8 @@ class ModelManager:
         msg = f"模型设置为了： {model_name}"
         logging.info(msg)
         model_type = ModelType.get_type(model_name)
+        if model_type != ModelType.OpenAI:
+            config.local_embedding = True
         if model_type == ModelType.OpenAI:
             model = OpenAIClient(
                 model_name=model_name,

requirements.txt

@@ -19,3 +19,4 @@ mpi4py
 icetk
 git+https://github.com/OptimalScale/LMFlow.git
 cpm-kernels
+sentence_transformers