<commit_before>import glob
import json
import logging
import math
import time

import os
from torch import nn, torch
from torch.autograd import Variable

from nmmt.torch_utils import torch_is_multi_gpu, torch_is_using_cuda
from onmt import Constants, Optim


class _Stats(object):
    def __init__(self):
        self.start_time = time.time()
        self.total_loss = 0
        self.src_words = 0
        self.tgt_words = 0
        self.num_correct = 0

    def update(self, loss, src_words, tgt_words, num_correct):
        self.total_loss += loss
        self.src_words += src_words
        self.tgt_words += tgt_words
        self.num_correct += num_correct

    @property
    def accuracy(self):
        return float(self.num_correct) / self.tgt_words

    @property
    def loss(self):
        return self.total_loss / self.tgt_words

    @property
    def perplexity(self):
        return math.exp(self.loss)

    def __str__(self):
        elapsed_time = time.time() - self.start_time

        return '[num_correct: %6.2f; %3.0f src tok; %3.0f tgt tok; ' \
               'acc: %6.2f; ppl: %6.2f; %3.0f src tok/s; %3.0f tgt tok/s]' % (
                   self.num_correct, self.src_words, self.tgt_words,
                   self.accuracy * 100, self.perplexity, self.src_words / elapsed_time, self.tgt_words / elapsed_time
               )


class NMTEngineTrainer:
    class Options(object):
        def __init__(self):
            self.log_level = logging.INFO

            self.bpe_symbols = 32000
            self.max_vocab_size = None  # Unlimited
            self.vocab_pruning_threshold = None  # Skip pruning

            self.batch_size = 64
            self.max_generator_batches = 32  # Maximum batches of words in a seq to run the generator on in parallel.

            self.report_steps = 100  # Log status every 'report_steps' steps
            self.validation_steps = 10000  # compute the validation score every 'validation_steps' steps
            self.checkpoint_steps = 10000  # Drop a checkpoint every 'checkpoint_steps' steps
            self.step_limit = None  # If set, run 'step_limit' steps at most

            self.optimizer = 'sgd'
            self.learning_rate = 1.
            self.max_grad_norm = 5
            self.lr_decay = 0.9
            self.lr_decay_steps = 10000  # decrease learning rate every 'lr_decay_steps' steps
            self.lr_decay_start_at = 50000  # start learning rate decay after 'start_decay_at' steps

            self.n_checkpoints = 20  # checkpoints saved during training and used for termination condition
            self.n_avg_checkpoints = 20  # number of checkpoints to merge at the end of training process

        def load_from_dict(self, d):
            for key in self.__dict__:
                if key in d:
                    self.__dict__[key] = d[key]

        def __str__(self):
            return str(self.__dict__)

        def __repr__(self):
            return str(self.__dict__)

    class State(object):
        def __init__(self, size):
            self.size = size
            self.checkpoint = None
            self.history = []

        def empty(self):
            return len(self.history) == 0

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

        @property
        def last_step(self):
            return self.checkpoint['step'] if self.checkpoint is not None else 0

        def average_perplexity(self):
            if self.empty():
                return 0

            s = 0
            for checkpoint in self.history:
                s += checkpoint['perplexity']

            return s / len(self.history)

        @staticmethod
        def _delete_checkpoint(checkpoint):
            for path in glob.glob(checkpoint['file'] + '.*'):
                os.remove(path)

        def add_checkpoint(self, step, file_path, perplexity):
            self.checkpoint = {
                'step': step,
                'file': file_path,
                'perplexity': perplexity
            }

            self.history.insert(0, self.checkpoint)

            if len(self.history) > self.size:
                self._delete_checkpoint(self.history.pop())

        def save_to_file(self, file_path):
            with open(file_path, 'w') as stream:
                stream.write(json.dumps(self.__dict__, indent=4))

        @staticmethod
        def load_from_file(file_path):
            state = NMTEngineTrainer.State(0)
            with open(file_path, 'r') as stream:
                state.__dict__ = json.loads(stream.read())
            return state

    def __init__(self, engine, options=None, optimizer=None, state=None):
        self._logger = logging.getLogger('nmmt.NMTEngineTrainer')
        self._engine = engine
        self.opts = options if options is not None else NMTEngineTrainer.Options()

        self.state = state if state is not None else NMTEngineTrainer.State(self.opts.n_checkpoints)

        if optimizer is None:
            optimizer = Optim(self.opts.optimizer, self.opts.learning_rate, max_grad_norm=self.opts.max_grad_norm,
                              lr_decay=self.opts.lr_decay, lr_start_decay_at=self.opts.lr_decay_start_at)
            optimizer.set_parameters(engine.model.parameters())
        self.optimizer = optimizer

    def reset_learning_rate(self, value):
        self.optimizer.lr = value
        self.optimizer.set_parameters(self._engine.model.parameters())

    def _log(self, message):
        if self.opts.log_level > logging.NOTSET:
            self._logger.log(self.opts.log_level, message)

    @staticmethod
    def _new_nmt_criterion(vocab_size):
        weight = torch.ones(vocab_size)
        weight[Constants.PAD] = 0
        criterion = nn.NLLLoss(weight, size_average=False)
        if torch_is_using_cuda():
            criterion.cuda()
        return criterion

    def _compute_memory_efficient_loss(self, outputs, targets, generator, criterion, evaluation=False):
        # compute generations one piece at a time
        num_correct, loss = 0, 0
        outputs = Variable(outputs.data, requires_grad=(not evaluation), volatile=evaluation)

        batch_size = outputs.size(1)
        outputs_split = torch.split(outputs, self.opts.max_generator_batches)
        targets_split = torch.split(targets, self.opts.max_generator_batches)

        for i, (out_t, targ_t) in enumerate(zip(outputs_split, targets_split)):
            out_t = out_t.view(-1, out_t.size(2))
            scores_t = generator(out_t)
            loss_t = criterion(scores_t, targ_t.view(-1))
            pred_t = scores_t.max(1)[1]
            num_correct_t = pred_t.data.eq(targ_t.data).masked_select(targ_t.ne(Constants.PAD).data).sum()
            num_correct += num_correct_t
            loss += loss_t.data[0]
            if not evaluation:
                loss_t.div(batch_size).backward()

        grad_output = None if outputs.grad is None else outputs.grad.data
        return loss, grad_output, num_correct

    def _evaluate(self, step, criterion, dataset):
        total_loss = 0
        total_words = 0
        total_num_correct = 0

        model = self._engine.model
        model.eval()<commit_msg>bug fixing: required linearization of the 2D tensor of targets<commit_after>import glob
import json
import logging
import math
import time

import os
from torch import nn, torch
from torch.autograd import Variable

from nmmt.torch_utils import torch_is_multi_gpu, torch_is_using_cuda
from onmt import Constants, Optim


class _Stats(object):
    def __init__(self):
        self.start_time = time.time()
        self.total_loss = 0
        self.src_words = 0
        self.tgt_words = 0
        self.num_correct = 0

    def update(self, loss, src_words, tgt_words, num_correct):
        self.total_loss += loss
        self.src_words += src_words
        self.tgt_words += tgt_words
        self.num_correct += num_correct

    @property
    def accuracy(self):
        return float(self.num_correct) / self.tgt_words

    @property
    def loss(self):
        return self.total_loss / self.tgt_words

    @property
    def perplexity(self):
        return math.exp(self.loss)

    def __str__(self):
        elapsed_time = time.time() - self.start_time

        return '[num_correct: %6.2f; %3.0f src tok; %3.0f tgt tok; ' \
               'acc: %6.2f; ppl: %6.2f; %3.0f src tok/s; %3.0f tgt tok/s]' % (
                   self.num_correct, self.src_words, self.tgt_words,
                   self.accuracy * 100, self.perplexity, self.src_words / elapsed_time, self.tgt_words / elapsed_time
               )


class NMTEngineTrainer:
    class Options(object):
        def __init__(self):
            self.log_level = logging.INFO

            self.bpe_symbols = 32000
            self.max_vocab_size = None  # Unlimited
            self.vocab_pruning_threshold = None  # Skip pruning

            self.batch_size = 64
            self.max_generator_batches = 32  # Maximum batches of words in a seq to run the generator on in parallel.

            self.report_steps = 100  # Log status every 'report_steps' steps
            self.validation_steps = 10000  # compute the validation score every 'validation_steps' steps
            self.checkpoint_steps = 10000  # Drop a checkpoint every 'checkpoint_steps' steps
            self.step_limit = None  # If set, run 'step_limit' steps at most

            self.optimizer = 'sgd'
            self.learning_rate = 1.
            self.max_grad_norm = 5
            self.lr_decay = 0.9
            self.lr_decay_steps = 10000  # decrease learning rate every 'lr_decay_steps' steps
            self.lr_decay_start_at = 50000  # start learning rate decay after 'start_decay_at' steps

            self.n_checkpoints = 20  # checkpoints saved during training and used for termination condition
            self.n_avg_checkpoints = 20  # number of checkpoints to merge at the end of training process

        def load_from_dict(self, d):
            for key in self.__dict__:
                if key in d:
                    self.__dict__[key] = d[key]

        def __str__(self):
            return str(self.__dict__)

        def __repr__(self):
            return str(self.__dict__)

    class State(object):
        def __init__(self, size):
            self.size = size
            self.checkpoint = None
            self.history = []

        def empty(self):
            return len(self.history) == 0

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

        @property
        def last_step(self):
            return self.checkpoint['step'] if self.checkpoint is not None else 0

        def average_perplexity(self):
            if self.empty():
                return 0

            s = 0
            for checkpoint in self.history:
                s += checkpoint['perplexity']

            return s / len(self.history)

        @staticmethod
        def _delete_checkpoint(checkpoint):
            for path in glob.glob(checkpoint['file'] + '.*'):
                os.remove(path)

        def add_checkpoint(self, step, file_path, perplexity):
            self.checkpoint = {
                'step': step,
                'file': file_path,
                'perplexity': perplexity
            }

            self.history.insert(0, self.checkpoint)

            if len(self.history) > self.size:
                self._delete_checkpoint(self.history.pop())

        def save_to_file(self, file_path):
            with open(file_path, 'w') as stream:
                stream.write(json.dumps(self.__dict__, indent=4))

        @staticmethod
        def load_from_file(file_path):
            state = NMTEngineTrainer.State(0)
            with open(file_path, 'r') as stream:
                state.__dict__ = json.loads(stream.read())
            return state

    def __init__(self, engine, options=None, optimizer=None, state=None):
        self._logger = logging.getLogger('nmmt.NMTEngineTrainer')
        self._engine = engine
        self.opts = options if options is not None else NMTEngineTrainer.Options()

        self.state = state if state is not None else NMTEngineTrainer.State(self.opts.n_checkpoints)

        if optimizer is None:
            optimizer = Optim(self.opts.optimizer, self.opts.learning_rate, max_grad_norm=self.opts.max_grad_norm,
                              lr_decay=self.opts.lr_decay, lr_start_decay_at=self.opts.lr_decay_start_at)
            optimizer.set_parameters(engine.model.parameters())
        self.optimizer = optimizer

    def reset_learning_rate(self, value):
        self.optimizer.lr = value
        self.optimizer.set_parameters(self._engine.model.parameters())

    def _log(self, message):
        if self.opts.log_level > logging.NOTSET:
            self._logger.log(self.opts.log_level, message)

    @staticmethod
    def _new_nmt_criterion(vocab_size):
        weight = torch.ones(vocab_size)
        weight[Constants.PAD] = 0
        criterion = nn.NLLLoss(weight, size_average=False)
        if torch_is_using_cuda():
            criterion.cuda()
        return criterion

    def _compute_memory_efficient_loss(self, outputs, targets, generator, criterion, evaluation=False):
        # compute generations one piece at a time
        num_correct, loss = 0, 0
        outputs = Variable(outputs.data, requires_grad=(not evaluation), volatile=evaluation)

        batch_size = outputs.size(1)
        outputs_split = torch.split(outputs, self.opts.max_generator_batches)
        targets_split = torch.split(targets, self.opts.max_generator_batches)

        for i, (out_t, targ_t) in enumerate(zip(outputs_split, targets_split)):
            out_t = out_t.view(-1, out_t.size(2))
            scores_t = generator(out_t)
            loss_t = criterion(scores_t, targ_t.view(-1))
            pred_t = scores_t.max(1)[1]
            num_correct_t = pred_t.data.eq(targ_t.view(-1).data).masked_select(targ_t.ne(Constants.PAD).view(-1).data).sum()
            num_correct += num_correct_t
            loss += loss_t.data[0]
            if not evaluation:
                loss_t.div(batch_size).backward()

        grad_output = None if outputs.grad is None else outputs.grad.data
        return loss, grad_output, num_correct

    def _evaluate(self, step, criterion, dataset):
        total_loss = 0
        total_words = 0
        total_num_correct = 0

        model = self._engine.model
        model.eval()
{'fraction_non_alphanumeric': 0.08886264156187035, 'fraction_numerical': 0.011463906065448481, 'mean_word_length': 2.152431011826544, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, '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': '372022', 'n_tokens_mistral': 4502, 'n_tokens_neox': 4216, 'n_words': 1008}