Show Code
import torch
import torch.nn as nn
import random
import re
from collections import CounterFirst approach to ATTENTION
We saw how the seq2seq approach for a translation task failed to recognize the meaning of individual words, and compleatly failed to translate even new sentances that combine sentances from the training. The bottleneck we descriped in the previous experiment is where the model (encoder) collapsed all information for the training sentance into a single vector. The new proposed approach by Bahdanau (Neural Machine Translation by Jointly Learning to Align and Translate) paper is to allow the decoder to go back and pick relative parts of the encoder, giving it a richer context to find the correct meaning.
Impotrs and data
Show Code
pairs = [
# greetings & common phrases
("hello", "مرحبا"),
("goodbye", "مع السلامة"),
("thank you", "شكرا لك"),
("good morning", "صباح الخير"),
("good night", "تصبح على خير"),
("welcome", "أهلا وسهلا"),
("how are you", "كيف حالك"),
("i am fine", "أنا بخير"),
("see you later", "أراك لاحقا"),
("please", "من فضلك"),
("excuse me", "عذرا"),
("congratulations", "مبروك"),
("happy birthday", "عيد ميلاد سعيد"),
("good luck", "حظا سعيدا"),
("no problem", "لا مشكلة"),
# i am ...
("i am happy", "أنا سعيد"),
("i am sad", "أنا حزين"),
("i am tired", "أنا متعب"),
("i am hungry", "أنا جائع"),
("i am thirsty", "أنا عطشان"),
("i am here", "أنا هنا"),
("i am a student", "أنا طالب"),
("i am a teacher", "أنا معلم"),
("i am ready", "أنا مستعد"),
("i am busy", "أنا مشغول"),
# i + verb + object
("i love you", "أحبك"),
("i love reading", "أحب القراءة"),
("i read a book", "أقرأ كتابا"),
("i drink water", "أشرب الماء"),
("i eat bread", "آكل الخبز"),
("i drink coffee", "أشرب القهوة"),
("i drink tea", "أشرب الشاي"),
("i write a letter", "أكتب رسالة"),
("i want water", "أريد ماء"),
("i have a car", "لدي سيارة"),
("i have a book", "لدي كتاب"),
("i see the moon", "أرى القمر"),
("i go to school", "أذهب إلى المدرسة"),
("i live in cairo", "أعيش في القاهرة"),
("i speak arabic", "أتحدث العربية"),
("i play football", "ألعب كرة القدم"),
("i like music", "أحب الموسيقى"),
("i need help", "أحتاج مساعدة"),
("i know the answer", "أعرف الإجابة"),
("i open the door", "أفتح الباب"),
# he / she
("he is a doctor", "هو طبيب"),
("she is a teacher", "هي معلمة"),
("he reads a book", "هو يقرأ كتابا"),
("she writes a letter", "هي تكتب رسالة"),
("he loves coffee", "هو يحب القهوة"),
("she drinks tea", "هي تشرب الشاي"),
("he goes to work", "هو يذهب إلى العمل"),
("she plays piano", "هي تعزف البيانو"),
("he is my friend", "هو صديقي"),
("she is my sister", "هي أختي"),
("he has a dog", "لديه كلب"),
("she has a cat", "لديها قطة"),
("he speaks english", "هو يتحدث الإنجليزية"),
("she lives in dubai", "هي تعيش في دبي"),
("he is tall", "هو طويل"),
# we / they / you
("we are friends", "نحن أصدقاء"),
("we love our country", "نحن نحب بلدنا"),
("we go home", "نذهب إلى البيت"),
("they are students", "هم طلاب"),
("they play together", "هم يلعبون معا"),
("you are welcome", "على الرحب والسعة"),
("you are kind", "أنت لطيف"),
("where are you", "أين أنت"),
("what is your name", "ما اسمك"),
("my name is ahmed", "اسمي أحمد"),
# questions & statements
("what time is it", "كم الساعة"),
("how much is this", "بكم هذا"),
("where is the bathroom", "أين الحمام"),
("i do not understand", "لا أفهم"),
("can you help me", "هل يمكنك مساعدتي"),
("i am sorry", "أنا آسف"),
("it is cold", "الجو بارد"),
("it is hot", "الجو حار"),
("the weather is nice", "الطقس جميل"),
("today is monday", "اليوم الاثنين"),
("the book is on the table", "الكتاب على الطاولة"),
("the food is delicious", "الطعام لذيذ"),
("the house is big", "البيت كبير"),
("the car is fast", "السيارة سريعة"),
("the water is cold", "الماء بارد"),
# simple subject-verb-object
("the boy eats an apple", "الولد يأكل تفاحة"),
("the girl reads a book", "البنت تقرأ كتابا"),
("the cat drinks milk", "القطة تشرب الحليب"),
("the dog runs fast", "الكلب يجري بسرعة"),
("the sun is bright", "الشمس ساطعة"),
("the children play in the garden", "الأطفال يلعبون في الحديقة"),
("the man drives a car", "الرجل يقود سيارة"),
("the woman cooks food", "المرأة تطبخ الطعام"),
("the bird flies", "الطائر يطير"),
("the teacher explains the lesson", "المعلم يشرح الدرس"),
("the student asks a question", "الطالب يسأل سؤالا"),
("mother makes tea", "الأم تصنع الشاي"),
("father reads the newspaper", "الأب يقرأ الجريدة"),
("the train is late", "القطار متأخر"),
("i will travel tomorrow", "سأسافر غدا"),
]
print(len(pairs), "pairs")100 pairsShow Code
PAD, SOS, EOS, UNK = "<pad>", "<sos>", "<eos>", "<unk>"
special_tokens = [PAD, SOS, EOS, UNK] # PAD=0, SOS=1, EOS=2, UNK=3
def tokenize(text):
text = text.lower().strip()
text = re.sub(r"([^\w\s])", r" \1 ", text) # space out punctuation
return text.split() # \w matches Arabic letters too
def build_vocab(tokenized_sentences, min_freq=1):
# count how many times each word appears → every word with its frequency
word_counts = Counter(tok for sent in tokenized_sentences for tok in sent)
# keep a word only if it appears at least min_freq times, dropping rare words
frequent_words = [w for w, c in word_counts.items() if c >= min_freq]
# prepend the special tokens (so PAD=0, SOS=1, EOS=2, UNK=3)
vocab = special_tokens + frequent_words
# map each word to an index, and the reverse
word_to_index = {word: idx for idx, word in enumerate(vocab)}
index_to_word = {idx: word for word, idx in word_to_index.items()}
return word_to_index, index_to_word
english_tokens = [tokenize(en) for en, ar in pairs]
arabic_tokens = [tokenize(ar) for en, ar in pairs]
english_word_to_index, english_index_to_word = build_vocab(english_tokens)
arabic_word_to_index, arabic_index_to_word = build_vocab(arabic_tokens)
print(f"English vocab: {len(english_word_to_index):,} words")
print(f"Arabic vocab: {len(arabic_word_to_index):,} words")English vocab: 163 words
Arabic vocab: 180 wordsThen below we prepare the input and target data for the decoder:
The function
encode_sourcemapp the sentance to indicies[18, 28], and set<unk>to unknown words, then append with the end of sentance token<eos>=2.The function
encode_target_wrappedwrapps the sentance in special chars,<sos>=1at start and<eos>=2at the end.Then
make_decoder_iowill output 2 lists, the first one is the input to the decoder which contains all chars without the first one (<sos>), and the second list contains all chars without the last char (<eos>).
Show Code
def encode_source(english_tokens):
# English words → indices (+ EOS), with <unk> fallback
ids = [english_word_to_index.get(
t, english_word_to_index[UNK]) for t in english_tokens]
return ids + [english_word_to_index[EOS]]
def encode_target_wrapped(arabic_tokens):
# Arabic words → indices, wrapped with <sos> at the start and <eos> at the end
ids = [arabic_word_to_index.get(t, arabic_word_to_index[UNK])
for t in arabic_tokens] # this gets the index for each word
# then wrapp it in special tokens
return [arabic_word_to_index[SOS]] + ids + [arabic_word_to_index[EOS]]
def make_decoder_io(arabic_tokens):
# [<sos>, w1, w2, ..., wn, <eos>]
wrapped = encode_target_wrapped(arabic_tokens)
# [<sos>, w1, w2, ..., wn] → FED to decoder
decoder_input = wrapped[:-1] # all but the last char <EOS>
# [w1, w2, ..., wn, <eos>] → what it PREDICTS
decoder_target = wrapped[1:] # all but the first char <SOS>
return decoder_input, decoder_target
make_decoder_io(["أنا", "سعيد"])([1, 18, 28], [18, 28, 2])Here we create a padding function that do the following: - First finds the longest sequence in the list –> max_len
make a tensor pre-filled with
pad_index(0) of shape(num_sequences, max_len)copy each real sequence into the start of its row, leaving the rest as padding
Then the build_full_batch function loops over all pairs to build 3 lists : - english_input : the input to the encoder (English text) - decoder_input : the input to the decoder (Arabic text) - decoder_target : the target (true vlues) to the decoder (Arabic text)
Becasue we have 100 pairs, the output of the above 3 lists is: (100, max_english_len), (100, max_arabic_len), (100, max_arabic_len)
Show Code
def pad_sequences(sequences, pad_index=0):
max_len = max(len(s) for s in sequences)
padded = torch.full((len(sequences), max_len), pad_index, dtype=torch.long)
for i, seq in enumerate(sequences):
# seq is a plain Python list of integers, e.g. [1, 15, 23, 2]
# :len(seq) fill columns from 0 to len(seq)
padded[i, :len(seq)] = torch.tensor(seq, dtype=torch.long)
return padded
def build_full_batch(pairs):
english_seqs, decoder_in_seqs, decoder_tgt_seqs = [], [], []
for english, arabic in pairs:
english_seqs.append(encode_source(tokenize(english)))
dec_in, dec_tgt = make_decoder_io(tokenize(arabic))
decoder_in_seqs.append(dec_in)
decoder_tgt_seqs.append(dec_tgt)
english_input = pad_sequences(english_seqs) # (100, max_english_len)
decoder_input = pad_sequences(decoder_in_seqs) # (100, max_arabic_len)
decoder_target = pad_sequences(decoder_tgt_seqs) # (100, max_arabic_len)
return english_input, decoder_input, decoder_target
english_input, decoder_input, decoder_target = build_full_batch(pairs)
print("english_input :", english_input.shape)
print("decoder_input :", decoder_input.shape)
print("decoder_target:", decoder_target.shape)english_input : torch.Size([100, 7])
decoder_input : torch.Size([100, 5])
decoder_target: torch.Size([100, 5])Building the model
Encoder
The encoder is responsable of reading the english sentances and create context vectors that represent the relationship between words in the sentance. In the previous experiment, the output of the encoder was a single vector hidden that got fed as input to the decoder.
In this encoder we will define a bidirectional RNN to so that each word in the sentance will contain information not only about the previous words, but also the words after. And we will take all the hidden vectors (not only the final squeeze).
In the bidirectional approach, we don’t need the model to start from the padding tokens then move on backward capturing the context along the way, this will contaminate the context vectors and it will be carried into every backward-direction annotation for that sentence, including the ones at real word positions. This is why we are using nn.utils.rnn.pack_padded_sequence in defining the encoder so the backward pass actually starts at each sequence’s real last token instead of at padding.
Show Code
import torch.nn.utils.rnn as rnn_utils
PAD_IDX = english_word_to_index[PAD]
class Encoder(nn.Module):
def __init__(self, input_vocab_size, embed_size, hidden_size):
super().__init__()
self.embedding = nn.Embedding(
input_vocab_size, embed_size, padding_idx=PAD_IDX)
self.lstm = nn.LSTM(embed_size, hidden_size,
batch_first=True, bidirectional=True)
def forward(self, x, lengths):
# x: (batch_size, seq_len)
embedded = self.embedding(x) # (batch_size, seq_len, embed_size)
packed = rnn_utils.pack_padded_sequence(
embedded, lengths.cpu(), batch_first=True, enforce_sorted=False
)
packed_outputs, (hidden, cell) = self.lstm(packed)
# unpack back to (batch_size, seq_len, hidden_size*2), re-padding with zeros
outputs, _ = rnn_utils.pad_packed_sequence(
packed_outputs, batch_first=True, total_length=x.size(1)
)
return outputs, (hidden, cell)
encoder = Encoder(input_vocab_size=len(english_word_to_index),
embed_size=32, hidden_size=64)
# real length of each sentence = count of non-pad tokens (PAD_IDX=0, no real word uses index 0)
english_lengths = (english_input != PAD_IDX).sum(dim=1)
outputs, (hidden, cell) = encoder(english_input, english_lengths)
print("outputs:", outputs.shape) # (100, 7, 128)
print("hidden :", hidden.shape) # (2, 100, 64)
print("cell :", cell.shape) # (2, 100, 64)outputs: torch.Size([100, 7, 128])
hidden : torch.Size([2, 100, 64])
cell : torch.Size([2, 100, 64])Decoder
Every decoding step needs:
- Decoder input vector, start with
<sos> - Decoder output vector, end with
<eos> - previous word embedding
- previous hidden state
- context vector \(c_i\)
We need to compute the \(c_i\):
\[ c_i = \sum_{j=1}^{T} \alpha_{ij} h_j \]
- first compute a score called alignment score \(e_{ij}\), the alignment between the current decoder state and all hidden states in the encoder. It is a liner model. \[ e_{ij} = a(s_{i-1}, h_j) \]
Note: both \(s_{i-1}\) and \(h_j\) live in differnet spaces, so we need first to translate them into the same attention_space. Then we need to find a score that is based on encoder output and decoder state so we add the two then pass the resutl into a tanh() to allow the model to learn nonlinear interactions. Finally, we need one score per word, so the energy_layer copress the attention vector to a single score.
- Pass the scores into a softmax to get the weights. \(\alpha_{ij}\)
- Then find the context vector using the above equation.
We will do the above 3 steps in a separate class Attention(nn.Module).
Show Code
class Attention(nn.Module):
def __init__(self, encoder_hidden_size, decoder_hidden_size):
super().__init__()
# Transform the encoder outputs and decoder hidden state to a common dimension
self.encoder_transform = nn.Linear(encoder_hidden_size, decoder_hidden_size, bias=False)
self.decoder_transform = nn.Linear(decoder_hidden_size, decoder_hidden_size, bias=False)
self.energy_layer = nn.Linear(decoder_hidden_size, 1, bias=False)
def forward(self, encoder_outputs, decoder_hidden):
'''
encoder_outputs:
(batch_size, seq_len, encoder_hidden_size*2) // the 2 because the encoder is bidirectional
decoder_hidden:
(batch_size, decoder_hidden_size)
Returns:
---------
attention_weights: (batch_size, seq_len)
context_vector: (batch_size, encoder_hidden_size*2)
'''
# compute allignment scores (energy) between decoder hidden state and encoder outputs
scores = self.energy_layer(torch.tanh(
self.encoder_transform(encoder_outputs) + self.decoder_transform(decoder_hidden).unsqueeze(1)
)).squeeze(-1) # (batch_size, seq_len)
attention_weights = torch.softmax(scores, dim=1) # (batch_size, seq_len)
# compute context vector as weighted sum of encoder outputs
context_vector = torch.bmm(
attention_weights.unsqueeze(1),
encoder_outputs
).squeeze(1) # (batch_size, encoder_hidden_size*2)
return attention_weights, context_vector
Show Code
# encoder is bidirectional with hidden_size=64 → encoder_outputs' last dim is 64*2=128
attention = Attention(encoder_hidden_size=128, decoder_hidden_size=64)
# outputs: (100, 7, 128) from the encoder cell — this is encoder_outputs
# stand in for a decoder hidden state until the decoder exists: take one direction's final hidden state
dummy_decoder_hidden = hidden[-1] # (100, 64)
attn_weights, context_vector = attention(outputs, dummy_decoder_hidden)
print("attn_weights :", attn_weights.shape) # expect (100, 7)
print("context_vector:", context_vector.shape) # expect (100, 128)
# sanity check: softmax weights should sum to 1 across seq_len for every example
print("weights sum to 1:", torch.allclose(
attn_weights.sum(dim=1), torch.ones(100)))attn_weights : torch.Size([100, 7])
context_vector: torch.Size([100, 128])
weights sum to 1: TruePrevious word
│
▼
Embedding
│
▼
Attention
│
▼
Context vector
│
▼
Concatenate
(Embedding + Context)
│
▼
LSTM
│
▼
Output layer
│
▼
Vocabulary logitsShow Code
class Decoder(nn.Module):
def __init__(self, output_vocab_size, embed_size, encoder_hidden_size, decoder_hidden_size):
super().__init__()
self.embedding = nn.Embedding(output_vocab_size, embed_size, padding_idx=PAD_IDX)
self.lstm = nn.LSTM(embed_size, decoder_hidden_size, batch_first=True)
self.attention = Attention(encoder_hidden_size, decoder_hidden_size)
self.fc = nn.Linear(decoder_hidden_size,output_vocab_size)
def forward(self, input_token, encoder_outputs, hidden, cell):
# ----------------------------
# 1) Embed previous word
# ----------------------------
embedded = self.embedding(input_token)
embedded = embedded.unsqueeze(1)
# (batch, embed_size)
decoder_hidden = hidden[-1]
# ----------------------------
# 2) Attention
# ----------------------------
attention_weights, context = self.attention(
encoder_outputs,
decoder_hidden
)
# (batch, encoder_hidden_size)
context = context.unsqueeze(1)
# ----------------------------
# 3) Concatenate
# ----------------------------
lstm_input = torch.cat(
(embedded, context),
dim=2
)
# ----------------------------
# 4) Decoder LSTM
# ----------------------------
output, (hidden, cell) = self.lstm(
lstm_input,
(hidden, cell)
)
# ----------------------------
# 5) Vocabulary prediction
# ----------------------------
prediction = self.fc(
output.squeeze(1)
)
return prediction, hidden, cell, attention_weights
Training
Show Code
import random
import torch.optim as optim
# ---- hyperparameters ----
# must match encoder's hidden_size for the state-bridge below
DECODER_HIDDEN_SIZE = 64
NUM_EPOCHS = 500
LEARNING_RATE = 1e-3
TEACHER_FORCING_RATIO = 0.9
PRINT_EVERY = 25
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# fresh encoder/decoder for training
encoder = Encoder(len(english_word_to_index),
embed_size=32, hidden_size=64).to(device)
decoder = Decoder(len(arabic_word_to_index), embed_size=32,
encoder_hidden_size=128, decoder_hidden_size=DECODER_HIDDEN_SIZE).to(device)
# fix: Decoder.lstm was declared with input_size=embed_size, but forward() feeds it
# cat(embedding, context) = embed_size + encoder_hidden_size (32 + 128). Patch it here,
# or better, correct the input_size directly in the Decoder class definition above.
decoder.lstm = nn.LSTM(32 + 128, DECODER_HIDDEN_SIZE,
batch_first=True).to(device)
optimizer = optim.Adam(list(encoder.parameters()) +
list(decoder.parameters()), lr=LEARNING_RATE)
criterion = nn.CrossEntropyLoss(ignore_index=PAD_IDX)
english_input_d = english_input.to(device)
decoder_input_d = decoder_input.to(device)
decoder_target_d = decoder_target.to(device)
english_lengths_d = english_lengths.to(device)
def bridge_bidirectional_state(hidden, cell):
"""Encoder is bidirectional -> hidden/cell are (2, batch, H).
Decoder LSTM is unidirectional -> needs (1, batch, H).
Sum forward + backward states into one (valid since both share hidden_size=64)."""
h = (hidden[0] + hidden[1]).unsqueeze(0)
c = (cell[0] + cell[1]).unsqueeze(0)
return h, c
for epoch in range(1, NUM_EPOCHS + 1):
encoder.train()
decoder.train()
optimizer.zero_grad()
encoder_outputs, (enc_hidden, enc_cell) = encoder(
english_input_d, english_lengths_d)
dec_hidden, dec_cell = bridge_bidirectional_state(enc_hidden, enc_cell)
batch_size, target_len = decoder_input_d.shape
input_token = decoder_input_d[:, 0] # <sos> for every example
loss = 0.0
for t in range(target_len):
logits, dec_hidden, dec_cell, _ = decoder(
input_token, encoder_outputs, dec_hidden, dec_cell)
loss += criterion(logits, decoder_target_d[:, t])
if t + 1 < target_len:
use_teacher_forcing = random.random() < TEACHER_FORCING_RATIO
input_token = decoder_input_d[:, t +
1] if use_teacher_forcing else logits.argmax(dim=1)
loss = loss / target_len
loss.backward()
torch.nn.utils.clip_grad_norm_(
list(encoder.parameters()) + list(decoder.parameters()), max_norm=1.0)
optimizer.step()
if epoch == 1 or epoch % PRINT_EVERY == 0:
print(f"epoch {epoch:4d}/{NUM_EPOCHS} | loss {loss.item():.4f}")epoch 1/500 | loss 5.2488
epoch 25/500 | loss 3.3559
epoch 50/500 | loss 2.4866
epoch 75/500 | loss 2.1090
epoch 100/500 | loss 1.7001
epoch 125/500 | loss 1.2025
epoch 150/500 | loss 0.7893
epoch 175/500 | loss 0.4836
epoch 200/500 | loss 0.3005
epoch 225/500 | loss 0.1967
epoch 250/500 | loss 0.1366
epoch 275/500 | loss 0.1003
epoch 300/500 | loss 0.0772
epoch 325/500 | loss 0.0616
epoch 350/500 | loss 0.0505
epoch 375/500 | loss 0.0424
epoch 400/500 | loss 0.0362
epoch 425/500 | loss 0.0314
epoch 450/500 | loss 0.0276
epoch 475/500 | loss 0.0244
epoch 500/500 | loss 0.0218Show Code
MAX_DECODE_LEN = 20
@torch.no_grad()
def translate(english_sentence):
encoder.eval()
decoder.eval()
# word indices + <eos>, <unk> fallback
ids = encode_source(tokenize(english_sentence))
src = torch.tensor([ids], dtype=torch.long,
device=device) # (1, seq_len)
length = torch.tensor([len(ids)], dtype=torch.long)
encoder_outputs, (enc_hidden, enc_cell) = encoder(src, length)
dec_hidden, dec_cell = bridge_bidirectional_state(enc_hidden, enc_cell)
input_token = torch.tensor(
[arabic_word_to_index[SOS]], dtype=torch.long, device=device)
output_words = []
for _ in range(MAX_DECODE_LEN):
logits, dec_hidden, dec_cell, _ = decoder(
input_token, encoder_outputs, dec_hidden, dec_cell)
next_id = logits.argmax(dim=1)
if next_id.item() == arabic_word_to_index[EOS]:
break
output_words.append(arabic_index_to_word[next_id.item()])
input_token = next_id
return " ".join(output_words)
for sentence in ["i love you", "the boy eats an apple", "he is a doctor"]:
print(f"{sentence:30s} -> {translate(sentence)}")i love you -> أحبك
the boy eats an apple -> الولد يأكل تفاحة
he is a doctor -> هو طبيبShow Code
for sentence in ["the teacher drink fast"]:
print(f"{sentence:30s} -> {translate(sentence)}")the teacher drink fast -> المعلم يقود
Comments