Decode a batch of token distributions. Args: token_dists: softmax probabilities over the token distribution. Shape: N, L, C raw: return unprocessed labels (will return list of list of strings) Returns: list of string labels (arbitrary length) and
(self,
token_dists: Tensor,
raw: bool = False)
| 73 | raise NotImplementedError |
| 74 | |
| 75 | def decode(self, |
| 76 | token_dists: Tensor, |
| 77 | raw: bool = False) -> Tuple[List[str], List[Tensor]]: |
| 78 | """Decode a batch of token distributions. |
| 79 | |
| 80 | Args: |
| 81 | token_dists: softmax probabilities over the token distribution. Shape: N, L, C |
| 82 | raw: return unprocessed labels (will return list of list of strings) |
| 83 | |
| 84 | Returns: |
| 85 | list of string labels (arbitrary length) and |
| 86 | their corresponding sequence probabilities as a list of Tensors |
| 87 | """ |
| 88 | batch_tokens = [] |
| 89 | batch_probs = [] |
| 90 | for dist in token_dists: |
| 91 | probs, ids = dist.max(-1) # greedy selection |
| 92 | if not raw: |
| 93 | probs, ids = self._filter(probs, ids) |
| 94 | tokens = self._ids2tok(ids, not raw) |
| 95 | batch_tokens.append(tokens) |
| 96 | batch_probs.append(probs) |
| 97 | return batch_tokens, batch_probs |
| 98 | |
| 99 | |
| 100 | class Tokenizer(BaseTokenizer): |