| """Greedy and beam-search decoding (no KV cache: simple and correct; fast enough for small models).""" |
| import math |
|
|
| import torch |
|
|
|
|
| def _encode_source(model, tok, text, device): |
| ids = tok.encode(text)[: model.max_seq_len - 2] |
| src = torch.tensor([[tok.BOS_ID] + ids + [tok.EOS_ID]], dtype=torch.long, device=device) |
| src_mask = model.make_src_mask(src) |
| return src, src_mask, model.encode(src, src_mask) |
|
|
|
|
| def _max_steps(model, src): |
| return max(1, min(model.max_seq_len - 1, int(src.size(1) * 2) + 10)) |
|
|
|
|
| @torch.no_grad() |
| def greedy_decode(model, tok, text, device, max_len=None): |
| model.eval() |
| if not text.strip(): |
| return "" |
| src, src_mask, memory = _encode_source(model, tok, text, device) |
| steps = min(max_len or 10 ** 9, _max_steps(model, src)) |
| ys = torch.tensor([[tok.BOS_ID]], dtype=torch.long, device=device) |
| for _ in range(steps): |
| logits = model.decode(ys, memory, src_mask)[:, -1, :].float() |
| logits[:, tok.PAD_ID] = float("-inf") |
| logits[:, tok.BOS_ID] = float("-inf") |
| nxt = int(logits.argmax(-1)) |
| if nxt == tok.EOS_ID: |
| break |
| ys = torch.cat([ys, torch.tensor([[nxt]], device=device)], dim=1) |
| return tok.decode(ys[0, 1:].tolist()) |
|
|
|
|
| def _length_penalty(length, alpha): |
| return ((5.0 + length) / 6.0) ** alpha |
|
|
|
|
| @torch.no_grad() |
| def beam_search_decode(model, tok, text, device, beam_size=5, max_len=None, alpha=0.6): |
| model.eval() |
| if not text.strip(): |
| return "" |
| src, src_mask, memory = _encode_source(model, tok, text, device) |
| steps = min(max_len or 10 ** 9, _max_steps(model, src)) |
|
|
| seqs = torch.tensor([[tok.BOS_ID]], dtype=torch.long, device=device) |
| scores = torch.zeros(1, device=device) |
| finished = [] |
|
|
| for step in range(steps): |
| k = seqs.size(0) |
| logits = model.decode(seqs, memory.expand(k, -1, -1), src_mask.expand(k, -1, -1, -1))[:, -1, :] |
| logp = torch.log_softmax(logits.float(), dim=-1) |
| logp[:, tok.PAD_ID] = float("-inf") |
| logp[:, tok.BOS_ID] = float("-inf") |
| V = logp.size(-1) |
| flat = (scores.unsqueeze(1) + logp).view(-1) |
| top_scores, top_idx = flat.topk(min(2 * beam_size, flat.numel())) |
|
|
| new_seqs, new_scores = [], [] |
| for s, idx in zip(top_scores.tolist(), top_idx.tolist()): |
| if not math.isfinite(s): |
| continue |
| b, t = divmod(idx, V) |
| if t == tok.EOS_ID: |
| finished.append((s / _length_penalty(step + 1, alpha), seqs[b, 1:].tolist())) |
| else: |
| new_seqs.append(torch.cat([seqs[b], torch.tensor([t], device=device)])) |
| new_scores.append(s) |
| if len(new_seqs) == beam_size: |
| break |
| if not new_seqs: |
| break |
| seqs = torch.stack(new_seqs) |
| scores = torch.tensor(new_scores, device=device) |
| if len(finished) >= beam_size: |
| |
| |
| |
| finished.sort(key=lambda x: x[0], reverse=True) |
| finished = finished[:beam_size] |
| best_alive = scores.max().item() / _length_penalty(step + 1, alpha) |
| if finished[-1][0] >= best_alive: |
| break |
|
|
| if not finished: |
| for s, seq in zip(scores.tolist(), seqs): |
| finished.append((s / _length_penalty(seqs.size(1) - 1, alpha), seq[1:].tolist())) |
| best = max(finished, key=lambda x: x[0])[1] |
| return tok.decode(best) |
|
|
|
|
| def translate(model, tok, text, device, beam_size=5, max_len=None): |
| """beam_size <= 1 -> greedy.""" |
| if beam_size and beam_size > 1: |
| return beam_search_decode(model, tok, text, device, beam_size=beam_size, max_len=max_len) |
| return greedy_decode(model, tok, text, device, max_len=max_len) |
|
|
