<jupyter_start><jupyter_text>Revised and fixed code from https://github.com/JayParks/transformer (MT) for LM Stripped the code from the JayParks repo for MT Transformer. Introduced a few updates and changes for speed, but it's still frustratingly slow. Possible improvement - speed it up.Another issue - hyperparameter search for language modelling (number of heads, number of self-attention layers, etc). Does not work well from the box. This might be of help https://arxiv.org/pdf/1804.00247.pdf.Also consider parallelizing.  TODO* Clean up* Add MoS  Random sequence length batchingThis version of Transformer LM usesrandom sequence length batching.**NB** Make sure the src code does not assuem the existence of PAD.<jupyter_code>import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="1"
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
from showprogress import showprogress

import torch
torch.cuda.device(0)
import torch.nn as nn
import torch.optim as optim

from torch.nn.utils.rnn import pack_padded_sequence as pack
from torch.nn.utils.rnn import pad_packed_sequence as pad
from torch.nn.utils import clip_grad_norm_ as clip
from torch.optim.lr_scheduler import StepLR

import const
from data import *
from transformermos import *
d_model = 512
voc_size = 10000
batch_size = 12
seq_len = 74
n_experts = 10
a = torch.ones([batch_size, seq_len, d_model])
l1 = torch.nn.Linear(d_model, n_experts*d_model)
l2 = torch.nn.Linear(d_model, voc_size)
l3 = torch.nn.Linear(d_model, n_experts, bias=False)
output = a
latent = l1(output)  # h  [batch_size x seq_len x n_experts * d_model]
logit = l2(latent.view(-1, d_model))  # HW [batch_size * seq_len * n_experts x d_model]
print(logit.shape)
prior_logit = l3(output).contiguous().view(-1, n_experts)
prior = torch.nn.functional.softmax(prior_logit, dim=1)  # pi
print(prior.shape)
prob = torch.nn.functional.softmax(logit.view(-1, voc_size), dim=1).view(-1, n_experts, voc_size)  # exp(hw) / sum(exp(hw))
prob = (prob * prior.unsqueeze(2).expand_as(prob)).sum(1)
print(prob.shape)
prob.size(-1)
def logging(s, print_=True, log_=True):
    if print_:
        print(s)
#     if log_:
#         with open(os.path.join(args.save, 'log.txt'), 'a+') as f_log:
#             f_log.write(s + '\n')
# ptb_datapath_train = 'data/penn/train.txt'
# ptb_datapath_valid = 'data/penn/valid.txt'
# ptb_datapath_test = 'data/penn/test.txt'

# batch_size = 128

# ptb_train = DataSet(ptb_datapath_train, batch_size, display_freq=0, max_len=90, trunc_len=90)
# ptb_valid = DataSet(ptb_datapath_valid, batch_size, display_freq=0, max_len=90, trunc_len=90)
# ptb_test = DataSet(ptb_datapath_test, batch_size, display_freq=0, max_len=90, trunc_len=90)
# ptb_train.build_dict()
# ptb_valid.change_dict(ptb_train.dictionary)
# ptb_test.change_dict(ptb_train.dictionary)
############ Optional: get data by tokens ###############
corpus = Corpus('data/penn')
eval_batch_size = 10
test_batch_size = 1
batch_size = 12
train_data = batchify(corpus.train, batch_size, )
val_data = batchify(corpus.valid, eval_batch_size)
test_data = batchify(corpus.test, test_batch_size)

#### how to take a batch ####
# the data is already splitten into batch_size(now we need to decide about seq length)
# batch_num = 2
# batch = get_batch(train_data, batch_num, seq_len=35)


#### TODO (if needed) ###
# 1) repackage hiddens for learning by tokens
# 2) learn not every step (depends on 1st point)
# 3) add grad clipping
val_data.device
voc_size = len(corpus.dictionary) #corpus.dictionary.total # ptb_train.num_vocb
n_tokens = voc_size
emb_dim = 512
d_k = 64
d_v = 64
n_layers = 2
n_heads = 4
d_ff = 2048
max_tgt_seq_len = 90
dropout = 0.1
weighted_model = False
share_proj_weight = True
lr = 1e-6
n_epochs = 1000
clip_grad = 5
warmup_steps = 2000
log_interval = 200
model = LMTransformer(n_layers, d_k, d_v, emb_dim, d_ff,
                      n_heads, max_tgt_seq_len, voc_size,
                      dropout, weighted_model, share_proj_weight)
criterion = nn.CrossEntropyLoss(ignore_index=const.PAD)

if torch.cuda.is_available():
    model = model.cuda()
    criterion = criterion.cuda()

#opt = optim.Adam(model.trainable_params(), lr=lr)
# lr_lambda = lambda epoch: 0.99 ** epoch
#lrsched = StepLR(opt, step_size=10, gamma=0.5)
bptt0 = 70
max_seq_len_delta = 40
# Evaluate model
# seq_len is strange parameter
def evaluate(data_source, model, ntokens, seq_len):
    model.eval()
    total_loss = 0
    batch = 0
    for i in range(0, data_source.size(0) - 1, seq_len):
        data, targets = get_batch(data_source, i, seq_len=seq_len)
        seq_len = data.shape[1]
        lengths = torch.ones(data.shape[0], device=device, dtype=torch.long) * seq_len

        log_prob, self_attn = model(data, lengths)
        loss = criterion(log_prob, targets.view(-1))

        total_loss += loss.item()
        batch += 1
    return total_loss / batch

n_epochs = 500
opt = optim.Adam(model.trainable_params(),betas=(0.9, 0.98), eps=1e-09, lr=lr, weight_decay=1e-5)
i=0
best_val_loss = []
try:
    for epoch in range(n_epochs):
        epoch_start_time = time.time()
        total_loss = 0
        print('Start epoch %d, learning rate %f '%(epoch + 1, opt.state_dict()['param_groups'][0]['lr']))
        start_time = time.time()
        model.train()
        batch, i = 0, 0
        while i < train_data.size(0) - 2:
            bptt = bptt0 if np.random.random() < 0.95 else bptt0 / 2.
            # Prevent excessively small or negative sequence lengths
            seq_len = max(5, int(np.random.normal(bptt, 5))) # loc 70, scale 5
            # There's a very small chance that it could select a very long sequence length resulting in OOM
            seq_len = min(seq_len, bptt + max_seq_len_delta)

            data, targets = get_batch(train_data, i, seq_len=seq_len)
            seq_len = data.shape[1]
            lengths = torch.ones(data.shape[0], device=device, dtype=torch.long) * seq_len

            opt.zero_grad()
            output, self_attn = model.forward(data, lengths)
#             break
#         break
            loss = criterion(output, targets.view(-1))
            
            loss.backward()
            opt.step()
            
            batch += 1
            i += seq_len
            
            new_lr = np.power(emb_dim, -0.5) * np.min([
                np.power((batch), -0.5),
                np.power(warmup_steps, -1.5) * (batch)])
            for param_group in opt.param_groups:
                param_group['lr'] = new_lr

            if batch % log_interval == 0 and batch > 0:
                cur_loss = loss.item()
                elapsed = time.time() - start_time
                logging('| epoch {:3d} | {}/{} batches | lr {:02.4f} | ms/batch {:5.2f} | '
                        'loss {:5.2f} | ppl {:8.2f}'.format(
                    epoch, batch, len(train_data) // bptt0, opt.param_groups[0]['lr'],
                    elapsed * 1000 / log_interval, cur_loss, math.exp(cur_loss)))
                total_loss = 0
                start_time = time.time()

        val_loss = evaluate(val_data, model, voc_size, bptt0)
        logging('-' * 89)
        logging('| end of epoch {:3d} | time: {:5.2f}s | valid loss {:5.2f} | '
                'valid ppl {:8.2f}'.format(epoch, (time.time() - epoch_start_time),
                                           val_loss, math.exp(val_loss)))
        logging('-' * 89)

        best_val_loss.append(val_loss)

except KeyboardInterrupt:
    logging('-' * 89)
    logging('Exiting from training early')<jupyter_output>Start epoch 1, learning rate 0.000001 
| epoch   0 | 200/1106 batches | lr 0.0001 | ms/batch 40.20 | loss 12.26 | ppl 211825.15
| epoch   0 | 400/1106 batches | lr 0.0002 | ms/batch 39.63 | loss 10.30 | ppl 29599.82
| epoch   0 | 600/1106 batches | lr 0.0003 | ms/batch 39.16 | loss  9.06 | ppl  8587.34
| epoch   0 | 800/1106 batches | lr 0.0004 | ms/batch 40.22 | loss  8.95 | ppl  7737.95
| epoch   0 | 1000/1106 batches | lr 0.0005 | ms/batch 39.71 | loss  8.75 | ppl  6329.87
-----------------------------------------------------------------------------------------
| end of epoch   0 | time: 47.99s | valid loss  8.11 | valid ppl  3324.46
-----------------------------------------------------------------------------------------
Start epoch 2, learning rate 0.000565 
| epoch   1 | 200/1106 batches | lr 0.0001 | ms/batch 39.83 | loss  7.89 | ppl  2680.51
| epoch   1 | 400/1106 batches | lr 0.0002 | ms/batch 39.62 | loss  7.73 | ppl  2275.75
| epoch   1 | 600/1106 batches | lr 0.0003 | ms/ba[...]
{'fraction_non_alphanumeric': 0.14005830903790087, 'fraction_numerical': 0.049912536443148685, 'mean_word_length': 2.768014059753954, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 5, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '499347', 'license': 'MIT', 'n_tokens_mistral': 3306, 'n_tokens_neox': 2965, 'n_words': 770, 'path': 'TransformerLM_mos.ipynb', 'repo_name': 'maremun/mix-attend-break-sticks', 'size': '8575'}