<jupyter_start><jupyter_text>Comparing LSH implementationsCompare performance for LSH between implementations on the syntethich task of the reformer paper.  Create dataset Copied from synthetic task notebook. Consider adding to `core`.<jupyter_code>class TwinSequence(Dataset):
    def __init__(self, sl=1024, len=100):
        assert sl%2 == 0
        self.sl = sl
        self.len = len
    def __getitem__(self, idx):
        seq = torch.randint(1,128,(self.sl//2,))             # w: [1-127] of len sl//2
        seq[0] = 0                                           # seq = 0w
        seq = torch.cat((seq,seq), -1)                       # seq = 0w0w
        target = torch.cat((seq[1:],torch.tensor([0])), -1)  # return offset target x:[0123], y:[1230]
        return (seq, target)                     
    def __len__(self):
        return self.len

class MaskTargCallback(Callback):
    def before_batch(self):
        self.y[:, :self.dls.train_ds.sl//2] = -100

def masked_accuracy(inp, targ, ignore=-100):
    pred = inp.argmax(dim=-1)
    mask = targ[0] != ignore
    return (pred[:,mask] == targ[:,mask]).float().mean()<jupyter_output><empty_output><jupyter_text>Long sequence modelling Sequence length of 1024, similar what is used in the paper.<jupyter_code>bs, sl = 32,1024
n_epochs = 5
train_sz = 50_000
valid_sz = 10_000
dls = DataLoaders.from_dsets(TwinSequence(sl, train_sz), TwinSequence(sl, valid_sz), bs=bs, shuffle=False, device='cuda')<jupyter_output><empty_output><jupyter_text>Total training steps:<jupyter_code>len(dls.train)*n_epochs<jupyter_output><empty_output><jupyter_text>Our implementation - LSHLM We set `n_hashes` to 1. The paper reach around 80% masked accuracy with this setup. Other hyperparameters aren't specified in the paper.<jupyter_code>n_hashes=1
bucket_size = 64
assert sl % (bucket_size * 2) == 0
model = LSHLM(vocab_sz=128, d_model=256, n_layers=1, n_heads=4, max_seq_len=sl,
              bucket_size=bucket_size, n_hashes=n_hashes, causal=True)
learn = Learner(dls, model, loss_func=CrossEntropyLossFlat(ignore_index=-100), 
                metrics=masked_accuracy, cbs=[MaskTargCallback()])
learn.fit_one_cycle(n_epochs, 1e-3)
learn.recorder.plot_loss()
del learn
torch.cuda.empty_cache()<jupyter_output><empty_output><jupyter_text>As expected per the paper results, the model struggles with learning the task.  Our implementation - Reformer LM Let's try our `ReformerLM`<jupyter_code>from reformer_fastai.reformer import ReformerLM as fastReformerLM
n_hashes=1
bucket_size = 64
assert sl % (bucket_size * 2) == 0
model = fastReformerLM(128, 256, d_ff=256, n_layers=1, n_heads=4, max_seq_len=sl,attn_dropout=0,
                   ff_dropout=0, bucket_size=bucket_size, n_hashes=n_hashes, causal=True)
learn = Learner(dls, model, loss_func=CrossEntropyLossFlat(ignore_index=-100), 
                metrics=masked_accuracy, cbs=[MaskTargCallback()])
learn.fit_one_cycle(n_epochs, 1e-3)
learn.recorder.plot_loss()
del learn
torch.cuda.empty_cache()<jupyter_output><empty_output><jupyter_text>> Similar performance as `LSHLM`.  Lucidrains implementation - Reformer LM Testing [lucidrains implementation](https://github.com/lucidrains/reformer-pytorch/).<jupyter_code>from reformer_pytorch import ReformerLM
n_hashes=1
bucket_size = 64
assert sl % (bucket_size * 2) == 0
model = ReformerLM(num_tokens=128, dim=256, depth=1, max_seq_len = sl, heads=4, bucket_size=bucket_size, 
          n_hashes=n_hashes, causal=True)
learn = Learner(dls, model, loss_func=CrossEntropyLossFlat(ignore_index=-100), 
                metrics=masked_accuracy, cbs=[MaskTargCallback()])
learn.fit_one_cycle(n_epochs, 1e-3)
learn.recorder.plot_loss()
del learn
torch.cuda.empty_cache()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11717495987158909, 'fraction_numerical': 0.03424291064740503, 'mean_word_length': 4.36441893830703, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 26, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '573989', 'license': 'MIT', 'n_tokens_mistral': 1403, 'n_tokens_neox': 1301, 'n_words': 275, 'path': 'nbs/10a_experiment.synthetic-task-comparison.ipynb', 'repo_name': 'tyoc213-contrib/reformer_fastai', 'size': '3738'}