Cohere

CohereEmbedding

Bases: BaseEmbedding

CohereEmbedding使用Cohere API为文本生成嵌入。

Source code in llama_index/embeddings/cohere/base.py

class CohereEmbedding(BaseEmbedding):
    """CohereEmbedding使用Cohere API为文本生成嵌入。"""

    # Instance variables initialized via Pydantic's mechanism
    cohere_client: cohere.Client = Field(description="CohereAI client")
    cohere_async_client: cohere.AsyncClient = Field(description="CohereAI Async client")
    truncate: str = Field(description="Truncation type - START/ END/ NONE")
    input_type: Optional[str] = Field(
        description="Model Input type. If not provided, search_document and search_query are used when needed."
    )
    embedding_type: str = Field(
        description="Embedding type. If not provided float embedding_type is used when needed."
    )

    def __init__(
        self,
        cohere_api_key: Optional[str] = None,
        model_name: str = "embed-english-v3.0",
        truncate: str = "END",
        input_type: Optional[str] = None,
        embedding_type: str = "float",
        embed_batch_size: int = DEFAULT_EMBED_BATCH_SIZE,
        callback_manager: Optional[CallbackManager] = None,
        base_url: Optional[str] = None,
        timeout: Optional[float] = None,
        httpx_client: Optional[httpx.Client] = None,
        httpx_async_client: Optional[httpx.AsyncClient] = None,
    ):
        """一个使用Cohere API生成嵌入的类表示。

Args:
    cohere_client (Any): Cohere客户端的实例，用于与Cohere API进行通信。
    truncate (str): 一个字符串，指示要应用于输入文本的截断策略。可能的取值为'START'、'END'或'NONE'。
    input_type (Optional[str]): 一个可选的字符串，指定提供给模型的输入类型。这取决于模型，可能是以下之一：'search_query'、'search_document'、'classification'或'clustering'。
    model_name (str): 用于生成嵌入的模型的名称。该类确保该模型得到支持，并且提供的输入类型与模型兼容。
"""
        # Validate model_name and input_type
        if model_name not in VALID_MODEL_INPUT_TYPES:
            raise ValueError(f"{model_name} is not a valid model name")

        if input_type not in VALID_MODEL_INPUT_TYPES[model_name]:
            raise ValueError(
                f"{input_type} is not a valid input type for the provided model."
            )
        if embedding_type not in VALID_MODEL_EMBEDDING_TYPES[model_name]:
            raise ValueError(
                f"{embedding_type} is not a embedding type for the provided model."
            )

        if truncate not in VALID_TRUNCATE_OPTIONS:
            raise ValueError(f"truncate must be one of {VALID_TRUNCATE_OPTIONS}")

        super().__init__(
            cohere_client=cohere.Client(
                cohere_api_key,
                client_name="llama_index",
                base_url=base_url,
                timeout=timeout,
                httpx_client=httpx_client,
            ),
            cohere_async_client=cohere.AsyncClient(
                cohere_api_key,
                client_name="llama_index",
                base_url=base_url,
                timeout=timeout,
                httpx_client=httpx_async_client,
            ),
            cohere_api_key=cohere_api_key,
            model_name=model_name,
            input_type=input_type,
            embedding_type=embedding_type,
            truncate=truncate,
            embed_batch_size=embed_batch_size,
            callback_manager=callback_manager,
        )

    @classmethod
    def class_name(cls) -> str:
        return "CohereEmbedding"

    def _embed(self, texts: List[str], input_type: str) -> List[List[float]]:
        """使用Cohere嵌入句子。"""
        if self.model_name in V3_MODELS:
            result = self.cohere_client.embed(
                texts=texts,
                input_type=self.input_type or input_type,
                embedding_types=[self.embedding_type],
                model=self.model_name,
                truncate=self.truncate,
            ).embeddings
        else:
            result = self.cohere_client.embed(
                texts=texts,
                model=self.model_name,
                embedding_types=[self.embedding_type],
                truncate=self.truncate,
            ).embeddings
        return getattr(result, self.embedding_type, None)

    async def _aembed(self, texts: List[str], input_type: str) -> List[List[float]]:
        """使用Cohere嵌入句子。"""
        if self.model_name in V3_MODELS:
            result = (
                await self.cohere_async_client.embed(
                    texts=texts,
                    input_type=self.input_type or input_type,
                    embedding_types=[self.embedding_type],
                    model=self.model_name,
                    truncate=self.truncate,
                )
            ).embeddings
        else:
            result = (
                await self.cohere_async_client.embed(
                    texts=texts,
                    model=self.model_name,
                    embedding_types=[self.embedding_type],
                    truncate=self.truncate,
                )
            ).embeddings
        return getattr(result, self.embedding_type, None)

    def _get_query_embedding(self, query: str) -> List[float]:
        """获取查询嵌入。对于查询嵌入，输入类型为'search_query'。"""
        return self._embed([query], input_type="search_query")[0]

    async def _aget_query_embedding(self, query: str) -> List[float]:
        """异步获取查询嵌入。对于查询嵌入，input_type='search_query'。"""
        return (await self._aembed([query], input_type="search_query"))[0]

    def _get_text_embedding(self, text: str) -> List[float]:
        """获取文本嵌入。"""
        return self._embed([text], input_type="search_document")[0]

    async def _aget_text_embedding(self, text: str) -> List[float]:
        """异步获取文本嵌入。"""
        return (await self._aembed([text], input_type="search_document"))[0]

    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
        """获取文本嵌入。"""
        return self._embed(texts, input_type="search_document")

    async def _aget_text_embeddings(self, texts: List[str]) -> List[List[float]]:
        """获取文本嵌入。"""
        return await self._aembed(texts, input_type="search_document")