| import torch |
|
|
|
|
| class DecodeStrategy(object): |
| def __init__(self, pad, bos, eos, batch_size, parallel_paths, min_length, max_length, |
| return_attention=False, return_hidden=False): |
| self.pad = pad |
| self.bos = bos |
| self.eos = eos |
|
|
| self.batch_size = batch_size |
| self.parallel_paths = parallel_paths |
| |
| self.predictions = [[] for _ in range(batch_size)] |
| self.scores = [[] for _ in range(batch_size)] |
| self.token_scores = [[] for _ in range(batch_size)] |
| self.attention = [[] for _ in range(batch_size)] |
| self.hidden = [[] for _ in range(batch_size)] |
|
|
| self.alive_attn = None |
| self.alive_hidden = None |
|
|
| self.min_length = min_length |
| self.max_length = max_length |
|
|
| n_paths = batch_size * parallel_paths |
| self.return_attention = return_attention |
| self.return_hidden = return_hidden |
|
|
| self.done = False |
|
|
| def initialize(self, memory_bank, device=None): |
| if device is None: |
| device = torch.device('cpu') |
| self.alive_seq = torch.full( |
| [self.batch_size * self.parallel_paths, 1], self.bos, |
| dtype=torch.long, device=device) |
| self.is_finished = torch.zeros( |
| [self.batch_size, self.parallel_paths], |
| dtype=torch.uint8, device=device) |
| self.alive_log_token_scores = torch.zeros( |
| [self.batch_size * self.parallel_paths, 0], |
| dtype=torch.float, device=device) |
|
|
| return None, memory_bank |
|
|
| def __len__(self): |
| return self.alive_seq.shape[1] |
|
|
| def ensure_min_length(self, log_probs): |
| if len(self) <= self.min_length: |
| log_probs[:, self.eos] = -1e20 |
|
|
| def ensure_max_length(self): |
| if len(self) == self.max_length + 1: |
| self.is_finished.fill_(1) |
|
|
| def advance(self, log_probs, attn): |
| raise NotImplementedError() |
|
|
| def update_finished(self): |
| raise NotImplementedError |
|
|
|
|
