<commit_before>""""""
import os
import re
import threading
import inspect
from functools import partial

import numpy as np
import torch
import torch.nn as nn

from .layers import ConvBlock
from .losses import CrossEntropyLoss
from ..utils import unpack_fn_from_config
from ... import Config


# TODO:
# layers/ConvBlock
# microbatch (rename/alias to virtual batch?)
# multi-device
# async
# 



LOSSES = {
    'mse': nn.MSELoss,
    'bce': nn.BCEWithLogitsLoss,
    'ce': CrossEntropyLoss,
    'crossentropy': CrossEntropyLoss,
    'absolutedifference': nn.L1Loss,
    'l1': nn.L1Loss,
    'cosine': nn.CosineSimilarity,
    'cos': nn.CosineSimilarity,
    'hinge': nn.HingeEmbeddingLoss,
    'huber': nn.SmoothL1Loss,
    'logloss': CrossEntropyLoss,
}


DECAYS = {
    'exp': torch.optim.lr_scheduler.ExponentialLR,
    'lambda': torch.optim.lr_scheduler.LambdaLR,
    'step': torch.optim.lr_scheduler.StepLR,
    'multistep': torch.optim.lr_scheduler.MultiStepLR,
    'cos': torch.optim.lr_scheduler.CosineAnnealingLR,
}


DECAYS_DEFAULTS = {
    torch.optim.lr_scheduler.ExponentialLR : dict(gamma=0.96),
    torch.optim.lr_scheduler.LambdaLR : dict(lr_lambda=lambda epoch: 0.96**epoch),
    torch.optim.lr_scheduler.StepLR: dict(step_size=30),
    torch.optim.lr_scheduler.MultiStepLR: dict(milestones=[30, 80]),
    torch.optim.lr_scheduler.CosineAnnealingLR: dict(T_max=None)
}


class EagerTorch:
    """ Eagerly! """

    def __init__(self, config=None, *args, **kwargs):
        self.config = Config(config)
        self.full_config = None
        self.train_lock = threading.Lock()

        self.model = None
        self.device = None
        self.microbatch = None

        self.n_iters = None
        self.current_iter = 0

        self.loss = None
        self.loss_fn = None
        self.lr_decay = None
        self.optimizer = None

        self.predictions = None


        load = self.config.get('load')
        build = self.config.get('build', default=load is None)
        if load:
            self.load(**load)
        if build:
            self.build()

    def reset(self):
        pass


    def build(self):
        """ Build the model """
        config = self.build_config()
        self.full_config = config

        self.device = self._get_device()

        # If the inputs were set in config with their shapes we can build rightaway
        if config.get('inputs'):
            print('_BUILD IN BUILD')
            self._build()


    def _get_device(self):
        device = self.config.get('device')
        if isinstance(device, torch.device) or device is None:
            _device = device
        elif isinstance(device, str):
            _device = device.split(':')
            unit, index = _device if len(_device) > 1 else (device, '0')
            if unit.lower() == 'gpu':
                _device = torch.device('cuda', int(index))
            elif unit.lower() == 'cpu':
                _device = torch.device('cpu')
            else:
                raise ValueError('Unknown device type: ', device)
        else:
            raise TypeError('Wrong device type: ', type(device))
        return _device


    @classmethod
    def default_config(cls):
        config = Config()
        config['inputs'] = {}
        config['common'] = {}
        config['initial_block'] = {}
        config['body'] = {}
        config['head'] = {}
        config['predictions'] = None
        config['output'] = None
        config['optimizer'] = ('Adam', dict())
        config['microbatch'] = None
        return config

    def build_config(self):
        config = self.default_config()
        config = config + self.config

        if config.get('inputs'):
            inputs = config.get('initial_block/inputs')
            if isinstance(inputs, str):
                config['common/data_format'] = config['inputs'][inputs].get('data_format')
        return config


    def _build(self, inputs=None):
        config = self.full_config

        inputs = inputs or self._placeholder_data()

        blocks = []
        for loc in ['initial_block', 'body', 'head']:
            inputs = inputs[0] if isinstance(inputs, (tuple, list)) and len(inputs) == 1 else inputs
            block = self._make_block(loc, config, inputs)
            if block is not None:
                inputs = block(inputs)
                blocks.append(block)

        self.model = nn.Sequential(*blocks)

        if self.loss_fn is None:
            self._make_loss(config)
        if self.optimizer is None:
            self._make_optimizer(config)

    def _placeholder_data(self, batch_size=2):
        config = self.full_config

        input_names = config.pop('initial_block/inputs', default=None) or list(config.get('inputs').keys())
        input_names = input_names if isinstance(input_names, (tuple, list)) else [input_names]
        shapes = [(batch_size, *config['inputs'][name]['shape']) for name in input_names]

        data = [np.zeros(shape, dtype=np.float32) for shape in shapes]
        data = self._fill_param(data)
        data = data[0] if len(data) == 1 else data
        return data

    def _make_block(self, name, config, inputs):
        config = {**config['common'], **config[name]}

        if 'module' in config:
            module = config['module']
            kwargs = config.get('module_kwargs')
            if 'inputs' in inspect.getfullargspec(module.__init__)[0]:
                kwargs = {**kwargs, **{'inputs': inputs}}
            block = module(*config.get('module_args', []), **kwargs)

        elif isinstance(config, dict):
            block = getattr(self, name)(inputs, **config)
        else:
            raise ValueError('Bad')
        return block


    def _make_loss(self, config):
        loss, args = unpack_fn_from_config('loss', config)

        if isinstance(loss, str):
            loss = LOSSES.get(re.sub('[-_ ]', '', loss).lower(), None)
        elif isinstance(loss, str) and hasattr(nn, loss):
            loss = getattr(nn, loss)<commit_msg>Add ability to chain blocks in the architecture<commit_after>""""""
import os
import re
import threading
import inspect
from collections import OrderedDict
from functools import partial

import numpy as np
import torch
import torch.nn as nn

from .layers import ConvBlock
from .losses import CrossEntropyLoss
from ..utils import unpack_fn_from_config
from ... import Config


# TODO:
# layers/ConvBlock
# microbatch (rename/alias to virtual batch?)
# multi-device
# async
# 



LOSSES = {
    'mse': nn.MSELoss,
    'bce': nn.BCEWithLogitsLoss,
    'ce': CrossEntropyLoss,
    'crossentropy': CrossEntropyLoss,
    'absolutedifference': nn.L1Loss,
    'l1': nn.L1Loss,
    'cosine': nn.CosineSimilarity,
    'cos': nn.CosineSimilarity,
    'hinge': nn.HingeEmbeddingLoss,
    'huber': nn.SmoothL1Loss,
    'logloss': CrossEntropyLoss,
}


DECAYS = {
    'exp': torch.optim.lr_scheduler.ExponentialLR,
    'lambda': torch.optim.lr_scheduler.LambdaLR,
    'step': torch.optim.lr_scheduler.StepLR,
    'multistep': torch.optim.lr_scheduler.MultiStepLR,
    'cos': torch.optim.lr_scheduler.CosineAnnealingLR,
}


DECAYS_DEFAULTS = {
    torch.optim.lr_scheduler.ExponentialLR : dict(gamma=0.96),
    torch.optim.lr_scheduler.LambdaLR : dict(lr_lambda=lambda epoch: 0.96**epoch),
    torch.optim.lr_scheduler.StepLR: dict(step_size=30),
    torch.optim.lr_scheduler.MultiStepLR: dict(milestones=[30, 80]),
    torch.optim.lr_scheduler.CosineAnnealingLR: dict(T_max=None)
}


class EagerTorch:
    """ Eagerly! """

    def __init__(self, config=None, *args, **kwargs):
        self.config = Config(config)
        self.full_config = None
        self.train_lock = threading.Lock()

        self.model = None
        self.device = None
        self.microbatch = None

        self.n_iters = None
        self.current_iter = 0

        self.loss = None
        self.loss_fn = None
        self.lr_decay = None
        self.optimizer = None

        self.predictions = None


        load = self.config.get('load')
        build = self.config.get('build', default=load is None)
        if load:
            self.load(**load)
        if build:
            self.build()

    def reset(self):
        pass


    def build(self):
        """ Build the model """
        config = self.build_config()
        self.full_config = config

        self.device = self._get_device()

        # If the inputs were set in config with their shapes we can build rightaway
        if config.get('inputs'):
            print('_BUILD IN BUILD')
            self._build()


    def _get_device(self):
        device = self.config.get('device')
        if isinstance(device, torch.device) or device is None:
            _device = device
        elif isinstance(device, str):
            _device = device.split(':')
            unit, index = _device if len(_device) > 1 else (device, '0')
            if unit.lower() == 'gpu':
                _device = torch.device('cuda', int(index))
            elif unit.lower() == 'cpu':
                _device = torch.device('cpu')
            else:
                raise ValueError('Unknown device type: ', device)
        else:
            raise TypeError('Wrong device type: ', type(device))
        return _device


    @classmethod
    def default_config(cls):
        config = Config()
        config['inputs'] = {}
        config['common'] = {}
        config['order'] = ['initial_block', 'body', 'head']
        config['initial_block'] = {}
        config['body'] = {}
        config['head'] = {}
        config['predictions'] = None
        config['output'] = None
        config['optimizer'] = ('Adam', dict())
        config['microbatch'] = None
        return config

    def build_config(self):
        config = self.default_config()
        config = config + self.config

        if config.get('inputs'):
            inputs = config.get('initial_block/inputs')
            if isinstance(inputs, str):
                config['common/data_format'] = config['inputs'][inputs].get('data_format')
        return config


    def _build(self, inputs=None):
        config = self.full_config
        order = config.get('order')

        inputs = inputs or self._placeholder_data()

        blocks = []
        for item in order:
            if isinstance(item, str):
                block_name = config_name = method = item
            elif isinstance(item, tuple):
                if len(item) == 2:
                    block_name, method = config_name, _ = item
                elif len(item) == 3:
                    block_name, config_name, method = item


            inputs = inputs[0] if isinstance(inputs, (tuple, list)) and len(inputs) == 1 else inputs
            block = self._make_block(config_name, method, config, inputs)
            if block is not None:
                inputs = block(inputs)
                blocks.append((block_name, block))

        self.model = nn.Sequential(OrderedDict(blocks))

        if self.loss_fn is None:
            self._make_loss(config)
        if self.optimizer is None:
            self._make_optimizer(config)

    def _placeholder_data(self, batch_size=2):
        config = self.full_config

        input_names = config.pop('initial_block/inputs', default=None) or list(config.get('inputs').keys())
        input_names = input_names if isinstance(input_names, (tuple, list)) else [input_names]
        shapes = [(batch_size, *config['inputs'][name]['shape']) for name in input_names]

        data = [np.zeros(shape, dtype=np.float32) for shape in shapes]
        data = self._fill_param(data)
        data = data[0] if len(data) == 1 else data
        return data

    def _make_block(self, name, method, config, inputs):
        config = {**config['common'], **config[name]}

        if 'module' in config:
            module = config['module']
            kwargs = config.get('module_kwargs')
            if 'inputs' in inspect.getfullargspec(module.__init__)[0]:
                kwargs = {**kwargs, **{'inputs': inputs}}
            block = module(*config.get('module_args', []), **kwargs)

        elif isinstance(config, dict):
            method = getattr(self, method) if isinstance(method, str) else method
            block = method(inputs, **config)
        else:
            raise ValueError('Bad')
        return block


    def _make_loss(self, config):
        loss, args = unpack_fn_from_config('loss', config)

        if isinstance(loss, str):
            loss = LOSSES.get(re.sub('[-_ ]', '', loss).lower(), None)
        elif isinstance(loss, str) and hasattr(nn, loss):
            loss = getattr(nn, loss)
{'fraction_non_alphanumeric': 0.12380045998889681, 'fraction_numerical': 0.004441272107224998, 'mean_word_length': 2.349269588313413, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5432752', 'n_tokens_mistral': 3860, 'n_tokens_neox': 3677, 'n_words': 894}