Spaces:
Runtime error
Runtime error
File size: 4,714 Bytes
dd0f067 | 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 102 103 104 105 | import math
import torch
import torch.nn as nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
class Attention(nn.Module):
"""
Scaled Dot Product Attention (Luong-Style) with Masking Support.
Args:
query: (N, Lq, Dq)
key: (N, Lk, Dk)
value: (N, Lv, Dv), optional. If None, value=key
mask: (N, Lk), optional. 1 for valid tokens, 0 for padding.
Returns:
context: (N, Lq, Dv)
"""
def __init__(self, use_scale=True):
super().__init__()
self.use_scale = use_scale
def forward(self, query, key, value=None, mask=None):
assert query.shape[-1] == key.shape[-1], "query & key must have same hidden dim"
if value is None:
value = key
else:
assert key.shape[1] == value.shape[1], "key & value must have same sequence length"
scale_factor = 1 / math.sqrt(key.size(-1)) if self.use_scale else 1 # 1 / sqer(Dk)
attention_scores = query @ key.transpose(-2, -1) * scale_factor # (N, Lq, Lk)
if mask is not None:
mask = mask.unsqueeze(1) # (N, 1, Lk)
attention_scores = attention_scores.masked_fill(mask == 0, float('-inf'))
attention_weights = torch.softmax(attention_scores, dim=-1) # (N, Lq, Lk)
return attention_weights @ value # (N, Lq, Dv)
class NMTModel(nn.Module):
"""
Neural Machine Translation (NMT) model with GRU encoder–decoder and Luong-style attention.
Args:
vocab_size (int): Size of the vocabulary.
embedding_dim (int, optional): Dimension of token embeddings. Default: 512.
hidden_dim (int, optional): Dimension of GRU hidden states. Default: 512.
gru_layers (int, optional): Number of GRU layers for encoder and decoder. Default: 2.
gru_dropout (float, optional): Dropout probability between GRU layers. Default: 0.1.
pad_token_id (int, optional): Index of the padding token. Default: 0.
Forward Inputs:
input_ids (LongTensor): Source token IDs, shape (N, L_src).
attention_mask (LongTensor): Source mask, shape (N, L_src), 1 for valid tokens.
decoder_input_ids (LongTensor): Target token IDs, shape (N, L_tgt).
decoder_attention_mask (LongTensor, optional): Target mask, shape (N, L_tgt).
Forward Returns:
logits (FloatTensor): Prediction scores, shape (N, vocab_size, L_tgt).
"""
def __init__(self, vocab_size, embedding_dim = 512,
hidden_dim = 512, gru_layers = 2,
gru_dropout = 0.1, pad_token_id = 0):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=pad_token_id)
self.encoder = nn.GRU(embedding_dim, hidden_dim,
num_layers = gru_layers, batch_first=True,
dropout = gru_dropout)
self.decoder = nn.GRU(embedding_dim, hidden_dim,
num_layers = gru_layers, batch_first=True,
dropout = gru_dropout)
self.attention = Attention()
self.output = nn.Linear(hidden_dim * 2, vocab_size)
def forward(self, input_ids, attention_mask,
decoder_input_ids, decoder_attention_mask=None):
# source and target embedding
src_embeddings = self.embedding(input_ids) # (N, L_src, emb_dim)
tgt_embeddings = self.embedding(decoder_input_ids) # (N, L_tgt, emb_dim)
# encoder
src_lenghts = attention_mask.sum(dim=1)
packed_encoder_inputs = pack_padded_sequence(src_embeddings,
src_lenghts.cpu(),
batch_first=True,
enforce_sorted=False)
encoder_outputs_packed , encoder_hidden = self.encoder(packed_encoder_inputs) # encoder_hidden: (num_layers, N, H_enc)
encoder_outputs, _ = pad_packed_sequence(encoder_outputs_packed, batch_first=True) # (N, L_src, H_enc)
# decoder
decoder_outputs, _ = self.decoder(tgt_embeddings, encoder_hidden) # decoder_outputs: (N, L_tgt, H_dec)
# attention
attention_outputs = self.attention(query = decoder_outputs,
key = encoder_outputs,
mask = attention_mask) # (N, L_tgt, H_enc)
# output
cat_outputs = torch.concat([decoder_outputs, attention_outputs], dim=-1) # (N, L_tgt, H_dec + H_enc)
logits = self.output(cat_outputs) # (N, L_tgt, vocab_size)
return logits.permute(0, 2, 1) # (N, vocab_size, L)
|