Spaces:
Runtime error
Runtime error
File size: 5,196 Bytes
ed4d993 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 |
"""written under MIT Licence, Michael Feil 2023."""
import asyncio
from logging import getLogger
from typing import Any, Dict, List, Optional
from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import BaseModel, Extra, root_validator
__all__ = ["InfinityEmbeddingsLocal"]
logger = getLogger(__name__)
class InfinityEmbeddingsLocal(BaseModel, Embeddings):
"""Optimized Infinity embedding models.
https://github.com/michaelfeil/infinity
This class deploys a local Infinity instance to embed text.
The class requires async usage.
Infinity is a class to interact with Embedding Models on https://github.com/michaelfeil/infinity
Example:
.. code-block:: python
from langchain_community.embeddings import InfinityEmbeddingsLocal
async with InfinityEmbeddingsLocal(
model="BAAI/bge-small-en-v1.5",
revision=None,
device="cpu",
) as embedder:
embeddings = await engine.aembed_documents(["text1", "text2"])
"""
model: str
"Underlying model id from huggingface, e.g. BAAI/bge-small-en-v1.5"
revision: Optional[str] = None
"Model version, the commit hash from huggingface"
batch_size: int = 32
"Internal batch size for inference, e.g. 32"
device: str = "auto"
"Device to use for inference, e.g. 'cpu' or 'cuda', or 'mps'"
backend: str = "torch"
"Backend for inference, e.g. 'torch' (recommended for ROCm/Nvidia)"
" or 'optimum' for onnx/tensorrt"
model_warmup: bool = True
"Warmup the model with the max batch size."
engine: Any = None #: :meta private:
"""Infinity's AsyncEmbeddingEngine."""
# LLM call kwargs
class Config:
"""Configuration for this pydantic object."""
extra = Extra.forbid
@root_validator(allow_reuse=True)
def validate_environment(cls, values: Dict) -> Dict:
"""Validate that api key and python package exists in environment."""
try:
from infinity_emb import AsyncEmbeddingEngine # type: ignore
except ImportError:
raise ImportError(
"Please install the "
"`pip install 'infinity_emb[optimum,torch]>=0.0.24'` "
"package to use the InfinityEmbeddingsLocal."
)
logger.debug(f"Using InfinityEmbeddingsLocal with kwargs {values}")
values["engine"] = AsyncEmbeddingEngine(
model_name_or_path=values["model"],
device=values["device"],
revision=values["revision"],
model_warmup=values["model_warmup"],
batch_size=values["batch_size"],
engine=values["backend"],
)
return values
async def __aenter__(self) -> None:
"""start the background worker.
recommended usage is with the async with statement.
async with InfinityEmbeddingsLocal(
model="BAAI/bge-small-en-v1.5",
revision=None,
device="cpu",
) as embedder:
embeddings = await engine.aembed_documents(["text1", "text2"])
"""
await self.engine.__aenter__()
async def __aexit__(self, *args: Any) -> None:
"""stop the background worker,
required to free references to the pytorch model."""
await self.engine.__aexit__(*args)
async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
"""Async call out to Infinity's embedding endpoint.
Args:
texts: The list of texts to embed.
Returns:
List of embeddings, one for each text.
"""
if not self.engine.running:
logger.warning(
"Starting Infinity engine on the fly. This is not recommended."
"Please start the engine before using it."
)
async with self:
# spawning threadpool for multithreaded encode, tokenization
embeddings, _ = await self.engine.embed(texts)
# stopping threadpool on exit
logger.warning("Stopped infinity engine after usage.")
else:
embeddings, _ = await self.engine.embed(texts)
return embeddings
async def aembed_query(self, text: str) -> List[float]:
"""Async call out to Infinity's embedding endpoint.
Args:
text: The text to embed.
Returns:
Embeddings for the text.
"""
embeddings = await self.aembed_documents([text])
return embeddings[0]
def embed_documents(self, texts: List[str]) -> List[List[float]]:
"""
This method is async only.
"""
logger.warning(
"This method is async only. "
"Please use the async version `await aembed_documents`."
)
return asyncio.run(self.aembed_documents(texts))
def embed_query(self, text: str) -> List[float]:
""" """
logger.warning(
"This method is async only."
" Please use the async version `await aembed_query`."
)
return asyncio.run(self.aembed_query(text))
|