7|hdZddlZddlmZddlmZmZmZddlm Z ddl m Z m Z m Z ddlmZdgZeeZGd de e Zy) z-written under MIT Licence, Michael Feil 2023.N) getLogger)AnyListOptional) Embeddings) BaseModel ConfigDictmodel_validator)SelfInfinityEmbeddingsLocalcFeZdZUdZeed< dZeeed< dZe ed< dZ eed< d Z eed < d Z e ed < dZeed < eddZeddefdZddZdeddfdZdeedeeefdZdedeefdZdeedeeefdZdedeefdZy)r aOptimized 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"]) modelNrevision batch_sizeautodevicetorchbackendT model_warmupengineforbid)extraprotected_namespacesafter)modereturnc ddlm}||j|j|j |j |j|j|_ |S#t$r tdwxYw)z?Validate that api key and python package exists in environment.r)AsyncEmbeddingEnginezrPlease install the `pip install 'infinity_emb[optimum,torch]>=0.0.24'` package to use the InfinityEmbeddingsLocal.)model_name_or_pathrrrrr) infinity_embr ImportErrorrrrrrrr)selfr s l/var/www/html/test/engine/venv/lib/python3.12/site-packages/langchain_community/embeddings/infinity_local.pyvalidate_environmentz,InfinityEmbeddingsLocal.validate_environmentBsn  9+#zz;;]]**<<     >  s AA.cTK|jjd{y7w)aQstart 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"]) N)r __aenter__)r$s r%r(z"InfinityEmbeddingsLocal.__aenter__Xskk$$&&&s (&(argscRK|jj|d{y7w)zUstop the background worker, required to free references to the pytorch model.N)r __aexit__)r$r)s r%r+z!InfinityEmbeddingsLocal.__aexit__es"$dkk##T***s '%'textscK|jjsmtjd|4d{|jj |d{\}}dddd{tjdS|jj |d{\}}|S7y7X7G#1d{7swYWxYw7$w)zAsync call out to Infinity's embedding endpoint. Args: texts: The list of texts to embed. Returns: List of embeddings, one for each text. zeStarting Infinity engine on the fly. This is not recommended.Please start the engine before using it.Nz$Stopped infinity engine after usage.)rrunningloggerwarningembed)r$r, embeddings_s r%aembed_documentsz(InfinityEmbeddingsLocal.aembed_documentsjs{{"" NN;  ? ?&*kk&7&7&> > A ? ? NNA B#'++"3"3E"::MJ ? > ? ? ? ? ;sb2C B.C B4B0B4 C *B2+9C $C % C 0B42C 4C:B= ;CC textcLK|j|gd{}|dS7 w)zAsync call out to Infinity's embedding endpoint. Args: text: The text to embed. Returns: Embeddings for the text. Nr)r4)r$r5r2s r% aembed_queryz$InfinityEmbeddingsLocal.aembed_querys, 00$88 !}9s $" $cttjdtj|j |S)z, This method is async only. zQThis method is async only. Please use the async version `await aembed_documents`.)r/r0asynciorunr4)r$r,s r%embed_documentsz'InfinityEmbeddingsLocal.embed_documentss2  E {{400788cttjdtj|j |S) zMThis method is async only. Please use the async version `await aembed_query`.)r/r0r9r:r7)r$r5s r% embed_queryz#InfinityEmbeddingsLocal.embed_querys0 B {{4,,T233r<)rN)__name__ __module__ __qualname____doc__str__annotations__rrrintrrrboolrrr model_configr r r&r(r+rfloatr4r7r;r?rr<r%r r s* JG"Hhsm"5J0FCAGSG%L$/FC*L '"d#* '+S+T+ DI$tE{:K. s tE{ 9T#Y94U 3D944U 4r<)rCr9loggingrtypingrrrlangchain_core.embeddingsrpydanticrr r typing_extensionsr __all__r@r/r rr<r%rPsB3&&0;;" $ % 8 M4iM4r<