# Search Module

<a id="geniml.search"></a>

<a id="geniml.search.HNSWBackend"></a>

## `HNSWBackend`

```python
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

- <a id="geniml.search.HNSWBackend.idx"></a>**`idx`**
- <a id="geniml.search.HNSWBackend.payloads"></a>**`payloads`**
- <a id="geniml.search.HNSWBackend.idx_path"></a>**`idx_path`**

### Methods

<a id="geniml.search.HNSWBackend.load"></a>

#### `load`

```python
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, <vector size>)
    ids: list of n point ids, or None to generate ids automatically
    payloads: optional list of n dictionaries that contain vector metadata

<a id="geniml.search.HNSWBackend.search"></a>

#### `search`

```python
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 (<dim>, ), returns a list of limit dictionaries in format:
    {"id": <id>, "distance": <distance>, "payload": {...}, "vector": [...]}
    If query shape is (n, <dim>), returns a 2D list where each sublist has limit dictionaries.

<a id="geniml.search.HNSWBackend.retrieve_info"></a>

#### `retrieve_info`

```python
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.

<a id="geniml.search.QdrantBackend"></a>

## `QdrantBackend`

```python
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

- <a id="geniml.search.QdrantBackend.collection"></a>**`collection`**
- <a id="geniml.search.QdrantBackend.config"></a>**`config`**
- <a id="geniml.search.QdrantBackend.qd_client"></a>**`qd_client`**

### Class methods

<a id="geniml.search.QdrantBackend.from_credentials"></a>

#### `from_credentials`

```python
@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

<a id="geniml.search.QdrantBackend.load"></a>

#### `load`

```python
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, <vector size>)
    ids: list of n point ids, or None to generate ids automatically
    payloads: optional list of n dictionaries that contain vector metadata

<a id="geniml.search.QdrantBackend.search"></a>

#### `search`

```python
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": <id>,
        "score": <score>,
        "payload": {<information of the vector>},
        "vector": [<the vector>]
    }

<a id="geniml.search.QdrantBackend.batch_search"></a>

#### `batch_search`

```python
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.

<a id="geniml.search.QdrantBackend.__len__"></a>

#### `__len__`

```python
def __len__() -> int
```

Return the number of embeddings in the backend.

Returns:
    The number of embeddings in the backend.

<a id="geniml.search.QdrantBackend.retrieve_info"></a>

#### `retrieve_info`

```python
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.

<a id="geniml.search.merge_backends"></a>

## `merge_backends`

```python
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.

<a id="geniml.search.BED2BEDSearchInterface"></a>

## `BED2BEDSearchInterface`

```python
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

- <a id="geniml.search.BED2BEDSearchInterface.query2vec"></a>**`query2vec`**
- <a id="geniml.search.BED2BEDSearchInterface.backend"></a>**`backend`**

### Methods

<a id="geniml.search.BED2BEDSearchInterface.query_search"></a>

#### `query_search`

```python
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.

<a id="geniml.search.Text2BEDSearchInterface"></a>

## `Text2BEDSearchInterface`

```python
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

- <a id="geniml.search.Text2BEDSearchInterface.query2vec"></a>**`query2vec`**
- <a id="geniml.search.Text2BEDSearchInterface.backend"></a>**`backend`**

### Methods

<a id="geniml.search.Text2BEDSearchInterface.query_search"></a>

#### `query_search`

```python
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.

<a id="geniml.search.Text2BEDSearchInterface.eval"></a>

#### `eval`

```python
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:
        {
            <query string>: [
                <store id in backend>,
                ...
            ],
            ...
        }

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

<a id="geniml.search.BED2Vec"></a>

## `BED2Vec`

```python
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

- <a id="geniml.search.BED2Vec.model"></a>**`model`**

### Methods

<a id="geniml.search.BED2Vec.forward"></a>

#### `forward`

```python
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.

<a id="geniml.search.Text2Vec"></a>

## `Text2Vec`

```python
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

- <a id="geniml.search.Text2Vec.text_embedder"></a>**`text_embedder`**
- <a id="geniml.search.Text2Vec.v2v"></a>**`v2v`**

### Methods

<a id="geniml.search.Text2Vec.forward"></a>

#### `forward`

```python
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.

<a id="geniml.search.anecdotal_search_from_hf_data"></a>

## `anecdotal_search_from_hf_data`

```python
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.

<a id="geniml.search.rand_eval"></a>

## `rand_eval`

```python
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.