Skip to content
BEDbase

Search Module

HNSWBackend

Section titled “HNSWBackend”
class HNSWBackend(local_index_path: str = DEFAULT_INDEX_PATH, payloads: Union[dict, str] = dict(), space: str = DEFAULT_HNSW_SPACE, dim: int = DEFAULT_DIM, ef: int = DEFAULT_EF, m: int = DEFAULT_M)

Bases: EmSearchBackend

A search backend that uses a local HNSW index to store and search embeddings

Initialize the HNSWBackend.

Args: local_index_path: local path where the index is saved to space: possible options are l2, cosine or ip dim: dimension of vectors that will be stored ef: defines a construction time/accuracy trade-off, higher ef -> more accurate but slower m: connected with internal dimensionality of the data, higher M -> higher accuracy/run_time when ef is fixed payloads: optional dictionary or path to payloads file

Attributes

Section titled “Attributes”
  • idx
  • payloads
  • idx_path

Methods

Section titled “Methods”

load

Section titled “load”
def load(vectors: np.ndarray, ids: Union[np.ndarray, None] = None, payloads: Union[List[Dict[str, str]], None] = None)

Upload embedding vectors into the HNSW index and store metadata.

Args: vectors: embedding vectors, a np.ndarray with shape of (n, ) ids: list of n point ids, or None to generate ids automatically payloads: optional list of n dictionaries that contain vector metadata

Section titled “search”
def search(query: np.ndarray, limit: int, with_payload: bool = True, with_vectors: bool = True, offset: int = 0) -> Union[List[Dict[str, Union[int, float, Dict[str, str], np.ndarray]]], List[List[Dict[str, Union[int, float, Dict[str, str], np.ndarray]]]]]

Get the limit nearest neighbors for query vector(s).

Args: query: the query vector, np.ndarray with shape of (1, dim) or (dim, ) limit: number of nearest neighbors to search for query vector with_payload: whether payload is included in the result with_vectors: whether the stored vector is included in the result offset: the offset of the search results

Returns: If query shape is (, ), returns a list of limit dictionaries in format: {"id": , "distance": , "payload": {...}, "vector": [...]} If query shape is (n, ), returns a 2D list where each sublist has limit dictionaries.

retrieve_info

Section titled “retrieve_info”
def retrieve_info(ids: Union[List[int], int], with_vectors: bool = False) -> Union[Dict[str, Union[int, List[float], Dict[str, str]]], List[Dict[str, Union[int, List[float], Dict[str, str]]]]]

Return the information of vectors for given storage ids.

Args: ids: storage id, or a list of ids with_vectors: whether the stored vector is included in the result

Returns: If single id, returns a dictionary. If list of ids, returns a list of dictionaries.

QdrantBackend

Section titled “QdrantBackend”
class QdrantBackend(qdrant_client: QdrantClient, dim: int = DEFAULT_DIM, dist: Distance = DEFAULT_QDRANT_DIST, collection: str = DEFAULT_COLLECTION_NAME)

Bases: EmSearchBackend

A search backend that uses a qdrant server to store and search embeddings

Initialize QdrantBackend connection.

Connect to Qdrant on commandline first: (Ubuntu Linux terminal) sudo docker run -p 6333:6333 -v $(pwd)/qdrant_storage:/qdrant/storage qdrant/qdrant

Parameters:

  • qdrant_client (QdrantClient): an instance of QdrantClient that is already connected to a Qdrant server
  • dim (int): dimension of the vectors to be stored (default: DEFAULT_DIM)
  • dist (Distance): distance metric used in the collection (default: DEFAULT_QDRANT_DIST)
  • collection (str): name of collection (default: DEFAULT_COLLECTION_NAME)

Attributes

Section titled “Attributes”
  • collection
  • config
  • qd_client

Class methods

Section titled “Class methods”

from_credentials

Section titled “from_credentials”
@classmethod
def from_credentials(dim: int = DEFAULT_DIM, dist: Distance = DEFAULT_QDRANT_DIST, collection: str = DEFAULT_COLLECTION_NAME, qdrant_host: str = DEFAULT_QDRANT_HOST, qdrant_port: int = DEFAULT_QDRANT_PORT, qdrant_api_key: str = None) -> QdrantBackend

Initialize QdrantBackend from connection parameters.

Args: dim: dimension of the vectors to be stored dist: distance metric used in the collection collection: collection name qdrant_host: host of the qdrant server qdrant_port: port of the qdrant server qdrant_api_key: api key for the qdrant server if needed

Returns: QdrantBackend instance

Methods

Section titled “Methods”

load

Section titled “load”
def load(vectors: np.ndarray, ids: Union[List[str], None] = None, payloads: Union[List[Dict[str, str]], None] = None)

Upload vectors and their labels into qdrant storage.

Args: vectors: embedding vectors, a np.ndarray with shape of (n, ) ids: list of n point ids, or None to generate ids automatically payloads: optional list of n dictionaries that contain vector metadata

search

Section titled “search”
def search(query: np.ndarray, limit: int, with_payload: bool = True, with_vectors: bool = True, offset: int = 0) -> Union[List[Dict[str, Union[int, float, Dict[str, str], List[float]]]], List[List[Dict[str, Union[int, float, Dict[str, str], List[float]]]]]]

Get k nearest neighbors from vectors in the collection.

Args: query: a vector to search limit: number of returned results with_payload: whether payload is included in the result with_vectors: whether the stored vector is included in the result offset: the offset of the search results

Returns: A list of dictionaries containing search results in this format: { "id": , "score": , "payload": {}, "vector": [] }

Section titled “batch_search”
def batch_search(queries: np.ndarray, limit: int, with_payload: bool = True, with_vectors: bool = True, offset: int = 0) -> List[List[Dict[str, Union[int, float, Dict[str, str], List[float]]]]]

Perform batch search with multiple query vectors.

Args: queries: multiple search vectors, np.ndarray with shape of (n, dim) limit: see docstring of search() with_payload: whether payload is included in the result with_vectors: whether the stored vector is included in the result offset: the offset of the search results

Returns: Results of all search requests, one list per query vector.

__len__

Section titled “__len__”
def __len__() -> int

Return the number of embeddings in the backend.

Returns: The number of embeddings in the backend.

retrieve_info

Section titled “retrieve_info”
def retrieve_info(ids: Union[List[int], int, List[str], str], with_vectors: bool = False) -> Union[Dict[str, Union[int, str, List[float], Dict[str, str]]], List[Dict[str, Union[int, str, List[float], Dict[str, str]]]]]

Return the information of vectors for given storage ids.

Args: ids: list of ids, or a single id with_vectors: whether the vectors themselves will also be returned in the output

Returns: If ids is one id, a dictionary similar to the output of search() without "score". If ids is a list, a list of dictionaries will be returned.

merge_backends

Section titled “merge_backends”
def merge_backends(backends_to_merge: List[HNSWBackend], local_index_path: str, dim: int) -> HNSWBackend

Merge multiple backends into a single backend.

Args: backends_to_merge: a list of HNSWBackend instances to merge local_index_path: the path to the local index file of the merged output HNSWBackend dim: the dimension of vectors stored in the HNSWBackend

Returns: A merged HNSWBackend containing all vectors from input backends.

BED2BEDSearchInterface

Section titled “BED2BEDSearchInterface”
class BED2BEDSearchInterface(backend: Union[QdrantBackend, HNSWBackend], query2vec: Union[str, BED2Vec])

Bases: BEDSearchInterface

Search interface for the query that is a region set

Initialize BED2BEDSearchInterface.

Args: backend: the backend where vectors are stored query2vec: a BED2Vec model or a Hugging Face model repository of Region2VecExModel

Attributes

Section titled “Attributes”
  • query2vec
  • backend

Methods

Section titled “Methods”

Section titled “query_search”
def query_search(query: Union[str, RegionSet, np.ndarray], limit: int, with_payload: bool = True, with_vectors: bool = True, offset: int = 0) -> List[Dict]

Search for BED files similar to a query region set.

Args: query: a region set, path to a BED file on disk, or a region set embedding vector limit: number of results to return with_payload: whether payload is included in the result with_vectors: whether the stored vector is included in the result offset: the offset of the search results

Returns: A list of dictionaries containing search results.

Text2BEDSearchInterface

Section titled “Text2BEDSearchInterface”
class Text2BEDSearchInterface(backend: Union[QdrantBackend, HNSWBackend], query2vec: Text2Vec)

Bases: BEDSearchInterface

Search interface for the query that is a natural language string

Initialize the Text2BEDSearchInterface.

Args: backend: the backend where vectors are stored query2vec: a Text2Vec instance (see geniml.search.query2vec.text2vec for details)

Attributes

Section titled “Attributes”
  • query2vec
  • backend

Methods

Section titled “Methods”

query_search

Section titled “query_search”
def query_search(query: Union[str, np.ndarray], limit: int, with_payload: bool = True, with_vectors: bool = True, offset: int = 0) -> List[Dict]

Search for BED files using natural language query or embedding vector.

Args: query: the natural language query string or a vector in the embedding space of region sets limit: number of results to return with_payload: whether payload is included in the result with_vectors: whether the stored vector is included in the result offset: the offset of the search results

Returns: A list of dictionaries containing search results.

eval

Section titled “eval”
def eval(query_dict: Dict[str, List[Union[int, np.int64]]]) -> Dict[str, float]

Evaluate retrieval performance using Mean Average Precision, AUC-ROC, and R-precision.

Args: query_dict: a dictionary containing queries and relevant results in this format: { : [ , ... ], ... }

Returns: A dictionary with keys: "Mean Average Precision", "Mean AUC-ROC", "Average R-Precision"

BED2Vec

Section titled “BED2Vec”
class BED2Vec(model: Union[str, Region2VecExModel])

Bases: Query2Vec

Embed a query region set into a vector

Initialize BED2Vec with a Region2VecExModel.

Args: model: a Region2VecExModel instance or a model repository on Hugging Face

Attributes

Section titled “Attributes”
  • model

Methods

Section titled “Methods”

forward

Section titled “forward”
def forward(query: Union[str, RegionSet]) -> np.ndarray

Embed the query region set into a vector.

Args: query: a RegionSet instance or the path to a BED file on disk

Returns: The region set embedding vector.

Text2Vec

Section titled “Text2Vec”
class Text2Vec(model: Union[str, TextEmbedding], v2v: Union[str, Vec2VecFNN, None])

Map a query string into a vector into the embedding space of region sets

Parameters:

  • model (Union[str, TextEmbedding]): a model repository on Hugging Face
  • v2v (Union[str, Vec2VecFNN, None]): a Vec2VecFNN (see geniml/text2bednn/text2bednn.py) or a model repository on Hugging Face

Attributes

Section titled “Attributes”
  • text_embedder
  • v2v

Methods

Section titled “Methods”

forward

Section titled “forward”
def forward(query: str) -> np.ndarray

Embed a natural language query string into a vector.

Args: query: a natural language string

Returns: The embedding vector of the query.

anecdotal_search_from_hf_data

Section titled “anecdotal_search_from_hf_data”
def anecdotal_search_from_hf_data(query: str, dataset_repo: str, search_model_repo: str, text_embed_model_repo: str, k: int = 10) -> List[Dict[str, Union[float, int, Dict[str, str]]]]

Test retrieval performance of a trained search model on a Hugging Face dataset.

Args: query: user input search term dataset_repo: Hugging Face repository of the dataset search_model_repo: Hugging Face repository of the search model text_embed_model_repo: Hugging Face repository of the text encoder model k: number of results to return

Returns: A list of dictionaries containing search scores (distance to the mapped query vector) and search result metadata.

rand_eval

Section titled “rand_eval”
def rand_eval(n: int, query_dict: Dict) -> Tuple[float, float, float]

Compute evaluation scores for random retrieval performance.

Args: n: total number of results

query_dict: a dictionary containing queries and relevant results in this format:
    {
        <query string>: [
            <store id in backend>,
            ...
        ],
        ...
    }

Returns: A tuple of (Average Precision, AUC-ROC, R-precision) scores.