Repository: deepset-ai/haystack

Issue:
`ExtractiveReader` (v2)
`ExtractiveReader` is a port of the former `Reader`s. If it's the only Reader we end up porting to v2, we may be able to stick with a generic name that clarifies it's an extractive component, not a generative one.

Depending on what we decide in the discussion, we may have another Reader, which means we have to find a better name, or we may add features to this component, which may also mean finding another name.

Draft API:

```python
@component
class ExtractiveReader:

    def __init__(self, model: str, ... params ...):
        # stores the params and the model name - no model initialization here!!

    def warm_up(self):
        # init the model

    @component.output_types(answers=List[List[Answer]])
    def run(self, queries: List[str], documents: List[List[Document]], top_k: Optional[int] = None, ...):
        # runs the model
        return {"answers": answers}
```

This issue is about creating a simple Reader component that leverages Transformers pipelines to the maximum, in an attempt to keep this component as simple as possible. More sophisticated features will be addressed in later issues.


Discussion:

Resolving patch:
diff --git a/haystack/preview/components/readers/__init__.py b/haystack/preview/components/readers/__init__.py
new file mode 100644
--- /dev/null
+++ b/haystack/preview/components/readers/__init__.py
@@ -0,0 +1 @@
+from haystack.preview.components.readers.extractive import ExtractiveReader
diff --git a/haystack/preview/components/readers/extractive.py b/haystack/preview/components/readers/extractive.py
new file mode 100644
--- /dev/null
+++ b/haystack/preview/components/readers/extractive.py
@@ -0,0 +1,322 @@
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Union
+import math
+import warnings
+
+from haystack.preview import component, default_from_dict, default_to_dict, ComponentError, Document, ExtractedAnswer
+from haystack.preview.lazy_imports import LazyImport
+
+with LazyImport(message="Run 'pip install farm-haystack[inference]'") as torch_and_transformers_import:
+    from transformers import AutoModelForQuestionAnswering, AutoTokenizer
+    from tokenizers import Encoding
+    import torch
+
+
+@component
+class ExtractiveReader:
+    """
+    A component for performing extractive QA.
+    Every possible answer span is assigned a confidence score independent of other answer spans. This fixes a common issue of other implementations which make comparisons across documents harder by normalising each document's answers independently.
+    """
+
+    def __init__(
+        self,
+        model_name_or_path: Union[Path, str] = "deepset/roberta-base-squad2-distilled",
+        device: Optional[str] = None,
+        top_k: int = 20,
+        confidence_threshold: Optional[float] = None,
+        max_seq_length: int = 384,
+        stride: int = 128,
+        max_batch_size: Optional[int] = None,
+        answers_per_seq: Optional[int] = None,
+        no_answer: bool = True,
+        calibration_factor: float = 0.1,
+    ) -> None:
+        """
+        Creates an ExtractiveReader
+        :param model: A HuggingFace transformers question answering model.
+            Can either be a path to a folder containing the model files or an identifier for the HF hub
+            Default: `'deepset/roberta-base-squad2-distilled'`
+        :param device: Pytorch device string. Uses GPU by default if available
+        :param top_k: Number of answers to return per query.
+            It is required even if confidence_threshold is set. Defaults to 20.
+        :param confidence_threshold: Answers with a confidence score below this value will not be returned
+        :param max_seq_length: Maximum number of tokens.
+            If exceeded by a sequence, the sequence will be split.
+            Default: 384
+        :param stride: Number of tokens that overlap when sequence is split because it exceeds max_seq_length
+            Default: 128
+        :param max_batch_size: Maximum number of samples that are fed through the model at the same time
+        :param answers_per_seq: Number of answer candidates to consider per sequence.
+            This is relevant when a document has been split into multiple sequence due to max_seq_length.
+        :param no_answer: Whether to return no answer scores
+        :param calibration_factor: Factor used for calibrating confidence scores
+        """
+        torch_and_transformers_import.check()
+        self.model_name_or_path = str(model_name_or_path)
+        self.model = None
+        self.device = device
+        self.max_seq_length = max_seq_length
+        self.top_k = top_k
+        self.confidence_threshold = confidence_threshold
+        self.stride = stride
+        self.max_batch_size = max_batch_size
+        self.answers_per_seq = answers_per_seq
+        self.no_answer = no_answer
+        self.calibration_factor = calibration_factor
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Serialize this component to a dictionary.
+        """
+        return default_to_dict(
+            self,
+            model_name_or_path=self.model_name_or_path,
+            device=self.device,
+            max_seq_length=self.max_seq_length,
+            top_k=self.top_k,
+            confidence_threshold=self.confidence_threshold,
+            stride=self.stride,
+            max_batch_size=self.max_batch_size,
+            answers_per_seq=self.answers_per_seq,
+            no_answer=self.no_answer,
+            calibration_factor=self.calibration_factor,
+        )
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ExtractiveReader":
+        """
+        Deserialize this component from a dictionary.
+        """
+        return default_from_dict(cls, data)
+
+    def warm_up(self):
+        if self.model is None:
+            if torch.cuda.is_available():
+                self.device = self.device or "cuda:0"
+            else:
+                self.device = self.device or "cpu:0"
+            self.model = AutoModelForQuestionAnswering.from_pretrained(self.model_name_or_path).to(self.device)
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path)
+
+    def _flatten_documents(
+        self, queries: List[str], documents: List[List[Document]]
+    ) -> Tuple[List[str], List[Document], List[int]]:
+        """
+        Flattens queries and documents so all query-document pairs are arranged along one batch axis.
+        """
+        flattened_queries = [query for documents_, query in zip(documents, queries) for _ in documents_]
+        flattened_documents = [document for documents_ in documents for document in documents_]
+        query_ids = [i for i, documents_ in enumerate(documents) for _ in documents_]
+        return flattened_queries, flattened_documents, query_ids
+
+    def _preprocess(
+        self, queries: List[str], documents: List[Document], max_seq_length: int, query_ids: List[int], stride: int
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[Encoding], List[int], List[int]]:
+        """
+        Split and tokenise documents and preserve structures by returning mappings to query and document ids.
+        """
+        texts = []
+        document_ids = []
+        for i, doc in enumerate(documents):
+            if doc.text is None:
+                warnings.warn(
+                    f"Document with id {doc.id} was passed to ExtractiveReader, but does not contain any text. It will be ignored."
+                )
+                continue
+            texts.append(doc.text)
+            document_ids.append(i)
+        encodings_pt = self.tokenizer(
+            queries,
+            [document.text for document in documents],
+            padding=True,
+            truncation=True,
+            max_length=max_seq_length,
+            return_tensors="pt",
+            return_overflowing_tokens=True,
+            stride=stride,
+        )
+
+        input_ids = encodings_pt.input_ids.to(self.device)
+        attention_mask = encodings_pt.attention_mask.to(self.device)
+
+        query_ids = [query_ids[index] for index in encodings_pt.overflow_to_sample_mapping]
+        document_ids = [document_ids[sample_id] for sample_id in encodings_pt.overflow_to_sample_mapping]
+
+        encodings = encodings_pt.encodings
+        sequence_ids = torch.tensor(
+            [[id_ if id_ is not None else -1 for id_ in encoding.sequence_ids] for encoding in encodings]
+        ).to(self.device)
+
+        return input_ids, attention_mask, sequence_ids, encodings, query_ids, document_ids
+
+    def _postprocess(
+        self,
+        start: torch.Tensor,
+        end: torch.Tensor,
+        sequence_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        answers_per_seq: int,
+        encodings: List[Encoding],
+    ) -> Tuple[List[List[int]], List[List[int]], torch.Tensor]:
+        """
+        Turn start and end logits into confidence scores for each answer span. Unlike most other implementations, there is no normalisation in each split to make the scores more comparable across different splits. The top k answer spans are returned.
+        """
+        mask = sequence_ids == 1
+        mask = torch.logical_and(mask, attention_mask == 1)
+        start = torch.where(mask, start, -torch.inf)
+        end = torch.where(mask, end, -torch.inf)
+        start = start.unsqueeze(-1)
+        end = end.unsqueeze(-2)
+
+        logits = start + end  # shape: (batch_size, seq_length (start), seq_length (end))
+        mask = torch.ones(logits.shape[-2:], dtype=torch.bool, device=self.device)
+        mask = torch.triu(mask)  # End shouldn't be before start
+        masked_logits = torch.where(mask, logits, -torch.inf)
+        probabilities = torch.sigmoid(masked_logits * self.calibration_factor)
+
+        flat_probabilities = probabilities.flatten(-2, -1)  # necessary for topk
+        candidates = torch.topk(flat_probabilities, answers_per_seq)
+        seq_length = logits.shape[-1]
+        start_candidates = candidates.indices // seq_length  # Recover indices from flattening
+        end_candidates = candidates.indices % seq_length
+        start_candidates = start_candidates.cpu()
+        end_candidates = end_candidates.cpu()
+
+        start_candidates = [
+            [encoding.token_to_chars(start)[0] for start in candidates]
+            for candidates, encoding in zip(start_candidates, encodings)
+        ]
+        end_candidates = [
+            [encoding.token_to_chars(end)[1] for end in candidates]
+            for candidates, encoding in zip(end_candidates, encodings)
+        ]
+        probabilities = candidates.values.cpu()
+
+        return start_candidates, end_candidates, probabilities
+
+    def _nest_answers(
+        self,
+        start: List[List[int]],
+        end: List[List[int]],
+        probabilities: torch.Tensor,
+        flattened_documents: List[Document],
+        queries: List[str],
+        answers_per_seq: int,
+        top_k: Optional[int],
+        confidence_threshold: Optional[float],
+        query_ids: List[int],
+        document_ids: List[int],
+        no_answer: bool,
+    ) -> List[List[ExtractedAnswer]]:
+        """
+        Reconstructs the nested structure that existed before flattening. Also computes a no answer probability. This probability is different from most other implementations because it does not consider the no answer logit introduced with SQuAD 2. Instead, it just computes the probability that the answer does not exist in the top k or top p.
+        """
+        flat_answers_without_queries = []
+        for document_id, start_candidates_, end_candidates_, probabilities_ in zip(
+            document_ids, start, end, probabilities
+        ):
+            for start_, end_, probability in zip(start_candidates_, end_candidates_, probabilities_):
+                doc = flattened_documents[document_id]
+                flat_answers_without_queries.append({"data": doc.text[start_:end_], "document": doc, "probability": probability.item(), "start": start_, "end": end_, "metadata": {}})  # type: ignore # doc.text cannot be None, because those documents are filtered when preprocessing. However, mypy doesn't know that.
+        i = 0
+        nested_answers = []
+        for query_id in range(query_ids[-1] + 1):
+            current_answers = []
+            while i < len(flat_answers_without_queries) and query_ids[i // answers_per_seq] == query_id:
+                answer = flat_answers_without_queries[i]
+                answer["query"] = queries[query_id]
+                current_answers.append(ExtractedAnswer(**answer))
+                i += 1
+            current_answers = sorted(current_answers, key=lambda answer: answer.probability, reverse=True)
+            current_answers = current_answers[:top_k]
+            if no_answer:
+                no_answer_probability = math.prod(1 - answer.probability for answer in current_answers)
+                answer_ = ExtractedAnswer(
+                    data=None, query=queries[query_id], metadata={}, document=None, probability=no_answer_probability
+                )
+                current_answers.append(answer_)
+            current_answers = sorted(current_answers, key=lambda answer: answer.probability, reverse=True)
+            if confidence_threshold is not None:
+                current_answers = [answer for answer in current_answers if answer.probability >= confidence_threshold]
+            nested_answers.append(current_answers)
+
+        return nested_answers
+
+    @component.output_types(answers=List[ExtractedAnswer])
+    def run(
+        self,
+        query: str,
+        document: List[Document],
+        top_k: Optional[int] = None,
+        confidence_threshold: Optional[float] = None,
+        max_seq_length: Optional[int] = None,
+        stride: Optional[int] = None,
+        max_batch_size: Optional[int] = None,
+        answers_per_seq: Optional[int] = None,
+        no_answer: Optional[bool] = None,
+    ):
+        """
+        Performs extractive QA on the given documents using the given query.
+        """
+        queries = [query]  # Temporary solution until we have decided what batching should look like in v2
+        documents = [document]
+        if self.model is None:
+            raise ComponentError("The component was not warmed up. Run 'warm_up()' before calling 'run()'.")
+
+        top_k = top_k or self.top_k
+        confidence_threshold = confidence_threshold or self.confidence_threshold
+        max_seq_length = max_seq_length or self.max_seq_length
+        stride = stride or self.stride
+        max_batch_size = max_batch_size or self.max_batch_size
+        answers_per_seq = answers_per_seq or self.answers_per_seq or top_k or 20
+        no_answer = no_answer if no_answer is not None else self.no_answer
+
+        flattened_queries, flattened_documents, query_ids = self._flatten_documents(queries, documents)
+        input_ids, attention_mask, sequence_ids, encodings, query_ids, document_ids = self._preprocess(
+            flattened_queries, flattened_documents, max_seq_length, query_ids, stride
+        )
+
+        num_batches = math.ceil(input_ids.shape[0] / max_batch_size) if max_batch_size else 1
+        batch_size = max_batch_size or input_ids.shape[0]
+
+        start_logits_list = []
+        end_logits_list = []
+
+        for i in range(num_batches):
+            start_index = i * batch_size
+            end_index = start_index + batch_size
+            cur_input_ids = input_ids[start_index:end_index]
+            cur_attention_mask = attention_mask[start_index:end_index]
+
+            output = self.model(input_ids=cur_input_ids, attention_mask=cur_attention_mask)
+            cur_start_logits = output.start_logits
+            cur_end_logits = output.end_logits
+            if num_batches != 1:
+                cur_start_logits = cur_start_logits.cpu()
+                cur_end_logits = cur_end_logits.cpu()
+            start_logits_list.append(cur_start_logits)
+            end_logits_list.append(cur_end_logits)
+
+        start_logits = torch.cat(start_logits_list)
+        end_logits = torch.cat(end_logits_list)
+
+        start, end, probabilities = self._postprocess(
+            start_logits, end_logits, sequence_ids, attention_mask, answers_per_seq, encodings
+        )
+
+        answers = self._nest_answers(
+            start,
+            end,
+            probabilities,
+            flattened_documents,
+            queries,
+            answers_per_seq,
+            top_k,
+            confidence_threshold,
+            query_ids,
+            document_ids,
+            no_answer,
+        )
+
+        return {"answers": answers[0]}  # same temporary batching fix as above
diff --git a/haystack/preview/dataclasses/answer.py b/haystack/preview/dataclasses/answer.py
--- a/haystack/preview/dataclasses/answer.py
+++ b/haystack/preview/dataclasses/answer.py
@@ -13,7 +13,7 @@ class Answer:
 @dataclass(frozen=True)
 class ExtractedAnswer(Answer):
     data: Optional[str]
-    document: Document
+    document: Optional[Document]
     probability: float
     start: Optional[int] = None
     end: Optional[int] = None


Test patch:
diff --git a/test/preview/components/readers/test_extractive.py b/test/preview/components/readers/test_extractive.py
new file mode 100644
--- /dev/null
+++ b/test/preview/components/readers/test_extractive.py
@@ -0,0 +1,281 @@
+from math import ceil, exp
+from typing import List
+from unittest.mock import patch, Mock
+import pytest
+
+import torch
+from transformers import pipeline
+
+from haystack.preview.components.readers import ExtractiveReader
+from haystack.preview import Document
+
+
+@pytest.fixture
+def mock_tokenizer():
+    def mock_tokenize(
+        texts: List[str],
+        text_pairs: List[str],
+        padding: bool,
+        truncation: bool,
+        max_length: int,
+        return_tensors: str,
+        return_overflowing_tokens: bool,
+        stride: int,
+    ):
+        assert padding
+        assert truncation
+        assert return_tensors == "pt"
+        assert return_overflowing_tokens
+
+        tokens = Mock()
+
+        num_splits = [ceil(len(text + pair) / max_length) for text, pair in zip(texts, text_pairs)]
+        tokens.overflow_to_sample_mapping = [i for i, num in enumerate(num_splits) for _ in range(num)]
+        num_samples = sum(num_splits)
+        tokens.encodings = [Mock() for _ in range(num_samples)]
+        sequence_ids = [0] * 16 + [1] * 16 + [None] * (max_length - 32)
+        for encoding in tokens.encodings:
+            encoding.sequence_ids = sequence_ids
+            encoding.token_to_chars = lambda i: (i - 16, i - 15)
+        tokens.input_ids = torch.zeros(num_samples, max_length, dtype=torch.int)
+        attention_mask = torch.zeros(num_samples, max_length, dtype=torch.int)
+        attention_mask[:32] = 1
+        tokens.attention_mask = attention_mask
+        return tokens
+
+    with patch("haystack.preview.components.readers.extractive.AutoTokenizer.from_pretrained") as tokenizer:
+        tokenizer.return_value = mock_tokenize
+        yield tokenizer
+
+
+@pytest.fixture()
+def mock_reader(mock_tokenizer):
+    class MockModel(torch.nn.Module):
+        def to(self, device):
+            assert device == "cpu:0"
+            self.device_set = True
+            return self
+
+        def forward(self, input_ids, attention_mask, *args, **kwargs):
+            assert input_ids.device == torch.device("cpu")
+            assert attention_mask.device == torch.device("cpu")
+            assert self.device_set
+            start = torch.zeros(input_ids.shape[:2])
+            end = torch.zeros(input_ids.shape[:2])
+            start[:, 27] = 1
+            end[:, 31] = 1
+            end[:, 32] = 1
+            prediction = Mock()
+            prediction.start_logits = start
+            prediction.end_logits = end
+            return prediction
+
+    with patch("haystack.preview.components.readers.extractive.AutoModelForQuestionAnswering.from_pretrained") as model:
+        model.return_value = MockModel()
+        reader = ExtractiveReader(model_name_or_path="mock-model", device="cpu:0")
+        reader.warm_up()
+        return reader
+
+
+example_queries = ["Who is the chancellor of Germany?", "Who is the head of the department?"]
+example_documents = [
+    [
+        Document(text="Angela Merkel was the chancellor of Germany."),
+        Document(text="Olaf Scholz is the chancellor of Germany"),
+        Document(text="Jerry is the head of the department."),
+    ]
+] * 2
+
+
+@pytest.mark.unit
+def test_output(mock_reader: ExtractiveReader):
+    answers = mock_reader.run(example_queries[0], example_documents[0], top_k=3)[
+        "answers"
+    ]  # [0] Uncomment and remove first two indices when batching support is reintroduced
+    doc_ids = set()
+    no_answer_prob = 1
+    for doc, answer in zip(example_documents[0], answers[:3]):
+        assert answer.start == 11
+        assert answer.end == 16
+        assert answer.data == doc.text[11:16]
+        assert answer.probability == pytest.approx(1 / (1 + exp(-2 * mock_reader.calibration_factor)))
+        no_answer_prob *= 1 - answer.probability
+        doc_ids.add(doc.id)
+    assert len(doc_ids) == 3
+    assert answers[-1].probability == pytest.approx(no_answer_prob)
+
+
+@pytest.mark.unit
+def test_flatten_documents(mock_reader: ExtractiveReader):
+    queries, docs, query_ids = mock_reader._flatten_documents(example_queries, example_documents)
+    i = 0
+    for j, query in enumerate(example_queries):
+        for doc in example_documents[j]:
+            assert queries[i] == query
+            assert docs[i] == doc
+            assert query_ids[i] == j
+            i += 1
+    assert len(docs) == len(queries) == len(query_ids) == i
+
+
+@pytest.mark.unit
+def test_preprocess(mock_reader: ExtractiveReader):
+    _, _, seq_ids, _, query_ids, doc_ids = mock_reader._preprocess(
+        example_queries * 3, example_documents[0], 384, [1, 1, 1], 0
+    )
+    expected_seq_ids = torch.full((3, 384), -1, dtype=torch.int)
+    expected_seq_ids[:, :16] = 0
+    expected_seq_ids[:, 16:32] = 1
+    assert torch.equal(seq_ids, expected_seq_ids)
+    assert query_ids == [1, 1, 1]
+    assert doc_ids == [0, 1, 2]
+
+
+def test_preprocess_splitting(mock_reader: ExtractiveReader):
+    _, _, seq_ids, _, query_ids, doc_ids = mock_reader._preprocess(
+        example_queries * 4, example_documents[0] + [Document(text="a" * 64)], 96, [1, 1, 1, 1], 0
+    )
+    assert seq_ids.shape[0] == 5
+    assert query_ids == [1, 1, 1, 1, 1]
+    assert doc_ids == [0, 1, 2, 3, 3]
+
+
+@pytest.mark.unit
+def test_postprocess(mock_reader: ExtractiveReader):
+    start = torch.zeros((2, 8))
+    start[0, 3] = 4
+    start[0, 1] = 5  # test attention_mask
+    start[0, 4] = 3
+    start[1, 2] = 1
+
+    end = torch.zeros((2, 8))
+    end[0, 1] = 5  # test attention_mask
+    end[0, 2] = 4  # test that end can't be before start
+    end[0, 3] = 3
+    end[0, 4] = 2
+    end[1, :] = -10
+    end[1, 4] = -1
+
+    sequence_ids = torch.ones((2, 8))
+    attention_mask = torch.ones((2, 8))
+    attention_mask[0, :2] = 0
+    encoding = Mock()
+    encoding.token_to_chars = lambda i: (int(i), int(i) + 1)
+
+    start_candidates, end_candidates, probs = mock_reader._postprocess(
+        start, end, sequence_ids, attention_mask, 3, [encoding, encoding]
+    )
+
+    assert len(start_candidates) == len(end_candidates) == len(probs) == 2
+    assert len(start_candidates[0]) == len(end_candidates[0]) == len(probs[0]) == 3
+    assert start_candidates[0][0] == 3
+    assert end_candidates[0][0] == 4
+    assert start_candidates[0][1] == 3
+    assert end_candidates[0][1] == 5
+    assert start_candidates[0][2] == 4
+    assert end_candidates[0][2] == 5
+    assert probs[0][0] == pytest.approx(1 / (1 + exp(-7 * mock_reader.calibration_factor)))
+    assert probs[0][1] == pytest.approx(1 / (1 + exp(-6 * mock_reader.calibration_factor)))
+    assert probs[0][2] == pytest.approx(1 / (1 + exp(-5 * mock_reader.calibration_factor)))
+    assert start_candidates[1][0] == 2
+    assert end_candidates[1][0] == 5
+    assert probs[1][0] == pytest.approx(1 / 2)
+
+
+@pytest.mark.unit
+def test_nest_answers(mock_reader: ExtractiveReader):
+    start = list(range(5))
+    end = [i + 5 for i in start]
+    start = [start] * 6
+    end = [end] * 6
+    probabilities = torch.arange(5).unsqueeze(0) / 5 + torch.arange(6).unsqueeze(-1) / 25
+    query_ids = [0] * 3 + [1] * 3
+    document_ids = list(range(3)) * 2
+    nested_answers = mock_reader._nest_answers(
+        start, end, probabilities, example_documents[0], example_queries, 5, 3, None, query_ids, document_ids, True
+    )
+    expected_no_answers = [0.2 * 0.16 * 0.12, 0]
+    for query, answers, expected_no_answer, probabilities in zip(
+        example_queries, nested_answers, expected_no_answers, [probabilities[:3, -1], probabilities[3:, -1]]
+    ):
+        assert len(answers) == 4
+        for doc, answer, probability in zip(example_documents[0], reversed(answers[:3]), probabilities):
+            assert answer.query == query
+            assert answer.document == doc
+            assert answer.probability == pytest.approx(probability)
+        no_answer = answers[-1]
+        assert no_answer.query == query
+        assert no_answer.document is None
+        assert no_answer.probability == pytest.approx(expected_no_answer)
+
+
+@pytest.mark.integration
+def test_t5():
+    reader = ExtractiveReader("TARUNBHATT/flan-t5-small-finetuned-squad")
+    reader.warm_up()
+    answers = reader.run(example_queries[0], example_documents[0], top_k=2)[
+        "answers"
+    ]  # remove indices when batching support is reintroduced
+    assert answers[0].data == "Angela Merkel"
+    assert answers[0].probability == pytest.approx(0.7764519453048706)
+    assert answers[1].data == "Olaf Scholz"
+    assert answers[1].probability == pytest.approx(0.7703777551651001)
+    assert answers[2].data is None
+    assert answers[2].probability == pytest.approx(0.051331606147570596)
+    # Uncomment assertions below when batching is reintroduced
+    # assert answers[0][2].probability == pytest.approx(0.051331606147570596)
+    # assert answers[1][0].data == "Jerry"
+    # assert answers[1][0].probability == pytest.approx(0.7413333654403687)
+    # assert answers[1][1].data == "Olaf Scholz"
+    # assert answers[1][1].probability == pytest.approx(0.7266613841056824)
+    # assert answers[1][2].data is None
+    # assert answers[1][2].probability == pytest.approx(0.0707035798685709)
+
+
+@pytest.mark.integration
+def test_roberta():
+    reader = ExtractiveReader("deepset/tinyroberta-squad2")
+    reader.warm_up()
+    answers = reader.run(example_queries[0], example_documents[0], top_k=2)[
+        "answers"
+    ]  # remove indices when batching is reintroduced
+    assert answers[0].data == "Olaf Scholz"
+    assert answers[0].probability == pytest.approx(0.8614975214004517)
+    assert answers[1].data == "Angela Merkel"
+    assert answers[1].probability == pytest.approx(0.857952892780304)
+    assert answers[2].data is None
+    assert answers[2].probability == pytest.approx(0.0196738764278237)
+    # uncomment assertions below when there is batching in v2
+    # assert answers[0][0].data == "Olaf Scholz"
+    # assert answers[0][0].probability == pytest.approx(0.8614975214004517)
+    # assert answers[0][1].data == "Angela Merkel"
+    # assert answers[0][1].probability == pytest.approx(0.857952892780304)
+    # assert answers[0][2].data is None
+    # assert answers[0][2].probability == pytest.approx(0.0196738764278237)
+    # assert answers[1][0].data == "Jerry"
+    # assert answers[1][0].probability == pytest.approx(0.7048940658569336)
+    # assert answers[1][1].data == "Olaf Scholz"
+    # assert answers[1][1].probability == pytest.approx(0.6604189872741699)
+    # assert answers[1][2].data is None
+    # assert answers[1][2].probability == pytest.approx(0.1002123719777046)
+
+
+@pytest.mark.integration
+def test_matches_hf_pipeline():
+    reader = ExtractiveReader("deepset/tinyroberta-squad2")
+    reader.warm_up()
+    answers = reader.run(example_queries[0], [[example_documents[0][0]]][0], top_k=20, no_answer=False)[
+        "answers"
+    ]  # [0] Remove first two indices when batching support is reintroduced
+    answers_hf = pipeline("question-answering", model=reader.model, tokenizer=reader.tokenizer, align_to_words=False)(
+        question=example_queries[0],
+        context=example_documents[0][0].text,
+        max_answer_len=1_000,
+        handle_impossible_answer=False,
+        top_k=20,
+    )  # We need to disable HF postprocessing features to make the results comparable. This is related to https://github.com/huggingface/transformers/issues/26286
+    assert len(answers) == len(answers_hf) == 20
+    for answer, answer_hf in zip(answers, answers_hf):
+        assert answer.start == answer_hf["start"]
+        assert answer.end == answer_hf["end"]
+        assert answer.data == answer_hf["answer"]