%reload_ext autoreload
%autoreload 2
%matplotlib inline

from fastai.io import *
from fastai.conv_learner import *
from fastai.column_data import *

"""
Explanation: Lesson 6 RNN (Code Along)

2018/7/22|8/10  –– Wayne H Nixalo
End of explanation
"""


# PATH = Path('data/nietzsche/')
PATH = 'data/nietzsche/'

get_data("https://s3.amazonaws.com/text-datasets/nietzsche.txt", f'{PATH}nietzsche.txt')
text = open(f'{PATH}nietzsche.txt').read()
print('corpus length:', len(text))

text[:400]

chars = sorted(list(set(text)))
vocab_size = len(chars) + 1
print('total chars', vocab_size)

"""
Explanation: 1. Setup

We're going to download the collected works of Nietzsche to use as our data for this class.
End of explanation
"""


chars.insert(0, '\0')
''.join(chars[1:-5])

"""
Explanation: Sometimes it's useful to have a zero value in the dataset, eg: for padding.
End of explanation
"""


char_indices = {c: i for i, c in enumerate(chars)}
indices_char = {i: c for i, c in enumerate(chars)}

"""
Explanation: Map from chars to indices and back again:
End of explanation
"""


idx = [char_indices[c] for c in text]
idx[:10]

''.join(indices_char[i] for i in idx[:70])

"""
Explanation: idx will be the data we use form now on – it simply converts all characters to their index (based on the mapping above).
End of explanation
"""


cs = 3
c1_dat = [idx[i]   for i in range(0, len(idx)-cs, cs)] # every 1st char
c2_dat = [idx[i+1] for i in range(0, len(idx)-cs, cs)] # every 2nd
c3_dat = [idx[i+2] for i in range(0, len(idx)-cs, cs)] # every 3rd
c4_dat = [idx[i+3] for i in range(0, len(idx)-cs, cs)] # every 4th

"""
Explanation: 2. Three char model
2.1 Create inputs

Create a list of every 4th character, starting at the 0th, 1st, 2nd, then 3rd characters.
End of explanation
"""


x1 = np.stack(c1_dat)
x2 = np.stack(c2_dat)
x3 = np.stack(c3_dat)

"""
Explanation: Our inputs:
End of explanation
"""


y = np.stack(c4_dat)

"""
Explanation: Our outputs:
End of explanation
"""


x1[:4], x2[:4], x3[:4]

y[:4]

x1.shape, y.shape

"""
Explanation: The first 4 inputs and outputs:
End of explanation
"""


n_hidden = 256

"""
Explanation: 2.2 Create and train model

Pick a size for our hidden state
End of explanation
"""


n_fac = 42 # about half the number of our characters

'0.3' in torch.__version__

class Char3Model(nn.Module):
    def __init__(self, vocab_size, n_fac):
        super().__init__()
        self.e = nn.Embedding(vocab_size, n_fac) # embedding
        
        # the 'green arrow' from our diagram – the layer operation from input to hidden
        self.l_in = nn.Linear(n_fac, n_hidden)
        
        # the 'orange arrow' from our diagram – the layer operation from hidden to hidden
        self.l_hidden = nn.Linear(n_hidden, n_hidden)
        
        # the 'blue arrow' from our diagram – the layer operation from hidden to output
        self.l_out = nn.Linear(n_hidden, vocab_size)
        
    def forward(self, c1, c2, c3):
        in1 = F.relu(self.l_in(self.e(c1)))
        in2 = F.relu(self.l_in(self.e(c2)))
        in3 = F.relu(self.l_in(self.e(c3)))
        
        if '0.3' in torch.__version__:
            h = V(torch.zeros(in1.size()).cuda())
            h = F.tanh(self.l_hidden(h+in1))
            h = F.tanh(self.l_hidden(h+in2))
            h = F.tanh(self.l_hidden(h+in3))
        else:
            h = torch.zeros(in1.size()).cuda() # I dont think I have to wrap as Variable since this is pytorch 0.4, no?
            h = torch.tanh(self.l_hidden(h + in1))
            h = torch.tanh(self.l_hidden(h + in2))
            h = torch.tanh(self.l_hidden(h + in3))
        
        return F.log_softmax(self.l_out(h))

mdata = ColumnarModelData.from_arrays('.', [-1], np.stack([x1,x2,x3], axis=1), y, bs=512)
model = Char3Model(vocab_size, n_fac).cuda()

it = iter(mdata.trn_dl)
*xs,yt = next(it)
# tensor = model(*xs)
tensor = model(*V(xs))

optimizer = optim.Adam(model.parameters(), 1e-2)

set_lrs(optimizer, 1e-3)
fit(model, mdata, 1, optimizer, F.nll_loss)

set_lrs(optimizer, 1e-3)
fit(model, mdata, 1, optimizer, F.nll_loss)

"""
Explanation: The number of latent factors to create (ie: size of the embedding matrix):
End of explanation
"""


def get_next(inp):
    """
    Takes a 3-char string. 
    Turns it into a Tensor of an array of the char index of the string.
    Passes that tensor to the model.
    Does an argmax to get the predicted char-number; then coverts to char.
    """
    idxs = T(np.array([char_indices[c] for c in inp]))
#     pred = model(*idxs)
    pred = model(*VV(idxs))
    i = np.argmax(to_np(pred))
    return chars[i]

get_next('y. '), get_next('ppl'), get_next(' th'), get_next('and')

"""
Explanation: 2.3 Test model
End of explanation
"""


cs = 8

"""
Explanation: 3. Our first RNN
Lecture 6
3.1 Create inputs

This is the size of our unrolled RNN:
End of explanation
"""


c_in_dat = [[idx[i + j] for i in range(cs)] for j in range(len(idx) - cs)]

c_out_dat = [idx[j + cs] for j in range(len(idx) - cs)]

xs = np.stack(c_in_dat, axis=0); xs.shape

y = np.stack(c_out_dat); y.shape

"""
Explanation: For each of 0 thru 8, create a list of every 8th character with that starting point. These will be the 8 inputs to our model.
End of explanation
"""


xs[:cs, :cs]

"""
Explanation: So each column below is one series of 8 characters from the text.
End of explanation
"""


y[:cs]

"""
Explanation: they're overlapping. So after '[42, 29, 30, 25, 27, 29,  1,  1]' comes '1', and after '[29, 30, 25, 27, 29,  1,  1,  1]' comes '43', and so on. The nth row is the same as the nth column.

...and this is the next character after each sequence
End of explanation
"""


val_idx = get_cv_idxs(len(idx) - cs - 1)

mdata = ColumnarModelData.from_arrays('.', val_idx, xs, y, bs=512)

class CharLoopModel(nn.Module):
    """This is an RNN."""
    def __init__(self, vocab_size, n_fac):
        super().__init__()
        self.e = nn.Embedding(vocab_size, n_fac)
        self.l_in = nn.Linear(n_fac, n_hidden)
        self.l_hidden = nn.Linear(n_hidden, n_hidden)
        self.l_out = nn.Linear(n_hidden, vocab_size)
        
    def forward(self, *cs):
        bs = cs[0].size(0)
#         h  = torch.zeros(bs, n_hidden).cuda()
        h = V(torch.zeros(bs, n_hidden).cuda())
        for c in cs:
#             inp = torch.tanh(self.l_in(self.e(c)))   # the torch.tanh vs F.tanh warning didnt pop
#             h   = torch.tanh(self.l_hidden(h + inp)) # up on Mac, but did on Linux-gpu. Odd.
            inp = F.relu(self.l_in(self.e(c)))
            h   = F.tanh(self.l_hidden(h+inp))
        
        return F.log_softmax(self.l_out(h), dim=-1)

model = CharLoopModel(vocab_size, n_fac).cuda()
optimizer = optim.Adam(model.parameters(), 1e-2)

fit(model, mdata, 1, optimizer, F.nll_loss)

set_lrs(optimizer, 1e-3)
fit(model, mdata, 1, optimizer, F.nll_loss)

"""
Explanation: 3.2 Create and train model
End of explanation
"""


class CharLoopConcatModel(nn.Module):
    def __init__(self, vocab_size, n_fac):
        super().__init__()
        self.e = nn.Embedding(vocab_size, n_fac)
        self.l_in = nn.Linear(n_fac + n_hidden, n_hidden)
        self.l_hidden = nn.Linear(n_hidden, n_hidden)
        self.l_out = nn.Linear(n_hidden, vocab_size)
        
    def forward(self, *cs):
        bs = cs[0].size(0)
#         h  = torch.zeros(bs, n_hidden).cuda()
        h = V(torch.zeros(bs, n_hidden).cuda())
        for c in cs:
            inp = torch.cat((h, self.e(c)), 1)
            inp = F.relu(self.l_in(inp))
#             h   = torch.tanh(self.l_hidden(inp))
            h   = F.tanh(self.l_hidden(inp))
            
        return F.log_softmax(self.l_out(inp), dim=-1)

model = CharLoopConcatModel(vocab_size, n_fac).cuda()
optimizer = optim.Adam(model.parameters(), 1e-3)

it = iter(mdata.trn_dl)
*xs,yt = next(it)
# t = model(*xs)
t = model(*V(xs))

xs[0].size(0)

t

fit(model, mdata, 1, optimizer, F.nll_loss)

set_lrs(optimizer, 1e-4)
fit(model, mdata, 1, optimizer, F.nll_loss)

"""
Explanation: The input and hidden states represent qualitatively different types of information, so adding them together can potentially lose information. Instead we can concatenate them together.
End of explanation
"""


if '0.3' in torch.__version__:
    def get_next(inp):
        idxs = T(np.array([char_indices[c] for c in inp]))
        p = model(*VV(idxs))
        i = np.argmax(to_np(p))
        return chars[i]
else:
    def get_next(inp):
    #     idxs = [T(np.array([char_indices[c] for c in inp]))]
        idxs = [T(np.array([char_indices[c]])) for c in inp]
        p = model(*idxs)
        i = np.argmax(to_np(p))
    #     pdb.set_trace()
        return chars[i]

get_next('for thos')

get_next('part of ')

get_next('queens a')

"""
Explanation: 3.3 Test Model
End of explanation
"""


class CharRNN(nn.Module):
    def __init__(self, vocab_size, n_fac):
        super().__init__()
        self.e = nn.Embedding(vocab_size, n_fac)
        self.rnn = nn.RNN(n_fac, n_hidden)
        self.l_out = nn.Linear(n_hidden, vocab_size)
        
    def forward(self, *cs):
        bs = cs[0].size(0)
#         h  = torch.zeros(1, bs, n_hidden)
        h = V(torch.zeros(1, bs, n_hidden))
        inp = self.e(torch.stack(cs))
        outp,h = self.rnn(inp, h)
        
        return F.log_softmax(self.l_out(outp[-1]))
#         return F.log_softmax(self.l_out(outp[-1]), dim=-1) # outp[-1] to get last hidden state

model = CharRNN(vocab_size, n_fac).cuda()
optimizer = optim.Adam(model.parameters(), 1e-3)

it = iter(mdata.trn_dl)
*xs,yt = next(it)

# tensor = model.e(V(torch.stack(xs))) # works w/o V(.). but takes longer when switching btwn w/wo V(.)?
# tensor = model.e(torch.stack(xs)) # these are ints so cannot require gradients
# tensor = model.e(T(torch.stack(xs)))
tensor = model.e(V(torch.stack(xs)))
tensor.size()

# htensor = V(torch.zeros(1, 512, n_hidden)) # V(.) required here, else: RuntimeError: CuDNN error: CUDNN_STATUS_EXECUTION_FAILED
# NOTE: does not work: htensor = torch.zeros(1, 512, n_hidden, requires_grad=True) # requires_grad=True accomplishes what V(.) did in 0.3.1 for 0.4.
# htensor = T(torch.zeros(1, 512, n_hidden))
htensor = V(torch.zeros(1, 512, n_hidden))

outp, hn = model.rnn(tensor, htensor)
outp.size(), hn.size()

"""
Explanation: 4. RNN with PyTorch
Lecture 6, 1:48:52
End of explanation
"""


# the error when using pytorch 0.4:
tensor = model(*V(xs)); tensor.size()

tensor = model(*V(xs)); tensor.size()

fit(model, mdata, 4, optimizer, F.nll_loss)

set_lrs(opt, 1e-4)
fit(model, mdata, 2, optimizer, F.nll_loss)

"""
Explanation: I'm able to get this far in pytorch 0.4, using T instead of V. The problem is the next line keeps giving me a:
<div class="alert alert-info" style="border: 3px solid #900C3F; color: #900C3F; background-color: #F08080">
RuntimeError: CuDNN error: CUDNN_STATUS_EXECUTION_FAILED
</div>

As per here, I'm going to use pytorch 0.3 from here to the end.
End of explanation
"""


def get_next(inp):
    idxs = T(np.array([char_indices[c] for c in inp]))
    p = model(*VV(idxs))
    i = np.argmax(to_np(p))
    return chars[i]

get_next('for thos')

def get_next_n(inp, n):
    res = inp
    for i in range(n):
        c = get_next(inp)
        res += c
        inp = inp[1:] + c
    return res

get_next_n('for thos', 40)

"""
Explanation: 4.1 Test model
End of explanation
"""


c_in_dat = [[idx[i+j] for i in range(cs)] for j in range(0, len(idx) - cs - 1, cs)]

"""
Explanation: 5. Multi-output model
5.1 Setup
Lecture 1:58:07

Let's take non-overlapping sets of characters this time.
End of explanation
"""


c_out_dat = [[idx[i+j] for i in range(cs)] for j in range(1, len(idx) - cs, cs)]

xs = np.stack(c_in_dat)
xs.shape

ys = np.stack(c_out_dat)
ys.shape

xs[:cs, :cs]

ys[:cs, :cs]

"""
Explanation: Then create the exact same thing, offset by 1, as our labels.
End of explanation
"""


val_idx = get_cv_idxs(len(xs) - cs - 1)

mdata = ColumnarModelData.from_arrays('.', val_idx, xs, ys, bs=512)

class CharSeqRNN(nn.Module):
    def __init__(self, vocab_size, n_fac):
        super().__init__()
        self.e = nn.Embedding(vocab_size, n_fac)
        self.rnn = nn.RNN(n_fac, n_hidden)
        self.l_out = nn.Linear(n_hidden, vocab_size)
        
    def forward(self, *cs):
        bs = cs[0].size(0)
        h = V(torch.zeros(1, bs, n_hidden))
        inp = self.e(torch.stack(cs))
        outp,h = self.rnn(inp, h)
        return F.log_softmax(self.l_out(outp), dim=-1)

model = CharSeqRNN(vocab_size, n_fac).cuda()
optimizer = optim.Adam(model.parameters(), 1e-3)

it = iter(mdata.trn_dl)
*xst, yt = next(it)

def nll_loss_seq(inp, targ):
    sl,bs,nh = inp.size() # 8 x 512 x nhidden 
    targ = targ.transpose(0,1).contiguous().view(-1)
    return F.nll_loss(inp.view(-1, nh), targ)

fit(model, mdata, 4, optimizer, nll_loss_seq)

set_lrs(opt, 1e-4)

fit(model, mdata, 1, optimizer, nll_loss_seq)

"""
Explanation: 5.2 Create and train model
End of explanation
"""


model     = CharSeqRNN(vocab_size, n_fac).cuda()
optimizer = optim.Adam(model.parameters(), 1e-2)

m.rnn.weight_hh_l0.data.copy_(torch.eye(n_hidden))

fit(model, mdata, 4, optimizer, nll_loss_seq)

set_lrs(optimizer, 1e-3)

fit(model, mdata, 4, opt, nll_loss_seq)

set_lrs(optimizer, 1e-4)

fit(model, mdata, 4, optimizer, nll_loss_seq)

"""
Explanation: 5.3 Identity init
End of explanation
"""


from torchtext import vocab, data

from fastai.nlp import *
from fastai.lm_rnn import *

PATH = 'data/nietzsche/'
TRN_PATH = 'trn/'
VAL_PATH = 'val/'
TRN = f'{PATH}{TRN_PATH}'
VAL = f'{PATH}{VAL_PATH}'

## line counting: https://stackoverflow.com/a/3137099
# $ wc -l nietzsche/nietzsche.txt
## splitting: https://stackoverflow.com/a/2016918
# $ split -l 7947 nietzsche/nietzsche.txt
# $ mv xaa nietzsche/trn.txt
# $ mv xab nietzsche/val.txt

%ls {PATH}

%ls {PATH}trn

TEXT = data.Field(lower=True, tokenize=list) # torchtext
bs = 64; bptt = 8; n_fac = 42; n_hidden = 256

FILES = dict(train=TRN_PATH, validation=VAL_PATH, test=VAL_PATH)
mdata = LanguageModelData.from_text_files(PATH, TEXT, **FILES, bs=bs, bptt=bptt, min_freq=3)

len(mdata.trn_dl), mdata.nt, len(mdata.trn_ds), len(mdata.trn_ds[0].text)

"""
Explanation: 6. Stateful model
Lecture 7
End of explanation
"""


class CharSeqStatefulRNN(nn.Module):
    def __init__(self, vocab_size, n_fac, bs):
        self.vocab_size = vocab_size
        super().__init__()
        self.e = nn.Embedding(vocab_size, n_fac)
        self.rnn = nn.RNN(n_fac, n_hidden)
        self.l_out = nn.Linear(n_hidden, vocab_size)
        self.init_hidden(bs)
        
    def forward(self, cs):
        bs = cs[0].size(0)
        if self.h.size(1) != bs: self.init_hidden(bs)
        outp,h = self.rnn(self.e(cs), self.h)
        self.h = repackage_var(h) # bptt here; throw away hidden state's history
        return F.log_softmax(self.l_out(outp), dim=-1).view(-1, self.vocab_size)
    
    def init_hidden(self, bs): self.h = V(torch.zeros(1, bs, n_hidden))

m = CharSeqStatefulRNN(mdata.nt, n_fac, 512).cuda()
opt = optim.Adam(m.parameters(), 1e-3)

fit(m, mdata, 4, opt, F.nll_loss)

set_lrs(opt, 1e-4)

fit(m, mdata, 4, opt, F.nll_loss)

"""
Explanation: 6.2 RNN
End of explanation
"""


# # From pytorch source:
# def RNNCell(input, hidden, w_ih, w_hh, b_ih, b_hh):
#     return F.tanh(F.linear(input, w_ih, b_ih) + F.linear(hidden, w_hh, b_hh))

class CharSeqStatefulRNN2(nn.Module):
    def __init__(self, vocab_size, n_fac, bs):
        super().__init__()
        self.vocab_size = vocab_size
        self.e = nn.Embedding(vocab_size, n_fac)
        self.rnn = nn.RNNCell(n_fac, n_hidden)
        self.l_out = nn.Linear(n_hidden, vocab_size)
        self.init_hidden(bs)
        
    def forward(self, cs):
        bs = cs[0].size(0)
        if self.h.size(1) != bs: self.init_hidden(bs)
        outp = []
        o = self.h
        for c in cs:
            o = self.rnn(self.e(c), o)
            outp.append(o)
        outp = self.l_out(torch.stack(outp))
        self.h = repackage_var(o)
        return F.log_softmax(outp, dim=-1).view(-1, self.vocab_size)
    
    def init_hidden(self, bs): self.h = V(torch.zeros(1, bs, n_hidden))

m = CharSeqStatefulRNN2(mdata.nt, n_fac, 512).cuda()
opt = optim.Adam(m.parameters(), 1e-3)

fit(m, mdata, 4, opt, F.nll_loss)

"""
Explanation: 6.3 RNN loop
End of explanation
"""


class CharSeqStatefulGRU(nn.Module):
    def __init__(self, vocab_size, n_fac, bs):
        super().__init__()
        self.vocab_size = vocab_size
        self.e = nn.Embedding(vocab_size, n_fac)
        self.rnn = nn.GRU(n_fac, n_hidden)
        self.l_out = nn.Linear(n_hidden, vocab_size)
        self.init_hidden(bs)
        
    def forward(self, cs):
        bs = cs[0].size(0)
        if self.h.size(1) != bs: self.init_hidden(bs)
        outp,h = self.rnn(self.e(cs), self.h)
        self.h = repackage_var(h)
        return F.log_softmax(self.l_out(outp), dim=-1).view(-1, self.vocab_size)
    
    def init_hidden(self, bs): self.h = V(torch.zeros(1, bs, n_hidden))

# # From pytorch source code – for reference

# def GRUCell(input, hidden, w_ih, w_hh, b_ih, b_hh):
#     gi = F.linear(input, w_ih, b_ih)
#     gh = F.linear(hidden, w_hh, b_hh)
#     i_r, i_i, i_n = gi.chunk(3, 1)
#     h_r, h_i, h_n = gh.chunk(3, 1)
    
#     resetgate = F.sigmoid(i_r + h_r)
#     inputgate = F.sigmoid(i_i + h_i)
#     newgate = F.tanh(i_h + resetgate * h_n)
#     return newgate + inputgate * (hidden - newgate)

m = CharSeqStatefulGRU(mdata.nt, n_fac, 512).cuda()

opt = optim.Adam(m.parameters(), 1e-3)

fit(m, mdata, 6, opt, F.nll_loss)

set_lrs(opt, 1e-4)

fit(m, mdata, 3, opt, F.nll_loss)

"""
Explanation: 6.4 GRU
Lecture 6 46:45
End of explanation
"""


from fastai import sgdr

n_hidden = 512

class CharSeqStatefulLSTM(nn.Module):
    def __init__(self, vocab_size, n_fac, bs, nl):
        super().__init__()
        self.vocab_size,self.nl = vocab_size,nl
        self.e = nn.Embedding(vocab_size, n_fac)
        self.rnn = nn.LSTM(n_fac, n_hidden, nl, dropout=0.5)
        self.l_out = nn.Linear(n_hidden, vocab_size)
        self.init_hidden(bs)
        
    def forward(self, cs):
        bs = cs[0].size(0)
        if self.h[0].size(1) != bs: self.init_hidden(bs)
        outp,h = self.rnn(self.e(cs), self.h)
        self.h = repackage_var(h)
        return F.log_softmax(self.l_out(outp), dim=-1).view(-1, self.vocab_size)
    
    def init_hidden(self, bs):
        self.h = (V(torch.zeros(self.nl, bs, n_hidden)),
                  V(torch.zeros(self.nl, bs, n_hidden)))

m = CharSeqStatefulLSTM(mdata.nt, n_fac, 512, 2).cuda()
lo = LayerOptimizer(optim.Adam, m, 1e-2, 1e-5)

os.makedirs(f'{PATH}models', exist_ok=True)

fit(m, mdata, 2, lo.opt, F.nll_loss)

on_end = lambda sched, cycle: save_model(m, f'{PATH}models/cyc_{cycle}')
cb = [CosAnneal(lo, len(mdata.trn_dl), cycle_mult=2, on_cycle_end=on_end)]
fit(m, mdata, 2**4-1, lo.opt, F.nll_loss, callbacks=cb)

on_end = lambda sched, cycle: save_model(m, f'{PATH}models/cyc_{cycle}')
cb = [CosAnneal(lo, len(mdata.trn_dl), cycle_mult=2, on_cycle_end=on_end)]
fit(m, mdata, 2**6-1, lo.opt, F.nll_loss, callbacks=cb)

"""
Explanation: 6.5 Putting it all together LSTM
End of explanation
"""


def get_next(inp):
    idxs = TEXT.numericalize(inp)
    p = m(VV(idxs.transpose(0,1)))
    r = torch.multinomial(p[-1].exp(), 1)
    return TEXT.vocab.itos[to_np(r)[0]]

get_next('for thos')

def get_next_n(inp, n):
    res = inp
    for i in range(n):
        c = get_next(inp)
        res += c
        inp = inp[1:] + c
    return res

print(get_next_n('for thos', 400))

print(get_next_n('the reason', 400))

"""
Explanation: 6.6 Test
End of explanation
"""


