File size: 1,148 Bytes
2ea9ced |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
# pylint: disable=abstract-method
from typing import * # pylint: disable=wildcard-import,unused-wildcard-import
from abc import abstractmethod
import torch
from .base import LMScorer
class BatchedLMScorer(LMScorer):
# @overrides
def _build(self, model_name: str, options: Dict[str, Any]) -> None:
super()._build(model_name, options)
batch_size = options.get("batch_size", 1)
if batch_size < 1:
raise ValueError("The batch_size option must be positive")
# pylint: disable=attribute-defined-outside-init
self.batch_size = batch_size
# @overrides
def _tokens_log_prob(
self, text: List[str]
) -> List[Tuple[torch.DoubleTensor, torch.LongTensor, List[str]]]:
outputs = []
for i in range(0, len(text), self.batch_size):
batch = text[i : i + self.batch_size]
outputs.extend(self._tokens_log_prob_for_batch(batch))
return outputs
@abstractmethod
def _tokens_log_prob_for_batch(
self, text: List[str]
) -> List[Tuple[torch.DoubleTensor, torch.LongTensor, List[str]]]:
... # pragma: no cover
|