Package: pytorch-beam-search
PyTorch Beam SearchThis library implements fully vectorized Beam Search, Greedy Search and sampling for sequence models written in PyTorch. This is specially useful for tasks in Natural Language Processing, but can also be used for anything that requires generating a sequence from a sequence model.UsageA GPT-like character-level language modelfrom pytorch_beam_search import autoregressive

# Create vocabulary and examples
# Tokenize the way you need
corpus = list("abcdefghijklmnopqrstwxyz ")
# len(corpus) == 25
# An Index object represents a mapping from the vocabulary
# to integers (indices) to feed into the models
index = autoregressive.Index(corpus)
n_gram_size = 17    # 16 with an offset of 1 
n_grams = [corpus[i:n_gram_size + i] for i in range(len(corpus))[:-n_gram_size + 1]]

# Create tensor
X = index.text2tensor(n_grams)
# X.shape == (n_examples, len_examples) == (25 - 17 + 1 = 9, 17)

# Create and train the model
model = autoregressive.TransformerEncoder(index)    # just a PyTorch model
model.fit(X)    # basic method included

# Generate new predictions
new_examples = [list("new first"), list("new second")]
X_new = index.text2tensor(new_examples)
loss, error_rate = model.evaluate(X_new)    # basic method included
predictions, log_probabilities = autoregressive.beam_search(model, X_new)
# every element in predictions is the list of candidates for each example
output = [index.tensor2text(p) for p in predictions]
outputA Transformer character sequence-to-sequence modelfrom pytorch_beam_search import seq2seq

# Create vocabularies
# Tokenize the way you need
source = [list("abcdefghijkl"), list("mnopqrstwxyz")]
target = [list("ABCDEFGHIJKL"), list("MNOPQRSTWXYZ")]
# An Index object represents a mapping from the vocabulary
# to integers (indices) to feed into the models
source_index = seq2seq.Index(source)
target_index = seq2seq.Index(target)

# Create tensors
X = source_index.text2tensor(source)
Y = target_index.text2tensor(target)
# X.shape == (n_source_examples, len_source_examples) == (2, 11)
# Y.shape == (n_target_examples, len_target_examples) == (2, 12)

# Create and train the model
model = seq2seq.Transformer(source_index, target_index)    # just a PyTorch model
model.fit(X, Y, epochs = 100)    # basic method included

# Generate new predictions
new_source = [list("new first in"), list("new second in")]
new_target = [list("new first out"), list("new second out")]
X_new = source_index.text2tensor(new_source)
Y_new = target_index.text2tensor(new_target)
loss, error_rate = model.evaluate(X_new, Y_new)    # basic method included
predictions, log_probabilities = seq2seq.beam_search(model, X_new) 
output = [target_index.tensor2text(p) for p in predictions]
outputFeaturesAlgorithmsThegreedy_searchfunction implements Greedy Search, which simply picks the most likely token at every step. This is the fastest and simplest algorithm, but can work well if the model is properly trained.Thesamplefunction implements sampling from a sequence model, using the learned distribution at every step to build the output token by token. This is very useful to inspect what the model learned.Thebeam_searchfunction implements Beam Search, a form of pruned Breadth-First Search that expands a fixed number of the best candidates at every step. This is the slowest algorithm, but usually outperforms Greedy Search.ModelsTheautoregressivemodule implements the search algorithms and some architectures for unsupervised models that learn to predict the next token in a sequence.LSTMis a simple baseline/sanity check.TransformerEncoderis aGPT-like model for state-of-the-art performance.Theseq2seqmodule implements the search algorithms and some architectures for supervised encoder-decoder models that learn how to map sequences to sequences.LSTMis a sequence-to-sequence unidirectional LSTM model similar to the one inCho et al., 2014, useful as a simple baseline/sanity check.ReversingLSTMis a sequence-to-sequence unidirectional LSTM model that reverses the order of the tokens in the input, similar to the one inSutskever et al., 2014. A bit more complex than LSTM but gives better performance.Transformeris a standardTransformermodel for state-of-the-art performance.Installationpip install pytorch_beam_searchContributeIssue TrackerSupportIf you are having issues, feel free to contact me atjarobyte91@gmail.comLicenseThe project is licensed under the MIT License.