使用Embedding API完成Embedding模型接入并进行Embedding封装
项目代码
该部分架构:
embedding
_init_.py
call_embedding.py
zhipuai_embedding.py
"""zhipiai_embedding.py"""
from __future__ import annotations
"""定义了一个名为 ZhipuAIEmbeddings 的类,用于生成文本的嵌入(embedding)。
这个类继承了 BaseModel 和 Embeddings,并使用了 pydantic 进行数据验证。"""
"""导入模块"""
import logging
from typing import Dict, List, Any
from langchain.embeddings.base import Embeddings
from langchain.pydantic_v1 import BaseModel, root_validator
"""这些导入提供了所需的模块和类型,logging 用于日志记录,typing 提供类型提示,
langchain 模块提供了 Embeddings 基类和 pydantic 数据验证。"""
"""设置日志记录"""
logger = logging.getLogger(__name__)
"""创建一个日志生成器,用于记录日志信息"""
"""定义 ZhipuAIEmbeddings 类:
用于表示 ZhipuAI 的嵌入模型。client 属性用于存储 ZhipuAI 的客户端实例。"""
class ZhipuAIEmbeddings(BaseModel, Embeddings):
"""`Zhipuai Embeddings` embedding models."""
client: Any
"""`zhipuai.ZhipuAI"""
"""环境验证:"""
@root_validator()
def validate_environment(cls, values: Dict) -> Dict:
"""
实例化ZhipuAI为values["client"]
Args:
values (Dict): 包含配置信息的字典,必须包含 client 的字段.
Returns:
values (Dict): 包含配置信息的字典。如果环境中有zhipuai库,则将返回实例化的ZhipuAI类;否则将报错 'ModuleNotFoundError: No module named 'zhipuai''.
"""
from zhipuai import ZhipuAI
values["client"] = ZhipuAI()
return values
"""在类实例化时调用,用于验证和初始化环境。
如果环境中有 zhipuai 库,则实例化 ZhipuAI 并赋值给 client 属性,
否则抛出模块未找到错误。"""
"""生成单个文本的嵌入:这个方法生成输入文本的嵌入,返回一个浮点数值列表"""
def embed_query(self, text: str) -> List[float]:
"""
生成输入文本的 embedding.
Args:
texts (str): 要生成 embedding 的文本.
Return:
embeddings (List[float]): 输入文本的 embedding,一个浮点数值列表.
"""
embeddings = self.client.embeddings.create(
model="embedding-2",
input=text
)
return embeddings.data[0].embedding
"""生成多个文本的嵌入:生成输入文本列表的嵌入,返回嵌入列表"""
def embed_documents(self, texts: List[str]) -> List[List[float]]:
"""
生成输入文本列表的 embedding.
Args:
texts (List[str]): 要生成 embedding 的文本列表.
Returns:
List[List[float]]: 输入列表中每个文档的 embedding 列表。每个 embedding 都表示为一个浮点值列表。
"""
return [self.embed_query(text) for text in texts]
"""异步方法(未实现)"""
async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
"""Asynchronous Embed search docs."""
raise NotImplementedError("Please use `embed_documents`. Official does not support asynchronous requests")
async def aembed_query(self, text: str) -> List[float]:
"""Asynchronous Embed query text."""
raise NotImplementedError("Please use `aembed_query`. Official does not support asynchronous requests")
"""这两个方法是异步版本的 embed_documents 和 embed_query,但未实现,并提示用户使用同步方法。"""