File size: 4,226 Bytes
599bb35 | 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 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """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) # [k, t]
scores = torch.zeros(1, device=device) # [k]
finished = [] # (normalised_score, ids)
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:
# Stop only when even the WORST of the best `beam_size` finished hypotheses already beats
# the best hypothesis still being extended. (Stopping as soon as `beam_size` hypotheses
# have finished would let early, low-quality endings crowd out the real best one.)
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: # ran out of steps: fall back to the best unfinished hypotheses
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)
|