txt = 'The U.S. dollar $1 is $1.00.'
spacy = WordTokenizer() #class ที่ทำหน้าที่ tokenizer ภาษาอังกฤษจาก spacy
tkn = Tokenizer(spacy) #สร้าง Tokenizer object ของ fastai
print(coll_repr(tkn(txt), 31)) #print 31 อันจากข้อความที่ถูก tokenize
#xxbos คือ beginning of sentence token (เพราะเราจะเอาข้อความาต่อกันยาวๆเวลาเทรน ต้องมีไว้เพื่อให้รู้ว่าตรงไหนคือเริ่มต้นประโยค)
#xxmaj คือ capitalization token บอกว่า token ต่อไปถูก capitalize ตัวแรก (ในที่นี้คือ The)
#xxup คือ all-caps token บอกว่า token ต่อไปเป็น uppercase ทั้งหมด (ในที่นี้คือ u.s)
#>> ['xxbos','xxmaj','the','xxup','u.s','.','dollar','$','1','is','$','1.00','.']

#วิธีการสร้าง fastai tokenizer ด้วย PythaiNLP (หรือจริงๆ tokenizer อะไรก็ได้)
#สร้าง class สำหรับ tokenizer ให้ผลเป็น list
from pythainlp.tokenize import word_tokenize
class NewMMTokenizer:
    def __init__(self, tok_func):
        self.tok_func = tok_func
    def __call__(self,items):
        return (self.tok_func(t) for t in items) 

#สร้าง newmm object จากคลาสที่เราสร้าง
newmm = NewMMTokenizer(word_tokenize)
#สร้าง tokenizer เหมือนภาษาอังกฤษ
tkn = Tokenizer(newmm)
tkn('สวัสดีครับพี่น้อง')
#>> ['xxbos', ' ', 'สวัสดี', 'ครับ', 'พี่น้อง']

def subword(sz):
    sp = SubwordTokenizer(vocab_sz=sz) #create sentencepiece tokenizer with vocab size of `sz`
    sp.setup(txts) #เทรน sentencepiece
    return '|'.join(first(sp([txt]))[:40]) #รีเทิร์น 40 subwords แรก

# ▁ ในที่นี้แทน space เพราะ subword จะรวบ space เข้ามาด้วย
subword(200)
# >> ▁The|▁|U|.|S|.|▁do|l|l|ar|▁|$|1|▁is|▁|$|1|.|0|0|.
subword(20000)
# >> ▁The|▁U|.|S|.|▁dollar|▁$1|▁is|▁$1|.00|.

toks = tkn(txt)
toks200 = txts[:200].map(tkn) #tokens

num = Numericalize()
num.setup(toks200) #set up using the tokens
nums = num(toks)[:20]; nums #numericalized; เปลี่ยนจาก token เป็น index
' '.join(num.vocab[o] for o in nums) #เปลี่ยนกลับจาก index เป็น token

#batch size = 3
#sequence length = 2
['A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R']

#mini-batches จะเป็น
#batch 0
[
'A','B'
'G','H'
'M','N'
]

#batch 1
[
'C','D'
'I','J'
'O','P'
]

#batch 2
[
'E','F'
'K','L'
'Q','R'
]

#batch 0; X
[
'A','B'
'G','H'
'M','N'
]

#batch 0; y
[
'B','C'
'H','I'
'N','O'
]

dls_clas = DataBlock(
    blocks=(TextBlock.from_folder(path, vocab=dls_lm.vocab),CategoryBlock),
    get_y = parent_label,
    get_items=partial(get_text_files, folders=['train', 'test']),
    splitter=GrandparentSplitter(valid_name='test')
).dataloaders(path, path=path, bs=128, seq_len=72)

#a long text with numbers in English from one to nine thousand nine hundred ninety nine
#"one . two . three . four . five . six . seven . eight . nine . ten . eleven . twelve . thirteen . ..."
text = ' . '.join([l.strip() for l in lines])
tokens = text.split(' ')

#get vocab list
vocab = L(*tokens).unique()
#['one','.','two','three','four','five','six','seven','eight','nine'...]

#numericalize with a dictionary
word2idx = {w:i for i,w in enumerate(vocab)}
nums = L(word2idx[i] for i in tokens)

#create X,y pairs with 3 tokens to predict the next token
seqs = L((tensor(nums[i:i+3]), nums[i+3]) for i in range(0,len(nums)-4,3))
#[(['one', '.', 'two'], '.'),(['.', 'three', '.'], 'four'),(['four', '.', 'five'], '.'),(['.', 'six', '.'], 'seven')]
#[(tensor([0, 1, 2]), 1),(tensor([1, 3, 1]), 4),(tensor([4, 1, 5]), 1),(tensor([1, 6, 1]), 7)]

#create dataloader with batch size = 64
#validation split at 80/20
bs = 64
cut = int(len(seqs) * 0.8)
dls = DataLoaders.from_dsets(seqs[:cut], seqs[cut:], bs=64, shuffle=False) #do not shuffle

#group_chunks ช่วยสร้าง mini-batch ตามวิธีข้างต้น
#index 0 เป็น (tensor([0,1,2]), 1)
#index 1 เป็น (tensor([11,  1,  2]), 28)
#index 64 (batch size) เป็น (tensor([1,3,1]), 4)
#วิธีเช็คว่าเรียง seq เป็น mini-batch ยังไง; หา seq[1] หรือ (tensor([1,3,1]),4)
#dls.train_ds.argwhere(lambda x: (x[1]==4)&(x[0][0].item()==1)&(x[0][1].item()==3)&(x[0][2].item()==1))

def group_chunks(ds, bs):
    m = len(ds) // bs
    new_ds = L()
    for i in range(m): new_ds += L(ds[i + m*j] for j in range(bs))
    return new_ds

cut = int(len(seqs) * 0.8)
dls = DataLoaders.from_dsets(
    group_chunks(seqs[:cut], bs), 
    group_chunks(seqs[cut:], bs), 
    bs=bs, drop_last=True, shuffle=False)

#RNN architecture for language model
class LMModel1(Module):
    def __init__(self, vocab_sz, n_hidden):
        self.i_h = nn.Embedding(vocab_sz, n_hidden) #input layer; takes (batch size, 1) returns (batch size, n_hidden)
        self.h_h = nn.Linear(n_hidden, n_hidden) #hidden layer; takes (batch size, n_hidden) returns (batch size, n_hidden)
        self.h_o = nn.Linear(n_hidden,vocab_sz) #output layer; takes (batch size, n_hidden) returns (batch size, vocab size)
        
    def forward(self, x):
        #x: batch size * 3 tokens
        h = F.relu(self.h_h(self.i_h(x[:,0]))) #put first token through i_h, h_h and ReLU
        h = h + self.i_h(x[:,1]) #put the second token through i_h, then plus with result from last step
        h = F.relu(self.h_h(h)) #put result of last step through ReLU
        h = h + self.i_h(x[:,2]) #put the third token through i_h, then plus with result from last step
        h = F.relu(self.h_h(h)) #put result of last step through ReLU
        return self.h_o(h) #put result of last step through h_o; predicts which token comes next

#Refactored RNN ด้วย loop
class LMModel2(Module):
    def __init__(self, vocab_sz, n_hidden):
        self.i_h = nn.Embedding(vocab_sz, n_hidden)  
        self.h_h = nn.Linear(n_hidden, n_hidden)     
        self.h_o = nn.Linear(n_hidden,vocab_sz)
        
    def forward(self, x):
        h = 0 #initialize h as 0
        for i in range(3):
            h = h + self.i_h(x[:,i])
            h = F.relu(self.h_h(h))
        return self.h_o(h)

#รักษา hidden state ไว้ในแต่ละครั้งที่ทำ forward()
class LMModel3(Module):
    def __init__(self, vocab_sz, n_hidden):
        self.i_h = nn.Embedding(vocab_sz, n_hidden)  
        self.h_h = nn.Linear(n_hidden, n_hidden)     
        self.h_o = nn.Linear(n_hidden,vocab_sz)
        self.h = 0
        
    def forward(self, x):
        for i in range(3):
            self.h = self.h + self.i_h(x[:,i])
            self.h = F.relu(self.h_h(self.h))
        out = self.h_o(self.h)
        #.detach() เพื่อให้ Pytorch คำนวณ gradients แยกเป็นครั้งๆที่ forward ถูกเรียก
        #ไม่งั้น pytorch จะคำนวร gradients ตั้งแต่ token แรกยัน token สุดท้ายใน dataset
        self.h = self.h.detach() 
        return out
    
    def reset(self): self.h = 0

sl = 16
seqs = L((tensor(nums[i:i+sl]), tensor(nums[i+1:i+sl+1]))
         for i in range(0,len(nums)-sl-1,sl))
cut = int(len(seqs) * 0.8)
dls = DataLoaders.from_dsets(group_chunks(seqs[:cut], bs),
                             group_chunks(seqs[cut:], bs),
                             bs=bs, drop_last=True, shuffle=False)

dls.train_ds[0]
# >>(tensor([0, 1, 2, 1, 3, 1, 4, 1, 5, 1, 6, 1, 7, 1, 8, 1]),
# >>tensor([1, 2, 1, 3, 1, 4, 1, 5, 1, 6, 1, 7, 1, 8, 1, 9]))

class LMModel4(Module):
    def __init__(self, vocab_sz, n_hidden):
        self.i_h = nn.Embedding(vocab_sz, n_hidden)  
        self.h_h = nn.Linear(n_hidden, n_hidden)     
        self.h_o = nn.Linear(n_hidden,vocab_sz)
        self.h = 0
        
    def forward(self, x):
        outs = [] #รับ output จากทุก token
        for i in range(sl):
            self.h = self.h + self.i_h(x[:,i])
            self.h = F.relu(self.h_h(self.h))
            outs.append(self.h_o(self.h))
        self.h = self.h.detach()
        return torch.stack(outs, dim=1) #stack output เข้าด้วยกัน
    
    def reset(self): self.h = 0

class LMModel5(Module):
    def __init__(self, vocab_sz, n_hidden, n_layers):
        self.i_h = nn.Embedding(vocab_sz, n_hidden)
        self.rnn = nn.RNN(n_hidden, n_hidden, n_layers, batch_first=True) #เลือกได้ว่าอยากได้กี่ layer ซ้อนกัน
        self.h_o = nn.Linear(n_hidden, vocab_sz)
        self.h = torch.zeros(n_layers, bs, n_hidden)
        
    def forward(self, x):
        res,h = self.rnn(self.i_h(x), self.h)
        self.h = h.detach()
        return self.h_o(res)
    
    def reset(self): self.h.zero_()

class LSTMCell(Module):
    def __init__(self, ni, nh):
        self.forget_gate = nn.Linear(ni + nh, nh)
        self.input_gate  = nn.Linear(ni + nh, nh)
        self.cell_gate   = nn.Linear(ni + nh, nh)
        self.output_gate = nn.Linear(ni + nh, nh)

    def forward(self, input, state):
        h,c = state
        h = torch.cat([h, input], dim=1)
        forget = torch.sigmoid(self.forget_gate(h))
        c = c * forget
        inp = torch.sigmoid(self.input_gate(h))
        cell = torch.tanh(self.cell_gate(h))
        c = c + inp * cell
        out = torch.sigmoid(self.output_gate(h))
        h = out * torch.tanh(c)
        return h, (h,c)

#refactored version
class LSTMCell(Module):
    def __init__(self, ni, nh):
        self.ih = nn.Linear(ni,4*nh)
        self.hh = nn.Linear(nh,4*nh)

    def forward(self, input, state):
        h,c = state
        # One big multiplication for all the gates is better than 4 smaller ones
        gates = (self.ih(input) + self.hh(h)).chunk(4, 1)
        ingate,forgetgate,outgate = map(torch.sigmoid, gates[:3])
        cellgate = gates[3].tanh()

        c = (forgetgate*c) + (ingate*cellgate)
        h = outgate * c.tanh()
        return h, (h,c)

class Dropout(Module):
    def __init__(self, p): self.p = p
    def forward(self, x):
        if not self.training: return x #ถ้าไม่ได้เทรนอยู่ ไม่ต้องทำ dropout
        mask = x.new(*x.shape).bernoulli_(1-p) #สุ่มว่าแต่ละ activation อันไหนที่เราจะเปลี่ยนเป็น 0
        return x * mask.div_(1-p)

loss += alpha * activations.pow(2).mean()

loss += beta * (activations[:,1:] - activations[:,:-1]).pow(2).mean()

class LMModel7(Module):
    def __init__(self, vocab_sz, n_hidden, n_layers, p):
        self.i_h = nn.Embedding(vocab_sz, n_hidden)
        self.rnn = nn.LSTM(n_hidden, n_hidden, n_layers, batch_first=True)
        self.drop = nn.Dropout(p)
        self.h_o = nn.Linear(n_hidden, vocab_sz)
        self.h_o.weight = self.i_h.weight #weight-tying
        self.h = [torch.zeros(n_layers, bs, n_hidden) for _ in range(2)]
        
    def forward(self, x):
        raw,h = self.rnn(self.i_h(x), self.h)
        out = self.drop(raw)
        self.h = [h_.detach() for h_ in h]
        return self.h_o(out),raw,out
    
    def reset(self): 
        for h in self.h: h.zero_()