codekingpro/portable-devtools

1import hashlib2from typing import List3 4import numpy as np5from langchain_core.embeddings import Embeddings6from pydantic import BaseModel7 8 9class FakeEmbeddings(Embeddings, BaseModel):10    """Fake embedding model."""11 12    size: int13    """The size of the embedding vector."""14 15    def _get_embedding(self) -> List[float]:16        return list(np.random.normal(size=self.size))17 18    def embed_documents(self, texts: List[str]) -> List[List[float]]:19        return [self._get_embedding() for _ in texts]20 21    def embed_query(self, text: str) -> List[float]:22        return self._get_embedding()23 24 25class DeterministicFakeEmbedding(Embeddings, BaseModel):26    """27    Fake embedding model that always returns28    the same embedding vector for the same text.29    """30 31    size: int32    """The size of the embedding vector."""33 34    def _get_embedding(self, seed: int) -> List[float]:35        # set the seed for the random generator36        np.random.seed(seed)37        return list(np.random.normal(size=self.size))38 39    @staticmethod40    def _get_seed(text: str) -> int:41        """42        Get a seed for the random generator, using the hash of the text.43        """44        return int(hashlib.sha256(text.encode("utf-8")).hexdigest(), 16) % 10**845 46    def embed_documents(self, texts: List[str]) -> List[List[float]]:47        return [self._get_embedding(seed=self._get_seed(_)) for _ in texts]48 49    def embed_query(self, text: str) -> List[float]:50        return self._get_embedding(seed=self._get_seed(text))51