codekingpro/portable-devtools

1from typing import Any, Dict, List, Mapping, Optional2 3import requests4from langchain_core.embeddings import Embeddings5from langchain_core.utils import get_from_dict_or_env, pre_init6from pydantic import BaseModel, ConfigDict7 8DEFAULT_MODEL_ID = "sentence-transformers/clip-ViT-B-32"9MAX_BATCH_SIZE = 102410 11 12class DeepInfraEmbeddings(BaseModel, Embeddings):13    """Deep Infra's embedding inference service.14 15    To use, you should have the16    environment variable ``DEEPINFRA_API_TOKEN`` set with your API token, or pass17    it as a named parameter to the constructor.18    There are multiple embeddings models available,19    see https://deepinfra.com/models?type=embeddings.20 21    Example:22        .. code-block:: python23 24            from langchain_community.embeddings import DeepInfraEmbeddings25            deepinfra_emb = DeepInfraEmbeddings(26                model_id="sentence-transformers/clip-ViT-B-32",27                deepinfra_api_token="my-api-key"28            )29            r1 = deepinfra_emb.embed_documents(30                [31                    "Alpha is the first letter of Greek alphabet",32                    "Beta is the second letter of Greek alphabet",33                ]34            )35            r2 = deepinfra_emb.embed_query(36                "What is the second letter of Greek alphabet"37            )38 39    """40 41    model_id: str = DEFAULT_MODEL_ID42    """Embeddings model to use."""43    normalize: bool = False44    """whether to normalize the computed embeddings"""45    embed_instruction: str = "passage: "46    """Instruction used to embed documents."""47    query_instruction: str = "query: "48    """Instruction used to embed the query."""49    model_kwargs: Optional[dict] = None50    """Other model keyword args"""51    deepinfra_api_token: Optional[str] = None52    """API token for Deep Infra. If not provided, the token is 53    fetched from the environment variable 'DEEPINFRA_API_TOKEN'."""54    batch_size: int = MAX_BATCH_SIZE55    """Batch size for embedding requests."""56 57    model_config = ConfigDict(extra="forbid", protected_namespaces=())58 59    @pre_init60    def validate_environment(cls, values: Dict) -> Dict:61        """Validate that api key and python package exists in environment."""62        deepinfra_api_token = get_from_dict_or_env(63            values, "deepinfra_api_token", "DEEPINFRA_API_TOKEN"64        )65        values["deepinfra_api_token"] = deepinfra_api_token66        return values67 68    @property69    def _identifying_params(self) -> Mapping[str, Any]:70        """Get the identifying parameters."""71        return {"model_id": self.model_id}72 73    def _embed(self, input: List[str]) -> List[List[float]]:74        _model_kwargs = self.model_kwargs or {}75        # HTTP headers for authorization76        headers = {77            "Authorization": f"bearer {self.deepinfra_api_token}",78            "Content-Type": "application/json",79        }80        # send request81        try:82            res = requests.post(83                f"https://api.deepinfra.com/v1/inference/{self.model_id}",84                headers=headers,85                json={"inputs": input, "normalize": self.normalize, **_model_kwargs},86            )87        except requests.exceptions.RequestException as e:88            raise ValueError(f"Error raised by inference endpoint: {e}")89 90        if res.status_code != 200:91            raise ValueError(92                "Error raised by inference API HTTP code: %s, %s"93                % (res.status_code, res.text)94            )95        try:96            t = res.json()97            embeddings = t["embeddings"]98        except requests.exceptions.JSONDecodeError as e:99            raise ValueError(100                f"Error raised by inference API: {e}.\nResponse: {res.text}"101            )102 103        return embeddings104 105    def embed_documents(self, texts: List[str]) -> List[List[float]]:106        """Embed documents using a Deep Infra deployed embedding model.107        For larger batches, the input list of texts is chunked into smaller108        batches to avoid exceeding the maximum request size.109 110        Args:111            texts: The list of texts to embed.112 113        Returns:114            List of embeddings, one for each text.115        """116 117        embeddings = []118        instruction_pairs = [f"{self.embed_instruction}{text}" for text in texts]119 120        chunks = [121            instruction_pairs[i : i + self.batch_size]122            for i in range(0, len(instruction_pairs), self.batch_size)123        ]124        for chunk in chunks:125            embeddings += self._embed(chunk)126 127        return embeddings128 129    def embed_query(self, text: str) -> List[float]:130        """Embed a query using a Deep Infra deployed embedding model.131 132        Args:133            text: The text to embed.134 135        Returns:136            Embeddings for the text.137        """138        instruction_pair = f"{self.query_instruction}{text}"139        embedding = self._embed([instruction_pair])[0]140        return embedding141