Module base

BaseRetriever Objects

class BaseRetriever(BaseComponent)

retrieve

| @abstractmethod
 | retrieve(query: str, filters: dict = None, top_k: Optional[int] = None, index: str = None) -> List[Document]

Scan through documents in DocumentStore and return a small number documents that are most relevant to the query.

Arguments:

query: The query
filters: A dictionary where the keys specify a metadata field and the value is a list of accepted values for that field
top_k: How many documents to return per query.
index: The name of the index in the DocumentStore from which to retrieve documents

timing

| timing(fn, attr_name)

Wrapper method used to time functions.

eval

| eval(label_index: str = "label", doc_index: str = "eval_document", label_origin: str = "gold_label", top_k: int = 10, open_domain: bool = False, return_preds: bool = False) -> dict

Performs evaluation on the Retriever. Retriever is evaluated based on whether it finds the correct document given the query string and at which position in the ranking of documents the correct document is.

| Returns a dict containing the following metrics:

- &quot;recall&quot;: Proportion of questions for which correct document is among retrieved documents
- &quot;mrr&quot;: Mean of reciprocal rank. Rewards retrievers that give relevant documents a higher rank.
  Only considers the highest ranked relevant document.
- &quot;map&quot;: Mean of average precision for each question. Rewards retrievers that give relevant
  documents a higher rank. Considers all retrieved relevant documents. If ``open_domain=True``,
  average precision is normalized by the number of retrieved relevant documents per query.
  If ``open_domain=False``, average precision is normalized by the number of all relevant documents
  per query.

Arguments:

label_index: Index/Table in DocumentStore where labeled questions are stored
doc_index: Index/Table in DocumentStore where documents that are used for evaluation are stored
top_k: How many documents to return per query
open_domain: If True, retrieval will be evaluated by checking if the answer string to a question is contained in the retrieved docs (common approach in open-domain QA). If False, retrieval uses a stricter evaluation that checks if the retrieved document ids are within ids explicitly stated in the labels.
return_preds: Whether to add predictions in the returned dictionary. If True, the returned dictionary contains the keys "predictions" and "metrics".

Module sparse

ElasticsearchRetriever Objects

class ElasticsearchRetriever(BaseRetriever)

init

| __init__(document_store: ElasticsearchDocumentStore, top_k: int = 10, custom_query: str = None)

Arguments:

document_store: an instance of a DocumentStore to retrieve documents from.

custom_query: query string as per Elasticsearch DSL with a mandatory query placeholder(query).

               Optionally, ES `filter` clause can be added where the values of `terms` are placeholders
               that get substituted during runtime. The placeholder(${filter_name_1}, ${filter_name_2}..)
               names must match with the filters dict supplied in self.retrieve().
               ::

                   **An example custom_query:**
                   ```python
                  |    {
                  |        &quot;size&quot;: 10,
                  |        &quot;query&quot;: {
                  |            &quot;bool&quot;: {
                  |                &quot;should&quot;: [{&quot;multi_match&quot;: {
                  |                    &quot;query&quot;: ${query},                 // mandatory query placeholder
                  |                    &quot;type&quot;: &quot;most_fields&quot;,
                  |                    &quot;fields&quot;: [&quot;text&quot;, &quot;title&quot;]}}],
                  |                &quot;filter&quot;: [                                 // optional custom filters
                  |                    {&quot;terms&quot;: {&quot;year&quot;: ${years}}},
                  |                    {&quot;terms&quot;: {&quot;quarter&quot;: ${quarters}}},
                  |                    {&quot;range&quot;: {&quot;date&quot;: {&quot;gte&quot;: ${date}}}}
                  |                    ],
                  |            }
                  |        },
                  |    }
                   ```

               **For this custom_query, a sample retrieve() could be:**
               ```python
              |    self.retrieve(query=&quot;Why did the revenue increase?&quot;,
              |                  filters={&quot;years&quot;: [&quot;2019&quot;], &quot;quarters&quot;: [&quot;Q1&quot;, &quot;Q2&quot;]})
              ```

top_k: How many documents to return per query.

retrieve

| retrieve(query: str, filters: dict = None, top_k: Optional[int] = None, index: str = None) -> List[Document]

Scan through documents in DocumentStore and return a small number documents that are most relevant to the query.

Arguments:

query: The query
filters: A dictionary where the keys specify a metadata field and the value is a list of accepted values for that field
top_k: How many documents to return per query.
index: The name of the index in the DocumentStore from which to retrieve documents

ElasticsearchFilterOnlyRetriever Objects

class ElasticsearchFilterOnlyRetriever(ElasticsearchRetriever)

Naive "Retriever" that returns all documents that match the given filters. No impact of query at all. Helpful for benchmarking, testing and if you want to do QA on small documents without an "active" retriever.

retrieve

| retrieve(query: str, filters: dict = None, top_k: Optional[int] = None, index: str = None) -> List[Document]

Scan through documents in DocumentStore and return a small number documents that are most relevant to the query.

Arguments:

query: The query
filters: A dictionary where the keys specify a metadata field and the value is a list of accepted values for that field
top_k: How many documents to return per query.
index: The name of the index in the DocumentStore from which to retrieve documents

TfidfRetriever Objects

class TfidfRetriever(BaseRetriever)

Read all documents from a SQL backend.

Split documents into smaller units (eg, paragraphs or pages) to reduce the computations when text is passed on to a Reader for QA.

It uses sklearn's TfidfVectorizer to compute a tf-idf matrix.

init

| __init__(document_store: BaseDocumentStore, top_k: int = 10)

Arguments:

document_store: an instance of a DocumentStore to retrieve documents from.
top_k: How many documents to return per query.

retrieve

| retrieve(query: str, filters: dict = None, top_k: Optional[int] = None, index: str = None) -> List[Document]

Scan through documents in DocumentStore and return a small number documents that are most relevant to the query.

Arguments:

query: The query
filters: A dictionary where the keys specify a metadata field and the value is a list of accepted values for that field
top_k: How many documents to return per query.
index: The name of the index in the DocumentStore from which to retrieve documents

fit

| fit()

Performing training on this class according to the TF-IDF algorithm.

Module dense

DensePassageRetriever Objects

class DensePassageRetriever(BaseRetriever)

Retriever that uses a bi-encoder (one transformer for query, one transformer for passage). See the original paper for more details: Karpukhin, Vladimir, et al. (2020): "Dense Passage Retrieval for Open-Domain Question Answering." (https://arxiv.org/abs/2004.04906).

init

| __init__(document_store: BaseDocumentStore, query_embedding_model: Union[Path, str] = "facebook/dpr-question_encoder-single-nq-base", passage_embedding_model: Union[Path, str] = "facebook/dpr-ctx_encoder-single-nq-base", model_version: Optional[str] = None, max_seq_len_query: int = 64, max_seq_len_passage: int = 256, top_k: int = 10, use_gpu: bool = True, batch_size: int = 16, embed_title: bool = True, use_fast_tokenizers: bool = True, infer_tokenizer_classes: bool = False, similarity_function: str = "dot_product", global_loss_buffer_size: int = 150000, progress_bar: bool = True, devices: Optional[List[Union[int, str, torch.device]]] = None)

Init the Retriever incl. the two encoder models from a local or remote model checkpoint. The checkpoint format matches huggingface transformers' model format

Example:

    ```python
    |    # remote model from FAIR
    |    DensePassageRetriever(document_store=your_doc_store,
    |                          query_embedding_model=&quot;facebook/dpr-question_encoder-single-nq-base&quot;,
    |                          passage_embedding_model=&quot;facebook/dpr-ctx_encoder-single-nq-base&quot;)
    |    # or from local path
    |    DensePassageRetriever(document_store=your_doc_store,
    |                          query_embedding_model=&quot;model_directory/question-encoder&quot;,
    |                          passage_embedding_model=&quot;model_directory/context-encoder&quot;)
    ```

Arguments:

document_store: An instance of DocumentStore from which to retrieve documents.
query_embedding_model: Local path or remote name of question encoder checkpoint. The format equals the one used by hugging-face transformers' modelhub models Currently available remote names: "facebook/dpr-question_encoder-single-nq-base"
passage_embedding_model: Local path or remote name of passage encoder checkpoint. The format equals the one used by hugging-face transformers' modelhub models Currently available remote names: "facebook/dpr-ctx_encoder-single-nq-base"
model_version: The version of model to use from the HuggingFace model hub. Can be tag name, branch name, or commit hash.
max_seq_len_query: Longest length of each query sequence. Maximum number of tokens for the query text. Longer ones will be cut down."
max_seq_len_passage: Longest length of each passage/context sequence. Maximum number of tokens for the passage text. Longer ones will be cut down."
top_k: How many documents to return per query.
use_gpu: Whether to use all available GPUs or the CPU. Falls back on CPU if no GPU is available.
batch_size: Number of questions or passages to encode at once. In case of multiple gpus, this will be the total batch size.
embed_title: Whether to concatenate title and passage to a text pair that is then used to create the embedding. This is the approach used in the original paper and is likely to improve performance if your titles contain meaningful information for retrieval (topic, entities etc.) . The title is expected to be present in doc.meta["name"] and can be supplied in the documents before writing them to the DocumentStore like this: {"text": "my text", "meta": {"name": "my title"}}.
use_fast_tokenizers: Whether to use fast Rust tokenizers
infer_tokenizer_classes: Whether to infer tokenizer class from the model config / name. If False, the class always loads DPRQuestionEncoderTokenizer and DPRContextEncoderTokenizer.
similarity_function: Which function to apply for calculating the similarity of query and passage embeddings during training. Options: dot_product (Default) or cosine
global_loss_buffer_size: Buffer size for all_gather() in DDP. Increase if errors like "encoded data exceeds max_size ..." come up
progress_bar: Whether to show a tqdm progress bar or not. Can be helpful to disable in production deployments to keep the logs clean.
devices: List of GPU devices to limit inference to certain GPUs and not use all available ones (e.g. ["cuda:0"]). As multi-GPU training is currently not implemented for DPR, training will only use the first device provided in this list.

retrieve

| retrieve(query: str, filters: dict = None, top_k: Optional[int] = None, index: str = None) -> List[Document]

Scan through documents in DocumentStore and return a small number documents that are most relevant to the query.

Arguments:

query: The query
filters: A dictionary where the keys specify a metadata field and the value is a list of accepted values for that field
top_k: How many documents to return per query.
index: The name of the index in the DocumentStore from which to retrieve documents

embed_queries

| embed_queries(texts: List[str]) -> List[np.ndarray]

Create embeddings for a list of queries using the query encoder

Arguments:

texts: Queries to embed

Returns:

Embeddings, one per input queries

embed_passages

| embed_passages(docs: List[Document]) -> List[np.ndarray]

Create embeddings for a list of passages using the passage encoder

Arguments:

docs: List of Document objects used to represent documents / passages in a standardized way within Haystack.

Returns:

Embeddings of documents / passages shape (batch_size, embedding_dim)

train

| train(data_dir: str, train_filename: str, dev_filename: str = None, test_filename: str = None, max_samples: int = None, max_processes: int = 128, dev_split: float = 0, batch_size: int = 2, embed_title: bool = True, num_hard_negatives: int = 1, num_positives: int = 1, n_epochs: int = 3, evaluate_every: int = 1000, n_gpu: int = 1, learning_rate: float = 1e-5, epsilon: float = 1e-08, weight_decay: float = 0.0, num_warmup_steps: int = 100, grad_acc_steps: int = 1, use_amp: str = None, optimizer_name: str = "TransformersAdamW", optimizer_correct_bias: bool = True, save_dir: str = "../saved_models/dpr", query_encoder_save_dir: str = "query_encoder", passage_encoder_save_dir: str = "passage_encoder")

train a DensePassageRetrieval model

Arguments:

data_dir: Directory where training file, dev file and test file are present
train_filename: training filename
dev_filename: development set filename, file to be used by model in eval step of training
test_filename: test set filename, file to be used by model in test step after training
max_samples: maximum number of input samples to convert. Can be used for debugging a smaller dataset.
max_processes: the maximum number of processes to spawn in the multiprocessing.Pool used in DataSilo. It can be set to 1 to disable the use of multiprocessing or make debugging easier.
dev_split: The proportion of the train set that will sliced. Only works if dev_filename is set to None
batch_size: total number of samples in 1 batch of data
embed_title: whether to concatenate passage title with each passage. The default setting in official DPR embeds passage title with the corresponding passage
num_hard_negatives: number of hard negative passages(passages which are very similar(high score by BM25) to query but do not contain the answer
num_positives: number of positive passages
n_epochs: number of epochs to train the model on
evaluate_every: number of training steps after evaluation is run
n_gpu: number of gpus to train on
learning_rate: learning rate of optimizer
epsilon: epsilon parameter of optimizer
weight_decay: weight decay parameter of optimizer
grad_acc_steps: number of steps to accumulate gradient over before back-propagation is done
use_amp: Whether to use automatic mixed precision (AMP) or not. The options are: "O0" (FP32) "O1" (Mixed Precision) "O2" (Almost FP16) "O3" (Pure FP16). For more information, refer to: https://nvidia.github.io/apex/amp.html
optimizer_name: what optimizer to use (default: TransformersAdamW)
num_warmup_steps: number of warmup steps
optimizer_correct_bias: Whether to correct bias in optimizer
save_dir: directory where models are saved
query_encoder_save_dir: directory inside save_dir where query_encoder model files are saved
passage_encoder_save_dir: directory inside save_dir where passage_encoder model files are saved

save

| save(save_dir: Union[Path, str], query_encoder_dir: str = "query_encoder", passage_encoder_dir: str = "passage_encoder")

Save DensePassageRetriever to the specified directory.

Arguments:

save_dir: Directory to save to.
query_encoder_dir: Directory in save_dir that contains query encoder model.
passage_encoder_dir: Directory in save_dir that contains passage encoder model.

Returns:

None

load

| @classmethod
 | load(cls, load_dir: Union[Path, str], document_store: BaseDocumentStore, max_seq_len_query: int = 64, max_seq_len_passage: int = 256, use_gpu: bool = True, batch_size: int = 16, embed_title: bool = True, use_fast_tokenizers: bool = True, similarity_function: str = "dot_product", query_encoder_dir: str = "query_encoder", passage_encoder_dir: str = "passage_encoder", infer_tokenizer_classes: bool = False)

Load DensePassageRetriever from the specified directory.

EmbeddingRetriever Objects

class EmbeddingRetriever(BaseRetriever)

init

| __init__(document_store: BaseDocumentStore, embedding_model: str, model_version: Optional[str] = None, use_gpu: bool = True, model_format: str = "farm", pooling_strategy: str = "reduce_mean", emb_extraction_layer: int = -1, top_k: int = 10, progress_bar: bool = True)

Arguments:

document_store: An instance of DocumentStore from which to retrieve documents.
embedding_model: Local path or name of model in Hugging Face's model hub such as 'deepset/sentence_bert'
model_version: The version of model to use from the HuggingFace model hub. Can be tag name, branch name, or commit hash.
use_gpu: Whether to use gpu or not

model_format: Name of framework that was used for saving the model. Options:

               - ``&#39;farm&#39;``
               - ``&#39;transformers&#39;``
               - ``&#39;sentence_transformers&#39;``

pooling_strategy: Strategy for combining the embeddings from the model (for farm / transformers models only). Options:

                   - ``&#39;cls_token&#39;`` (sentence vector)
                   - ``&#39;reduce_mean&#39;`` (sentence vector)
                   - ``&#39;reduce_max&#39;`` (sentence vector)
                   - ``&#39;per_token&#39;`` (individual token vectors)

emb_extraction_layer: Number of layer from which the embeddings shall be extracted (for farm / transformers models only). Default: -1 (very last layer).
top_k: How many documents to return per query.
progress_bar: If true displays progress bar during embedding.

retrieve

| retrieve(query: str, filters: dict = None, top_k: Optional[int] = None, index: str = None) -> List[Document]

Scan through documents in DocumentStore and return a small number documents that are most relevant to the query.

Arguments:

query: The query
filters: A dictionary where the keys specify a metadata field and the value is a list of accepted values for that field
top_k: How many documents to return per query.
index: The name of the index in the DocumentStore from which to retrieve documents

embed_queries

| embed_queries(texts: List[str]) -> List[np.ndarray]

Create embeddings for a list of queries.

Arguments:

texts: Queries to embed

Returns:

Embeddings, one per input queries

embed_passages

| embed_passages(docs: List[Document]) -> List[np.ndarray]

Create embeddings for a list of passages.

Arguments:

docs: List of documents to embed

Returns:

Embeddings, one per input passage

_EmbeddingEncoder Objects

class _EmbeddingEncoder()

embed_queries

| @abstractmethod
 | embed_queries(texts: List[str]) -> List[np.ndarray]

Create embeddings for a list of queries.

Arguments:

texts: Queries to embed

Returns:

Embeddings, one per input queries

embed_passages

| @abstractmethod
 | embed_passages(docs: List[Document]) -> List[np.ndarray]

Create embeddings for a list of passages.

Arguments:

docs: List of documents to embed

Returns:

Embeddings, one per input passage

Stars

2464

Edit on GitHub

BaseRetriever Objects
ElasticsearchRetriever Objects
ElasticsearchFilterOnlyRetriever Objects
TfidfRetriever Objects
DensePassageRetriever Objects
EmbeddingRetriever Objects
_EmbeddingEncoder Objects