pytorch-transformers

This repository aims at providing the main variations of the transformer model in PyTorch. Currently it includes the initial model based on "Attention Is All You Need" (Vaswani et al. 2017) and the OpenAI GPT2 model based on Radford et al. 2018 and Radford et al. 2019.

Installation

Install via pip:

pip install git+https://github.com/butyr/pytorch-transformer.git

Usage

from transformer import Transformer


model = Transformer(
        vocab_size=25_000,
        model_dim=512,
        hidden_dim=2048,
        nheads=8,
        max_len=512,
        depth=6,
    )

src_batch = # Tensor with shape (batch_size, src_sentence_length)
tgt_batch = # Tensor with shape (batch_size, tgt_sentence_length)

outputs = model(src_batch, tgt_batch)