# Copyright (c) RenChu Wang - All Rights Reserved
import dataclasses as dcls
from collections.abc import Callable, Iterable, Mapping, Sequence
from typing import Any
from numpy.typing import NDArray
from bocoel.corpora.embedders import Embedder
from bocoel.corpora.indices import Index
from bocoel.corpora.storages import Storage
[docs]
@dcls.dataclass(frozen=True)
class Corpus:
"""
Corpus is the entry point to handling the data in this library.
A corpus has 3 main components:
- Index: Searches one particular column in the storage.Provides fast retrival.
- Storage: Used to store the questions / answers / texts.
- Embedder: Embeds the text into vectors for faster access.
An index only corresponds to one key. If search over multiple keys is desired,
a new column or a new corpus (with shared storage) should be created.
"""
index: Index
storage: Storage
[docs]
@classmethod
def index_storage(
cls,
storage: Storage,
embedder: Embedder,
keys: Sequence[str],
index_backend: type[Index],
concat: Callable[[Iterable[Any]], str] = " [SEP] ".join,
**index_kwargs: Any,
) -> "Corpus":
"""
Creates a corpus from the given storage, embedder, key and index class,
where storage entries would be mapped to strings,
Parameters:
storage: The storage to index.
embedder: The embedder to use.
keys: The keys to use for the index.
index_backend: The index class to use.
concat: The function to use to concatenate the keys.
**index_kwargs: Additional arguments to pass to the index class.
Returns:
The created corpus.
"""
def transform(mapping: Mapping[str, Sequence[Any]]) -> Sequence[str]:
data = [mapping[k] for k in keys]
return [concat(datum) for datum in zip(*data)]
return cls.index_mapped(
storage=storage,
embedder=embedder,
transform=transform,
index_backend=index_backend,
**index_kwargs,
)
[docs]
@classmethod
def index_mapped(
cls,
storage: Storage,
embedder: Embedder,
transform: Callable[[Mapping[str, Sequence[Any]]], Sequence[str]],
index_backend: type[Index],
**index_kwargs: Any,
) -> "Corpus":
"""
Creates a corpus from the given storage, embedder, key and index class,
where storage entries would be mapped to strings,
using the specified batched transform function.
Parameters:
storage: The storage to index.
embedder: The embedder to use.
transform: The function to use to transform the storage entries.
index_backend: The index class to use.
**index_kwargs: Additional arguments to pass to the index class.
Returns:
The created corpus.
"""
embeddings = embedder.encode_storage(storage, transform=transform)
return cls.index_embeddings(
embeddings=embeddings,
storage=storage,
index_backend=index_backend,
**index_kwargs,
)
[docs]
@classmethod
def index_embeddings(
cls,
storage: Storage,
embeddings: NDArray,
index_backend: type[Index],
**index_kwargs: Any,
) -> "Corpus":
"""
Create the corpus with the given embeddings.
This can be used to save time by encoding once and caching embeddings.
Parameters:
storage: The storage to use.
embeddings: The embeddings to use.
index_backend: The index class to use.
**index_kwargs: Additional arguments to pass to the index class.
Returns:
The created corpus.
"""
index = index_backend(embeddings, **index_kwargs)
return cls(index=index, storage=storage)
