# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/19_text.callbacks.ipynb.

# %% auto 0
__all__ = ['LabelForcing', 'RunvalCallback']

# %% ../../nbs/19_text.callbacks.ipynb 1
from fastai.torch_imports import *
from fastai.torch_core import *
from fastai.callback.core import *
from fastcore.all import *
from ..imports import *
from ..metrics import *

# %% ../../nbs/19_text.callbacks.ipynb 6
class LabelForcing(Callback):
    def __init__(self, end_epoch):
        self.end_epoch = end_epoch
    
    @torch.no_grad()
    def after_pred(self):
        if self.training and self.epoch <= self.end_epoch:
            pred, attn, wgts = self.learn.pred
            pos = Tensor(self.y) == 1
            # pred[pos] += 3*pred.std()
            # pred[~pos] -= 3*pred.std()
            attn[pos] += attn[pos].std(dim=-1).unsqueeze(-1)
            attn[~pos] -= attn[~pos].std(dim=-1).unsqueeze(-1)

# %% ../../nbs/19_text.callbacks.ipynb 8
from fastai.callback.progress import *
from fastai.learner import Recorder

# %% ../../nbs/19_text.callbacks.ipynb 9
class _FakeLearner: 
    def to_detach(self,b,cpu=True,gather=True):
        return to_detach(b,cpu,gather)

# %% ../../nbs/19_text.callbacks.ipynb 10
def _cpupy(t): return t.cpu().numpy() if isinstance(t, Tensor) else t

# %% ../../nbs/19_text.callbacks.ipynb 11
class RunvalCallback(Callback):
    order=ProgressCallback.order-1
    
    def __init__(self, mets):
        self.mets = mets
        
    def before_train(self): 
        self.val_cyclit = itertools.cycle(self.dls.valid) 
        self._fake_l = _FakeLearner()
        self.mets.map(Self.reset())
        self.len = len(self.dls.valid)
        self._launch_tqbar()
        
    def _launch_tqbar(self):
        if hasattr(self, 'tqbar'): self.tqbar.close()
        self.mets.map(Self.reset())
        self.counter = 0
        self.tqbar = tqdm(total=self.len, leave=False)
        self.tqbar.set_description('Running validation')
        
    def after_batch(self):
        if not self.training: return
        self.model.eval()
        self.learn.training=False
        self._batch_val()
        self.model.train()
        self.learn.training=True
        
    @torch.no_grad()
    def _batch_val(self):
        xb, yb = next(self.val_cyclit)
        self.tqbar.update(1)
        if self.counter == self.len:
            self._launch_tqbar()
        else: self.counter += 1
        self._fake_l.yb = (yb,)
        self._fake_l.pred, *_ = self.model(xb) 
        self._fake_l.loss = Tensor(self.loss_func(self._fake_l.pred, yb))
        for met in self.mets: met.accumulate(self._fake_l)
        # pdb.set_trace()
        comment = dict(self.mets.attrgot('name').zipwith(self.mets.attrgot('value').map(_cpupy)))
        self.tqbar.set_postfix(comment)
    
    def after_train(self):
        if hasattr(self, 'val_cyclit'): delattr(self, 'val_cyclit')
        if hasattr(self, 'tqbar'): self.tqbar.close(); delattr(self, 'tqbar')
