[iddlmZmZmZmZmZmZddlmZddl m Z ddl m Z ddl mZddlmZddlmZmZddlmZdd lZdd lmZGd d eZy ) )AnyIterableOptionalSequenceTypeUnion)asdict) OnnxProvider)Bm25)Bm42)MiniCOIL)SparseEmbeddingSparseTextEmbeddingBase)SpladePPN)SparseModelDescriptioncleZdZUeeeegZee e e d<e dee eeffdZe deefdZ ddedeedeedeeed ed eeed ed effd Z ddeeeefdedeed edeef dZdeeeefd edeefdZ ddeeeefded edefdZxZS)SparseTextEmbeddingEMBEDDINGS_REGISTRYreturncZ|jDcgc] }t|c}Scc}w)a Lists the supported models. Returns: list[dict[str, Any]]: A list of dictionaries containing the model information. Example: ``` [ { "model": "prithvida/SPLADE_PP_en_v1", "vocab_size": 30522, "description": "Independent Implementation of SPLADE++ Model for English", "license": "apache-2.0", "size_in_GB": 0.532, "sources": { "hf": "qdrant/SPLADE_PP_en_v1", }, } ] ``` )_list_supported_modelsr )clsmodels b/mnt/e/genesis-system/.venv/lib/python3.12/site-packages/fastembed/sparse/sparse_text_embedding.pylist_supported_modelsz)SparseTextEmbedding.list_supported_modelss%0,/+E+E+GH%u HHHs(cjg}|jD]!}|j|j#|SN)rextendr)rresult embeddings rrz*SparseTextEmbedding._list_supported_models.s6/100 >I MM)::< = > model_name cache_dirthreads providerscuda device_ids lazy_loadkwargsc ^t |||fi|jdjk(rtjdt dd|j D]=} | j} tfd| Ds(| |f|||||d||_ ytdd ) Nzprithvida/Splade_PP_en_v1zvThe right spelling is prithivida/Splade_PP_en_v1. Support of this name will be removed soon, please fix the model_name) stacklevelzprithivida/Splade_PP_en_v1c3tK|]/}j|jjk(1ywr)lowerr).0rr"s r z/SparseTextEmbedding.__init__..Ls,[:##%):):)<<[s58)r$r%r&r'r(zModel z~ is not supported in SparseTextEmbedding.Please check the supported models using `SparseTextEmbedding.list_supported_models()`) super__init__r.warningswarnDeprecationWarningrranyr ValueError) selfr"r#r$r%r&r'r(r)EMBEDDING_MODEL_TYPEsupported_models __class__s ` rr2zSparseTextEmbedding.__init__5s YB6B    ! 1, data-parallel encoding will be used, recommended for offline encoding of large datasets. If 0, use all available cores. If None, don't use data-parallel processing, use default onnxruntime threading instead. Returns: List of embeddings, one per document N)rembed)r8r<r=r>r)s rr@zSparseTextEmbedding.embed^s+,$4::##Iz8NvNNNs #-+-queryc+ZK|jj|fi|Ed{y7w)z Embeds queries Args: query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries. Returns: Iterable[SparseEmbedding]: The sparse embeddings. N)r query_embed)r8rAr)s rrCzSparseTextEmbedding.query_embedvs'*4::))%:6:::s !+)+textsc @|jj|fd|i|S)aReturns the number of tokens in the texts. Args: texts (str | Iterable[str]): The list of texts to embed. batch_size (int): Batch size for encoding Returns: int: Sum of number of tokens in the texts. r=)r token_count)r8rDr=r)s rrFzSparseTextEmbedding.token_counts%&tzz%%eM MfMMr!)NNNFNF)N)i) __name__ __module__ __qualname__rr r r rlistrr__annotations__ classmethoddictstrrrrrrintrr boolr2rrrr@rCrF __classcell__)r;s@rrrs@H$PTV^?_d#:;<_Id4S>&:II2t,B'C$(!%6:*.' ' C=' # ' H\23 '  ' T#Y'' ' ' X"& Ohsm+,OO3- O  O / " O0 ;3 -. ;:= ; / " ;CG N3 -. Nr]s;AA)&&.0EN1Nr!