Hugging Face's logo

Spaces:
ME22B191
/
Neural_Translation
Sleeping

App Files Community
Neural_Translation / model.py
Soham Shah
phase1:serving the translator
3d658ec
Raw History Blame Contribute Delete
28.3 kB
"""
model.py β€” Transformer Architecture Skeleton
DA6401 Assignment 3: "Attention Is All You Need"
AUTOGRADER CONTRACT (DO NOT MODIFY SIGNATURES):
β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
β”‚ scaled_dot_product_attention(Q, K, V, mask) β†’ (out, weights) β”‚
β”‚ MultiHeadAttention.forward(q, k, v, mask) β†’ Tensor β”‚
β”‚ PositionalEncoding.forward(x) β†’ Tensor β”‚
β”‚ make_src_mask(src, pad_idx) β†’ BoolTensor β”‚
β”‚ make_tgt_mask(tgt, pad_idx) β†’ BoolTensor β”‚
β”‚ Transformer.encode(src, src_mask) β†’ Tensor β”‚
β”‚ Transformer.decode(memory,src_m,tgt,tgt_m) β†’ Tensor β”‚
β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
"""
import math
import copy
import os
from typing import Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
# ══════════════════════════════════════════════════════════════════════
# STANDALONE ATTENTION FUNCTION
# Exposed at module level so the autograder can import and test it
# independently of MultiHeadAttention.
# ══════════════════════════════════════════════════════════════════════
def scaled_dot_product_attention(
Q: torch.Tensor,
K: torch.Tensor,
V: torch.Tensor,
mask: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Compute Scaled Dot-Product Attention.
Attention(Q, K, V) = softmax( QΒ·Kα΅€ / √dβ‚– ) Β· V
Args:
Q : Query tensor, shape (..., seq_q, d_k)
K : Key tensor, shape (..., seq_k, d_k)
V : Value tensor, shape (..., seq_k, d_v)
mask : Optional Boolean mask, shape broadcastable to
(..., seq_q, seq_k).
Positions where mask is True are MASKED OUT
(set to -inf before softmax).
Returns:
output : Attended output, shape (..., seq_q, d_v)
attn_w : Attention weights, shape (..., seq_q, seq_k)
"""
d_k = Q.size(-1)
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
# Q : [..., seq_q, d_k]
# K.transpose(-2, -1) : [..., d_k, seq_k]
# matmul result : [..., seq_q, seq_k]
# great observation: A plain K.T would not work here β€” .T reverses all dims,
# which would scramble the batch/head axes.
# transpose(-2, -1) surgically swaps only the two we want.
if mask is not None:
scores = scores.masked_fill(mask, float("-inf"))
attn_w = F.softmax(scores, dim=-1)
output = torch.matmul(attn_w, V)
return output, attn_w
# ══════════════════════════════════════════════════════════════════════
# ❷ MASK HELPERS
# Exposed at module level so they can be tested independently and
# reused inside Transformer.forward.
# ══════════════════════════════════════════════════════════════════════
def make_src_mask(
src: torch.Tensor,
pad_idx: int = 1,
) -> torch.Tensor:
"""
Build a padding mask for the encoder (source sequence).
Args:
src : Source token-index tensor, shape [batch, src_len]
pad_idx : Vocabulary index of the <pad> token (default 1)
Returns:
Boolean mask, shape [batch, 1, 1, src_len]
True β†’ position is a PAD token (will be masked out)
False β†’ real token
"""
return (src == pad_idx).unsqueeze(1).unsqueeze(2)
def make_tgt_mask(
tgt: torch.Tensor,
pad_idx: int = 1,
) -> torch.Tensor:
"""
Build a combined padding + causal (look-ahead) mask for the decoder.
Args:
tgt : Target token-index tensor, shape [batch, tgt_len]
pad_idx : Vocabulary index of the <pad> token (default 1)
Returns:
Boolean mask, shape [batch, 1, tgt_len, tgt_len]
True β†’ position is masked out (PAD or future token)
"""
T = tgt.size(1)
pad_mask = (tgt == pad_idx).unsqueeze(1).unsqueeze(2) # [B, 1, 1, T]
causal = torch.triu(
torch.ones(T, T, device=tgt.device, dtype=torch.bool),
diagonal=1,
) # [T, T]
return pad_mask | causal # [B, 1, T, T]
# ══════════════════════════════════════════════════════════════════════
# MULTI-HEAD ATTENTION
# ══════════════════════════════════════════════════════════════════════
class MultiHeadAttention(nn.Module):
"""
Multi-Head Attention as in "Attention Is All You Need", Β§3.2.2.
MultiHead(Q,K,V) = Concat(head_1,...,head_h) Β· W_O
head_i = Attention(QΒ·W_Qi, KΒ·W_Ki, VΒ·W_Vi)
You are NOT allowed to use torch.nn.MultiheadAttention.
Args:
d_model (int) : Total model dimensionality. Must be divisible by num_heads.
num_heads (int) : Number of parallel attention heads h.
dropout (float): Dropout probability applied to attention weights.
"""
def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1) -> None:
super().__init__()
assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
self.d_model = d_model
self.num_heads = num_heads
self.d_k = d_model // num_heads # depth per head
self.W_q = nn.Linear(d_model, d_model) # matrix multiplication of W_q to X equivalently means passing through MLP with weight W_q and no bias.
# Same for W_k, W_v, W_o.
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
self.W_o = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(p=dropout)
self.attn = None # cached attention weights (for visualization, expt 2.3)
def forward(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""
Args:
query : shape [batch, seq_q, d_model]
key : shape [batch, seq_k, d_model]
value : shape [batch, seq_k, d_model]
mask : Optional BoolTensor broadcastable to
[batch, num_heads, seq_q, seq_k]
True β†’ masked out (attend nowhere)
Returns:
output : shape [batch, seq_q, d_model]
"""
B = query.size(0)
Q = self.W_q(query) # [B, seq_q, d_model]
K = self.W_k(key) # [B, seq_k, d_model]
V = self.W_v(value) # [B, seq_k, d_model]
#
# split d_model into (num_heads, d_k); bring head axis up to batch position
Q = Q.view(B, -1, self.num_heads, self.d_k).transpose(1, 2) # [B, h, seq_q, d_k]
K = K.view(B, -1, self.num_heads, self.d_k).transpose(1, 2) # [B, h, seq_k, d_k]
V = V.view(B, -1, self.num_heads, self.d_k).transpose(1, 2) # [B, h, seq_k, d_k]
#
out, attn = scaled_dot_product_attention(Q, K, V, mask=mask)
self.attn = attn # cache for expt 2.3
out = self.dropout(out) # [B, h, seq_q, d_k]
#
# concat heads back: [B, h, seq_q, d_k] -> [B, seq_q, d_model]
out = out.transpose(1, 2).contiguous().view(B, -1, self.d_model)
return self.W_o(out)
# ══════════════════════════════════════════════════════════════════════
# POSITIONAL ENCODING
# ══════════════════════════════════════════════════════════════════════
class PositionalEncoding(nn.Module):
"""
Sinusoidal Positional Encoding as in "Attention Is All You Need", Β§3.5.
Args:
d_model (int) : Embedding dimensionality.
dropout (float): Dropout applied after adding encodings.
max_len (int) : Maximum sequence length to pre-compute (default 5000).
"""
def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000) -> None:
super().__init__()
self.dropout = nn.Dropout(p=dropout)
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len).unsqueeze(1).float() # [max_len, 1]
div_term = torch.exp(
torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
) # [d_model/2]
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0) # [1, max_len, d_model]
self.register_buffer("pe", pe)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x : Input embeddings, shape [batch, seq_len, d_model]
Returns:
Tensor of same shape [batch, seq_len, d_model]
= x + PE[:, :seq_len, :]
"""
x = x + self.pe[:, : x.size(1)]
return self.dropout(x)
# ══════════════════════════════════════════════════════════════════════
# FEED-FORWARD NETWORK
# ══════════════════════════════════════════════════════════════════════
class PositionwiseFeedForward(nn.Module):
"""
Position-wise Feed-Forward Network, Β§3.3:
FFN(x) = max(0, xΒ·W₁ + b₁)Β·Wβ‚‚ + bβ‚‚
Args:
d_model (int) : Input / output dimensionality (e.g. 512).
d_ff (int) : Inner-layer dimensionality (e.g. 2048).
dropout (float): Dropout applied between the two linears.
"""
def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1) -> None:
super().__init__()
self.linear1 = nn.Linear(d_model, d_ff) # W₁, b₁ : d_model -> d_ff
self.linear2 = nn.Linear(d_ff, d_model) # Wβ‚‚, bβ‚‚ : d_ff -> d_model
self.dropout = nn.Dropout(p=dropout)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x : shape [batch, seq_len, d_model]
Returns:
shape [batch, seq_len, d_model]
"""
return self.linear2(self.dropout(F.relu(self.linear1(x))))
# ══════════════════════════════════════════════════════════════════════
# ENCODER LAYER
# ══════════════════════════════════════════════════════════════════════
class EncoderLayer(nn.Module):
"""
Single Transformer encoder sub-layer:
x β†’ [Self-Attention β†’ Add & Norm] β†’ [FFN β†’ Add & Norm]
Args:
d_model (int) : Model dimensionality.
num_heads (int) : Number of attention heads.
d_ff (int) : FFN inner dimensionality.
dropout (float): Dropout probability.
"""
def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1) -> None:
super().__init__()
self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
self.ff = PositionwiseFeedForward(d_model, d_ff, dropout)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor, src_mask: torch.Tensor) -> torch.Tensor:
"""
Args:
x : shape [batch, src_len, d_model]
src_mask : shape [batch, 1, 1, src_len]
Returns:
shape [batch, src_len, d_model]
"""
# Post-LN: every residual addition is immediately followed by LayerNorm.
# LayerNorm's Jacobian rescales gradients. Stacking N=6 of these means the
# gradient is repeatedly rescaled on its way back. Early in training, before
# weights settle, this makes the gradient magnitude at the bottom layers large
# and high-variance β€” the model can diverge in the first few hundred steps.
# This is exactly why the paper needs the Noam warmup: the LR must start
# near-zero and ramp up slowly, letting the layers stabilize before taking
# big steps. Post-LN without warmup typically blows up.
# sub-layer 1: self-attention + Add & Norm (Post-LN)
attn = self.self_attn(x, x, x, mask=src_mask)
x = self.norm1(x + self.dropout(attn))
# sub-layer 2: position-wise FFN + Add & Norm
ff = self.ff(x)
x = self.norm2(x + self.dropout(ff))
return x
# ══════════════════════════════════════════════════════════════════════
# DECODER LAYER
# ══════════════════════════════════════════════════════════════════════
class DecoderLayer(nn.Module):
"""
Single Transformer decoder sub-layer:
x β†’ [Masked Self-Attn β†’ Add & Norm]
β†’ [Cross-Attn(memory) β†’ Add & Norm]
β†’ [FFN β†’ Add & Norm]
Args:
d_model (int) : Model dimensionality.
num_heads (int) : Number of attention heads.
d_ff (int) : FFN inner dimensionality.
dropout (float): Dropout probability.
"""
def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1) -> None:
super().__init__()
self.self_attn = MultiHeadAttention(d_model, num_heads, dropout) # masked self-attn
self.cross_attn = MultiHeadAttention(d_model, num_heads, dropout) # attends to encoder
self.ff = PositionwiseFeedForward(d_model, d_ff, dropout)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.norm3 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(
self,
x: torch.Tensor,
memory: torch.Tensor,
src_mask: torch.Tensor,
tgt_mask: torch.Tensor,
) -> torch.Tensor:
"""
Args:
x : shape [batch, tgt_len, d_model]
memory : Encoder output, shape [batch, src_len, d_model]
src_mask : shape [batch, 1, 1, src_len]
tgt_mask : shape [batch, 1, tgt_len, tgt_len]
Returns:
shape [batch, tgt_len, d_model]
"""
# 1. masked self-attention over decoder inputs (causal + pad)
sa = self.self_attn(x, x, x, mask=tgt_mask)
x = self.norm1(x + self.dropout(sa))
# 2. cross-attention: Q from decoder, K/V from encoder memory
ca = self.cross_attn(x, memory, memory, mask=src_mask)
x = self.norm2(x + self.dropout(ca))
# 3. position-wise FFN
ff = self.ff(x)
x = self.norm3(x + self.dropout(ff))
return x
# ══════════════════════════════════════════════════════════════════════
# ENCODER & DECODER STACKS
# ══════════════════════════════════════════════════════════════════════
class Encoder(nn.Module):
"""Stack of N identical EncoderLayer modules with final LayerNorm."""
def __init__(self, layer: EncoderLayer, N: int) -> None:
super().__init__()
self.layers = nn.ModuleList([copy.deepcopy(layer) for _ in range(N)])
self.norm = nn.LayerNorm(layer.norm1.normalized_shape[0])
def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
"""
Args:
x : shape [batch, src_len, d_model]
mask : shape [batch, 1, 1, src_len]
Returns:
shape [batch, src_len, d_model]
"""
for layer in self.layers:
x = layer(x, mask)
return self.norm(x)
class Decoder(nn.Module):
"""Stack of N identical DecoderLayer modules with final LayerNorm."""
def __init__(self, layer: DecoderLayer, N: int) -> None:
super().__init__()
self.layers = nn.ModuleList([copy.deepcopy(layer) for _ in range(N)])
self.norm = nn.LayerNorm(layer.norm1.normalized_shape[0])
def forward(
self,
x: torch.Tensor,
memory: torch.Tensor,
src_mask: torch.Tensor,
tgt_mask: torch.Tensor,
) -> torch.Tensor:
"""
Args:
x : shape [batch, tgt_len, d_model]
memory : shape [batch, src_len, d_model]
src_mask : shape [batch, 1, 1, src_len]
tgt_mask : shape [batch, 1, tgt_len, tgt_len]
Returns:
shape [batch, tgt_len, d_model]
"""
for layer in self.layers:
x = layer(x, memory, src_mask, tgt_mask)
return self.norm(x)
# ══════════════════════════════════════════════════════════════════════
# FULL TRANSFORMER
# ══════════════════════════════════════════════════════════════════════
# cache so the checkpoint + vocab are fetched/built only once even if
# Transformer() is constructed many times
_PRETRAINED_CACHE: dict = {}
class Transformer(nn.Module):
"""
Full Encoder-Decoder Transformer for sequence-to-sequence tasks.
Args:
src_vocab_size (int) : Source vocabulary size.
tgt_vocab_size (int) : Target vocabulary size.
d_model (int) : Model dimensionality (default 512).
N (int) : Number of encoder/decoder layers (default 6).
num_heads (int) : Number of attention heads (default 8).
d_ff (int) : FFN inner dimensionality (default 2048).
dropout (float): Dropout probability (default 0.1).
"""
# trained weights live on Drive; gdown fetches them in __init__
CHECKPOINT_DRIVE_ID = "15X7of0L3awcux6m00UzXlhruuEaRzPUx"
CHECKPOINT_PATH = "checkpoint_baseline.pt"
def __init__(
self,
src_vocab_size: int = 1,
tgt_vocab_size: int = 1,
d_model: int = 512,
N: int = 6,
num_heads: int = 8,
d_ff: int = 2048,
dropout: float = 0.1,
checkpoint_path: str = None,
) -> None:
super().__init__()
self._build(src_vocab_size, tgt_vocab_size, d_model, N,
num_heads, d_ff, dropout)
self._weights_loaded = False
# Transformer() with default vocab sizes (or an explicit
# checkpoint_path) means inference: load weights/vocab/tokenizer
# now so infer() does no I/O and stays under the time limit.
inference_mode = (checkpoint_path is not None) or (
src_vocab_size == 1 and tgt_vocab_size == 1
)
if inference_mode:
self._load_pretrained(checkpoint_path or self.CHECKPOINT_PATH)
def _build(self, src_vocab_size, tgt_vocab_size, d_model, N,
num_heads, d_ff, dropout) -> None:
"""(Re)construct all submodules. Called from __init__, and again from
_load_pretrained once the checkpoint's true config is known."""
self.d_model = d_model
self.config = {
"src_vocab_size": src_vocab_size,
"tgt_vocab_size": tgt_vocab_size,
"d_model": d_model,
"N": N,
"num_heads": num_heads,
"d_ff": d_ff,
"dropout": dropout,
}
self.src_embed = nn.Embedding(src_vocab_size, d_model)
self.tgt_embed = nn.Embedding(tgt_vocab_size, d_model)
self.pos_enc = PositionalEncoding(d_model, dropout)
self.encoder = Encoder(EncoderLayer(d_model, num_heads, d_ff, dropout), N)
self.decoder = Decoder(DecoderLayer(d_model, num_heads, d_ff, dropout), N)
self.generator = nn.Linear(d_model, tgt_vocab_size)
# Xavier-uniform init on every matrix (Transformer is init-sensitive)
for p in self.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
def _load_pretrained(self, path: str) -> None:
"""Fill the cache once (download weights, build vocab/tokenizer),
then rebuild this instance from it. Called from __init__."""
cache = _PRETRAINED_CACHE
if not cache:
if not os.path.exists(path):
import gdown # lazy: not in requirements.txt
gdown.download(id=self.CHECKPOINT_DRIVE_ID, output=path,
quiet=True)
ckpt = torch.load(path, map_location="cpu")
cache["config"] = ckpt.get("model_config") if isinstance(ckpt, dict) else None
cache["state"] = (ckpt.get("model_state_dict", ckpt)
if isinstance(ckpt, dict) else ckpt)
import dataset as ds
base = ds.Multi30kDataset("train") # rebuild the same vocab
cache["src_vocab"] = base.src_vocab
cache["tgt_vocab"] = base.tgt_vocab
cache["nlp_de"] = ds.Multi30kDataset._nlp_de
cache["ds"] = ds
if cache["config"]:
device = next(self.parameters()).device
self._build(**cache["config"])
self.to(device)
self.load_state_dict(cache["state"])
self.src_vocab = cache["src_vocab"]
self.tgt_vocab = cache["tgt_vocab"]
self._nlp_de = cache["nlp_de"]
self._ds = cache["ds"]
self.eval()
self._weights_loaded = True
self._warmup()
def _warmup(self) -> None:
"""Run a few GPU decode steps now so the first real infer() call
doesn't pay the CUDA cold-start. The caller moves the model to its
device after __init__, so move to cuda, warm up, then back to cpu."""
if not torch.cuda.is_available():
return
ds = self._ds
try:
self.to("cuda")
with torch.no_grad():
src = torch.tensor([[ds.SOS_IDX, 5, ds.EOS_IDX]], device="cuda")
sm = make_src_mask(src, ds.PAD_IDX)
mem = self.encode(src, sm)
ys = torch.tensor([[ds.SOS_IDX]], device="cuda")
for _ in range(4):
self.decode(mem, sm, ys, make_tgt_mask(ys, ds.PAD_IDX))
ys = torch.cat([ys, ys[:, :1]], dim=1)
torch.cuda.synchronize()
except Exception:
pass
finally:
self.to("cpu") # autograder will .to(device) itself
# ── AUTOGRADER HOOKS ── keep these signatures exactly ─────────────
def encode(
self,
src: torch.Tensor,
src_mask: torch.Tensor,
) -> torch.Tensor:
"""
Run the full encoder stack.
Args:
src : Token indices, shape [batch, src_len]
src_mask : shape [batch, 1, 1, src_len]
Returns:
memory : Encoder output, shape [batch, src_len, d_model]
"""
x = self.pos_enc(self.src_embed(src) * math.sqrt(self.d_model))
return self.encoder(x, src_mask)
def decode(
self,
memory: torch.Tensor,
src_mask: torch.Tensor,
tgt: torch.Tensor,
tgt_mask: torch.Tensor,
) -> torch.Tensor:
"""
Run the full decoder stack and project to vocabulary logits.
Args:
memory : Encoder output, shape [batch, src_len, d_model]
src_mask : shape [batch, 1, 1, src_len]
tgt : Token indices, shape [batch, tgt_len]
tgt_mask : shape [batch, 1, tgt_len, tgt_len]
Returns:
logits : shape [batch, tgt_len, tgt_vocab_size]
"""
x = self.pos_enc(self.tgt_embed(tgt) * math.sqrt(self.d_model))
x = self.decoder(x, memory, src_mask, tgt_mask)
return self.generator(x)
def forward(
self,
src: torch.Tensor,
tgt: torch.Tensor,
src_mask: torch.Tensor,
tgt_mask: torch.Tensor,
) -> torch.Tensor:
"""
Full encoder-decoder forward pass.
Args:
src : shape [batch, src_len]
tgt : shape [batch, tgt_len]
src_mask : shape [batch, 1, 1, src_len]
tgt_mask : shape [batch, 1, tgt_len, tgt_len]
Returns:
logits : shape [batch, tgt_len, tgt_vocab_size]
"""
memory = self.encode(src, src_mask)
return self.decode(memory, src_mask, tgt, tgt_mask)
def infer(self, src_sentence: str) -> str:
"""
Translate a German sentence to English with greedy decoding.
Args:
src_sentence: raw German text.
Returns:
The translated English string.
"""
# weights/vocab/tokenizer already loaded in __init__, so this is
# just compute (no downloads here)
ds = self._ds
device = next(self.parameters()).device
self.eval()
with torch.no_grad():
tokens = [t.text for t in self._nlp_de.tokenizer(src_sentence.lower())]
ids = [ds.SOS_IDX] + self.src_vocab.encode(tokens) + [ds.EOS_IDX]
src = torch.tensor([ids], dtype=torch.long, device=device)
src_mask = make_src_mask(src, ds.PAD_IDX).to(device)
memory = self.encode(src, src_mask)
ys = torch.full((1, 1), ds.SOS_IDX, dtype=torch.long, device=device)
for _ in range(50): # Multi30k targets are short
tgt_mask = make_tgt_mask(ys, ds.PAD_IDX)
logits = self.decode(memory, src_mask, ys, tgt_mask)
nxt = logits[:, -1, :].argmax(-1, keepdim=True)
ys = torch.cat([ys, nxt], dim=1)
if nxt.item() == ds.EOS_IDX:
break
out = []
for i in ys[0].tolist()[1:]: # skip <sos>
if i == ds.EOS_IDX:
break
out.append(self.tgt_vocab.lookup_token(i))
return " ".join(out)