# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/02_text.models.core.ipynb.

# %% auto 0
__all__ = ['SequentialRNN', 'SentenceEncoder', 'AttentiveSentenceEncoder', 'masked_concat_pool', 'PoolingLinearClassifier',
           'OurPoolingLinearClassifier', 'LabelAttentionClassifier', 'LabelAttentionClassifier2',
           'LabelAttentionClassifier3', 'get_text_classifier']

# %% ../../../nbs/02_text.models.core.ipynb 1
from fastai.data.all import *
from fastai.text.models.core import *
from fastai.text.models.awdlstm import *
from ...layers import *

# %% ../../../nbs/02_text.models.core.ipynb 4
_model_meta = {AWD_LSTM: {'hid_name':'emb_sz', 'url':URLs.WT103_FWD, 'url_bwd':URLs.WT103_BWD,
                          'config_lm':awd_lstm_lm_config, 'split_lm': awd_lstm_lm_split,
                          'config_clas':awd_lstm_clas_config, 'split_clas': awd_lstm_clas_split},}

# %% ../../../nbs/02_text.models.core.ipynb 8
class SequentialRNN(nn.Sequential):
    "A sequential pytorch module that passes the reset call to its children."
    def reset(self):
        for c in self.children(): getattr(c, 'reset', noop)()

# %% ../../../nbs/02_text.models.core.ipynb 10
def _pad_tensor(t, bs):
    if t.size(0) < bs: return torch.cat([t, t.new_zeros(bs-t.size(0), *t.shape[1:])])
    return t

# %% ../../../nbs/02_text.models.core.ipynb 11
class SentenceEncoder(Module):
    "Create an encoder over `module` that can process a full sentence."
    def __init__(self, bptt, module, pad_idx=1, max_len=None): store_attr('bptt,module,pad_idx,max_len')
    def reset(self): getattr(self.module, 'reset', noop)()

    def forward(self, input):
        bs,sl = input.size()
        self.reset()
        mask = input == self.pad_idx
        outs,masks = [],[]
        for i in range(0, sl, self.bptt):
            #Note: this expects that sequence really begins on a round multiple of bptt
            real_bs = (input[:,i] != self.pad_idx).long().sum()
            o = self.module(input[:real_bs,i: min(i+self.bptt, sl)])
            if self.max_len is None or sl-i <= self.max_len:
                outs.append(o)
                masks.append(mask[:,i: min(i+self.bptt, sl)])
        outs = torch.cat([_pad_tensor(o, bs) for o in outs], dim=1)
        mask = torch.cat(masks, dim=1)
        return outs,mask

# %% ../../../nbs/02_text.models.core.ipynb 14
class AttentiveSentenceEncoder(Module):
    "Create an encoder over `module` that can process a full sentence."
    def __init__(self, bptt, module, pad_idx=1, max_len=None): store_attr('bptt,module,pad_idx,max_len')
    def reset(self): getattr(self.module, 'reset', noop)()

    def forward(self, input):
        bs,sl = input.size()
        self.reset()
        mask = input == self.pad_idx
        outs,masks = [],[]
        for i in range(0, sl, self.bptt):
            #Note: this expects that sequence really begins on a round multiple of bptt
            real_bs = (input[:,i] != self.pad_idx).long().sum()
            o = self.module(input[:real_bs,i: min(i+self.bptt, sl)])
            if self.max_len is None or sl-i <= self.max_len:
                outs.append(o)
                masks.append(mask[:,i: min(i+self.bptt, sl)])
        outs = torch.cat([_pad_tensor(o, bs) for o in outs], dim=1)
        mask = torch.cat(masks, dim=1)
        return outs,mask

# %% ../../../nbs/02_text.models.core.ipynb 17
def masked_concat_pool(output, mask, bptt):
    "Pool `MultiBatchEncoder` outputs into one vector [last_hidden, max_pool, avg_pool]"
    lens = output.shape[1] - mask.long().sum(dim=1)
    last_lens = mask[:,-bptt:].long().sum(dim=1)
    avg_pool = output.masked_fill(mask[:, :, None], 0).sum(dim=1)
    avg_pool.div_(lens.type(avg_pool.dtype)[:,None])
    max_pool = output.masked_fill(mask[:,:,None], -float('inf')).max(dim=1)[0]
    x = torch.cat([output[torch.arange(0, output.size(0)),-last_lens-1], max_pool, avg_pool], 1) #Concat pooling.
    return x

# %% ../../../nbs/02_text.models.core.ipynb 19
class PoolingLinearClassifier(Module):
    "Create a linear classifier with pooling"
    def __init__(self, dims, ps, bptt, y_range=None):
        if len(ps) != len(dims)-1: raise ValueError("Number of layers and dropout values do not match.")
        acts = [nn.ReLU(inplace=True)] * (len(dims) - 2) + [None]
        layers = [LinBnDrop(i, o, p=p, act=a) for i,o,p,a in zip(dims[:-1], dims[1:], ps, acts)]
        if y_range is not None: layers.append(SigmoidRange(*y_range))
        self.layers = nn.Sequential(*layers)
        self.bptt = bptt

    def forward(self, input):
        out,mask = input
        x = masked_concat_pool(out, mask, self.bptt)
        x = self.layers(x)
        return x, out, out

# %% ../../../nbs/02_text.models.core.ipynb 26
class OurPoolingLinearClassifier(Module):
    def __init__(self, dims, ps, bptt, y_range=None):
        self.layer = LinBnDrop(dims[0], dims[1], p=ps, act=None)
        self.bptt = bptt

    def forward(self, input):
        out, mask = input
        x = masked_concat_pool(out, mask, self.bptt)
        x = self.layer(x)
        return x, out, out

# %% ../../../nbs/02_text.models.core.ipynb 30
class LabelAttentionClassifier(Module):
    def __init__(self, dims, ps, bptt, y_range=None):
        self.fts = dims[0]
        self.lbs = dims[-1] 
        self.layers = LinBnDrop(self.lbs, ln=False, p=ps, act=None) # deb
        self.bptt = bptt
        self.emb_label = Embedding(self.lbs, self.fts) # deb: note that size of the label embeddings need not be same as nh 
        self.final_lin = nn.Linear(self.fts, self.lbs) 

    def forward(self, input):
        out, _ = input
        attn_wgts = out @ self.emb_label.weight.transpose(0, 1) # deb
        attn_wgts = F.softmax(attn_wgts, 1) # deb
        ctx = attn_wgts.transpose(1,2) @ out # deb
        x = self.layers(ctx)
        x = (self.final_lin.weight * x).sum(dim=2)
        return x, out, out

# %% ../../../nbs/02_text.models.core.ipynb 31
class LabelAttentionClassifier2(Module):
    initrange=0.1
    def __init__(self, dims, ps, bptt, y_range=None):
        self.fts = dims[0]
        self.lbs = dims[-1] 
        
        # ps = 0.1 # deb
        self.layers = LinBnDrop(self.lbs, ln=False, p=ps, act=None) # deb
        self.bptt = bptt
        # self.emb_label = Embedding(self.lbs, self.fts) # deb: note that size of the label embeddings need not be same as nh 
        self.emb_label = self._init_param(self.lbs, self.fts) # deb: note that size of the label embeddings need not be same as nh 
        self.final_lin = nn.Linear(self.fts, self.lbs) 
        self.final_lin.weight.data.uniform_(-self.initrange, self.initrange)
        self.final_lin.bias.data.zero_()
    
    def _init_param(self, *sz): return nn.Parameter(torch.zeros(sz).normal_(0, 0.01))

    def forward(self, input):
        out, _ = input
        # x = masked_concat_pool(out, mask, self.bptt)
        
        # bs = out.shape[0]
        # ctx = out.new_zeros((bs, self.lbs, self.fts))
        # for out_split in torch.split(out, 1, dim=1):
        # self.emb_label = nn.Parameter(self.emb_label * self.m1)
        attn_wgts = out @ self.emb_label.transpose(0, 1) # deb
        # attn_wgts = sigmoid_range(attn_wgts, 0, 5.5) # did not help
        attn_wgts = F.softmax(attn_wgts, 1) # deb
        # attn_wgts = torch.nn.functional.log_softmax(attn_wgts, 1) # deb
        # attn_wgts = torch.log(attn_wgts)/(attn_wgts.sum(dim=1, keepdim=True) + 1e-12)
        # attn_wgts[torch.isnan(attn_wgts)] = tensor(0.)
        # attn_wgts = torch.nn.functional.normalize(torch.log(attn_wgts), dim=1)
        ctx = attn_wgts.transpose(1,2) @ out # deb
        

        x = self.layers(ctx)
        # x = self.final_lin.weight.mul(x).sum(dim=2).add(self.final_lin.bias) #missed_deb
        x = (self.final_lin.weight * x).sum(dim=2) + self.final_lin.bias
        # x = (self.final_lin.weight * x + self.final_lin.bias.unsqueeze(1)).sum(dim=2)
        
        # x = x.view(x.shape[0], x.shape[1])
        return x, out, out

# %% ../../../nbs/02_text.models.core.ipynb 32
class LabelAttentionClassifier3(Module):
    initrange=0.1
    def __init__(self, dims, ps, bptt, y_range=None):
        self.fts = dims[0]
        self.lbs = dims[-1] 
        
        # ps = 0.1 # deb
        self.layers = LinBnDrop(self.lbs, ln=False, p=ps, act=None) # deb
        self.attn = XMLAttention(self.lbs, self.fts, 0.0)
        self.final_lin = nn.Linear(self.fts, self.lbs) 
        init_default(self.final_lin, 
                     func=partial(torch.nn.init.uniform_, a=-self.initrange, b=self.initrange))
    
    def forward(self, input):
        out, _ = input
        ctx = self.attn(out)
        x = self.layers(ctx)
        x = (self.final_lin.weight * ctx).sum(dim=2) + self.final_lin.bias
        
        return x, out, out

# %% ../../../nbs/02_text.models.core.ipynb 36
def get_text_classifier(arch, vocab_sz, n_class, seq_len=72, config=None, drop_mult=1., lin_ftrs=None,
                       ps=None, pad_idx=1, max_len=72*20, y_range=None):
    "Create a text classifier from `arch` and its `config`, maybe `pretrained`"
    meta = _model_meta[arch]
    config = ifnone(config, meta['config_clas']).copy()
    for k in config.keys():
        if k.endswith('_p'): config[k] *= drop_mult
    if lin_ftrs is None: lin_ftrs = [50]
    if ps is None: ps = [0.1]*len(lin_ftrs) # not required if not using OurPoolingLinearClasifier
#     layers = [config[meta['hid_name']] * 3] + lin_ftrs + [n_class]  # required if using fastai's PoolingLinearClassifier
    layers = [config[meta['hid_name']]] + [n_class]
#     ps = [config.pop('output_p')] + ps
    ps = config.pop('output_p')
    init = config.pop('init') if 'init' in config else None
    encoder = AttentiveSentenceEncoder(seq_len, arch(vocab_sz, **config), pad_idx=pad_idx, max_len=max_len)
    # decoder = OurPoolingLinearClassifier(layers, ps, bptt=seq_len, y_range=y_range)
    decoder = LabelAttentionClassifier3(layers, ps, bptt=seq_len, y_range=y_range)
    model = SequentialRNN(encoder, decoder)
    return model if init is None else model.apply(init)
