def _setup_infer_dataloader(self, queries: List[str], batch_size: int) -> 'torch.utils.data.DataLoader':
"""
Setup function for a infer data loader.
Args:
queries: lower cased text without punctuation
batch_size: batch size to use during inference
Returns:
A pytorch DataLoader.
"""
dataset = BertPunctuationCapitalizationInferDataset(
tokenizer=self.tokenizer, queries=queries, max_seq_length=self._cfg.dataset.max_seq_length
)
return torch.utils.data.DataLoader(
dataset=dataset,
collate_fn=dataset.collate_fn,
batch_size=batch_size,
shuffle=False,
num_workers=self._cfg.dataset.num_workers,
pin_memory=self._cfg.dataset.pin_memory,
drop_last=False,
)
def _setup_infer_dataloader(
self,
queries: List[str],
batch_size: int,
max_seq_length: int,
step: int,
margin: int,
) -> torch.utils.data.DataLoader:
"""
Setup function for a infer data loader.
Args:
model: a ``PunctuationCapitalizationModel`` instance for which data loader is created.
queries: lower cased text without punctuation
batch_size: batch size to use during inference
max_seq_length: length of segments into which queries are split. ``max_seq_length`` includes ``[CLS]`` and
``[SEP]`` so every segment contains at most ``max_seq_length-2`` tokens from input a query.
step: number of tokens by which a segment is offset to a previous segment. Parameter ``step`` cannot be greater
than ``max_seq_length-2``.
margin: number of tokens near the edge of a segment which label probabilities are not used in final prediction
computation.
Returns:
A pytorch DataLoader.
"""
if max_seq_length is None:
max_seq_length = self._cfg.dataset.max_seq_length
if step is None:
step = self._cfg.dataset.step
if margin is None:
margin = self._cfg.dataset.margin
dataset = BertPunctuationCapitalizationInferDataset(
tokenizer=self.tokenizer,
queries=queries,
max_seq_length=max_seq_length,
step=step,
margin=margin)
return torch.utils.data.DataLoader(
dataset=dataset,
collate_fn=dataset.collate_fn,
batch_size=batch_size,
shuffle=False,
num_workers=self._cfg.dataset.num_workers,
pin_memory=self._cfg.dataset.pin_memory,
drop_last=False,
)