qianfan_bot_retriever.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163
  1. # Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
  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. import json
  15. from typing import Dict, List
  16. import requests
  17. from langchain_core.embeddings import Embeddings
  18. from paddlex.utils import logging
  19. from .base import BaseRetriever
  20. class QianFanBotRetriever(BaseRetriever):
  21. """QianFan Bot Retriever"""
  22. entities = [
  23. "qianfan",
  24. ]
  25. MODELS = [
  26. "tao-8k",
  27. "embedding-v1",
  28. "bge-large-zh",
  29. "bge-large-en",
  30. ]
  31. def __init__(self, config: Dict) -> None:
  32. """
  33. Initializes the ErnieBotRetriever instance with the provided configuration.
  34. Args:
  35. config (Dict): A dictionary containing configuration settings.
  36. - model_name (str): The name of the model to use.
  37. - api_type (str): The type of API to use ('qianfan' or 'openai').
  38. - api_key (str): The API key for 'qianfan' API.
  39. - base_url (str): The base URL for 'qianfan' API.
  40. Raises:
  41. ValueError: If api_type is not one of ['qianfan','openai'],
  42. base_url is None for api_type is qianfan,
  43. api_key is None for api_type is qianfan.
  44. """
  45. super().__init__()
  46. model_name = config.get("model_name", None)
  47. api_key = config.get("api_key", None)
  48. base_url = config.get("base_url", None)
  49. if model_name not in self.MODELS:
  50. raise ValueError(
  51. f"model_name must be in {self.MODELS} of QianFanBotRetriever."
  52. )
  53. if api_key is None:
  54. raise ValueError("api_key cannot be empty when api_type is qianfan.")
  55. if base_url is None:
  56. raise ValueError("base_url cannot be empty when api_type is qianfan.")
  57. self.embedding = QianfanEmbeddings(
  58. model=model_name,
  59. base_url=base_url,
  60. api_key=api_key,
  61. )
  62. self.model_name = model_name
  63. self.config = config
  64. class QianfanEmbeddings(Embeddings):
  65. """`Baidu Qianfan Embeddings` embedding models."""
  66. def __init__(
  67. self,
  68. api_key: str,
  69. base_url: str = "https://qianfan.baidubce.com/v2",
  70. model: str = "embedding-v1",
  71. **kwargs,
  72. ):
  73. """
  74. Initialize the Baidu Qianfan Embeddings class.
  75. Args:
  76. api_key (str): The Qianfan API key.
  77. base_url (str): The base URL for 'qianfan' API.
  78. model (str): Model name. Default is "embedding-v1",select in ["tao-8k","embedding-v1","bge-large-en","bge-large-zh"].
  79. kwargs (dict): Additional keyword arguments passed to the base Embeddings class.
  80. """
  81. super().__init__(**kwargs)
  82. chunk_size_map = {
  83. "tao-8k": 1,
  84. "embedding-v1": 16,
  85. "bge-large-en": 16,
  86. "bge-large-zh": 16,
  87. }
  88. self.api_key = api_key
  89. self.base_url = base_url
  90. self.model = model
  91. self.chunk_size = chunk_size_map.get(model, 1)
  92. def embed(self, texts: str, **kwargs) -> List[float]:
  93. url = f"{self.base_url}/embeddings"
  94. payload = json.dumps(
  95. {"model": kwargs.get("model", self.model), "input": [f"{texts}"]}
  96. )
  97. headers = {
  98. "Content-Type": "application/json",
  99. "Authorization": f"Bearer {self.api_key}",
  100. }
  101. response = requests.request("POST", url, headers=headers, data=payload)
  102. if response.status_code != 200:
  103. logging.error(
  104. f"Failed to call Qianfan API. Status code: {response.status_code}, Response content: {response}"
  105. )
  106. return response.json()
  107. def embed_query(self, text: str) -> List[float]:
  108. resp = self.embed_documents([text])
  109. return resp[0]
  110. def embed_documents(self, texts: List[str]) -> List[List[float]]:
  111. """
  112. Embeds a list of text documents using the AutoVOT algorithm.
  113. Args:
  114. texts (List[str]): A list of text documents to embed.
  115. Returns:
  116. List[List[float]]: A list of embeddings for each document in the input list.
  117. Each embedding is represented as a list of float values.
  118. """
  119. lst = []
  120. for chunk in texts:
  121. resp = self.embed(texts=chunk)
  122. lst.extend([res["embedding"] for res in resp["data"]])
  123. return lst
  124. async def aembed_query(self, text: str) -> List[float]:
  125. embeddings = await self.aembed_documents([text])
  126. return embeddings[0]
  127. async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
  128. lst = []
  129. for chunk in texts:
  130. resp = await self.embed(texts=chunk)
  131. for res in resp["data"]:
  132. lst.extend([res["embedding"]])
  133. return lst