This repository implements the original "Attention is all you need" paper from scratch, and explores the complete transformer architecture, implementing scaled dot product attention, multi-headed attention, etc. Furthermore, this work also applies the transformer architecture for the task of machine translation by training on the iwslt2017-en-de dataset, thus equipping the model for bidirectional English <---> German translation.
The model was trained on the iwslt2017-en-de dataset for 20 epochs, with a learning rate of 1e-4.
Run the follwing command to train the model.
python train.py
Run the following command for inference.
python predict.py --source_sentence <src-sentence> --source_lang {en, de} --target_lang {en, de} --model_path <path-to-ckpt>
Feel free to play around with the translations on this kaggle notebook!
Check out the results below. They are quite interesting - small sentences which are very common are translated accurately, but slightly larger sentences tend to get messed up or "lost in translation" by the model.
Given that this model was trained in a toy setting for only 20 epochs, without even a learning rate warm up or any other hyperparameter tuning, these results are expected.
Figure 1: Correct translation of a simple, common sentence.
Figure 2: Another correct translation example.
Figure 3: Example where the model “loses” the meaning in a longer sentence.
@inproceedings{vaswani2017attention, title={Attention is all you need}, author={Vaswani, Ashish and Shazeer, Noam and Parmar, Niki and Uszkoreit, Jakob and Jones, Llion and Gomez, Aidan N. and Kaiser, {\L}ukasz and Polosukhin, Illia}, booktitle={Advances in Neural Information Processing Systems}, volume={30}, year={2017} }