text: # AUTOGENERATED! DO NOT EDIT! File to edit: model.ipynb (unless otherwise specified).
  
  __all__ = ['Model']
  
  # Cell
  import numpy as np
  import matplotlib.pyplot as plt
  from functools import partial
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  from fastai.basics import Learner, LearnerCallback, Callback
  
  from .blocks.blocks import DefaultModel, DefaultQModel
  from .data import create_db, create_tensors, get_moments
  from .losses import q_loss
  from .callbacks.callbacks import *
  
  # Cell
  class Model(nn.Module):
      """
      Default Model Wrapper
      parameters:
      - df: dataset used in training dataset.
      - model (optional): how to model time series. Default: DefaultModel.
      - loss (optional): loss function: Default l1 loss.
      - bs (optional): batchsize
      - alpha (optional): l2 weight decay
      - beta (optional): l1 weight decay
      - silent: whether to silence output or not
  
      Usage:
      ```
      model = Model(train_df)
      model.fit(train_df)
      y = model.predict(test_df)
      ```
      """
      def __init__(self, model_args=None, quantiles=[0.05, 0.5, 0.95],
                   bs=128, epochs=100, lr=0.3, alpha=0.1, silent=True):
          super().__init__()
          self.loss = partial(q_loss, quantiles=quantiles)
          self.model_args = model_args
          self.quantiles = quantiles
          self.lr = lr
          self.alpha = alpha
  
          self.bs = bs
          self.epochs = epochs
          self.silent = silent
  
      def fit(self, df):
          self.moments = get_moments(df.copy())
          self.model = DefaultQModel(self.moments,
                                     **self.model_args,
                                     quantiles=self.quantiles)
          # self.find_appropriate_lr(df)
          learner = self.create_learner(df)
          cb = [PrintLoss(learner)]
          # breakpoint()
  
          learner.fit_one_cycle(self.epochs, self.lr, callbacks=cb)
  
      def predict(self, df):
          x, _ = create_tensors(df.copy(), self.moments, predict=True)
          mean, sd = self.moments['y']
          y = sd * self.model(**x) + mean
          return y.detach().cpu().numpy()
  
      def forward(self, *args):
          return self.model(*args)
  
      def create_learner(self, df):
          db = create_db(df, bs=self.bs, moments=self.moments)
          learner = Learner(db,
                            self.model,
                            loss_func=self.loss,
                            wd=self.alpha,
                            silent=True)
          return learner
  
      def plot_components(self):
          idx = self.quantiles.index(0.5)
          median = self.model.models[idx]
          median.seasonal.plot_components()