ernie_bot_retriever.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from .base import BaseRetriever
  15. import os
  16. from langchain.docstore.document import Document
  17. from langchain.text_splitter import RecursiveCharacterTextSplitter
  18. from langchain_community.embeddings import QianfanEmbeddingsEndpoint
  19. from langchain_community.vectorstores import FAISS
  20. from langchain_community import vectorstores
  21. from erniebot_agent.extensions.langchain.embeddings import ErnieEmbeddings
  22. import time
  23. class ErnieBotRetriever(BaseRetriever):
  24. """Ernie Bot Retriever"""
  25. entities = [
  26. "ernie-4.0",
  27. "ernie-3.5",
  28. "ernie-3.5-8k",
  29. "ernie-lite",
  30. "ernie-tiny-8k",
  31. "ernie-speed",
  32. "ernie-speed-128k",
  33. "ernie-char-8k",
  34. ]
  35. def __init__(self, config):
  36. super().__init__()
  37. model_name = config.get('model_name', None)
  38. api_type = config.get('api_type', None)
  39. ak = config.get('ak', None)
  40. sk = config.get('sk', None)
  41. access_token = config.get('access_token', None)
  42. if model_name not in self.entities:
  43. raise ValueError(f"model_name must be in {self.entities} of ErnieBotChat.")
  44. if api_type not in ["aistudio", "qianfan"]:
  45. raise ValueError("api_type must be one of ['aistudio', 'qianfan']")
  46. if api_type == "aistudio" and access_token is None:
  47. raise ValueError("access_token cannot be empty when api_type is aistudio.")
  48. if api_type == "qianfan" and (ak is None or sk is None):
  49. raise ValueError("ak and sk cannot be empty when api_type is qianfan.")
  50. self.model_name = model_name
  51. self.config = config
  52. def generate_vector_database(self, text_list,
  53. block_size=300,
  54. separators=["\t", "\n", "。", "\n\n", ""],
  55. sleep_time=0.5):
  56. """
  57. args:
  58. return:
  59. """
  60. text_splitter = RecursiveCharacterTextSplitter(
  61. chunk_size=block_size, chunk_overlap=20, separators=separators
  62. )
  63. texts = text_splitter.split_text("\t".join(text_list))
  64. all_splits = [Document(page_content=text) for text in texts]
  65. api_type = self.config["api_type"]
  66. if api_type == "qianfan":
  67. os.environ["QIANFAN_AK"] = os.environ.get("EB_AK", self.config["ak"])
  68. os.environ["QIANFAN_SK"] = os.environ.get("EB_SK", self.config["sk"])
  69. user_ak = os.environ.get("EB_AK", self.config["ak"])
  70. user_id = hash(user_ak)
  71. vectorstore = FAISS.from_documents(
  72. documents=all_splits, embedding=QianfanEmbeddingsEndpoint()
  73. )
  74. elif api_type == "aistudio":
  75. token = self.config["access_token"]
  76. vectorstore = FAISS.from_documents(
  77. documents=all_splits[0:1],
  78. embedding=ErnieEmbeddings(aistudio_access_token=token),
  79. )
  80. #### ErnieEmbeddings.chunk_size = 16
  81. step = min(16, len(all_splits) - 1)
  82. for shot_splits in [
  83. all_splits[i : i + step] for i in range(1, len(all_splits), step)
  84. ]:
  85. time.sleep(sleep_time)
  86. vectorstore_slice = FAISS.from_documents(
  87. documents=shot_splits,
  88. embedding=ErnieEmbeddings(aistudio_access_token=token),
  89. )
  90. vectorstore.merge_from(vectorstore_slice)
  91. else:
  92. raise ValueError(f"Unsupported api_type: {api_type}")
  93. return vectorstore
  94. def encode_vector_store_to_bytes(self, vectorstore):
  95. vectorstore = self.encode_vector_store(vectorstore.serialize_to_bytes())
  96. return vectorstore
  97. def decode_vector_store_from_bytes(self, vectorstore):
  98. if not self.is_vector_store(vectorstore):
  99. raise ValueError("The retrieved vectorstore is not for PaddleX.")
  100. api_type = self.config["api_type"]
  101. if api_type == "aistudio":
  102. access_token = self.config["access_token"]
  103. embeddings = ErnieEmbeddings(aistudio_access_token=access_token)
  104. elif api_type == "qianfan":
  105. ak = self.config["ak"]
  106. sk = self.config["sk"]
  107. embeddings = QianfanEmbeddingsEndpoint(qianfan_ak=ak, qianfan_sk=sk)
  108. else:
  109. raise ValueError(f"Unsupported api_type: {api_type}")
  110. vectorstore = vectorstores.FAISS.deserialize_from_bytes(
  111. self.decode_vector_store(vector), embeddings
  112. )
  113. return vectorstore
  114. def similarity_retrieval(self, query_text_list, vectorstore, sleep_time=0.5):
  115. # 根据提问匹配上下文
  116. C = []
  117. for query_text in query_text_list:
  118. QUESTION = query_text
  119. time.sleep(sleep_time)
  120. docs = vectorstore.similarity_search_with_relevance_scores(QUESTION, k=2)
  121. context = [(document.page_content, score) for document, score in docs]
  122. context = sorted(context, key=lambda x: x[1])
  123. C.extend([x[0] for x in context[::-1]])
  124. C = list(set(C))
  125. all_C = " ".join(C)
  126. return all_C