Repository: deepset-ai/haystack
Pull request: feat: initial ExtractiveReader implementation
### Related Issues

- fixes #5431

### Proposed Changes:

It implements an `ExtractiveReader` similarly to the inference functionality of `FARMReader`, but without all the farm dependencies

### How did you test it?
I still need to add tests, but I tried it with the following code:
```py
docs = [[Document(content="Angela Merkel is the chancellor of Germany."), Document(content="Olaf Scholz is the chancellor of Germany")], [Document(content="Jerry is the head of the department.")]]
queries = ["Who is the chancellor of Germany?", "What is Jerry's role?"]
reader = ExtractiveReader("deepset/tinyroberta-squad2")
p = Pipeline()
p.add_component("reader", reader)

print(p.run({"reader": {"documents": docs, "queries": queries}}))
```
### Notes for the reviewer
It's currently lacking two features for inference compared to FARMReader:
- confidence score calibration (out of scope for this PR)
- filtering overlapping predictions
Also see questions below

### 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` (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.

Hints from issue discussion:

Commit message: initial ExtractiveReader implementation

Patch to review:
diff --git a/haystack/preview/__init__.py b/haystack/preview/__init__.py
index 66f27add11..f3a5804acf 100644
--- a/haystack/preview/__init__.py
+++ b/haystack/preview/__init__.py
@@ -1,4 +1,4 @@
 from canals.component import component
 from haystack.preview.document_stores.decorator import store
-from haystack.preview.dataclasses import Document
+from haystack.preview.dataclasses import Document, Answer
 from haystack.preview.pipeline import Pipeline, PipelineError, NoSuchStoreError, load_pipelines, save_pipelines
diff --git a/haystack/preview/components/reader.py b/haystack/preview/components/reader.py
new file mode 100644
index 0000000000..196482b403
--- /dev/null
+++ b/haystack/preview/components/reader.py
@@ -0,0 +1,151 @@
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+from haystack.preview import component, Document, Answer
+from haystack.preview.pipeline import Pipeline
+from haystack.lazy_imports import LazyImport
+
+with LazyImport(message="Run 'pip install farm-haystack[inference]'") as torch_and_transformers_import:
+    from transformers import AutoModelForQuestionAnswering
+    from tokenizers import Tokenizer, Encoding
+    import torch
+
+
+@component
+class ExtractiveReader:
+    def __init__(
+        self, reader: Union[Path, str], device: Optional[str] = None, max_seq_length: int = 384, top_k: int = 10
+    ) -> None:
+        self.reader = reader
+        self.device = device
+        self.loaded = False
+        self.max_seq_length = max_seq_length
+        self.top_k = top_k
+
+    def warm_up(self):
+        if not self.loaded:
+            torch_and_transformers_import.check()
+            if torch.cuda.is_available():
+                self.device = "cuda:0"
+            else:
+                self.device = "cpu:0"
+            self.model = AutoModelForQuestionAnswering.from_pretrained(self.reader).to(self.device)
+            self.tokenizer = Tokenizer.from_pretrained(self.reader)
+
+    def _flatten(self, queries: List[str], documents: List[List[Document]]) -> Tuple[List[str], List[str]]:
+        flattened_queries = [query for documents_, query in zip(documents, queries) for _ in documents_]
+        flattened_documents = [document.content for documents_ in documents for document in documents_]
+        return flattened_queries, flattened_documents
+
+    def _preprocess(
+        self, queries: List[str], documents: List[str], max_seq_length: int
+    ) -> Tuple[torch.tensor, torch.tensor, torch.tensor, List[Encoding]]:
+        self.tokenizer.enable_padding(length=max_seq_length)
+        self.tokenizer.enable_truncation(max_seq_length)
+
+        encodings = self.tokenizer.encode_batch(list(zip(queries, documents)))
+
+        input_ids = torch.tensor([encoding.ids for encoding in encodings]).to(self.device)
+        attention_mask = torch.tensor([encoding.attention_mask for encoding in encodings]).to(self.device)
+        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
+
+    def _postprocess(
+        self,
+        output,
+        sequence_ids: torch.tensor,
+        attention_mask: torch.tensor,
+        top_k: int,
+        encodings: List[Encoding],
+        max_seq_length: int,
+    ) -> Tuple[List[List[int]], List[List[int]], List[List[float]]]:
+        start = output.start_logits
+        end = output.end_logits
+        mask = sequence_ids == 1
+        mask[..., 0] = True
+        mask = torch.logical_and(mask, attention_mask == 1)
+        start = torch.where(mask, start, -torch.inf)
+        end = torch.where(mask, end, -torch.inf)
+        start = torch.softmax(start, -1)
+        end = torch.softmax(end, -1)
+        start = start.unsqueeze(-1)
+        end = end.unsqueeze(-2)
+
+        probabilities = start * end  # shape: (batch_size, seq_length (start), seq_length (end))
+        masked_probabilities = torch.triu(probabilities, diagonal=1)
+        masked_probabilities[..., 0, 0] = probabilities[..., 0, 0]
+        masked_probabilities[..., 0, 1:] = 0
+        flat_probabilities = masked_probabilities.flatten(-2, -1)  # necessary for topk
+        candidates = torch.topk(flat_probabilities, top_k)
+        start_candidates = candidates.indices // max_seq_length
+        end_candidates = candidates.indices % max_seq_length
+        start_candidates = start_candidates.cpu()
+        end_candidates = end_candidates.cpu()
+
+        start_candidates = [
+            [encoding.token_to_chars(start)[0] if start != 0 else 0 for start in candidates]
+            for candidates, encoding in zip(start_candidates, encodings)
+        ]
+        end_candidates = [
+            [encoding.token_to_chars(end)[1] if end != 0 else 0 for end in candidates]
+            for candidates, encoding in zip(end_candidates, encodings)
+        ]
+        probabilities = candidates.values.cpu().tolist()
+
+        return start_candidates, end_candidates, probabilities
+
+    def _unflatten(
+        self,
+        start: List[List[int]],
+        end: List[List[int]],
+        probabilities: List[List[float]],
+        flattened_documents: List[str],
+        documents: List[List[Document]],
+        top_k: int,
+    ) -> List[List[Answer]]:
+        answers = [
+            [document[start:end] for start, end in zip(start_candidates_, end_candidates_)]
+            for document, start_candidates_, end_candidates_ in zip(flattened_documents, start, end)
+        ]
+
+        i = 0
+        nested_answers = []
+        for documents_ in documents:
+            current_answers = []
+            j = i + len(documents_)
+            for answers_, probabilities_ in zip(answers[i:j], probabilities[i:j]):
+                for answer, probability in zip(answers_, probabilities_):
+                    current_answers.append(Answer(answer=answer, probability=probability))
+            i = j
+            current_answers = sorted(current_answers, key=lambda answer: answer.probability, reverse=True)[:top_k]
+            nested_answers.append(current_answers)
+
+        return nested_answers
+
+    @component.output_types(answers=List[List[Answer]])
+    def run(
+        self,
+        queries: List[str],
+        documents: List[List[Document]],
+        top_k: Optional[int] = None,
+        max_seq_length: Optional[int] = None,
+    ):
+        top_k = top_k or self.top_k
+        max_seq_length = max_seq_length or self.max_seq_length
+
+        flattened_queries, flattened_documents = self._flatten(queries, documents)
+        input_ids, attention_mask, sequence_ids, encodings = self._preprocess(
+            flattened_queries, flattened_documents, max_seq_length
+        )
+
+        output = self.model(input_ids=input_ids, attention_mask=attention_mask)
+
+        start, end, probabilities = self._postprocess(
+            output, sequence_ids, attention_mask, top_k, encodings, max_seq_length
+        )
+
+        answers = self._unflatten(start, end, probabilities, flattened_documents, documents, top_k)
+
+        return {"answers": answers}
diff --git a/haystack/preview/dataclasses/__init__.py b/haystack/preview/dataclasses/__init__.py
index 6432e492bb..514a70c120 100644
--- a/haystack/preview/dataclasses/__init__.py
+++ b/haystack/preview/dataclasses/__init__.py
@@ -1,3 +1,4 @@
 from haystack.preview.dataclasses.document import Document, ContentType
+from haystack.preview.dataclasses.answer import Answer
 
-__all__ = ["Document", "ContentType"]
+__all__ = ["Document", "ContentType", "Answer"]
diff --git a/haystack/preview/dataclasses/answer.py b/haystack/preview/dataclasses/answer.py
new file mode 100644
index 0000000000..b73af71533
--- /dev/null
+++ b/haystack/preview/dataclasses/answer.py
@@ -0,0 +1,8 @@
+from typing import Optional
+from dataclasses import dataclass
+
+
+@dataclass
+class Answer:
+    answer: Optional[str]
+    probability: float


Code review comments:
diff_hunk: @@ -0,0 +1,151 @@
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+from haystack.preview import component, Document, Answer
+from haystack.preview.pipeline import Pipeline
+from haystack.lazy_imports import LazyImport
+
+with LazyImport(message="Run 'pip install farm-haystack[inference]'") as torch_and_transformers_import:
+    from transformers import AutoModelForQuestionAnswering
+    from tokenizers import Tokenizer, Encoding
+    import torch
+
+
+@component
+class ExtractiveReader:
+    def __init__(
+        self, reader: Union[Path, str], device: Optional[str] = None, max_seq_length: int = 384, top_k: int = 10
path: haystack/preview/components/reader.py
text: @author:
Should there be a default reader?

@author:
Also, is there a better parameter name? This follows the HF convention where it may either be a HF hub id or a path to a local folder. HF always calls this `model_name_or_path`.

@user1:
> Should there be a default reader?

I'd say yes, let's put a default value.

> Also, is there a better parameter name?

I'd vote for calling `reader` just `model`. I'd like to have `model` as the default replacement for `model_name_or_path` in all new components (so far, this is the first).

@author:
> I'd say yes, let's put a default value.
Do you know which one you'd like to put? I feel like this is trade-off between not frustrating users because they have to download 1.7GB (deberta-large) or because the model isn't working well (tinyroberta-squad2).

@user1:
Uh good question... Maybe a compromise, medium model? There is a "base" roberta, which is still quite large but I'd imagine most regular laptops can still handle it. If so I think that a good choice.

@author:
I have set it to "deepset/roberta-base-squad2-distilled" which is our best performing roberta base model.
In general, I would think that download time might be a bigger consideration than the speed of the laptop. At least, on my 2020 MacBook Air even the best extractive model available (deepset/deberta-v3-large-squad2), is sill pretty fast (I haven't properly benchmarked it yet). But I don't know what kind of hardware haystack users typically have.

@user1:
Alright I think that's a good compromise. I also have no clue what hardware to expect, so smaller models look like a safer bet for now

@user2:
I wouldn't call the parameter just `model`. I would expect an actual `model` instance then, it's almost like a convention in data science/machine learning circles. `model_name_or_path` makes it clear that this is a string / path.

diff_hunk: @@ -0,0 +1,151 @@
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+from haystack.preview import component, Document, Answer
+from haystack.preview.pipeline import Pipeline
+from haystack.lazy_imports import LazyImport
+
+with LazyImport(message="Run 'pip install farm-haystack[inference]'") as torch_and_transformers_import:
+    from transformers import AutoModelForQuestionAnswering
+    from tokenizers import Tokenizer, Encoding
+    import torch
+
+
+@component
+class ExtractiveReader:
+    def __init__(
+        self, reader: Union[Path, str], device: Optional[str] = None, max_seq_length: int = 384, top_k: int = 10
+    ) -> None:
+        self.reader = reader
+        self.device = device
+        self.loaded = False
+        self.max_seq_length = max_seq_length
+        self.top_k = top_k
+
+    def warm_up(self):
+        if not self.loaded:
+            torch_and_transformers_import.check()
path: haystack/preview/components/reader.py
text: @author:
Should this check already be performed in __init__?

@user1:
Yeah this can go in init I'd say

diff_hunk: @@ -0,0 +1,8 @@
+from typing import Optional
+from dataclasses import dataclass
+
+
+@dataclass
+class Answer:
+    answer: Optional[str]
+    probability: float
path: haystack/preview/dataclasses/answer.py
text: @author:
Should this contain question, document, start and end?

@user2:
That is actually a very good question. I'd say all of them are useful, but then we might need an extra field that specifies whether it's generative, so it does not contain start and end, or extractive, so it does.

Maybe let's do the following:

```
@user1
class Answer:
    answer: Optional[str]
    probability: float
    question: str
    documents: List[Document]   # might be more than one if it's a generative answer
    positions = List[Span]    # we might need to also include something different for tables
```

@user2:
However I think the `Answer` object is worth its own small PR

@user2:
Superseded by https://github.com/deepset-ai/haystack/pull/5563

diff_hunk: @@ -0,0 +1,151 @@
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+from haystack.preview import component, Document, Answer
+from haystack.preview.pipeline import Pipeline
+from haystack.lazy_imports import LazyImport
+
+with LazyImport(message="Run 'pip install farm-haystack[inference]'") as torch_and_transformers_import:
+    from transformers import AutoModelForQuestionAnswering
+    from tokenizers import Tokenizer, Encoding
+    import torch
+
+
+@component
+class ExtractiveReader:
+    def __init__(
+        self, reader: Union[Path, str], device: Optional[str] = None, max_seq_length: int = 384, top_k: int = 10
+    ) -> None:
+        self.reader = reader
+        self.device = device
+        self.loaded = False
+        self.max_seq_length = max_seq_length
+        self.top_k = top_k
+
+    def warm_up(self):
+        if not self.loaded:
+            torch_and_transformers_import.check()
+            if torch.cuda.is_available():
+                self.device = "cuda:0"
+            else:
+                self.device = "cpu:0"
+            self.model = AutoModelForQuestionAnswering.from_pretrained(self.reader).to(self.device)
+            self.tokenizer = Tokenizer.from_pretrained(self.reader)
+
+    def _flatten(self, queries: List[str], documents: List[List[Document]]) -> Tuple[List[str], List[str]]:
+        flattened_queries = [query for documents_, query in zip(documents, queries) for _ in documents_]
+        flattened_documents = [document.content for documents_ in documents for document in documents_]
+        return flattened_queries, flattened_documents
+
+    def _preprocess(
+        self, queries: List[str], documents: List[str], max_seq_length: int
+    ) -> Tuple[torch.tensor, torch.tensor, torch.tensor, List[Encoding]]:
+        self.tokenizer.enable_padding(length=max_seq_length)
+        self.tokenizer.enable_truncation(max_seq_length)
+
+        encodings = self.tokenizer.encode_batch(list(zip(queries, documents)))
+
+        input_ids = torch.tensor([encoding.ids for encoding in encodings]).to(self.device)
+        attention_mask = torch.tensor([encoding.attention_mask for encoding in encodings]).to(self.device)
+        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
+
+    def _postprocess(
+        self,
+        output,
+        sequence_ids: torch.tensor,
+        attention_mask: torch.tensor,
+        top_k: int,
+        encodings: List[Encoding],
+        max_seq_length: int,
+    ) -> Tuple[List[List[int]], List[List[int]], List[List[float]]]:
+        start = output.start_logits
+        end = output.end_logits
+        mask = sequence_ids == 1
+        mask[..., 0] = True
+        mask = torch.logical_and(mask, attention_mask == 1)
+        start = torch.where(mask, start, -torch.inf)
+        end = torch.where(mask, end, -torch.inf)
+        start = torch.softmax(start, -1)
+        end = torch.softmax(end, -1)
+        start = start.unsqueeze(-1)
+        end = end.unsqueeze(-2)
+
+        probabilities = start * end  # shape: (batch_size, seq_length (start), seq_length (end))
+        masked_probabilities = torch.triu(probabilities, diagonal=1)
+        masked_probabilities[..., 0, 0] = probabilities[..., 0, 0]
+        masked_probabilities[..., 0, 1:] = 0
path: haystack/preview/components/reader.py
text: @author:
It's slightly dodgy to do this after applying softmax because then all possible answers will never add up to 1. The only solution I see is to repeat start and end and apply softmax separately. However, this would require more memory and FARMReader also didn't go that far.

@user1:
If this is the same approach FARMReades has taken, then I'd add a comment above this code explaining the issue and I'll leave it as it is :+1:

diff_hunk: @@ -0,0 +1,151 @@
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+from haystack.preview import component, Document, Answer
+from haystack.preview.pipeline import Pipeline
+from haystack.lazy_imports import LazyImport
+
+with LazyImport(message="Run 'pip install farm-haystack[inference]'") as torch_and_transformers_import:
+    from transformers import AutoModelForQuestionAnswering
+    from tokenizers import Tokenizer, Encoding
+    import torch
+
+
+@component
+class ExtractiveReader:
+    def __init__(
+        self, reader: Union[Path, str], device: Optional[str] = None, max_seq_length: int = 384, top_k: int = 10
+    ) -> None:
+        self.reader = reader
path: haystack/preview/components/reader.py
text: @author:
Should the reader already be validated in __init__? We can't do this completely because it would be equivalent to loading it already, but we could still do some things.

@user1:
If there's some lightweight validation you can do, go for it! We can always review later how heavy the validation is and decide.

diff_hunk: @@ -0,0 +1,151 @@
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+from haystack.preview import component, Document, Answer
+from haystack.preview.pipeline import Pipeline
+from haystack.lazy_imports import LazyImport
+
+with LazyImport(message="Run 'pip install farm-haystack[inference]'") as torch_and_transformers_import:
+    from transformers import AutoModelForQuestionAnswering
+    from tokenizers import Tokenizer, Encoding
+    import torch
+
+
+@component
+class ExtractiveReader:
+    def __init__(
+        self, reader: Union[Path, str], device: Optional[str] = None, max_seq_length: int = 384, top_k: int = 10
+    ) -> None:
+        self.reader = reader
+        self.device = device
+        self.loaded = False
+        self.max_seq_length = max_seq_length
+        self.top_k = top_k
+
+    def warm_up(self):
+        if not self.loaded:
+            torch_and_transformers_import.check()
+            if torch.cuda.is_available():
+                self.device = "cuda:0"
+            else:
+                self.device = "cpu:0"
+            self.model = AutoModelForQuestionAnswering.from_pretrained(self.reader).to(self.device)
+            self.tokenizer = Tokenizer.from_pretrained(self.reader)
+
+    def _flatten(self, queries: List[str], documents: List[List[Document]]) -> Tuple[List[str], List[str]]:
+        flattened_queries = [query for documents_, query in zip(documents, queries) for _ in documents_]
+        flattened_documents = [document.content for documents_ in documents for document in documents_]
+        return flattened_queries, flattened_documents
+
+    def _preprocess(
+        self, queries: List[str], documents: List[str], max_seq_length: int
+    ) -> Tuple[torch.tensor, torch.tensor, torch.tensor, List[Encoding]]:
+        self.tokenizer.enable_padding(length=max_seq_length)
+        self.tokenizer.enable_truncation(max_seq_length)
+
+        encodings = self.tokenizer.encode_batch(list(zip(queries, documents)))
+
+        input_ids = torch.tensor([encoding.ids for encoding in encodings]).to(self.device)
+        attention_mask = torch.tensor([encoding.attention_mask for encoding in encodings]).to(self.device)
+        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
+
+    def _postprocess(
+        self,
+        output,
+        sequence_ids: torch.tensor,
+        attention_mask: torch.tensor,
+        top_k: int,
+        encodings: List[Encoding],
+        max_seq_length: int,
+    ) -> Tuple[List[List[int]], List[List[int]], List[List[float]]]:
+        start = output.start_logits
+        end = output.end_logits
+        mask = sequence_ids == 1
+        mask[..., 0] = True
+        mask = torch.logical_and(mask, attention_mask == 1)
+        start = torch.where(mask, start, -torch.inf)
+        end = torch.where(mask, end, -torch.inf)
+        start = torch.softmax(start, -1)
+        end = torch.softmax(end, -1)
+        start = start.unsqueeze(-1)
+        end = end.unsqueeze(-2)
+
+        probabilities = start * end  # shape: (batch_size, seq_length (start), seq_length (end))
+        masked_probabilities = torch.triu(probabilities, diagonal=1)
+        masked_probabilities[..., 0, 0] = probabilities[..., 0, 0]
+        masked_probabilities[..., 0, 1:] = 0
+        flat_probabilities = masked_probabilities.flatten(-2, -1)  # necessary for topk
+        candidates = torch.topk(flat_probabilities, top_k)
+        start_candidates = candidates.indices // max_seq_length
+        end_candidates = candidates.indices % max_seq_length
+        start_candidates = start_candidates.cpu()
+        end_candidates = end_candidates.cpu()
+
+        start_candidates = [
+            [encoding.token_to_chars(start)[0] if start != 0 else 0 for start in candidates]
+            for candidates, encoding in zip(start_candidates, encodings)
+        ]
+        end_candidates = [
+            [encoding.token_to_chars(end)[1] if end != 0 else 0 for end in candidates]
+            for candidates, encoding in zip(end_candidates, encodings)
+        ]
+        probabilities = candidates.values.cpu().tolist()
+
+        return start_candidates, end_candidates, probabilities
+
+    def _unflatten(
+        self,
+        start: List[List[int]],
+        end: List[List[int]],
+        probabilities: List[List[float]],
+        flattened_documents: List[str],
+        documents: List[List[Document]],
+        top_k: int,
+    ) -> List[List[Answer]]:
+        answers = [
+            [document[start:end] for start, end in zip(start_candidates_, end_candidates_)]
+            for document, start_candidates_, end_candidates_ in zip(flattened_documents, start, end)
+        ]
+
+        i = 0
+        nested_answers = []
+        for documents_ in documents:
+            current_answers = []
+            j = i + len(documents_)
+            for answers_, probabilities_ in zip(answers[i:j], probabilities[i:j]):
+                for answer, probability in zip(answers_, probabilities_):
+                    current_answers.append(Answer(answer=answer, probability=probability))
+            i = j
+            current_answers = sorted(current_answers, key=lambda answer: answer.probability, reverse=True)[:top_k]
+            nested_answers.append(current_answers)
+
+        return nested_answers
+
+    @component.output_types(answers=List[List[Answer]])
+    def run(
+        self,
+        queries: List[str],
+        documents: List[List[Document]],
path: haystack/preview/components/reader.py
text: @author:
Is this already a batch or should a batch size be specified separately so the method can perform additional batching?

@user1:
My idea is that this is already batched. Are there constraints on the batch size that would make this approach not viable? Do you need to control the batch size in any way?

@author:
If the batch size is too big the GPU will OOM. With the same GPU, you would probably want different batch sizes depending on the model size.

@user1:
Maybe then we can have a "max batch size": if the incoming batch is larger, we split it to fit. It's a bit of a mess to juggle all these batches but it should be possible. Does that sound alright?

@author:
That sounds good and it wasn't too complex to add.