<reponame>aethersoft/textkit-learn
import abc
import collections
import copy
import json
import os
import typing

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset
from torch.utils.data.dataloader import DataLoader, default_collate

from tklearn.neural_network.callbacks import CallbackCollection
from tklearn.utils.multiclass import type_of_target as get_type_of_target
from tklearn.utils import progress

__all__ = [
    'TorchTrainer'
]


def format_error_message(type, *args, **kwargs):
    if type == 'not_supported_type':
        return 'Supported target types are: (\'multilabel-indicator\', ' \
               '\'binary\', \'multiclass\', \'multiclass-multioutput\'). ' \
               'Got \'{}\' instead.'.format(*args, **kwargs)
    else:
        return ''


def _get_item(v, idx) -> typing.Union[collections.Sequence, collections.Mapping, np.ndarray]:
    if isinstance(v, tuple):
        return tuple(_get_item(vi, idx) for vi in v)
    elif isinstance(v, collections.Mapping):
        return {key: _get_item(value, idx) for key, value in v.items()}
    elif isinstance(v, pd.Series):
        return v.iloc[idx]
    elif isinstance(v, pd.DataFrame):
        return v.iloc[idx].to_numpy()
    else:
        return v[idx]


class SequenceLoss(object):
    """Generic loss for sequence targets.

    The function is built on top of classic loss functions.
    """

    def __init__(self, criterion=None):
        """Constructor for loss function.

        Parameters
        ----------
        criterion
            Loss function
        """
        if criterion is None:
            criterion = nn.BCEWithLogitsLoss()
        self.criterion = criterion

    def __call__(self, output, target):
        """Calculates and returns loss on calling loss callable object.

        Parameters
        ----------
        output
            Model prediction labels
        target
            Ground-truth labels
        Returns
        -------

        """
        # reshape labels to give a flat vector of length batch_size*seq_len x
        target = target.view(-1, target.shape[-1])
        output = output.view(-1, output.shape[-1])
        mask = target >= 0  # mask out 'PAD' tokens
        masked_target = target[mask[:, 0], :]
        masked_output = output[mask[:, 0], :]
        return self.criterion(masked_output, masked_target)


def custom_collate(batch):
    if len(batch) < 1:
        return default_collate(batch)
    elem = batch[0]
    if isinstance(elem, tuple):
        return tuple(custom_collate([item[k] for item in batch]) for k in range(len(elem)))
    elif isinstance(elem, collections.Mapping):
        return {k: custom_collate([item[k] for item in batch]) for k in elem.keys()}
    elif isinstance(elem, collections.Sequence) and (len(elem) > 0):
        # if batch is a list of strings
        if isinstance(elem[0], int):
            return torch.tensor(batch)
        elif isinstance(elem[0], float):
            return torch.tensor(batch, dtype=torch.float64)
        else:
            raise TypeError('Only sequence of numeric values are supported.')
    return default_collate(batch)


def _get_class_weights(y):
    if isinstance(y, tuple):
        return tuple(_get_class_weights(yi) for yi in y)
    y = np.reshape(y, (-1, y.shape[-1]))
    num_pos = np.sum(y > 0, axis=0)
    class_weight = (y.shape[0] - num_pos) / num_pos
    return torch.as_tensor(class_weight, dtype=torch.float)


class TorchDataset(Dataset):
    """Auto-torch dataset from X, y used in Trainer.

    """

    def __init__(self, **kwargs):
        item_0 = next(iter(kwargs.values()))
        while isinstance(item_0, dict):
            item_0 = next(iter(item_0.values()))
        length = len(item_0)
        samples = []
        for idx in range(length):
            record = dict()
            for k, v in kwargs.items():
                record[k] = _get_item(v, idx)
            samples.append(record)
        self.samples = samples

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        return self.samples[idx]


class TorchTrainer(abc.ABC):
    def __init__(self, callbacks=None, device=None):
        super(TorchTrainer, self).__init__()
        if callbacks is None:
            callbacks = []
        self.callbacks = callbacks
        self.device = device
        self.model_ = None
        self.history_ = None
        self.type_of_target_ = None

    def _fit(self, X, y=None, epochs=1, batch_size=32, class_weights=None, shuffle=True, clip_grad_value=None,
             verbose=1):
        if self.type_of_target_ is None:
            self.type_of_target_ = get_type_of_target(y)
        self.model_ = self._to_device(self.build_model())
        dataset = TorchDataset(X=X, y=y)
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, collate_fn=custom_collate)
        num_batches = len(dataloader)
        num_training_steps = num_batches * epochs
        # calculate weights for positive class of each category
        if class_weights is not None:
            if class_weights == 'auto':
                class_weights = _get_class_weights(y)
            elif isinstance(class_weights, (collections.Sequence, np.ndarray)):
                class_weights = torch.as_tensor(class_weights, dtype=torch.float)
        if class_weights is not None:
            class_weights = self._to_device(class_weights)
        # create model and send to device
        criterion = self._get_criterion(type_of_target=self.type_of_target, class_weights=class_weights)
        optimizer = self._get_optimizer()
        scheduler = self._get_scheduler(optimizer=optimizer, num_training_steps=num_training_steps)
        callbacks = self._get_callbacks()
        self.history_ = []
        for epoch in range(epochs):  # loop over the dataset multiple times
            callbacks.on_epoch_start(epoch + 1, self.history)
            running_loss = 0.0
            batch_history = []
            progress_bar = progress() if verbose > 0 else None
            for i, data in (enumerate if progress_bar is None else progress_bar.enumerate)(dataloader, 0):
                callbacks.on_train_batch_start(i)
                batch_x, batch_y = data['X'], data['y']
                # move data to device
                batch_x = self._to_device(batch_x)
                batch_y = self._to_device(batch_y)
                # zero the parameter gradients - clear out the gradients of all Variables
                #   in this optimizer (i.e. W, b)
                optimizer.zero_grad()
                # forward + backward + optimize
                outputs = self.model(batch_x)
                if isinstance(outputs, tuple):
                    losses = []
                    for criterion_i, output, batch_y_i in zip(criterion, outputs, batch_y):
                        batch_y_i = batch_y_i.type_as(output)
                        loss_i = criterion_i(output, batch_y_i)
                        losses.append(loss_i)
                    _loss = sum(losses)
                else:
                    batch_y = batch_y.type_as(outputs)
                    _loss = criterion(outputs, batch_y)
                _loss.backward()
                if clip_grad_value is not None:
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), clip_grad_value)
                optimizer.step()
                if scheduler is not None:
                    # noinspection PyUnresolvedReferences
                    scheduler.step()
                # print statistics
                running_loss += _loss.item()
                batch_log = {
                    'epoch': epoch + 1,
                    'batch': i + 1,
                    'loss': running_loss / (i + 1),
                }
                if progress_bar is not None:
                    progress_bar.set_description(
                        'Epoch: {epoch:3d}/{epochs:<3d}\tBatch: {batch:4d}/{num_batches:<4d}\tLoss:{loss:0.4f}'.format(
                            num_batches=num_batches, epochs=epochs, **batch_log
                        )
                    )
                batch_history.append(batch_log)
                callbacks.on_train_batch_end(i, batch_log)
            callbacks.on_epoch_end(epoch + 1, batch_history)
            self.history_ += batch_history
        return self

    def fit(self, X, y=None):
        return self._fit(X, y)

    def _predict(self, X, batch_size=8, threshold=None, verbose=0):
        p = self._predict_proba(X, batch_size=batch_size, verbose=verbose)
        if self.type_of_target in ['multilabel-indicator', 'binary', 'multilabel-sequence-output']:
            if threshold is None:
                threshold = 0.5
            return (p > threshold).astype(int)
        elif self.type_of_target in ['multiclass', 'multiclass-multioutput', 'multiclass-sequence-output']:
            return (p == p.max(axis=-1, keepdims=True)).astype(int)
        else:
            raise ValueError(format_error_message('not_supported_type', self.type_of_target))

    def predict(self, X):
        return self._predict(X)

    def _predict_proba(self, X, batch_size=8, verbose=0):
        # update predict model device if required
        self.model_ = self._to_device(self.model)
        self.model_.eval()
        # create torch dataset and loader from provided data
        dataset = TorchDataset(x=X)
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=custom_collate)
        results = []
        progress_bar = progress() if verbose > 0 else None
        for i, data in (enumerate if progress_bar is None else progress_bar.enumerate)(dataloader, 0):
            batch_x = data['x']
            # move data to device
            batch_x = self._to_device(batch_x)
            output = self.model(batch_x)
            if isinstance(self.type_of_target, tuple):
                outputs = []
                for output_i, target_type in zip(output, self.type_of_target):
                    if target_type in ['multilabel-indicator', 'binary', 'multilabel-multioutput',
                                       'multilabel-sequence-output']:
                        output_i = torch.sigmoid(output_i).cpu().detach().numpy()
                    elif target_type in ['multiclass', 'multiclass-multioutput', 'multiclass-sequence-output']:
                        output_i = torch.softmax(output_i, dim=-1).cpu().detach().numpy()
                    else:
                        output_i = output_i.cpu().detach().numpy()
                    outputs.append(output_i)
                output = tuple(outputs)
            else:
                if self.type_of_target in ['multilabel-indicator', 'binary', 'multilabel-multioutput',
                                           'multilabel-sequence-output']:
                    output = torch.sigmoid(output).cpu().detach().numpy()
                elif self.type_of_target in ['multiclass', 'multiclass-multioutput', 'multiclass-sequence-output']:
                    output = torch.softmax(output, dim=-1).cpu().detach().numpy()
                else:
                    output = output.cpu().detach().numpy()
            results.append(output)
        if isinstance(self.type_of_target, tuple):
            return tuple(np.concatenate(r) for r in zip(*results))
        else:
            return np.concatenate(results)

    def predict_proba(self, X):
        return self._predict_proba(X)

    def _get_criterion(self, type_of_target=None, class_weights=None):
        if type_of_target is None:
            type_of_target = self.type_of_target
        if isinstance(type_of_target, tuple):
            if class_weights is None:
                class_weights = tuple()
            class_weights = dict(enumerate(class_weights))
            return tuple(self._get_criterion(t, class_weights.get(i, None)) for i, t in enumerate(type_of_target))
        if type_of_target in ['multiclass', 'multiclass-multioutput']:
            # This criterion combines LogSoftmax and NLLLoss in one single class.
            loss = nn.CrossEntropyLoss()
        elif type_of_target in ['multilabel-indicator', 'binary']:
            # This loss combines a Sigmoid layer and the BCELoss in one single class
            loss = nn.BCEWithLogitsLoss(pos_weight=class_weights)
        elif type_of_target in ['multiclass-sequence-output']:
            # loss = SequenceLoss(criterion=nn.CrossEntropyLoss())
            raise NotImplementedError('multi-class sequence loss is not implemented.')
        elif type_of_target in ['multilabel-sequence-output']:
            loss = SequenceLoss(criterion=nn.BCEWithLogitsLoss(pos_weight=class_weights))
        else:
            raise ValueError(format_error_message('not_supported_type', type_of_target))
        return loss

    def _get_optimizer(self):
        return optim.SGD(self.model_.parameters(), lr=0.001, momentum=0.9)

    def _get_scheduler(self, **kwargs):
        return None

    def _get_callbacks(self):
        callbacks = []
        for x in self.callbacks:
            x = copy.copy(x)
            x.model = self
            callbacks.append(x)
        callbacks = CallbackCollection(callbacks)
        callbacks.model = self
        return callbacks

    def _get_device(self):
        if self.device is None:
            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
            return device
        return torch.device(self.device)

    def _to_device(self, obj, device=None):
        if device is None:
            device = self._get_device()
        if isinstance(obj, tuple):
            return tuple(self._to_device(item, device=device) for item in obj)
        if isinstance(obj, collections.Mapping):
            return {k: self._to_device(v, device) for k, v in obj.items()}
        return obj.to(device)

    def build_model(self):
        raise NotImplementedError

    @property
    def model(self):
        if self.model_ is None:
            raise AttributeError('Model (\'model\') is not available yet.')
        return self.model_

    @property
    def type_of_target(self):
        if self.type_of_target_ is None:
            raise AttributeError('Type of target attribute (\'type_of_target\') is not available yet.')
        return self.type_of_target_

    @property
    def history(self):
        return self.history_

    def save(self, path):
        if not os.path.exists(path):
            os.makedirs(path)
        torch.save(self.model_.state_dict(), os.path.join(path, 'model_.pk'))
        params = self.get_params()
        with open(os.path.join(path, 'params.json'), 'w', encoding='utf-8') as fp:
            json.dump(params, fp)

    @classmethod
    def load(cls, path):
        self = cls()
        with open(os.path.join(path, 'params.json'), 'r', encoding='utf-8') as fp:
            params = json.load(fp)
        self.set_params(**params)
        self.model_ = self.build_model()
        state_dict = torch.load(os.path.join(path, 'model_.pk'), map_location=torch.device('cpu'))
        self.model_.load_state_dict(state_dict)
        return self

    def get_params(self, deep=True):
        params = {}
        for k, v in self.__dict__.items():
            if k not in {'callbacks', 'model_'}:
                params[k] = v
        return params

    def set_params(self, **parameters):
        for parameter, value in parameters.items():
            setattr(self, parameter, value)
        return self

import abc
import collections
import copy
import json
import os
import typing

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset
from torch.utils.data.dataloader import DataLoader, default_collate

from tklearn.neural_network.callbacks import CallbackCollection
from tklearn.utils.multiclass import type_of_target as get_type_of_target
from tklearn.utils import progress

__all__ = [
    'TorchTrainer'
]


def format_error_message(type, *args, **kwargs):
    if type == 'not_supported_type':
        return 'Supported target types are: (\'multilabel-indicator\', ' \
               '\'binary\', \'multiclass\', \'multiclass-multioutput\'). ' \
               'Got \'{}\' instead.'.format(*args, **kwargs)
    else:
        return ''


def _get_item(v, idx) -> typing.Union[collections.Sequence, collections.Mapping, np.ndarray]:
    if isinstance(v, tuple):
        return tuple(_get_item(vi, idx) for vi in v)
    elif isinstance(v, collections.Mapping):
        return {key: _get_item(value, idx) for key, value in v.items()}
    elif isinstance(v, pd.Series):
        return v.iloc[idx]
    elif isinstance(v, pd.DataFrame):
        return v.iloc[idx].to_numpy()
    else:
        return v[idx]


class SequenceLoss(object):
    """Generic loss for sequence targets.

    The function is built on top of classic loss functions.
    """

    def __init__(self, criterion=None):
        """Constructor for loss function.

        Parameters
        ----------
        criterion
            Loss function
        """
        if criterion is None:
            criterion = nn.BCEWithLogitsLoss()
        self.criterion = criterion

    def __call__(self, output, target):
        """Calculates and returns loss on calling loss callable object.

        Parameters
        ----------
        output
            Model prediction labels
        target
            Ground-truth labels
        Returns
        -------

        """
        # reshape labels to give a flat vector of length batch_size*seq_len x
        target = target.view(-1, target.shape[-1])
        output = output.view(-1, output.shape[-1])
        mask = target >= 0  # mask out 'PAD' tokens
        masked_target = target[mask[:, 0], :]
        masked_output = output[mask[:, 0], :]
        return self.criterion(masked_output, masked_target)


def custom_collate(batch):
    if len(batch) < 1:
        return default_collate(batch)
    elem = batch[0]
    if isinstance(elem, tuple):
        return tuple(custom_collate([item[k] for item in batch]) for k in range(len(elem)))
    elif isinstance(elem, collections.Mapping):
        return {k: custom_collate([item[k] for item in batch]) for k in elem.keys()}
    elif isinstance(elem, collections.Sequence) and (len(elem) > 0):
        # if batch is a list of strings
        if isinstance(elem[0], int):
            return torch.tensor(batch)
        elif isinstance(elem[0], float):
            return torch.tensor(batch, dtype=torch.float64)
        else:
            raise TypeError('Only sequence of numeric values are supported.')
    return default_collate(batch)


def _get_class_weights(y):
    if isinstance(y, tuple):
        return tuple(_get_class_weights(yi) for yi in y)
    y = np.reshape(y, (-1, y.shape[-1]))
    num_pos = np.sum(y > 0, axis=0)
    class_weight = (y.shape[0] - num_pos) / num_pos
    return torch.as_tensor(class_weight, dtype=torch.float)


class TorchDataset(Dataset):
    """Auto-torch dataset from X, y used in Trainer.

    """

    def __init__(self, **kwargs):
        item_0 = next(iter(kwargs.values()))
        while isinstance(item_0, dict):
            item_0 = next(iter(item_0.values()))
        length = len(item_0)
        samples = []
        for idx in range(length):
            record = dict()
            for k, v in kwargs.items():
                record[k] = _get_item(v, idx)
            samples.append(record)
        self.samples = samples

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        return self.samples[idx]


class TorchTrainer(abc.ABC):
    def __init__(self, callbacks=None, device=None):
        super(TorchTrainer, self).__init__()
        if callbacks is None:
            callbacks = []
        self.callbacks = callbacks
        self.device = device
        self.model_ = None
        self.history_ = None
        self.type_of_target_ = None

    def _fit(self, X, y=None, epochs=1, batch_size=32, class_weights=None, shuffle=True, clip_grad_value=None,
             verbose=1):
        if self.type_of_target_ is None:
            self.type_of_target_ = get_type_of_target(y)
        self.model_ = self._to_device(self.build_model())
        dataset = TorchDataset(X=X, y=y)
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, collate_fn=custom_collate)
        num_batches = len(dataloader)
        num_training_steps = num_batches * epochs
        # calculate weights for positive class of each category
        if class_weights is not None:
            if class_weights == 'auto':
                class_weights = _get_class_weights(y)
            elif isinstance(class_weights, (collections.Sequence, np.ndarray)):
                class_weights = torch.as_tensor(class_weights, dtype=torch.float)
        if class_weights is not None:
            class_weights = self._to_device(class_weights)
        # create model and send to device
        criterion = self._get_criterion(type_of_target=self.type_of_target, class_weights=class_weights)
        optimizer = self._get_optimizer()
        scheduler = self._get_scheduler(optimizer=optimizer, num_training_steps=num_training_steps)
        callbacks = self._get_callbacks()
        self.history_ = []
        for epoch in range(epochs):  # loop over the dataset multiple times
            callbacks.on_epoch_start(epoch + 1, self.history)
            running_loss = 0.0
            batch_history = []
            progress_bar = progress() if verbose > 0 else None
            for i, data in (enumerate if progress_bar is None else progress_bar.enumerate)(dataloader, 0):
                callbacks.on_train_batch_start(i)
                batch_x, batch_y = data['X'], data['y']
                # move data to device
                batch_x = self._to_device(batch_x)
                batch_y = self._to_device(batch_y)
                # zero the parameter gradients - clear out the gradients of all Variables
                #   in this optimizer (i.e. W, b)
                optimizer.zero_grad()
                # forward + backward + optimize
                outputs = self.model(batch_x)
                if isinstance(outputs, tuple):
                    losses = []
                    for criterion_i, output, batch_y_i in zip(criterion, outputs, batch_y):
                        batch_y_i = batch_y_i.type_as(output)
                        loss_i = criterion_i(output, batch_y_i)
                        losses.append(loss_i)
                    _loss = sum(losses)
                else:
                    batch_y = batch_y.type_as(outputs)
                    _loss = criterion(outputs, batch_y)
                _loss.backward()
                if clip_grad_value is not None:
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), clip_grad_value)
                optimizer.step()
                if scheduler is not None:
                    # noinspection PyUnresolvedReferences
                    scheduler.step()
                # print statistics
                running_loss += _loss.item()
                batch_log = {
                    'epoch': epoch + 1,
                    'batch': i + 1,
                    'loss': running_loss / (i + 1),
                }
                if progress_bar is not None:
                    progress_bar.set_description(
                        'Epoch: {epoch:3d}/{epochs:<3d}\tBatch: {batch:4d}/{num_batches:<4d}\tLoss:{loss:0.4f}'.format(
                            num_batches=num_batches, epochs=epochs, **batch_log
                        )
                    )
                batch_history.append(batch_log)
                callbacks.on_train_batch_end(i, batch_log)
            callbacks.on_epoch_end(epoch + 1, batch_history)
            self.history_ += batch_history
        return self

    def fit(self, X, y=None):
        return self._fit(X, y)

    def _predict(self, X, batch_size=8, threshold=None, verbose=0):
        p = self._predict_proba(X, batch_size=batch_size, verbose=verbose)
        if self.type_of_target in ['multilabel-indicator', 'binary', 'multilabel-sequence-output']:
            if threshold is None:
                threshold = 0.5
            return (p > threshold).astype(int)
        elif self.type_of_target in ['multiclass', 'multiclass-multioutput', 'multiclass-sequence-output']:
            return (p == p.max(axis=-1, keepdims=True)).astype(int)
        else:
            raise ValueError(format_error_message('not_supported_type', self.type_of_target))

    def predict(self, X):
        return self._predict(X)

    def _predict_proba(self, X, batch_size=8, verbose=0):
        # update predict model device if required
        self.model_ = self._to_device(self.model)
        self.model_.eval()
        # create torch dataset and loader from provided data
        dataset = TorchDataset(x=X)
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=custom_collate)
        results = []
        progress_bar = progress() if verbose > 0 else None
        for i, data in (enumerate if progress_bar is None else progress_bar.enumerate)(dataloader, 0):
            batch_x = data['x']
            # move data to device
            batch_x = self._to_device(batch_x)
            output = self.model(batch_x)
            if isinstance(self.type_of_target, tuple):
                outputs = []
                for output_i, target_type in zip(output, self.type_of_target):
                    if target_type in ['multilabel-indicator', 'binary', 'multilabel-multioutput',
                                       'multilabel-sequence-output']:
                        output_i = torch.sigmoid(output_i).cpu().detach().numpy()
                    elif target_type in ['multiclass', 'multiclass-multioutput', 'multiclass-sequence-output']:
                        output_i = torch.softmax(output_i, dim=-1).cpu().detach().numpy()
                    else:
                        output_i = output_i.cpu().detach().numpy()
                    outputs.append(output_i)
                output = tuple(outputs)
            else:
                if self.type_of_target in ['multilabel-indicator', 'binary', 'multilabel-multioutput',
                                           'multilabel-sequence-output']:
                    output = torch.sigmoid(output).cpu().detach().numpy()
                elif self.type_of_target in ['multiclass', 'multiclass-multioutput', 'multiclass-sequence-output']:
                    output = torch.softmax(output, dim=-1).cpu().detach().numpy()
                else:
                    output = output.cpu().detach().numpy()
            results.append(output)
        if isinstance(self.type_of_target, tuple):
            return tuple(np.concatenate(r) for r in zip(*results))
        else:
            return np.concatenate(results)

    def predict_proba(self, X):
        return self._predict_proba(X)

    def _get_criterion(self, type_of_target=None, class_weights=None):
        if type_of_target is None:
            type_of_target = self.type_of_target
        if isinstance(type_of_target, tuple):
            if class_weights is None:
                class_weights = tuple()
            class_weights = dict(enumerate(class_weights))
            return tuple(self._get_criterion(t, class_weights.get(i, None)) for i, t in enumerate(type_of_target))
        if type_of_target in ['multiclass', 'multiclass-multioutput']:
            # This criterion combines LogSoftmax and NLLLoss in one single class.
            loss = nn.CrossEntropyLoss()
        elif type_of_target in ['multilabel-indicator', 'binary']:
            # This loss combines a Sigmoid layer and the BCELoss in one single class
            loss = nn.BCEWithLogitsLoss(pos_weight=class_weights)
        elif type_of_target in ['multiclass-sequence-output']:
            # loss = SequenceLoss(criterion=nn.CrossEntropyLoss())
            raise NotImplementedError('multi-class sequence loss is not implemented.')
        elif type_of_target in ['multilabel-sequence-output']:
            loss = SequenceLoss(criterion=nn.BCEWithLogitsLoss(pos_weight=class_weights))
        else:
            raise ValueError(format_error_message('not_supported_type', type_of_target))
        return loss

    def _get_optimizer(self):
        return optim.SGD(self.model_.parameters(), lr=0.001, momentum=0.9)

    def _get_scheduler(self, **kwargs):
        return None

    def _get_callbacks(self):
        callbacks = []
        for x in self.callbacks:
            x = copy.copy(x)
            x.model = self
            callbacks.append(x)
        callbacks = CallbackCollection(callbacks)
        callbacks.model = self
        return callbacks

    def _get_device(self):
        if self.device is None:
            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
            return device
        return torch.device(self.device)

    def _to_device(self, obj, device=None):
        if device is None:
            device = self._get_device()
        if isinstance(obj, tuple):
            return tuple(self._to_device(item, device=device) for item in obj)
        if isinstance(obj, collections.Mapping):
            return {k: self._to_device(v, device) for k, v in obj.items()}
        return obj.to(device)

    def build_model(self):
        raise NotImplementedError

    @property
    def model(self):
        if self.model_ is None:
            raise AttributeError('Model (\'model\') is not available yet.')
        return self.model_

    @property
    def type_of_target(self):
        if self.type_of_target_ is None:
            raise AttributeError('Type of target attribute (\'type_of_target\') is not available yet.')
        return self.type_of_target_

    @property
    def history(self):
        return self.history_

    def save(self, path):
        if not os.path.exists(path):
            os.makedirs(path)
        torch.save(self.model_.state_dict(), os.path.join(path, 'model_.pk'))
        params = self.get_params()
        with open(os.path.join(path, 'params.json'), 'w', encoding='utf-8') as fp:
            json.dump(params, fp)

    @classmethod
    def load(cls, path):
        self = cls()
        with open(os.path.join(path, 'params.json'), 'r', encoding='utf-8') as fp:
            params = json.load(fp)
        self.set_params(**params)
        self.model_ = self.build_model()
        state_dict = torch.load(os.path.join(path, 'model_.pk'), map_location=torch.device('cpu'))
        self.model_.load_state_dict(state_dict)
        return self

    def get_params(self, deep=True):
        params = {}
        for k, v in self.__dict__.items():
            if k not in {'callbacks', 'model_'}:
                params[k] = v
        return params

    def set_params(self, **parameters):
        for parameter, value in parameters.items():
            setattr(self, parameter, value)
        return self
Generic loss for sequence targets.      The function is built on top of classic loss functions. Constructor for loss function.          Parameters         ----------         criterion             Loss function Calculates and returns loss on calling loss callable object.          Parameters         ----------         output             Model prediction labels         target             Ground-truth labels         Returns         ------- # reshape labels to give a flat vector of length batch_size*seq_len x # mask out 'PAD' tokens # if batch is a list of strings Auto-torch dataset from X, y used in Trainer. # calculate weights for positive class of each category # create model and send to device # loop over the dataset multiple times # move data to device # zero the parameter gradients - clear out the gradients of all Variables #   in this optimizer (i.e. W, b) # forward + backward + optimize # noinspection PyUnresolvedReferences # print statistics # update predict model device if required # create torch dataset and loader from provided data # move data to device # This criterion combines LogSoftmax and NLLLoss in one single class. # This loss combines a Sigmoid layer and the BCELoss in one single class # loss = SequenceLoss(criterion=nn.CrossEntropyLoss())