Neural Machine Translation with Python
Contents
- Introduction to Neural Machine Translation
- Encoder-Decoder Architecture
- Attention Mechanism
- Seq2Seq Model with Attention
- Training the NMT Model
- Evaluation and Inference
- Data Preprocessing
- Handling Out-of-Vocabulary Words
- Beam Search Decoding
- Handling Long Sentences
- Model Evaluation Metrics
- Fine-tuning and Transfer Learning
- Handling Low-Resource Languages
- Deployment and Serving
- Additional Resources
Introduction to Neural Machine Translation
Neural Machine Translation (NMT) is an advanced approach to machine translation that uses artificial neural networks to predict the likelihood of a sequence of words. This technique has revolutionized language translation by learning to align and translate simultaneously, resulting in more fluent and contextually accurate translations.
import torch
import torch.nn as nn
class NeuralMachineTranslation(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(NeuralMachineTranslation, self).__init__()
self.encoder = nn.LSTM(input_dim, hidden_dim, batch_first=True)
self.decoder = nn.LSTM(hidden_dim, hidden_dim, batch_first=True)
self.fc = nn.Linear(hidden_dim, output_dim)
def forward(self, src, trg):
encoder_outputs, (hidden, cell) = self.encoder(src)
decoder_outputs, _ = self.decoder(trg, (hidden, cell))
predictions = self.fc(decoder_outputs)
return predictions
Encoder-Decoder Architecture
The core of NMT systems is the encoder-decoder architecture. The encoder processes the input sequence, while the decoder generates the output sequence. This architecture allows the model to handle variable-length input and output sequences effectively.
class Encoder(nn.Module):
def __init__(self, input_dim, emb_dim, hidden_dim, n_layers, dropout):
super().__init__()
self.embedding = nn.Embedding(input_dim, emb_dim)
self.rnn = nn.LSTM(emb_dim, hidden_dim, n_layers, dropout=dropout)
self.dropout = nn.Dropout(dropout)
def forward(self, src):
embedded = self.dropout(self.embedding(src))
outputs, (hidden, cell) = self.rnn(embedded)
return hidden, cell
class Decoder(nn.Module):
def __init__(self, output_dim, emb_dim, hidden_dim, n_layers, dropout):
super().__init__()
self.output_dim = output_dim
self.embedding = nn.Embedding(output_dim, emb_dim)
self.rnn = nn.LSTM(emb_dim, hidden_dim, n_layers, dropout=dropout)
self.fc_out = nn.Linear(hidden_dim, output_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, input, hidden, cell):
input = input.unsqueeze(0)
embedded = self.dropout(self.embedding(input))
output, (hidden, cell) = self.rnn(embedded, (hidden, cell))
prediction = self.fc_out(output.squeeze(0))
return prediction, hidden, cell
Attention Mechanism
The attention mechanism allows the model to focus on different parts of the input sequence when generating each word in the output sequence. This greatly improves the model's ability to handle long sentences and maintain context throughout the translation process.
class Attention(nn.Module):
def __init__(self, hidden_dim):
super().__init__()
self.attn = nn.Linear(hidden_dim * 2, hidden_dim)
self.v = nn.Linear(hidden_dim, 1, bias=False)
def forward(self, hidden, encoder_outputs):
batch_size = encoder_outputs.shape[1]
src_len = encoder_outputs.shape[0]
hidden = hidden.unsqueeze(1).repeat(1, src_len, 1)
encoder_outputs = encoder_outputs.permute(1, 0, 2)
energy = torch.tanh(self.attn(torch.cat((hidden, encoder_outputs), dim=2)))
attention = self.v(energy).squeeze(2)
return F.softmax(attention, dim=1)
Seq2Seq Model with Attention
The Sequence-to-Sequence (Seq2Seq) model with attention combines the encoder-decoder architecture with the attention mechanism. This allows the model to selectively focus on relevant parts of the input sequence when generating each word of the output sequence.
class Seq2Seq(nn.Module):
def __init__(self, encoder, decoder, attention, device):
super().__init__()
self.encoder = encoder
self.decoder = decoder
self.attention = attention
self.device = device
def forward(self, src, trg, teacher_forcing_ratio=0.5):
batch_size = src.shape[1]
trg_len = trg.shape[0]
trg_vocab_size = self.decoder.output_dim
outputs = torch.zeros(trg_len, batch_size, trg_vocab_size).to(self.device)
hidden, cell = self.encoder(src)
input = trg[0,:]
for t in range(1, trg_len):
output, hidden, cell = self.decoder(input, hidden, cell)
outputs[t] = output
teacher_force = random.random() < teacher_forcing_ratio
top1 = output.argmax(1)
input = trg[t] if teacher_force else top1
return outputs
Training the NMT Model
Training an NMT model involves minimizing the difference between the predicted translations and the actual translations. This is typically done using cross-entropy loss and optimization algorithms like Adam.
def train(model, iterator, optimizer, criterion, clip):
model.train()
epoch_loss = 0
for i, batch in enumerate(iterator):
src = batch.src
trg = batch.trg
optimizer.zero_grad()
output = model(src, trg)
output_dim = output.shape[-1]
output = output[1:].view(-1, output_dim)
trg = trg[1:].view(-1)
loss = criterion(output, trg)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
optimizer.step()
epoch_loss += loss.item()
return epoch_loss / len(iterator)
# Example usage
LEARNING_RATE = 0.001
model = Seq2Seq(encoder, decoder, attention, device).to(device)
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
criterion = nn.CrossEntropyLoss(ignore_index=TRG_PAD_IDX)
for epoch in range(N_EPOCHS):
train_loss = train(model, train_iterator, optimizer, criterion, CLIP)
print(f'Epoch: {epoch+1:02} | Train Loss: {train_loss:.3f}')
Evaluation and Inference
After training, the model needs to be evaluated on a separate test set to assess its performance. During inference, the model generates translations for new input sentences.
def evaluate(model, iterator, criterion):
model.eval()
epoch_loss = 0
with torch.no_grad():
for i, batch in enumerate(iterator):
src = batch.src
trg = batch.trg
output = model(src, trg, 0) # turn off teacher forcing
output_dim = output.shape[-1]
output = output[1:].view(-1, output_dim)
trg = trg[1:].view(-1)
loss = criterion(output, trg)
epoch_loss += loss.item()
return epoch_loss / len(iterator)
def translate_sentence(sentence, src_field, trg_field, model, device, max_len=50):
model.eval()
tokens = [token.lower() for token in sentence.split()]
tokens = [src_field.init_token] + tokens + [src_field.eos_token]
src_indexes = [src_field.vocab.stoi[token] for token in tokens]
src_tensor = torch.LongTensor(src_indexes).unsqueeze(1).to(device)
with torch.no_grad():
encoder_outputs, hidden = model.encoder(src_tensor)
trg_indexes = [trg_field.vocab.stoi[trg_field.init_token]]
for i in range(max_len):
trg_tensor = torch.LongTensor([trg_indexes[-1]]).to(device)
with torch.no_grad():
output, hidden, _ = model.decoder(trg_tensor, hidden, encoder_outputs)
pred_token = output.argmax(1).item()
trg_indexes.append(pred_token)
if pred_token == trg_field.vocab.stoi[trg_field.eos_token]:
break
trg_tokens = [trg_field.vocab.itos[i] for i in trg_indexes]
return trg_tokens[1:] # Exclude the <sos> token
# Example usage
test_loss = evaluate(model, test_iterator, criterion)
print(f'Test Loss: {test_loss:.3f}')
example_sentence = "The cat sat on the mat."
translation = translate_sentence(example_sentence, SRC, TRG, model, device)
print(f'Source: {example_sentence}')
print(f'Translation: {" ".join(translation)}')
Data Preprocessing
Before training the NMT model, the input data needs to be preprocessed. This includes tokenization, lowercasing, and converting words to indices.
from torchtext.data import Field, BucketIterator
from torchtext.datasets import Multi30k
def tokenize_de(text):
return [tok.lower() for tok in spacy_de.tokenizer(text)]
def tokenize_en(text):
return [tok.lower() for tok in spacy_en.tokenizer(text)]
SRC = Field(tokenize=tokenize_de, init_token='<sos>', eos_token='<eos>', lower=True)
TRG = Field(tokenize=tokenize_en, init_token='<sos>', eos_token='<eos>', lower=True)
train_data, valid_data, test_data = Multi30k.splits(exts=('.de', '.en'), fields=(SRC, TRG))
SRC.build_vocab(train_data, min_freq=2)
TRG.build_vocab(train_data, min_freq=2)
BATCH_SIZE = 32
train_iterator, valid_iterator, test_iterator = BucketIterator.splits(
(train_data, valid_data, test_data),
batch_size=BATCH_SIZE,
device=device)
Handling Out-of-Vocabulary Words
One challenge in NMT is handling words that are not in the vocabulary. Techniques like subword tokenization (e.g., BPE) can help mitigate this issue.
from torchtext.data import Field
from subword_nmt import apply_bpe
def bpe_tokenize(text):
return apply_bpe.encode(text.lower())
SRC = Field(tokenize=bpe_tokenize, init_token='<sos>', eos_token='<eos>', lower=True)
TRG = Field(tokenize=bpe_tokenize, init_token='<sos>', eos_token='<eos>', lower=True)
# Example usage
input_sentence = "The quick brown fox jumps over the lazy dog"
tokenized = bpe_tokenize(input_sentence)
print(f"Original: {input_sentence}")
print(f"Tokenized: {tokenized}")
Beam Search Decoding
During inference, beam search can be used instead of greedy decoding to generate better translations by exploring multiple possible translation paths.
def beam_search(model, src_sentence, beam_width=3, max_len=50):
model.eval()
src_tensor = torch.LongTensor([SRC.vocab.stoi[token] for token in src_sentence]).unsqueeze(1).to(device)
with torch.no_grad():
encoder_outputs, hidden = model.encoder(src_tensor)
beams = [(hidden, [TRG.vocab.stoi[TRG.init_token]], 0)]
completed_beams = []
for _ in range(max_len):
new_beams = []
for hidden, sequence, score in beams:
if sequence[-1] == TRG.vocab.stoi[TRG.eos_token]:
completed_beams.append((sequence, score))
else:
trg_tensor = torch.LongTensor([sequence[-1]]).to(device)
with torch.no_grad():
output, hidden, _ = model.decoder(trg_tensor, hidden, encoder_outputs)
top_scores, top_tokens = output.topk(beam_width)
for token, token_score in zip(top_tokens[0], top_scores[0]):
new_score = score + token_score.item()
new_sequence = sequence + [token.item()]
new_beams.append((hidden, new_sequence, new_score))
beams = sorted(new_beams, key=lambda x: x[2], reverse=True)[:beam_width]
if completed_beams:
best_beam = max(completed_beams, key=lambda x: x[1])
else:
best_beam = max(beams, key=lambda x: x[2])
return [TRG.vocab.itos[idx] for idx in best_beam[0][1:-1]] # Exclude <sos> and <eos>
# Example usage
src_sentence = ["the", "cat", "sat", "on", "the", "mat"]
translation = beam_search(model, src_sentence)
print(f"Source: {' '.join(src_sentence)}")
print(f"Translation: {' '.join(translation)}")
Handling Long Sentences
Long sentences can be challenging for NMT models. Techniques like splitting long sentences or using models designed for long-range dependencies (e.g., Transformers) can help.
def split_long_sentence(sentence, max_length=50):
words = sentence.split()
chunks = []
current_chunk = []
for word in words:
if len(current_chunk) < max_length:
current_chunk.append(word)
else:
chunks.append(" ".join(current_chunk))
current_chunk = [word]
if current_chunk:
chunks.append(" ".join(current_chunk))
return chunks
# Example usage
long_sentence = "This is a very long sentence that exceeds the maximum length limit and needs to be split into multiple smaller chunks for better translation quality and to avoid potential issues with the neural machine translation model's ability to handle long-range dependencies effectively."
chunks = split_long_sentence(long_sentence)
for i, chunk in enumerate(chunks):
print(f"Chunk {i+1}: {chunk}")
translation = translate_sentence(chunk, SRC, TRG, model, device)
print(f"Translation: {' '.join(translation)}\n")
Model Evaluation Metrics
Various metrics can be used to evaluate the quality of machine translations. The BLEU (Bilingual Evaluation Understudy) score is one of the most common, measuring the similarity between the model's output and reference translations.
from nltk.translate.bleu_score import corpus_bleu
def calculate_bleu(model, test_data, src_field, trg_field, device):
model.eval()
references = []
hypotheses = []
for example in test_data:
src = vars(example)['src']
trg = vars(example)['trg']
translation = translate_sentence(src, src_field, trg_field, model, device)
references.append([trg])
hypotheses.append(translation)
return corpus_bleu(references, hypotheses)
# Example usage
bleu_score = calculate_bleu(model, test_data, SRC, TRG, device)
print(f'BLEU score: {bleu_score:.4f}')
Fine-tuning and Transfer Learning
Pre-trained models can be fine-tuned for specific language pairs or domains, leveraging transfer learning to improve performance and reduce training time.
def load_pretrained_model(model_path):
model = torch.load(model_path)
return model
def fine_tune(model, train_data, valid_data, optimizer, criterion, num_epochs):
best_valid_loss = float('inf')
for epoch in range(num_epochs):
model.train()
for batch in train_data:
optimizer.zero_grad()
output = model(batch.src, batch.trg)
loss = criterion(output, batch.trg)
loss.backward()
optimizer.step()
valid_loss = evaluate(model, valid_data, criterion)
if valid_loss < best_valid_loss:
best_valid_loss = valid_loss
torch.save(model.state_dict(), 'best_model.pt')
print(f'Epoch: {epoch+1}, Valid Loss: {valid_loss:.3f}')
# Example usage
pretrained_model = load_pretrained_model('pretrained_nmt_model.pt')
optimizer = optim.Adam(pretrained_model.parameters())
criterion = nn.CrossEntropyLoss(ignore_index=TRG_PAD_IDX)
fine_tune(pretrained_model, train_iterator, valid_iterator, optimizer, criterion, num_epochs=5)
Handling Low-Resource Languages
For language pairs with limited parallel data, techniques like data augmentation and multilingual training can be employed to improve translation quality.
def augment_data(sentence_pairs):
augmented_pairs = []
for src, trg in sentence_pairs:
# Original pair
augmented_pairs.append((src, trg))
# Back-translation
back_translated_src = translate(trg, 'target', 'source')
augmented_pairs.append((back_translated_src, trg))
# Synonym replacement
src_synonyms = replace_with_synonyms(src)
trg_synonyms = replace_with_synonyms(trg)
augmented_pairs.append((src_synonyms, trg_synonyms))
return augmented_pairs
def train_multilingual(models, train_data, valid_data, optimizer, criterion, num_epochs):
for epoch in range(num_epochs):
for model in models:
model.train()
for batch in train_data:
optimizer.zero_grad()
output = model(batch.src, batch.trg)
loss = criterion(output, batch.trg)
loss.backward()
optimizer.step()
valid_loss = evaluate(model, valid_data, criterion)
print(f'Model: {model.name}, Epoch: {epoch+1}, Valid Loss: {valid_loss:.3f}')
# Example usage
low_resource_data = load_low_resource_data()
augmented_data = augment_data(low_resource_data)
train_data = create_dataset(augmented_data)
multilingual_models = [NMTModel('en-fr'), NMTModel('en-es'), NMTModel('fr-es')]
optimizer = optim.Adam(sum([list(model.parameters()) for model in multilingual_models], []))
criterion = nn.CrossEntropyLoss(ignore_index=TRG_PAD_IDX)
train_multilingual(multilingual_models, train_data, valid_data, optimizer, criterion, num_epochs=10)
Deployment and Serving
Once trained, the NMT model needs to be deployed for real-world use. This involves setting up an API and potentially optimizing the model for inference speed.
from flask import Flask, request, jsonify
import torch
app = Flask(__name__)
model = NMTModel.load_from_checkpoint('best_model.pt')
model.eval()
@app.route('/translate', methods=['POST'])
def translate():
data = request.json
source_text = data['text']
source_lang = data['source_lang']
target_lang = data['target_lang']
# Preprocess input
tokenized_source = tokenize(source_text, source_lang)
# Translate
with torch.no_grad():
translation = model.translate(tokenized_source, target_lang)
# Postprocess output
translated_text = detokenize(translation, target_lang)
return jsonify({'translation': translated_text})
if __name__ == '__main__':
app.run(host='0.0.0.0', port=5000)
Additional Resources
For further exploration of Neural Machine Translation, consider the following resources:
- "Neural Machine Translation by Jointly Learning to Align and Translate" by Bahdanau et al. (2014): https://arxiv.org/abs/1409.0473
- "Attention Is All You Need" by Vaswani et al. (2017): https://arxiv.org/abs/1706.03762
- "Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation" by Wu et al. (2016): https://arxiv.org/abs/1609.08144
- "Convolutional Sequence to Sequence Learning" by Gehring et al. (2017): https://arxiv.org/abs/1705.03122
- "Achieving Human Parity on Automatic Chinese to English News Translation" by Hassan et al. (2018): https://arxiv.org/abs/1803.05567
These papers provide in-depth insights into various aspects of Neural Machine Translation, from fundamental concepts to advanced techniques.