<jupyter_start><jupyter_text>Prelims<jupyter_code>%matplotlib inline
%reload_ext autoreload
%autoreload 2
from fastai import *
from fastai.vision import *
from fastai.text import *
from fastai.callbacks.tracker import *
PATH = Path('drive/MyDrive/Thesis/data/iam')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device<jupyter_output><empty_output><jupyter_text>Helpers<jupyter_code>def show_img(im, figsize=None, ax=None, alpha=None, title=None):
    if not ax: fig,ax = plt.subplots(figsize=figsize)
    ax.imshow(image2np(im.data), alpha=alpha)
    if title: ax.set_title(title)
    return ax
def rshift(tgt, bos_token=1):
    "Shift y to the right by prepending token"
    bos = torch.zeros((tgt.size(0),1), device=device).type_as(tgt) + bos_token
    return torch.cat((bos, tgt[:,:-1]), dim=-1)

def subsequent_mask(size):
    return torch.tril(torch.ones((1,size,size), device=device).byte())<jupyter_output><empty_output><jupyter_text>Loss, Metrics, Callbacks<jupyter_code>class LabelSmoothing(nn.Module):
    def __init__(self, smoothing=0.1):
        super(LabelSmoothing, self).__init__()
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing
        
    def forward(self, pred, target):
        pred,targ = self.loss_prep(pred, target)
        pred = F.log_softmax(pred, dim=-1)        # necessary for KLDivLoss
        true_dist = pred.data.clone()
        true_dist.fill_(self.smoothing / pred.size(1)) 
        #print(true_dist)  
        #print(targ.shape)      
        true_dist.scatter_(1, targ.data.unsqueeze(1), self.confidence)  # [0.0012, 0.0012, 0.90, 0.0012]
        #print(F.kl_div(pred, true_dist, reduction='sum')/bs)
        return F.kl_div(pred, true_dist, reduction='sum')/bs
    
    def loss_prep(self, pred, target):
        "equalize input/target sl; combine bs/sl dimensions"
        bs,tsl = target.shape
        _ ,sl,vocab = pred.shape

        # F.pad( front,back for dimensions: 1,0,2 )
        if sl>tsl: target = F.pad(target, (0,sl-tsl))
        if tsl>sl: pred = F.pad(pred, (0,0,0,tsl-sl))

        targ = target.contiguous().view(-1).long()
        pred = pred.contiguous().view(-1, vocab)
        return pred, targ
!pip install Levenshtein
import Levenshtein as Lev

class CER(Callback):
    def __init__(self, fn, dump_fn=None):
        super().__init__()
        self.name = 'cer'
        self.fn = fn
        self.dump_fn = dump_fn
        if dump_fn:
          init_csv_dump(dump_fn)

    def on_epoch_begin(self, **kwargs):
        self.errors, self.total = 0, 0
    
    def on_batch_end(self, last_output, last_target, **kwargs):
        error,size = cer(last_output.cpu(), last_target.cpu(), self.fn, self.dump_fn)
        self.errors += error
        self.total += size
    
    def on_epoch_end(self, last_metrics, **kwargs):
        return add_metrics(last_metrics, self.errors/self.total)
class TKER(Callback):
    def __init__(self, fn):
        super().__init__()
        self.name = 'tker'
        self.fn = fn

    def on_epoch_begin(self, **kwargs):
        self.errors, self.total = 0, 0
    
    def on_batch_end(self, last_output, last_target, **kwargs):
        error,size = tker(last_output.cpu(), last_target.cpu(), self.fn)
        self.errors += error
        self.total += size
    
    def on_epoch_end(self, last_metrics, **kwargs):
        return add_metrics(last_metrics, self.errors/self.total)
class WER(Callback):
    def __init__(self, fn):
        super().__init__()
        self.name = 'wer'
        self.fn = fn

    def on_epoch_begin(self, **kwargs):
        self.errors, self.total = 0, 0
    
    def on_batch_end(self, last_output, last_target, **kwargs):
        error,size = wer(last_output.cpu(), last_target.cpu(), self.fn)
        self.errors += error
        self.total += size
    
    def on_epoch_end(self, last_metrics, **kwargs):
        return add_metrics(last_metrics, self.errors/self.total)
def cer(preds, targs, fn, export_fn=None):
    bs = targs.size(0)
    #print("Target tensor: ")
    #print(targs)
    res = torch.argmax(preds, dim=-1)
    error = 0
    for i in range(bs):
        #print(fn(targs[i]))
        p_re = fn(res[i])
        t_re = fn(targs[i])

        p = str(p_re)
        t = str(t_re)

        #print("Char prediction: ", p)
        #print("Char GT: ", t)

        if export_fn:
          with open(export_fn, 'a') as csvfile:
            csv_writer = csv.writer(csvfile, delimiter=' ')
            csv_writer.writerow([p.replace(' ', '|'), t.replace(' ', '|')])

        
        #print("Character Distance: ", Lev.distance(t, p))
        error += Lev.distance(t, p)/(len(t) or 1)
    return error, bs

def init_csv_dump(filename):
    with open(filename, 'w') as csvfile:
        csv_writer = csv.writer(csvfile, delimiter=' ')
        csv_writer.writerow(['prediction', 'ground_truth'])
pip install editdistance
import editdistance

def tker(preds, targs, fn):
    bs = targs.size(0)
    #print("Target tensor: ")
    #print(targs)
    res = torch.argmax(preds, dim=-1)
    error = 0
    for i in range(bs):
        p = fn(res[i])
        t = fn(targs[i])

        #print("Token Prediction: ", p)
        #print("Token Ground truth: ", t)
        
        dist = editdistance.eval(p, t)
        #print("Token Distance: ", dist)
        error += dist / (len(t) or 1)

    return error, bs
import string

def wer(preds, targs, fn):
    bs = targs.size(0)
    #print("Target tensor: ")
    #print(targs)
    res = torch.argmax(preds, dim=-1)
    error = 0
    for i in range(bs):
        #print(fn(targs[i]))
        p_re = fn(res[i])
        t_re = fn(targs[i])

        p = str(p_re)
        t = str(t_re)

        p = p.translate(str.maketrans("", "", string.punctuation))
        t = t.translate(str.maketrans("", "", string.punctuation))

        w_p = [w for w in p.split(' ') if w != '']
        w_t = [w for w in t.split(' ') if w != '']

        #print("Word predictions: ", w_p)
        #print("Word targets: ", w_t)

        dist = editdistance.eval(w_p, w_t)
        #print("Word Distance: ", dist)
        error += dist / (len(w_t) or 1)

    return error, bs
class TeacherForce(LearnerCallback):
    def __init__(self, learn:Learner):
        super().__init__(learn)
        
    def on_batch_begin(self, last_input, last_target, **kwargs):
        return {'last_input':(last_input, last_target), 'last_target':last_target}<jupyter_output><empty_output><jupyter_text>Data  handwriting<jupyter_code>from google.colab import drive
drive.mount('/content/drive')
#handwriting only
fname = 'dataset_short.csv'
FOLDER = 'lines'
def logits(x):
  x[x == 0.0] += 1.0e-06
  x[x == 1.0] -= 1.0e-06

  x = x / (1.0 - x)
  x = np.log(x)

  return x
CSV = PATH/fname
df = pd.read_csv(CSV, delimiter = ' ')

sz,bs = 100,15
seq_len,word_len = 750,300
len(df)<jupyter_output><empty_output><jupyter_text>ModelData<jupyter_code>class MatrixBase(ItemBase):
  def __init__(self, px:Tensor):
      self.data = px
      self.obj = px
      self._logit_px = None

  def clone(self):
      return self.__class__(self.data.clone())

  @property
  def logit_px(self)->LogitTensorImage:
      "Get logit(image.px)."
      if self._logit_px is None: self._logit_px = logit_(self.px)
      return self._logit_px

  def get_logits(self):
    self.data = self.obj - self.logit_px()

  def __str__(self):
    return str(self.data)
class MatrixList(ItemList):
  def get(self, i):
    fn = super().get(i)
    mat = np.genfromtxt(self.path/fn, delimiter=';', names=None)
    res = torch.from_numpy(mat.astype(np.float32, copy=False))
    return MatrixBase(res)<jupyter_output><empty_output><jupyter_text>SentencePiece<jupyter_code>!pip install sentencepiece
import sentencepiece as spm

sp = spm.SentencePieceProcessor()
sp.Load(str(PATH/'spm_iam_corpus.model'))
sp.SetEncodeExtraOptions("eos")
sp.SetDecodeExtraOptions("bos:eos")
def add_cap_tokens(text):  # before encode
    #print(text)
    text = text.replace('|', ' ')
    re_caps = re.compile(r'[A-Z]+')
    return re_caps.sub(_replace_caps, text)
    
def _replace_caps(m):
    tok = '[UP]' if m.end()-m.start() > 1 else '[MAJ]'
    return tok + m.group().lower()

def remove_cap_tokens(text):  # after decode
    text = text.replace(' \u2047 MAJ \u2047 ', '[MAJ]')
    text = text.replace(' \u2047 UP \u2047 ', '[UP]')
    text = text.replace(' \u2047 ', '')
    
    text = re.sub(r'\[UP\]\w+', lambda m: m.group()[4:].upper(), text)  #cap entire word
    text = re.sub(r'\[MAJ\]\w?', lambda m: m.group()[5:].upper(), text) #cap first letter
    
    return text

def remove_cap_pieces(pieces):
    #print("Before decap:", pieces)
    no_cap = []
    #open_tag = False
    for piece in pieces:
      if piece != '<unk>': #and not open_tag:
        no_cap.append(piece)
      #if piece == '<unk>':
      #  open_tag = not open_tag

    
    #print("After decap:", no_cap)
    return no_cap<jupyter_output><empty_output><jupyter_text>Word only (development)<jupyter_code>bs=20
def label_collater(samples:BatchSamples, pad_idx:int=0):
    "Function that collect samples and pads ends of labels."
    data = to_data(samples)
    ims, lbls = zip(*data)
    imgs = torch.stack(list(ims))
    if len(data) is 1 and lbls[0] is 0:   #predict
        labels = torch.zeros(1,1).long()
        return imgs, labels    
    max_len = max([len(s) for s in lbls])
    labels = torch.zeros(len(data), max_len+1).long() + pad_idx  # add 1 to max_len to account for bos token
    for i,lbl in enumerate(lbls):
        labels[i,:len(lbl)] = torch.from_numpy(lbl)  #padding end    
    return imgs, labels
class SPMTokenizer(BaseTokenizer):
    def tokenizer(self, t:str) -> List[str]: 
      #print("From tokenizer: ", t)
      #print("Encoded: ", sp.EncodeAsIds(t))
      return sp.EncodeAsIds(t)[1:]

class SPMProcessor(PreProcessor):
    def __init__(self, ds:ItemList=None, chunksize:int=10000):
        self.toknizr = Tokenizer(tok_func=SPMTokenizer, pre_rules=[rm_useless_spaces, add_cap_tokens],
                                 post_rules=[], special_cases=[])
        self.chunksize = chunksize
        
    def process_one(self, item):
        raise Exception("This isn't implemented!  I didn't think it was necessary...")
    
    def process(self, ds):
        tokens = []
        for i in progress_bar(range(0,len(ds),self.chunksize), leave=False):
            #print(tokens)
            #ds.items = [item.replace('|', ' ') for item in ds.items]
            #print(ds.items)
            tokens += self.toknizr.process_all(ds.items[i:i+self.chunksize])
        ds.items = tokens

class SPMList(ItemList):
    _processor = [SPMProcessor]

    def __init__(self, items:Iterator, sp, **kwargs):
        super().__init__(items, **kwargs)
        self.vocab = sp
        self.pad_idx = 0
        self.copy_new += ['vocab']

    def get(self, i):
        o = super().get(i)
        return Text(o, self.textify(o))

    def reconstruct(self, t:Tensor):
        return Text(t, self.textify(t))
    
    def textify(self, ids):
        if isinstance(ids, torch.Tensor): ids = ids.tolist()
        st = self.vocab.DecodeIds(ids)
        st = remove_cap_tokens(st)
        return st

    def get_pieces(self, ids):
        if isinstance(ids, torch.Tensor): ids = ids.tolist()
        
        pieces = [ self.vocab.IdToPiece(id) for id in ids ]
        pieces = remove_cap_pieces(pieces)

        return pieces
data = (MatrixList.from_df(df, path=PATH/FOLDER)
        .split_by_rand_pct(valid_pct=0.20, seed=42)
        .label_from_df(label_cls=SPMList, sp=sp)
        .databunch(bs=bs, device=device, collate_fn=label_collater)
        )
data<jupyter_output><empty_output><jupyter_text>Transformer Modules<jupyter_code># LayerNorm = nn.LayerNorm
LayerNorm = partial(nn.LayerNorm, eps=1e-4)  # eps: 1e-4 accomodates mixed precision training
class SublayerConnection(nn.Module):
    "A residual connection followed by a layer norm.  Note: (for code simplicity) norm is first."
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x)))
def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([deepcopy(module) for _ in range(N)])
class Encoder(nn.Module):
    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)
        
    def forward(self, x):
        #print("Shape of tensor passed to the encoding layers: ")
        #print(x.shape)
        for layer in self.layers:
            x = layer(x)
        return self.norm(x)
class EncoderLayer(nn.Module):
    "Encoder: self-attn and feed forward"
    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 2)

    def forward(self, x, mask=None):
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)
class Decoder(nn.Module):
    "Generic N layer decoder with masking."
    def __init__(self, layer, N):
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)
        
    def forward(self, x, src, tgt_mask=None):
        for layer in self.layers:
            x = layer(x, src, tgt_mask)
        return self.norm(x)
class DecoderLayer(nn.Module):
    "Decoder: self-attn, src-attn, and feed forward"
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 3)  # wraps layer in residual,dropout,norm
 
    def forward(self, x, src, tgt_mask=None):
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))  # acts as a weak LM
        x = self.sublayer[1](x, lambda x: self.src_attn(x, src, src))
        return self.sublayer[2](x, self.feed_forward)
def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    depth = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(depth)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e4)  #changed from: -1e9 to accomodate mixed precision  
    p_attn = F.softmax(scores, dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn
class SingleHeadedAttention(nn.Module):
    def __init__(self, d_model, dropout=0.2):
        super(SingleHeadedAttention, self).__init__()
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, query, key, value, mask=None):        
        query, key, value = [l(x) for l, x in zip(self.linears, (query, key, value))]
        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)
        return self.linears[-1](x)
class MultiHeadedAttention(nn.Module):
    def __init__(self, d_model, h=8, dropout=0.2):
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        self.d_k = d_model // h        # assume d_v always equals d_k
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, q, k, v, mask=None):
        if mask is not None: mask = mask.unsqueeze(1)
        bs = q.size(0)
        
        # 1) Do all the linear projections in batch from d_model => h x d_k 
        q, k, v = [l(x).view(bs, -1, self.h, self.d_k).transpose(1,2) for l, x in zip(self.linears, (q, k, v))]
        
        # 2) Apply attention on all the projected vectors in batch. 
        x, self.attn = attention(q, k, v, mask=mask, dropout=self.dropout)
        
        # 3) "Concat" using a view and apply a final linear. 
        x = x.transpose(1, 2).contiguous().view(bs, -1, self.h * self.d_k)
        return self.linears[-1](x)
class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, dropout=0.2):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_model*4)
        self.w_2 = nn.Linear(d_model*4, d_model)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x):
        return self.w_2(self.dropout(F.gelu(self.w_1(x))))
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=2000):
        super(PositionalEncoding, self).__init__()
        
        self.dropout = nn.Dropout(p=dropout)
        
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0.0, max_len).unsqueeze(1)
        log_increment = math.log(1e4) / d_model
        div_term = torch.exp(torch.arange(0.0, d_model, 2) * -log_increment)  
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe.unsqueeze_(0)

        self.register_buffer('pe', pe)    #(1,max_len,d_model)
        # registered buffers are Tensors (not Variables)
        # not a parameter but still want in the state_dict

    def forward(self, x):
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)
class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)<jupyter_output><empty_output><jupyter_text>SentencePiece Architectures  Word Arch<jupyter_code>class LearnedPositionalEmbeddings(nn.Module):
    def __init__(self, d_model, vocab, dropout=0.1):
        super(LearnedPositionalEmbeddings, self).__init__()
        self.nl_tok  = 4
        self.d_model = d_model

        self.embed = nn.Embedding(vocab, d_model, 0)
        self.rows = nn.Embedding(15, d_model//2, 0)
        self.w_cols = nn.Embedding(60, d_model//2, 0)
        
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, x):
        rows,cols = self.encode_spatial_positions(x)
        
        row_t = self.rows(rows)            
        col_t = self.w_cols(torch.clamp(cols, max=self.w_cols.num_embeddings-1))  # clamp to max column value
        pos_enc = torch.cat((row_t, col_t), dim=-1)
                
        x = self.embed(x)
        x = (x + pos_enc) * math.sqrt(self.d_model)
        return self.dropout(x)
    
    def encode_spatial_positions(self, x):
        rows,cols = torch.zeros_like(x),torch.zeros_like(x)
        for ii,batch in enumerate(x.unbind()):
            nls = torch.nonzero(batch==self.nl_tok).flatten()
            last = torch.nonzero(batch).flatten()[-1][None]
            splits = torch.cat([nls,last])

            p=0
            for i,n in enumerate(splits, start=1):
                rows[ii,p:n+1] = i
                cols[ii,p:n+1] = torch.arange(1,n-p+2)
                p = n+1
        return rows,cols
class ResnetBase(nn.Module):
    def __init__(self, em_sz):
        super().__init__()
        
        slices = {128: -4, 256: -3, 512: -2}
        s = slices[em_sz]

        net = models.resnet18(True)
        modules = list(net.children())[:s]
        self.base = nn.Sequential(*modules)                  #32x32 : 256
        
    def forward(self, x):
        # IVO
        return x
        return self.base(x)
class Adaptor(nn.Module):
    def forward(self, x):
        # IVO check out this too
        # NB! It might be to adapt the img data to a sequentially processed one
        # It might help you with passing it from the matrices
        #print("Before adapting")
        #print(x.shape)
        return x
        x = x.flatten(2,3).permute(0,2,1)

        #assert False
        return x.mul(8)
        ## IVO
        # return x OR
        # return logit(x)
class WordTransformer(nn.Module):
    def __init__(self, encoder, decoder, embeddings, generator):
        super(WordTransformer, self).__init__()
        self.encoder = encoder
        self.w_decoder = decoder
        self.embed = embeddings
        self.generator = generator
            
    def forward(self, src, tgt):
        #print("Source to the transformer")
        #print(src)
        #print(src.dtype)
        tgt = rshift(tgt, 1).long()
        mask = subsequent_mask(tgt.size(-1))
        return self.w_decoder(self.embed(tgt), self.encoder(src), mask)

    def generate(self, outs):
        return self.generator(outs)
def make_full_model(vocab, d_model, N=4, drops=0, heads=8):
    c = deepcopy
    attn = MultiHeadedAttention(d_model, heads)
    ff = PositionwiseFeedForward(d_model, drops)
    
    model = WordTransformer(
        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),
        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),
        LearnedPositionalEmbeddings(d_model, vocab, drops),  #word
        nn.Linear(d_model, vocab),
    )
        
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
                    
    return model
class Img2Seq(nn.Module):
    def __init__(self, adaptor, transformer): #img_encoder, 
        super(Img2Seq, self).__init__()
        #self.img_enc = img_encoder
        self.adaptor = adaptor
        self.transformer = transformer
        
    def forward(self, src, tgt=None, seq_len=300):
        if tgt is not None:   #train
            feats = self.adaptor(src)#self.img_enc(
            #print("Shape of input: ")
            #print(src)
            #print(src.shape)

            #print("Shape of the features after adapting: ")
            #print(feats)
            #print(feats.shape)
            # IVO Here is where you trace the data
            outs = self.transformer(feats, tgt)
            return self.transformer.generate(outs)
        else:                 #predict
            self.eval()
            with torch.no_grad():
                feats = self.transformer.encoder(self.adaptor(src))
                tgt = torch.ones((src.size(0),1), dtype=torch.long, device=device)

                res = []
                for i in progress_bar(range(seq_len)):
                    emb = self.transformer.embed(tgt)
                    #mask = subsequent_mask(tgt.size(-1))
                    dec_outs = self.transformer.w_decoder(emb, feats)#, mask)
                    prob = self.transformer.generate(dec_outs[:,-1])
                    res.append(prob)
                    pred = torch.argmax(prob, dim=-1, keepdim=True)
                    if (pred==0).all(): break
                    tgt = torch.cat([tgt,pred], dim=-1)
                return torch.stack(res).transpose(1,0).contiguous()
def make_learner(data, d_model, em_sz, N=4, drops=0.1, heads=8, smoothing=0.1, dump_name=None):
    #img_encoder = ResnetBase(em_sz)
    ## IVO 
    ## em_sz is the output of the CNN or the dimension of the features
    adaptor = Adaptor()
    transformer = make_full_model(10000, d_model, N, drops, heads)
    net = Img2Seq(adaptor, transformer)
    dump_file = None if not dump_name else PATH/dump_name
    learn = Learner(data, net, loss_func=LabelSmoothing(smoothing), 
                    metrics=[CER(data.y.reconstruct, dump_file),
                             TKER(data.y.get_pieces),
                             WER(data.y.reconstruct)], 
                    callback_fns=[TeacherForce])  #remove for export
    return learn
bs = 20
emb_size = 128
model_size = 128
learn = make_learner(data, model_size, emb_size, N=4, drops=0.1, heads=1, dump_name='dump/dump_1.0.csv')
# number of trainable parameters
sum(p.numel() for p in learn.model.parameters() if p.requires_grad)
# 50,873,424
# sd = torch.load(PATH/'models/font_sp10k.pth', map_location=device)
# learn.model.load_state_dict(sd['model'], strict=False)<jupyter_output><empty_output><jupyter_text>Train<jupyter_code>lrs = 1e-2
learn.fit_one_cycle(2, max_lr=lrs, start_epoch=0, callbacks=[SaveModelCallback(learn, every='epoch', name='full_model_1.0')])

# arch: (512, 512, N=4, drops=0.1, heads=8); subsequent_masks; lpe
# sentence_piece10k

# data: small dataset, sz:256, bs:50

# combo, 67.7M, 5cycle(1e-3)   
# 10.024531	10.917119	0.059539	0.100420	08:22
# chars:    1.71376   .10433
# words:    1.56690   .06228


# data: font generated, sz:512, bs:15

# combo, 67.7M, 3cycle(1e-3)
# 31.972498	28.244387	0.017639	0.037072	1:45:40   'font_sp10k'
# chars:    32.2701   .01937
# words:    41.7109   .02917


# data: handwriting, sz:512, bs:15

# combo, 67.7M, 5cycle(1e-3)
# 14.077939	11.943507	0.006205	0.003816	21:38   'hw_sp10k'
# chars:    50.1363   .00435
# words:    5.17031   .00043
# test pg:
# chars:    144.698   .03923
# words:    80.1309   .04630
# test upload:
# chars:    129.044   .26124
# words:    61.2923   .30572


# combo, 2cycle(1e-4), preload hw_sp10k
# 12.602509	11.120279	0.005010	0.002936	20:04    'hw_sp10k5'
# test pg:
# chars:    152.536   .03506
# words:    86.5458   .04597
# test upload:
# chars:    131.504   .35796
# words:    59.7730   .27973

# word only, 3cycle(1e-4), preload hw_sp10k
# 6.081270	6.041916	0.000911	14:51    'word_hw_sp10k' ****
#     pg:    82.7050   .04591
#     pg:    83.1463   .04556  no mask
# upload:    64.2214   .31966
# upload:    61.1524   .30548  no mask
curr_version = '0.1'
learn.save('transformer_ctc_' + curr_version)
learn.recorder.plot_losses()
assert False
learn.lr_find()
learn.recorder.plot(suggestion=True)
learn.recorder.plot(suggestion=True)
#learn.load('hw_sp10k'); None<jupyter_output><empty_output><jupyter_text>K-fold validation<jupyter_code>num_folds = 2
val_split = 0.2
part_size = math.floor(len(df) * val_split)
num_epochs = 3
import random

idxs = [*range(len(df))]
random.shuffle(idxs)

err_rates = []

err_rates_log = PATH/'err_rates_1.0.csv'

with open(err_rates_log, 'w') as csvfile:
  csv_writer = csv.writer(csvfile, delimiter=',')
  csv_writer.writerow(['fold', 'loss', 'cer', 'tkerr', 'wer'])

for k in range(num_folds):
  print("Fold #", k+1)
  
  start_int = k * part_size
  end_int = start_int + part_size

  valid_idx = idxs[start_int:end_int]
  train_idx = idxs[:start_int] + idxs[end_int:]

  data = (MatrixList.from_df(df, path=PATH/FOLDER)
        .split_by_idxs(train_idx=train_idx, valid_idx=valid_idx)
        .label_from_df(label_cls=SPMList, sp=sp)
        .databunch(bs=bs, device=device, collate_fn=label_collater)
        )
  
  dump_fn = 'dump/dump_k_fold_' + str(k) + '.csv'
  model_name = 'model_k_fold_' + str(k)
  
  learn = make_learner(data, model_size, emb_size, 
                       N=4, 
                       drops=0.1, 
                       heads=1, 
                       dump_name=dump_fn)
  
  learn.fit_one_cycle(num_epochs, max_lr=lrs, start_epoch=0, 
                      callbacks=[SaveModelCallback(
                          learn, every='epoch', name=model_name)])

  acc = learn.validate(data.valid_dl)
  print("Loss, CER, tkERR and WER for fold ", k, ": ", acc)

  err_rates.append(acc)

  with open(err_rates_log, 'a') as csvfile:
            csv_writer = csv.writer(csvfile, delimiter=',')
            csv_writer.writerow([k] + acc)

#for train_idx, val_idx in skf.split(df['label'],df['name']):
#  print(train_idx, val_idx)<jupyter_output><empty_output><jupyter_text>View Model Telemetry<jupyter_code>class FullStats(HookCallback):
    def on_train_begin(self, **kwargs):
        self.modules = [m for m in flatten_model(self.learn.model) if hasattr(m, 'weight')]
        self.g_hooks = Hooks(self.modules, self.g_hook, is_forward=False)
        self.a_hooks = Hooks(self.modules, self.a_hook)
        self.grads,self.acts = [],[]

    def g_hook(self, m:nn.Module, i:Tensors, o:Tensors)->Tuple[Rank0Tensor,Rank0Tensor]:
        oo = next(o)
        return oo.mean().item(),oo.std().item()
    
    def a_hook(self, m:nn.Module, i:Tensors, o:Tensors)->Tuple[Rank0Tensor,Rank0Tensor]:
        return o.mean().item(),o.std().item()

    def on_batch_end(self, train, **kwargs):
        if train:
            self.acts.append(self.a_hooks.stored)
            self.grads.append(self.g_hooks.stored)
            
    def on_train_end(self, **kwargs):
        self.a_hooks.remove()
        self.g_hooks.remove()
        self.acts = tensor(self.acts).permute(2,1,0)
        self.grads = tensor(self.grads).permute(2,1,0)
learn.fit(1, 1e-5, callbacks=[FullStats(learn)])#, StopAfterNBatches(n_batches=2)])
acts,grads = learn.full_stats.acts, learn.full_stats.grads
acts.shape,grads.shape
names=[]
for name, param in learn.model.named_parameters():
    if name.endswith('weight'):
        names.append(name)

names.insert(193, 'transformer.w_embed.embed.weight')
names.insert(194, 'transformer.w_embed.rows.weight')
len(names)
# :64     img_enc
# 64:67   adaptor
# 67:84   encoder
# 84:137  c_decoder
# 137:190 w_decoder
# 190:    embeddings/generator

plt.figure(figsize=(20,10))
for l in acts[1,137:190]:
    plt.plot(l)
plt.legend(names[137:190])
avg_act_stds_by_layer = acts[1,:].mean(-1)
avg_grad_stds_by_layer = grads[1,:].mean(-1)
plt.plot(avg_act_stds_by_layer)
plt.plot(avg_grad_stds_by_layer)
for (i,mod),a,g in zip(enumerate(names), avg_act_stds_by_layer, avg_grad_stds_by_layer):
    mod_name = str(mod).split('(')[0]
    print(f"{str(i).ljust(3)} {mod_name.ljust(60)} \
            {str(round(a.item(),5)).ljust(6)} {str(round(g.item(),5)).ljust(6)}")
# Last batch activations by layer

for (i,mod),m,s in zip(enumerate(names), acts[0,:,-1], acts[1,:,-1]):
    mod_name = str(mod).split('(')[0]
    print(f"{str(i).ljust(3)} {mod_name.ljust(50)} \
            {str(round(m.item(),5)).ljust(6)}  {str(round(s.item(),5)).ljust(6)}")<jupyter_output><empty_output><jupyter_text>Char/Word Greedy results<jupyter_code>def greedy_decode(src, model, seq_len, kind='char', bos_tok=1):
    model.eval()
    tfmr = model.transformer
    img_enc = model.img_enc
    adaptor = model.adaptor
    decoder = tfmr.c_decoder if kind=='char' else tfmr.w_decoder
    
    with torch.no_grad():
        feats = tfmr.encoder(adaptor(img_enc(src)))
        bs = src.size(0)
        tgt = torch.zeros((bs,1), dtype=torch.long, device=device) + bos_tok

        res = []
        for i in progress_bar(range(seq_len)):
            mask = subsequent_mask(tgt.size(-1))
            emb = tfmr.embed.encode_one(tgt, kind)
#             emb = tfmr.embed(tgt)
            
            dec_outs = decoder(emb, feats, mask)
            prob = tfmr.generator(dec_outs[:,-1])
            res.append(prob)
            pred = torch.argmax(prob, dim=-1, keepdim=True)
            if (pred==0).all(): break
            tgt = torch.cat([tgt,pred], dim=-1)
        out = torch.stack(res).transpose(1,0).contiguous()
        return out
vdl = iter(learn.data.valid_dl)
x,y = next(vdl)<jupyter_output><empty_output><jupyter_text>Single Word<jupyter_code># g_preds = greedy_decode(x, learn.model, word_len, 'word', 1)
g_preds = learn.model(x)
g_res = torch.argmax(g_preds, dim=-1)
g = [learn.loss_func(g_preds, y).item()/bs, cer(g_preds, y, data.y.reconstruct)[0]/bs]
print(f'greedy:    {str(g[0])[:7]}   {str(g[1])[1:7]}')
#greedy
fig, axes = plt.subplots(2,2, gridspec_kw={'hspace': 0.4}, figsize=(18, 20))
for i,ax in enumerate(axes.flat):
    #i += 8
    p = data.y.reconstruct(g_res[i])
    ax=show_img(x[i], ax=ax, title=p)<jupyter_output><empty_output><jupyter_text>Single Char<jupyter_code>g_preds = greedy_decode(x, learn.model, seq_len, 'char')
g_res = torch.argmax(g_preds, dim=-1)
g = [learn.loss_func(g_preds, y).item()/bs, cer(g_preds, y, data.y.reconstruct)[0]/bs]
print(f'greedy:    {str(g[0])[:7]}   {str(g[1])[1:7]}')
#greedy
fig, axes = plt.subplots(2,2, gridspec_kw={'hspace': 0.4}, figsize=(18, 20))
for i,ax in enumerate(axes.flat):
    p = data.y.reconstruct(g_res[i])
    ax=show_img(x[i], ax=ax, title=p)<jupyter_output><empty_output><jupyter_text>Combo Chars<jupyter_code>g_preds = greedy_decode(x, learn.model, seq_len, 'char', 1)
g_res = torch.argmax(g_preds, dim=-1)
loss_func = LabelSmoothing()
g = [loss_func(g_preds, y[0]).item()/bs, cer(g_preds, y[0], data.y.reconstruct_one)[0]/bs]
print(f'greedy:    {str(g[0])[:7]}   {str(g[1])[1:7]}')
#greedy
fig, axes = plt.subplots(3,2, gridspec_kw={'hspace': 0.4}, figsize=(18, 20))
for i,ax in enumerate(axes.flat):
    p = data.y.reconstruct_one(g_res[i])
    ax=show_img(x[i], ax=ax, title=p)<jupyter_output><empty_output><jupyter_text>Combo Words<jupyter_code>g_preds = greedy_decode(x, learn.model, word_len, 'word', 1, lm=False)
g_res = torch.argmax(g_preds, dim=-1)
loss_func = LabelSmoothing()
g = [loss_func(g_preds, y[1]).item()/bs, cer(g_preds, y[1], data.y.reconstruct_one)[0]/bs]
print(f'greedy:    {str(g[0])[:7]}   {str(g[1])[1:7]}')
#greedy
fig, axes = plt.subplots(3,2, gridspec_kw={'hspace': 0.4}, figsize=(18, 20))
for i,ax in enumerate(axes.flat):
    p = data.y.reconstruct_one(g_res[i])
    ax=show_img(x[i], ax=ax, title=p)<jupyter_output><empty_output><jupyter_text>Test Datasets<jupyter_code>FOLDER = 'uploads'
df = pd.read_csv(PATH/'uploads.csv')
len(df)

sz,bs = 512,14
seq_len = 700
FOLDER = 'paragraphs'
df = pd.read_csv(PATH/'test_pg.csv')
len(df)

sz,bs = 512,15
seq_len = 700
# sentencepiece combo
test_data = (ImageMultiList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)
        .split_none()
        .label_from_df(label_cls=SPMMultiList, sp=sp)
        .transform([], size=sz, resize_method=ResizeMethod.SQUISH)
        .databunch(bs=bs, device=device, collate_fn=multi_label_collater)
       )
# sentencepiece word
test_data = (ImageList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)
        .split_none()
        .label_from_df(label_cls=SPList)
        .transform([], size=sz, resize_method=ResizeMethod.SQUISH)
        .databunch(bs=bs, device=device, collate_fn=label_collater)
       )
x,y = next(iter(test_data.train_dl))<jupyter_output><empty_output><jupyter_text>Export Learner<jupyter_code>learn.export()
learn = load_learner(PATH)
img = open_image(PATH/'test/a03-014.png');
img
res = learn.predict(img)[0]
print(str(res))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14375163057657186, 'fraction_numerical': 0.0272487462677914, 'mean_word_length': 2.917111388668105, 'pattern_counts': {'":': 1, '<': 79, '<?xml version=': 0, '>': 88, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '19', 'dir': 'jupyter-structured-clean-dedup', 'id': '51633', 'license': 'Apache-2.0', 'n_tokens_mistral': 13287, 'n_tokens_neox': 12570, 'n_words': 2504, 'path': 'Transformer_CTC.ipynb', 'repo_name': 'ivaylojelev/fastai1', 'size': '34497'}