Repository: deepset-ai/haystack
Pull request: fix: Prevent invalid answer from being selected in ExtractiveReader
### Related Issues

- Fixes #6098 more fully

### Proposed Changes:

Issue: `ExtractiveReader` picked `answers_per_seq` number of answers per entry even if there weren't as many valid candidates. So it picked up invalid candidates too even if their probabilities were masked out to be 0. This meant the candidates could have special tokens which then broke with `token_to_chars` which returns `None` for them. 

Fix: Explicitly only keep candidates with probabilities > 0. 

Note: for now, the implementation loops over the instances which is necessary because different instances might have different number of final valid candidates.

### How did you test it?

Manual test with:
```python
from haystack.components.readers import ExtractiveReader
from haystack.dataclasses import Document

import traceback

docs = [
    Document(content="Mark lives in Turin"),
    Document(content="John lives in Boston"),
]

reader = ExtractiveReader()
reader.warm_up()

try:
    res = reader.run(query="Who lives in Turin?", documents=docs)
    print(res)
except Exception as e:
    print(traceback.format_exc())
```

### Checklist

- I have read the [contributors guidelines](https://github.com/deepset-ai/haystack/blob/main/CONTRIBUTING.md) and the [code of conduct](https://github.com/deepset-ai/haystack/blob/main/code_of_conduct.txt)
- I have updated the related issue with new insights and changes
- I added unit tests and updated the docstrings
- I've used one of the [conventional commit types](https://www.conventionalcommits.org/en/v1.0.0/) for my PR title: `fix:`, `feat:`, `build:`, `chore:`, `ci:`, `docs:`, `style:`, `refactor:`, `perf:`, `test:`.
- I documented my code
- I ran [pre-commit hooks](https://github.com/deepset-ai/haystack/blob/main/CONTRIBUTING.md#installation) and fixed any issue


Problem statement:
`ExtractiveReader` fails with proper input Documents
**Describe the bug**
`ExtractiveReader` seems to be very sensitive to the nature/format of input Documents

**Expected behavior**
It should be more robust and less error-prone.

**To Reproduce**
```python
from haystack.preview.components.readers import ExtractiveReader
from haystack.preview.dataclasses import Document

works = [Document(content="Mark lives in Turin.")]
does_not_work = [Document(content="Mark lives in Turin")]
does_not_work_2 = [Document(content="Mark lives in Turin."), Document(content="John lives in Boston.")]

#the error seems irrespective of the model used
reader = ExtractiveReader()
reader.warm_up()

for docs in [works, does_not_work, does_not_work_2]:
  try:
    res=reader.run(query="Who lives in Turin?", documents=docs)
  except Exception as e:
    print(traceback.format_exc())
```

**Error message**
The last 2 cases raise this error:

> Traceback (most recent call last):
  File "<ipython-input-15-3bfebcc12c00>", line 5, in <cell line: 1>
    res=reader.run(query="Who lives in Turin?", document=docs)
  File "/usr/local/lib/python3.10/dist-packages/canals/component/component.py", line 206, in wrapper
    return run_method(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/haystack/preview/components/readers/extractive.py", line 312, in run
    start, end, probabilities = self._postprocess(
  File "/usr/local/lib/python3.10/dist-packages/haystack/preview/components/readers/extractive.py", line 198, in _postprocess
    end_candidates = [
  File "/usr/local/lib/python3.10/dist-packages/haystack/preview/components/readers/extractive.py", line 199, in <listcomp>
    [encoding.token_to_chars(end)[1] for end in candidates]
  File "/usr/local/lib/python3.10/dist-packages/haystack/preview/components/readers/extractive.py", line 199, in <listcomp>
    [encoding.token_to_chars(end)[1] for end in candidates]
TypeError: 'NoneType' object is not subscriptable

**System:**
 - Haystack version (commit or version number): preview

Hints from issue discussion:

Commit message: Fix invalid answer being selected issue on ExtractiveReader

Patch to review:
diff --git a/haystack/components/readers/extractive.py b/haystack/components/readers/extractive.py
index 6cbda5f9f8..ebedc6d1d8 100644
--- a/haystack/components/readers/extractive.py
+++ b/haystack/components/readers/extractive.py
@@ -206,62 +206,50 @@ def _postprocess(
         implementations, it doesn't normalize the scores to make them easier to compare across different
         splits. Returns the top k answer spans.
         """
-        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)
+        mask = sequence_ids == 1  # Only keep tokens from the context (should ignore special tokens)
+        mask = torch.logical_and(mask, attention_mask == 1)  # Definitely remove special tokens
+        start = torch.where(mask, start, -torch.inf)  # Apply the mask on the start logits
+        end = torch.where(mask, end, -torch.inf)  # Apply the mask on the end logits
         start = start.unsqueeze(-1)
         end = end.unsqueeze(-2)
 
         logits = start + end  # shape: (batch_size, seq_length (start), seq_length (end))
+
+        # The mask here onwards is the same for all instances in the batch
+        # As such we do away with the batch dimension
         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
+
+        # topk can return invalid candidates as well if answers_per_seq > num_valid_candidates
+        # We only keep probability > 0 candidates later on
         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
+        candidates_values = candidates.values.cpu()
         start_candidates = start_candidates.cpu()
         end_candidates = end_candidates.cpu()
 
-        start_candidates_tokens_to_chars = [
-            [encoding.token_to_chars(start) for start in candidates]
-            for candidates, encoding in zip(start_candidates, encodings)
-        ]
-        if missing_start_tokens := [
-            (batch, index)
-            for batch, token_to_chars in enumerate(start_candidates_tokens_to_chars)
-            for index, pair in enumerate(token_to_chars)
-            if pair is None
-        ]:
-            logger.warning("Some start tokens could not be found in the context: %s", missing_start_tokens)
-        start_candidates_char_indices = [
-            [token_to_chars[0] if token_to_chars else None for token_to_chars in candidates]
-            for candidates in start_candidates_tokens_to_chars
-        ]
-
-        end_candidates_tokens_to_chars = [
-            [encoding.token_to_chars(end) for end in candidates]
-            for candidates, encoding in zip(end_candidates, encodings)
-        ]
-        if missing_end_tokens := [
-            (batch, index)
-            for batch, token_to_chars in enumerate(end_candidates_tokens_to_chars)
-            for index, pair in enumerate(token_to_chars)
-            if pair is None
-        ]:
-            logger.warning("Some end tokens could not be found in the context: %s", missing_end_tokens)
-        end_candidates_char_indices = [
-            [token_to_chars[1] if token_to_chars else None for token_to_chars in candidates]
-            for candidates in end_candidates_tokens_to_chars
-        ]
-
-        probabilities = candidates.values.cpu()
-
-        return start_candidates_char_indices, end_candidates_char_indices, probabilities
+        start_candidates_tokens_to_chars = []
+        end_candidates_tokens_to_chars = []
+        for i, (s_candidates, e_candidates, encoding) in enumerate(zip(start_candidates, end_candidates, encodings)):
+            # Those with probabilities > 0 are valid
+            valid = candidates_values[i] > 0
+            s_char_spans = []
+            e_char_spans = []
+            for start, end in zip(s_candidates[valid], e_candidates[valid]):
+                # token_to_chars returns `None` for special tokens
+                # But we shouldn't have special tokens in the answers at this point
+                s_char_spans.append(encoding.token_to_chars(start)[0])
+                e_char_spans.append(encoding.token_to_chars(end)[1])
+            start_candidates_tokens_to_chars.append(s_char_spans)
+            end_candidates_tokens_to_chars.append(e_char_spans)
+
+        return start_candidates_tokens_to_chars, end_candidates_tokens_to_chars, candidates_values
 
     def _nest_answers(
         self,
diff --git a/releasenotes/notes/extractive-reader-invalid-ans-a88e6b1d1ee897aa.yaml b/releasenotes/notes/extractive-reader-invalid-ans-a88e6b1d1ee897aa.yaml
new file mode 100644
index 0000000000..edec36fc43
--- /dev/null
+++ b/releasenotes/notes/extractive-reader-invalid-ans-a88e6b1d1ee897aa.yaml
@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    Fix issue with `ExtractiveReader` picking invalid answers when `answers_per_seq` > num of valid answers
diff --git a/test/components/readers/test_extractive.py b/test/components/readers/test_extractive.py
index a61cbe36bb..a48c984c4e 100644
--- a/test/components/readers/test_extractive.py
+++ b/test/components/readers/test_extractive.py
@@ -261,78 +261,6 @@ def test_warm_up_use_hf_token(mocked_automodel, mocked_autotokenizer):
     mocked_autotokenizer.assert_called_once_with("deepset/roberta-base-squad2", token="fake-token")
 
 
-def test_missing_token_to_chars_values():
-    # See https://github.com/deepset-ai/haystack/issues/6098
-
-    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: None
-        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
-
-    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.components.readers.extractive.AutoTokenizer.from_pretrained") as tokenizer, patch(
-        "haystack.components.readers.extractive.AutoModelForQuestionAnswering.from_pretrained"
-    ) as model:
-        tokenizer.return_value = mock_tokenize
-        model.return_value = MockModel()
-        reader = ExtractiveReader(model_name_or_path="mock-model", device="cpu:0")
-        reader.warm_up()
-
-        answers = reader.run(example_queries[0], example_documents[0], top_k=3)[
-            "answers"
-        ]  # [0] Uncomment and remove first two indices when batching support is reintroduced
-        for doc, answer in zip(example_documents[0], answers[:3]):
-            assert answer.start is None
-            assert answer.end is None
-            assert doc.content is not None
-            assert answer.data == doc.content
-
-
 @pytest.mark.integration
 def test_t5():
     reader = ExtractiveReader("TARUNBHATT/flan-t5-small-finetuned-squad")


Code review comments:
diff_hunk: @@ -206,62 +206,50 @@ def _postprocess(
         implementations, it doesn't normalize the scores to make them easier to compare across different
         splits. Returns the top k answer spans.
         """
-        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)
+        mask = sequence_ids == 1  # Only keep tokens from the context (should ignore special tokens)
+        mask = torch.logical_and(mask, attention_mask == 1)  # Definitely remove special tokens
+        start = torch.where(mask, start, -torch.inf)  # Apply the mask on the start logits
+        end = torch.where(mask, end, -torch.inf)  # Apply the mask on the end logits
         start = start.unsqueeze(-1)
         end = end.unsqueeze(-2)
 
         logits = start + end  # shape: (batch_size, seq_length (start), seq_length (end))
+
+        # The mask here onwards is the same for all instances in the batch
+        # As such we do away with the batch dimension
         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
+
+        # topk can return invalid candidates as well if answers_per_seq > num_valid_candidates
+        # We only keep probability > 0 candidates later on
         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
+        candidates_values = candidates.values.cpu()
         start_candidates = start_candidates.cpu()
         end_candidates = end_candidates.cpu()
 
-        start_candidates_tokens_to_chars = [
-            [encoding.token_to_chars(start) for start in candidates]
-            for candidates, encoding in zip(start_candidates, encodings)
-        ]
-        if missing_start_tokens := [
-            (batch, index)
-            for batch, token_to_chars in enumerate(start_candidates_tokens_to_chars)
-            for index, pair in enumerate(token_to_chars)
-            if pair is None
-        ]:
-            logger.warning("Some start tokens could not be found in the context: %s", missing_start_tokens)
-        start_candidates_char_indices = [
-            [token_to_chars[0] if token_to_chars else None for token_to_chars in candidates]
-            for candidates in start_candidates_tokens_to_chars
-        ]
-
-        end_candidates_tokens_to_chars = [
-            [encoding.token_to_chars(end) for end in candidates]
-            for candidates, encoding in zip(end_candidates, encodings)
-        ]
-        if missing_end_tokens := [
-            (batch, index)
-            for batch, token_to_chars in enumerate(end_candidates_tokens_to_chars)
-            for index, pair in enumerate(token_to_chars)
-            if pair is None
-        ]:
-            logger.warning("Some end tokens could not be found in the context: %s", missing_end_tokens)
-        end_candidates_char_indices = [
-            [token_to_chars[1] if token_to_chars else None for token_to_chars in candidates]
-            for candidates in end_candidates_tokens_to_chars
-        ]
-
-        probabilities = candidates.values.cpu()
-
-        return start_candidates_char_indices, end_candidates_char_indices, probabilities
+        start_candidates_tokens_to_chars = []
+        end_candidates_tokens_to_chars = []
+        for i, (s_candidates, e_candidates, encoding) in enumerate(zip(start_candidates, end_candidates, encodings)):
+            # Those with probabilities > 0 are valid
+            valid = candidates_values[i] > 0
+            s_char_spans = []
+            e_char_spans = []
+            for start, end in zip(s_candidates[valid], e_candidates[valid]):
path: haystack/components/readers/extractive.py
text: @user1:
can we rename the looping variables `start` and `end`? Pylint is complaining since these variables are already defined in the outer scope.