<gh_stars>1-10
import logging
import sys

sys.path.append('../../sr2018')
from components.data.morph_align.aligner import mcmc_align
from components.data.vocabulary.morph_vocab import MorphVocabulary
from components.constants import ALIGN_SYMBOL, STEP
from components.data.base_data import pad_seq, add_bos_eos, pad_attn
import copy
from components.utils.tensors import cuda_if_gpu
from torch.autograd import Variable
import torch
from components.data.base_morph_data import BaseMorphData
from components.utils.serialization import save_txt
from components.utils.readers import read_conll_data_file
from collections import Counter

logger = logging.getLogger('main')

class MorphRNNHardData(BaseMorphData):

    def training_setup(self, train_data_fname, dev_data_fname):

        assert train_data_fname is not None
        assert dev_data_fname is not None

        train_raw = read_conll_data_file(train_data_fname)
        dev_raw = read_conll_data_file(dev_data_fname)

        # extract from TRAIN lemmas, forms, feats if in training mode
        train_lemmas_l, train_forms_l, train_feat_d = BaseMorphData.extract_lemmas_forms_feature_dicts(train_raw)
        dev_lemmas_l, dev_forms_l, dev_feat_d = BaseMorphData.extract_lemmas_forms_feature_dicts(dev_raw)
        self.set_char_freqlist(train_lemmas_l)

        self.max_src_len = self.config['max_src_len']
        self.max_tgt_len = self.config['max_tgt_len']

        logger.info('Max lemma/form lenghts: %d/%d', self.max_src_len, self.max_tgt_len)

        train_lemmas_l, train_forms_l, train_feat_d = self.filter_data_outliers(train_lemmas_l,
                                                                                train_forms_l,
                                                                                train_feat_d)

        dev_lemmas_l, dev_forms_l, dev_feat_d = self.filter_data_outliers(dev_lemmas_l,
                                                                          dev_forms_l,
                                                                          dev_feat_d)

        # align forms and lemmas
        all_aligned_pairs = mcmc_align(wordpairs=list(zip(train_lemmas_l + dev_lemmas_l,
                                                          train_forms_l + dev_forms_l)),
                                       align_symbol=ALIGN_SYMBOL)

        # splitting aligned pairs
        num_train_pairs = len(train_lemmas_l)
        train_aligned_pairs = all_aligned_pairs[:num_train_pairs]
        dev_aligned_pairs = all_aligned_pairs[num_train_pairs:]

        # vectorize the data
        self.vocab.setup(vocab_path=self.vocab_fn,
                         data=(train_lemmas_l, train_forms_l, train_feat_d),
                         lower=self.lower,
                         source='lemma_form_feat')

        self.train = self.train_data_to_ids(data=list(zip(
            train_lemmas_l,
            train_feat_d,
            train_aligned_pairs,
            train_forms_l)))

        self.dev = self.train_data_to_ids(data=list(zip(
            dev_lemmas_l,
            dev_feat_d,
            dev_aligned_pairs,
            dev_forms_l)))

    def prediction_setup(self):
        self.max_src_len = self.config['max_src_len']
        self.max_tgt_len = self.config['max_tgt_len']

        # else load vocabulary file from the filesystem
        assert 'vocab_fn' in self.config, logger.error('Specify "vocab_fn" in the config file')
        self.vocab.setup(vocab_path=self.vocab_fn,
                         data=None,
                         lower=self.lower,
                         source='file')

    def set_vocabulary_class(self):
        self.vocab = MorphVocabulary()

    @staticmethod
    def extract_gold_labels(lemma, aligned_pair):

        y_dec_targets = []
        x_attn_ids = []

        padded_lemma = add_bos_eos(lemma, bos=BaseMorphData.BOS, eos=BaseMorphData.EOS)

        i_input_index = 0

        # go through alignments, progress j when new output is introduced, progress i when new char is seen on lemma (no ~)
        aligned_lemma, aligned_word = aligned_pair
        aligned_lemma += BaseMorphData.EOS
        aligned_word += BaseMorphData.EOS

        # run through the alignments
        for align_index, (x_char, y_char) in enumerate(zip(aligned_lemma, aligned_word)):

            # if reached the end word symbol
            if y_char == BaseMorphData.EOS:
                x_attn_ids.append(i_input_index)
                y_dec_targets.append(BaseMorphData.EOS)
                continue

            # initially, if there is no prefix in the output (shouldn't delay on current input), step forward
            if padded_lemma[i_input_index] == BaseMorphData.BOS and aligned_lemma[align_index] != ALIGN_SYMBOL:
                x_attn_ids.append(i_input_index)
                y_dec_targets.append(STEP)
                # prepare for the next iteration - "feedback"
                i_input_index += 1

            # if 0-to-1 or 1-to-1 alignment, compute loss for predicting the output symbol
            if aligned_word[align_index] != ALIGN_SYMBOL:
                x_attn_ids.append(i_input_index)
                y_dec_targets.append(aligned_word[align_index])

            # if the input's not done and the next is not a 0-to-1 alignment, perform step
            if i_input_index < len(padded_lemma) - 1 and aligned_lemma[align_index + 1] != ALIGN_SYMBOL:
                # dec_prev_outputs.append(prev_output_char_idx)
                x_attn_ids.append(i_input_index)
                y_dec_targets.append(STEP)
                i_input_index += 1

        # self.sanity_check_oov_classes(dec_targets)
        return y_dec_targets, x_attn_ids

    def train_data_to_ids(self, data):

        logger.debug('Vectorizing data')

        assert self.max_src_len is not None
        assert self.max_tgt_len is not None

        data_split_enc_x = []
        data_split_attn_ids = []
        data_split_feat_x = []
        data_split_y = []
        forms = []

        for idx, datum in enumerate(data):

            lemma, feat_dict, aligned_pair, form = datum
            x_enc_inputs = BaseMorphData.lemma_chars_to_ids(lemma, self.vocab)
            x_feature_ids = BaseMorphData.feat_d_to_ids(feat_dict, self.vocab)
            y_dec_targets, x_attn_ids = MorphRNNHardData.extract_gold_labels(lemma, aligned_pair)
            y_dec_target_ids = MorphRNNHardData.y_labels_to_ids(y_dec_targets, self.vocab)

            data_split_enc_x.append(x_enc_inputs)
            data_split_feat_x.append(x_feature_ids)
            data_split_attn_ids.append(x_attn_ids)
            data_split_y.append(y_dec_target_ids)
            forms.append(form)

        return (data_split_enc_x, data_split_feat_x, data_split_attn_ids, data_split_y, forms)

    def batchify_training_data(self, xy_ids_and_targets, batch_size, is_dev_data):

        logger.debug('Batchifying data')

        # sort according to the src lengths -> xy pairs
        sorted_data = sorted(zip(*xy_ids_and_targets), key=lambda p: len(p[0]), reverse=True)
        data_size = len(sorted_data)
        num_batches = data_size // batch_size
        data_indices = self.index_data(data_size, mode='no_shuffling')

        batch_pairs = []
        batch_forms_all = []
        for bi in range(num_batches+1): # including the last (smaller) batch
            batch_data_tuple, batch_forms = self.make_one_batch(sorted_data, data_indices, bi, batch_size)
            batch_pairs.append(batch_data_tuple)
            batch_forms_all.extend(batch_forms)

        if is_dev_data:
            self.dev_references = batch_forms_all
            logger.info('Saving dev (training) references to --> %s', self.fnames.dev_ref_fn)
            save_txt(self.dev_references, self.fnames.dev_ref_fn)

        assert data_size == len(batch_forms_all)

        return batch_pairs

    def make_one_batch(self, sorted_data, data_indices, bi, batch_size):

        batch_x_enc = []
        batch_x_feat = []
        batch_x_attn = []
        batch_y = []
        batch_forms = []

        curr_batch_indices = data_indices[bi * batch_size: (bi + 1) * batch_size]

        for idx in curr_batch_indices:
            x_enc_ids, x_feature_ids, x_attn_ids, y_ids, forms = sorted_data[idx]

            x_enc_ids_copy = copy.deepcopy(x_enc_ids)
            batch_x_enc.append(x_enc_ids_copy)

            x_attn_ids_copy = copy.deepcopy(x_attn_ids)
            batch_x_attn.append(x_attn_ids_copy)

            x_feature_ids_copy = copy.deepcopy(x_feature_ids)
            batch_x_feat.append(x_feature_ids_copy)

            y_ids_copy = copy.deepcopy(y_ids)
            batch_y.append(y_ids_copy)

            batch_forms.append(forms)

        # skipping features, since the num does not change
        batch_x_enc_lens = [len(s) for s in batch_x_enc]
        batch_x_attn_lens = [len(s) for s in batch_x_attn]
        batch_y_lens = [len(s) for s in batch_y]
        assert batch_y_lens == batch_x_attn_lens

        max_enc_seq_len = max(batch_x_enc_lens)
        max_dec_seq_len = max(batch_y_lens)

        batch_enc_x_padded = [pad_seq(x, max_enc_seq_len, pad_id=self.vocab.PAD_ID) for x in batch_x_enc]
        batch_dec_y_padded = [pad_seq(y, max_dec_seq_len, pad_id=self.vocab.PAD_ID) for y in batch_y]
        # hard alignment for padded sequences
        batch_attn_x_padded = [pad_attn(x, max_dec_seq_len) for x in batch_x_attn]

        batch_enc_x_var = cuda_if_gpu(Variable(torch.LongTensor(batch_enc_x_padded)).transpose(0, 1))
        batch_dec_y_var = cuda_if_gpu(Variable(torch.LongTensor(batch_dec_y_padded)).transpose(0, 1))
        batch_x_feat_var = cuda_if_gpu(Variable(torch.LongTensor(batch_x_feat)))
        batch_attn_x = cuda_if_gpu(torch.LongTensor(batch_attn_x_padded).transpose(0, 1))  # not a Variable!

        return (batch_enc_x_var, batch_x_feat_var, batch_attn_x, batch_dec_y_var), batch_forms

    @staticmethod
    def y_labels_to_ids(y_labels, vocab):
        return [vocab.lookup_tok(label) for label in y_labels]


component = MorphRNNHardData

import logging
import sys

sys.path.append('../../sr2018')
from components.data.morph_align.aligner import mcmc_align
from components.data.vocabulary.morph_vocab import MorphVocabulary
from components.constants import ALIGN_SYMBOL, STEP
from components.data.base_data import pad_seq, add_bos_eos, pad_attn
import copy
from components.utils.tensors import cuda_if_gpu
from torch.autograd import Variable
import torch
from components.data.base_morph_data import BaseMorphData
from components.utils.serialization import save_txt
from components.utils.readers import read_conll_data_file
from collections import Counter

logger = logging.getLogger('main')

class MorphRNNHardData(BaseMorphData):

    def training_setup(self, train_data_fname, dev_data_fname):

        assert train_data_fname is not None
        assert dev_data_fname is not None

        train_raw = read_conll_data_file(train_data_fname)
        dev_raw = read_conll_data_file(dev_data_fname)

        # extract from TRAIN lemmas, forms, feats if in training mode
        train_lemmas_l, train_forms_l, train_feat_d = BaseMorphData.extract_lemmas_forms_feature_dicts(train_raw)
        dev_lemmas_l, dev_forms_l, dev_feat_d = BaseMorphData.extract_lemmas_forms_feature_dicts(dev_raw)
        self.set_char_freqlist(train_lemmas_l)

        self.max_src_len = self.config['max_src_len']
        self.max_tgt_len = self.config['max_tgt_len']

        logger.info('Max lemma/form lenghts: %d/%d', self.max_src_len, self.max_tgt_len)

        train_lemmas_l, train_forms_l, train_feat_d = self.filter_data_outliers(train_lemmas_l,
                                                                                train_forms_l,
                                                                                train_feat_d)

        dev_lemmas_l, dev_forms_l, dev_feat_d = self.filter_data_outliers(dev_lemmas_l,
                                                                          dev_forms_l,
                                                                          dev_feat_d)

        # align forms and lemmas
        all_aligned_pairs = mcmc_align(wordpairs=list(zip(train_lemmas_l + dev_lemmas_l,
                                                          train_forms_l + dev_forms_l)),
                                       align_symbol=ALIGN_SYMBOL)

        # splitting aligned pairs
        num_train_pairs = len(train_lemmas_l)
        train_aligned_pairs = all_aligned_pairs[:num_train_pairs]
        dev_aligned_pairs = all_aligned_pairs[num_train_pairs:]

        # vectorize the data
        self.vocab.setup(vocab_path=self.vocab_fn,
                         data=(train_lemmas_l, train_forms_l, train_feat_d),
                         lower=self.lower,
                         source='lemma_form_feat')

        self.train = self.train_data_to_ids(data=list(zip(
            train_lemmas_l,
            train_feat_d,
            train_aligned_pairs,
            train_forms_l)))

        self.dev = self.train_data_to_ids(data=list(zip(
            dev_lemmas_l,
            dev_feat_d,
            dev_aligned_pairs,
            dev_forms_l)))

    def prediction_setup(self):
        self.max_src_len = self.config['max_src_len']
        self.max_tgt_len = self.config['max_tgt_len']

        # else load vocabulary file from the filesystem
        assert 'vocab_fn' in self.config, logger.error('Specify "vocab_fn" in the config file')
        self.vocab.setup(vocab_path=self.vocab_fn,
                         data=None,
                         lower=self.lower,
                         source='file')

    def set_vocabulary_class(self):
        self.vocab = MorphVocabulary()

    @staticmethod
    def extract_gold_labels(lemma, aligned_pair):

        y_dec_targets = []
        x_attn_ids = []

        padded_lemma = add_bos_eos(lemma, bos=BaseMorphData.BOS, eos=BaseMorphData.EOS)

        i_input_index = 0

        # go through alignments, progress j when new output is introduced, progress i when new char is seen on lemma (no ~)
        aligned_lemma, aligned_word = aligned_pair
        aligned_lemma += BaseMorphData.EOS
        aligned_word += BaseMorphData.EOS

        # run through the alignments
        for align_index, (x_char, y_char) in enumerate(zip(aligned_lemma, aligned_word)):

            # if reached the end word symbol
            if y_char == BaseMorphData.EOS:
                x_attn_ids.append(i_input_index)
                y_dec_targets.append(BaseMorphData.EOS)
                continue

            # initially, if there is no prefix in the output (shouldn't delay on current input), step forward
            if padded_lemma[i_input_index] == BaseMorphData.BOS and aligned_lemma[align_index] != ALIGN_SYMBOL:
                x_attn_ids.append(i_input_index)
                y_dec_targets.append(STEP)
                # prepare for the next iteration - "feedback"
                i_input_index += 1

            # if 0-to-1 or 1-to-1 alignment, compute loss for predicting the output symbol
            if aligned_word[align_index] != ALIGN_SYMBOL:
                x_attn_ids.append(i_input_index)
                y_dec_targets.append(aligned_word[align_index])

            # if the input's not done and the next is not a 0-to-1 alignment, perform step
            if i_input_index < len(padded_lemma) - 1 and aligned_lemma[align_index + 1] != ALIGN_SYMBOL:
                # dec_prev_outputs.append(prev_output_char_idx)
                x_attn_ids.append(i_input_index)
                y_dec_targets.append(STEP)
                i_input_index += 1

        # self.sanity_check_oov_classes(dec_targets)
        return y_dec_targets, x_attn_ids

    def train_data_to_ids(self, data):

        logger.debug('Vectorizing data')

        assert self.max_src_len is not None
        assert self.max_tgt_len is not None

        data_split_enc_x = []
        data_split_attn_ids = []
        data_split_feat_x = []
        data_split_y = []
        forms = []

        for idx, datum in enumerate(data):

            lemma, feat_dict, aligned_pair, form = datum
            x_enc_inputs = BaseMorphData.lemma_chars_to_ids(lemma, self.vocab)
            x_feature_ids = BaseMorphData.feat_d_to_ids(feat_dict, self.vocab)
            y_dec_targets, x_attn_ids = MorphRNNHardData.extract_gold_labels(lemma, aligned_pair)
            y_dec_target_ids = MorphRNNHardData.y_labels_to_ids(y_dec_targets, self.vocab)

            data_split_enc_x.append(x_enc_inputs)
            data_split_feat_x.append(x_feature_ids)
            data_split_attn_ids.append(x_attn_ids)
            data_split_y.append(y_dec_target_ids)
            forms.append(form)

        return (data_split_enc_x, data_split_feat_x, data_split_attn_ids, data_split_y, forms)

    def batchify_training_data(self, xy_ids_and_targets, batch_size, is_dev_data):

        logger.debug('Batchifying data')

        # sort according to the src lengths -> xy pairs
        sorted_data = sorted(zip(*xy_ids_and_targets), key=lambda p: len(p[0]), reverse=True)
        data_size = len(sorted_data)
        num_batches = data_size // batch_size
        data_indices = self.index_data(data_size, mode='no_shuffling')

        batch_pairs = []
        batch_forms_all = []
        for bi in range(num_batches+1): # including the last (smaller) batch
            batch_data_tuple, batch_forms = self.make_one_batch(sorted_data, data_indices, bi, batch_size)
            batch_pairs.append(batch_data_tuple)
            batch_forms_all.extend(batch_forms)

        if is_dev_data:
            self.dev_references = batch_forms_all
            logger.info('Saving dev (training) references to --> %s', self.fnames.dev_ref_fn)
            save_txt(self.dev_references, self.fnames.dev_ref_fn)

        assert data_size == len(batch_forms_all)

        return batch_pairs

    def make_one_batch(self, sorted_data, data_indices, bi, batch_size):

        batch_x_enc = []
        batch_x_feat = []
        batch_x_attn = []
        batch_y = []
        batch_forms = []

        curr_batch_indices = data_indices[bi * batch_size: (bi + 1) * batch_size]

        for idx in curr_batch_indices:
            x_enc_ids, x_feature_ids, x_attn_ids, y_ids, forms = sorted_data[idx]

            x_enc_ids_copy = copy.deepcopy(x_enc_ids)
            batch_x_enc.append(x_enc_ids_copy)

            x_attn_ids_copy = copy.deepcopy(x_attn_ids)
            batch_x_attn.append(x_attn_ids_copy)

            x_feature_ids_copy = copy.deepcopy(x_feature_ids)
            batch_x_feat.append(x_feature_ids_copy)

            y_ids_copy = copy.deepcopy(y_ids)
            batch_y.append(y_ids_copy)

            batch_forms.append(forms)

        # skipping features, since the num does not change
        batch_x_enc_lens = [len(s) for s in batch_x_enc]
        batch_x_attn_lens = [len(s) for s in batch_x_attn]
        batch_y_lens = [len(s) for s in batch_y]
        assert batch_y_lens == batch_x_attn_lens

        max_enc_seq_len = max(batch_x_enc_lens)
        max_dec_seq_len = max(batch_y_lens)

        batch_enc_x_padded = [pad_seq(x, max_enc_seq_len, pad_id=self.vocab.PAD_ID) for x in batch_x_enc]
        batch_dec_y_padded = [pad_seq(y, max_dec_seq_len, pad_id=self.vocab.PAD_ID) for y in batch_y]
        # hard alignment for padded sequences
        batch_attn_x_padded = [pad_attn(x, max_dec_seq_len) for x in batch_x_attn]

        batch_enc_x_var = cuda_if_gpu(Variable(torch.LongTensor(batch_enc_x_padded)).transpose(0, 1))
        batch_dec_y_var = cuda_if_gpu(Variable(torch.LongTensor(batch_dec_y_padded)).transpose(0, 1))
        batch_x_feat_var = cuda_if_gpu(Variable(torch.LongTensor(batch_x_feat)))
        batch_attn_x = cuda_if_gpu(torch.LongTensor(batch_attn_x_padded).transpose(0, 1))  # not a Variable!

        return (batch_enc_x_var, batch_x_feat_var, batch_attn_x, batch_dec_y_var), batch_forms

    @staticmethod
    def y_labels_to_ids(y_labels, vocab):
        return [vocab.lookup_tok(label) for label in y_labels]


component = MorphRNNHardData
# extract from TRAIN lemmas, forms, feats if in training mode # align forms and lemmas # splitting aligned pairs # vectorize the data # else load vocabulary file from the filesystem # go through alignments, progress j when new output is introduced, progress i when new char is seen on lemma (no ~) # run through the alignments # if reached the end word symbol # initially, if there is no prefix in the output (shouldn't delay on current input), step forward # prepare for the next iteration - "feedback" # if 0-to-1 or 1-to-1 alignment, compute loss for predicting the output symbol # if the input's not done and the next is not a 0-to-1 alignment, perform step # dec_prev_outputs.append(prev_output_char_idx) # self.sanity_check_oov_classes(dec_targets) # sort according to the src lengths -> xy pairs # including the last (smaller) batch # skipping features, since the num does not change # hard alignment for padded sequences # not a Variable!