# Hoplite

Hoplite is a database for working with large numbers of embeddings.

We develop it with bioacoustic audio applications in mind, where we have a large number of audio files and (usually) a small number of human annotations on those files. Each audio file is typically split up into (possibly overlapping) audio windows, and we have an embedding for each audio window.

Structurally, however, the database is not limited to audio, and may be useful for other domains (image embeddings, satellite image embeddings, etc).

The database handles the embeddings, source metadata for each embedding, labels (annotations) for individual embeddings, and a graph structure (edges) on the embeddings (used for fast approximate nearest neighbor search).

## Backends

The Hoplite database is accessed via the API in `interface.py`. It is possible to reimplement the API on different backends. As of this writing, two backends are implemented: an SQLite+usearch database and an in-memory implementation using a collection of Numpy arrays and dicts.

The `db_loader.py` file contains some tooling for copying a Hoplite databases, including into a new backend. (eg, persisting an in-memory database to SQLite+usearch.)

## Database structure

The tables within the database are as follows:

The **embeddings** table contains the embedding vectors, produced by some pretrained model. Each embedding has a unique ID, as well as a *source* and *offset* within the source. The offset is an arbitrary vector, indicating what region of the source the embedding refers to. For example, when working with bioacoustic audio and embeddings of fixed-length audio windows, the offset can be simply the time offset within the source audio file for the embedded audio. For image applications, the offset could be empty (indicating the whole image), or give the corners of a bounding box.

The **source** table holds information for joining embeddings with the source data.  Each *source* consists of a unique ID and a string identifier, for example a relative file path under some root directory, and a *dataset name*, which allows defining groups of related sources.

The **labels** table contains any labels or annotations on the embeddings. Each label consists of an *embedding id*, the *label* string, a *type* (eg, positive or negative), and a *provenance* string (eg, annotator's name, or id for the model which produced a pseudo-label). You can have many labels per embedding, for example indicating different entities represented in the audio, or different annotator's decisions for the same label+embedding pair.

Finally, the **metadata** table is a string-to-json key-value store for handling arbitrary additional metadata. This is used, for example, to indicate which embedding model was used to produce the embeddings in the database.

## Usage Examples

For up-to-date usage, see `tests/hoplite_test.py`. We give a brief overview of functionality here, however.

### Creating a Database

The in-memory database uses fixed-size numpy arrays for storing embeddings and sparse lists of graph edges. The embedding dimension should be known from the embedding model. The `max_size` is the maximum number of embeddings the database will contain. The `degree_bound` determines the maximum number of outgoing edges per embedding.

```
from perch_hoplite.db import in_mem_impl

db = in_mem_impl.InMemoryGraphSearchDB.create(
    embedding_dim=1280,
    max_size=2000,
    degree_bound=256,
)
db.setup()
```

The SQLite+usearch backend has no particular limit on the number of embeddings or number of edges per vertex. You can create an SQLite backend like so:

```
from perch_hoplite.db import sqlite_usearch_impl

db = sqlite_usearch_impl.SQLiteUsearchDB.create(
    db_path=db_file_path,
    usearch_cfg=sqlite_usearch_impl.get_default_usearch_config(
      embedding_dim=1280),
)
db.setup()
```

### Working with Embeddings

Adding embeddings to the database is simple. Describe the embedding source, and insert it using the API. The insertion call will return the unique ID for this embedding.

```
from perch_hoplite.db import interface

dataset_name = 'my_data'
source_id = 'some_file.wav'
offsets = np.array([100.0])
embedding = np.random.normal(size=[1280])

source = interface.EmbeddingSource(dataset_name, source_id, offsets)
uid = db.insert_embedding(embedding, source)
```

We can retrieve collections of embeddings with an array of UIDs. Note that backends with threading (eg, SQLite) may shuffle the order of the embeddings returned, so `get_embeddings` returns both the possibly-shuffled UIDs and associated embeddings in the matching order.
```
uids = np.array([uid1, uid2])
uids, embeddings = db.get_embeddings(uids)
```

### Working with Labels

Now to add a couple labels to the embedding:

```
label = interface.Label(
  embedding_id=uid,
  label='birb',
  type=interface.LabelType.POSITIVE,
  provenance='me')
db.insert_label(label)

other_label = interface.Label(
  embedding_id=uid,
  label='birb',
  type=interface.LabelType.NEGATIVE,
  provenance='you')
db.insert_label(other_label)
```

We can now get all the labels assigned to the embedding:
```
labels = db.get_labels(uid)
```

To retrieve embeddings with particular labels, we use `get_embeddings_by_label`.
By default, we get all POSITIVE embeddings annotated with the target label string, regardless of provenance.

```
got = db.get_embeddings_by_label('birb')
```

We can add `type` and `provenance` arguments to select embeddings with a particular label type or provenance, however. Specifying `None` for type or provenance will give unconstrained results.

```
got = db.get_embeddings_by_label('birb', provenance='me')
got = db.get_embeddings_by_label('birb', type=interface.LabelType.POSITIVE)

# Returns all embeddings with POSITIVE or NEGATIVE labels for 'birb'.
got = db.get_embeddings_by_label('birb', type=None)
```

### Working with Sources

To get all embeddings with a particular `dataset_name`:
```
got = db.get_embeddings_by_source(dataset_name='my_data')
```

To get all embeddings form a particular source:
```
got = db.get_embeddings_by_source(
  dataset_name='my_data', source_id='some_file.wav')
```

And to find out the source for some particular embedding:
```
source = db.get_embedding_source(uid)
```

### Utility functions

Some core info can be accessed through the API:

```
embedding_count = db.count_embeddings()

# Number of distinct labels in the database.
classes_count = db.count_classes()
```

We can get a random (in the [xkcd](https://xkcd.com/221/) sense) embedding ID from the database:
```
uid = db.get_one_embedding_id()
```

## Vector Similarity Search

We provide utilities for both brute force and indexed vector search.

Vector search requires a 'scoring function.' In this library, higher scores must indicate greater similarity. So, good scoring functions include cosine similarity, inner product, or *negative* Euclidean distance (nearest neighbor).

For bioacoustic embeddings, we recommend *inner product*, as it handles mixture embeddings for overlapping calls more gracefully than Euclidean distance or cosine similarity. However, the user is invited to experiment!

### Brute-force search

Brute force search capabilities are handled by `brutalism.py`. Here's some example usage:

```
from perch_hoplite.db import brutalism

query = np.random.normal(size=[1280])
score_fn = np.dot
results, scores = brutalism.threaded_brute_search(db, query, score_fn)

# The `results` are a search_results.TopKSearchResults object.
# This wraps a list of search results kept in heap order, but iterating over the
# TopKSearchResults object will yield results in descending order.
for r in results:
  print(r.embedding_id, r.sort_score)
```

Brute force search is appropriate in the following situations:

* Perfect recall is needed.
* The dataset contains less than about a million embeddings.
* A flexible scoring function is needed (eg, for margin sampling).

We can also use sampled search to apply brute force search on a random subset of data:

```
# Find elements whose dot product with q is approximately 2.
weird_score_fn = lambda q, t: -np.abs(2.0 - np.dot(q, t))
results, scores = brutalism.threaded_brute_search(
  db, query, weird_score_fn,
  sample_size=0.1)  # Search a random 10% of the embeddings.

results, scores = brutalism.threaded_brute_search(
  db, query, weird_score_fn,
  sample_size=100_000)  # Search up to 100k randomly selected embeddings.
```

In short, you should only move to indexed search when you're really annoyed with how slow brute force search is, don't mind missing some results, and have a specific scoring function to work with.

### Indexed Vector Search

*Under Construction*

For very large datasets, approximate nearest neighbor search can greatly speed up analysis. For this, we define a graph on the embeddings and explore using a greedy nearest-neighbor search. The algorithm is an adaptation of the Vamana algorithm, detailed in the [DiskANN paper](https://proceedings.neurips.cc/paper_files/paper/2019/file/09853c7fb1d3f8ee67a61b6bf4a7f8e6-Paper.pdf).

The `SQLiteUsearchDB` stores all vectors in a usearch vector database, allowing very fast approximate nearest neighbor search.
