PyTEI package

The main `TEIClient` class

The TEIClient is used to interface with a Text Embeddings Inference instance and is thus the main user-facing class.

Submodules

The pytei.store module

Use EmbeddingStore`s for caching embeddings. The `InMemoryEmbeddingStore serves as default in-memory cache. Use the DuckDBEmbeddingStore for persistent caching. Custom EmbeddingStore`s can be implemented by extending the abstract `EmbeddingStore base class.

class pytei.store.EmbeddingStore[source]

Bases: ABC

Abstract interface for a key-value store.

abstract get(key: str) → ndarray[source]: Get the embedding associated with the specified key. Raises KeyError if the key is not found. :param key: The key to get the embedding for. :type key: str :return: The embedding associated with the specified key. :rtype: numpy.ndarray

abstract get_all(keys: Collection[str]) → Dict[str, ndarray][source]: Get the embeddings associated with a set of keys. Returns a dictionary of all found key-value pairs. :param keys: The keys to get embeddings for. :type keys: Collection[str] :return: The embedding associated with the specified key. :rtype: numpy.ndarray

abstract put(key: str, value: ndarray) → None[source]: Store the embedding associated with the specified key. :param key: Identifier of the embedding. :type key: str :param value: Embedding to store. :type value: numpy.ndarray

abstract put_all(keys: List[str], values: List[ndarray]) → None[source]: Store the embeddings associated with the specified keys. :param keys: Identifiers of the embeddings to store. :type keys: List[str] :param values: Embeddings to store. :type values: List[numpy.ndarray]

abstract remove(key: str) → None[source]: Remove the embedding associated with the specified key. :param key: Identifier of the embedding to remove. :type key: str

class pytei.store.InMemoryEmbeddingStore[source]

Bases: EmbeddingStore

In-memory key-value store for embeddings.

get(key: str) → ndarray[source]: Get the embedding associated with the specified key. Raises KeyError if the key is not found. :param key: The key to get the embedding for. :type key: str :return: The embedding associated with the specified key. :rtype: numpy.ndarray

get_all(keys: Collection[str]) → Dict[str, ndarray][source]: Get the embeddings associated with a set of keys. Returns a dictionary of all found key-value pairs. :param keys: The keys to get embeddings for. :type keys: Collection[str] :return: The embedding associated with the specified key. :rtype: numpy.ndarray

put(key: str, value: ndarray)[source]: Store the embedding associated with the specified key. :param key: Identifier of the embedding. :type key: str :param value: Embedding to store. :type value: numpy.ndarray

put_all(keys: List[str], values: List[ndarray]) → None[source]: Store the embeddings associated with the specified keys. :param keys: Identifiers of the embeddings to store. :type keys: List[str] :param values: Embeddings to store. :type values: List[numpy.ndarray]

remove(key: str)[source]: Remove the embedding associated with the specified key. :param key: Identifier of the embedding to remove. :type key: str

class pytei.store.DuckDBEmbeddingStore(db_path: str = 'datastore.duckdb')[source]

Bases: EmbeddingStore

Persistent key-value store using DuckDB as backend.

get(key: str) → ndarray[source]: Get the embedding associated with the specified key. Raises KeyError if the key is not found. :param key: The key to get the embedding for. :type key: str :return: The embedding associated with the specified key. :rtype: numpy.ndarray

get_all(keys: Collection[str]) → Dict[str, ndarray][source]: Get the embeddings associated with a set of keys. Returns a dictionary of all found key-value pairs. :param keys: The keys to get embeddings for. :type keys: Collection[str] :return: The embedding associated with the specified key. :rtype: numpy.ndarray

put(key: str, value: ndarray)[source]: Store the embedding associated with the specified key. :param key: Identifier of the embedding. :type key: str :param value: Embedding to store. :type value: numpy.ndarray

put_all(keys: List[str], values: List[ndarray]) → None[source]: Store the embeddings associated with the specified keys. :param keys: Identifiers of the embeddings to store. :type keys: List[str] :param values: Embeddings to store. :type values: List[numpy.ndarray]

remove(key: str)[source]: Remove the embedding associated with the specified key. :param key: Identifier of the embedding to remove. :type key: str

The pytei.model module

The model defines the structure of the data returned by the TEIClient.

class pytei.model.PredictionResult(label: str, score: float)[source]

Bases: object

label: str

score: float

class pytei.model.Rank(index: int, score: float, text: str | None = None)[source]