<commit_before>
import numpy as np

# PyTorch
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F

from spinn.util.blocks import Embed, to_gpu, MLP, Linear, HeKaimingInitializer, gumbel_sample
from spinn.util.misc import Args, Vocab
from spinn.util.blocks import SimpleTreeLSTM
from spinn.util.sparks import sparks


def build_model(data_manager, initial_embeddings, vocab_size,
                num_classes, FLAGS, context_args, composition_args, logger=None):
    use_sentence_pair = data_manager.SENTENCE_PAIR_DATA
    model_cls = Pyramid

    return model_cls(model_dim=FLAGS.model_dim,
                     word_embedding_dim=FLAGS.word_embedding_dim,
                     vocab_size=vocab_size,
                     initial_embeddings=initial_embeddings,
                     num_classes=num_classes,
                     embedding_keep_rate=FLAGS.embedding_keep_rate,
                     use_sentence_pair=use_sentence_pair,
                     use_difference_feature=FLAGS.use_difference_feature,
                     use_product_feature=FLAGS.use_product_feature,
                     classifier_keep_rate=FLAGS.semantic_classifier_keep_rate,
                     mlp_dim=FLAGS.mlp_dim,
                     num_mlp_layers=FLAGS.num_mlp_layers,
                     mlp_ln=FLAGS.mlp_ln,
                     context_args=context_args,
                     gated=FLAGS.pyramid_gated,
                     trainable_temperature=FLAGS.pyramid_trainable_temperature,
                     test_temperature_mulitplier=FLAGS.pyramid_test_time_temperature_multiplier,
                     logger=logger,
                     gumbel=FLAGS.pyramid_gumbel,
                     )


class Pyramid(nn.Module):

    def __init__(self, model_dim=None,
                 word_embedding_dim=None,
                 vocab_size=None,
                 initial_embeddings=None,
                 num_classes=None,
                 embedding_keep_rate=None,
                 use_sentence_pair=False,
                 classifier_keep_rate=None,
                 mlp_dim=None,
                 num_mlp_layers=None,
                 mlp_ln=None,
                 context_args=None,
                 gated=None,
                 trainable_temperature=None,
                 test_temperature_mulitplier=None,
                 logger=None,
                 gumbel=None,
                 **kwargs
                 ):
        super(Pyramid, self).__init__()

        self.use_sentence_pair = use_sentence_pair
        self.model_dim = model_dim
        self.gated = gated
        self.test_temperature_mulitplier = test_temperature_mulitplier
        self.trainable_temperature = trainable_temperature
        self.logger = logger
        self.gumbel = gumbel

        classifier_dropout_rate = 1. - classifier_keep_rate

        args = Args()
        args.size = model_dim
        args.input_dropout_rate = 1. - embedding_keep_rate

        vocab = Vocab()
        vocab.size = initial_embeddings.shape[0] if initial_embeddings is not None else vocab_size
        vocab.vectors = initial_embeddings

        self.embed = Embed(word_embedding_dim, vocab.size, vectors=vocab.vectors)

        self.composition_fn = SimpleTreeLSTM(model_dim / 2,
                                             composition_ln=False)
        self.selection_fn = Linear(initializer=HeKaimingInitializer)(model_dim, 1)

        # TODO: Set up layer norm.

        mlp_input_dim = model_dim * 2 if use_sentence_pair else model_dim

        self.mlp = MLP(mlp_input_dim, mlp_dim, num_classes,
                       num_mlp_layers, mlp_ln, classifier_dropout_rate)

        if self.trainable_temperature:
            self.temperature = nn.Parameter(torch.ones(1, 1), requires_grad=True)

        self.encode = context_args.encoder
        self.reshape_input = context_args.reshape_input
        self.reshape_context = context_args.reshape_context

    def run_hard_pyramid(self, x):
        batch_size, seq_len, model_dim = x.data.size()
        # If we "merge left", that means words before the merge point stay put.
        # Example:
        # 012345678
        #     ++ : 4,5 -> X
        # 0123X678

        # This is supposed to be done for evaluation only, so I'm not
        # going to even bother with variables.
        words = torch.chunk(x, seq_len, 1) # A list of word vectors
        words = [torch.squeeze(word).data for word in words]

        # Used in phase 2.
        copy_left_buffer = torch.zeros(batch_size, model_dim)
        copy_right_buffer = torch.zeros(batch_size, model_dim)
        merge_buffer = torch.zeros(batch_size, model_dim)

        while len(words) > 1:
            print(len(words))
            composition_results = [] # Variable
            selection_logits_list = [] # Variable

            # Phase 1 - evaluate all possible mergers
            for position in range(len(words) - 1):
                left = Variable(words[position])
                right = Variable(words[position + 1])
                composed = self.composition_fn(left, right)
                selection_score = self.selection_fn(composed)
                composition_results.append(composed)
                selection_logits_list.append(selection_score)

            selection_logits = torch.cat(selection_logits_list, 1) # B x L
            merge_points = selection_logits.data.max(1)[1] # B

            # Phase 2 - apply mergers
            new_words = []
            for position in range(len(words) - 1): # destination position
                left = words[position]
                right = words[position + 1]

                # If merge at 4, positions 0-3 is copy left, 5-... is copy right
                is_copy_left = merge_points.gt(position).float()
                is_copy_right = merge_points.lt(position).float()
                is_merged = (1 - is_copy_left) * (1 - is_copy_right)

                # Expression equivalent to:
                #         is_copy_left.expand_as(left) * left + \
                #         is_copy_right.expand_as(right) * right + \
                #         is_merged.expand_as(left) * composition_results[position]
                copy_left_buffer.fill_(0.0).add_(left).mul_(is_copy_left.expand_as(left))
                copy_right_buffer.fill_(0.0).add_(right).mul_(is_copy_right.expand_as(left))
                merge_buffer.fill_(0.0).add_(composition_results[position].data).mul_(is_merged.expand_as(left))
                words[position].fill_(0.0).add_(copy_left_buffer).add_(copy_right_buffer).add_(merge_buffer)

            words.pop()
<commit_msg>buffers<commit_after>
import numpy as np

# PyTorch
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F

from spinn.util.blocks import Embed, to_gpu, MLP, Linear, HeKaimingInitializer, gumbel_sample
from spinn.util.misc import Args, Vocab
from spinn.util.blocks import SimpleTreeLSTM
from spinn.util.sparks import sparks


def build_model(data_manager, initial_embeddings, vocab_size,
                num_classes, FLAGS, context_args, composition_args, logger=None):
    use_sentence_pair = data_manager.SENTENCE_PAIR_DATA
    model_cls = Pyramid

    return model_cls(model_dim=FLAGS.model_dim,
                     word_embedding_dim=FLAGS.word_embedding_dim,
                     vocab_size=vocab_size,
                     initial_embeddings=initial_embeddings,
                     num_classes=num_classes,
                     embedding_keep_rate=FLAGS.embedding_keep_rate,
                     use_sentence_pair=use_sentence_pair,
                     use_difference_feature=FLAGS.use_difference_feature,
                     use_product_feature=FLAGS.use_product_feature,
                     classifier_keep_rate=FLAGS.semantic_classifier_keep_rate,
                     mlp_dim=FLAGS.mlp_dim,
                     num_mlp_layers=FLAGS.num_mlp_layers,
                     mlp_ln=FLAGS.mlp_ln,
                     context_args=context_args,
                     gated=FLAGS.pyramid_gated,
                     trainable_temperature=FLAGS.pyramid_trainable_temperature,
                     test_temperature_mulitplier=FLAGS.pyramid_test_time_temperature_multiplier,
                     logger=logger,
                     gumbel=FLAGS.pyramid_gumbel,
                     )


class Pyramid(nn.Module):

    def __init__(self, model_dim=None,
                 word_embedding_dim=None,
                 vocab_size=None,
                 initial_embeddings=None,
                 num_classes=None,
                 embedding_keep_rate=None,
                 use_sentence_pair=False,
                 classifier_keep_rate=None,
                 mlp_dim=None,
                 num_mlp_layers=None,
                 mlp_ln=None,
                 context_args=None,
                 gated=None,
                 trainable_temperature=None,
                 test_temperature_mulitplier=None,
                 logger=None,
                 gumbel=None,
                 **kwargs
                 ):
        super(Pyramid, self).__init__()

        self.use_sentence_pair = use_sentence_pair
        self.model_dim = model_dim
        self.gated = gated
        self.test_temperature_mulitplier = test_temperature_mulitplier
        self.trainable_temperature = trainable_temperature
        self.logger = logger
        self.gumbel = gumbel

        classifier_dropout_rate = 1. - classifier_keep_rate

        args = Args()
        args.size = model_dim
        args.input_dropout_rate = 1. - embedding_keep_rate

        vocab = Vocab()
        vocab.size = initial_embeddings.shape[0] if initial_embeddings is not None else vocab_size
        vocab.vectors = initial_embeddings

        self.embed = Embed(word_embedding_dim, vocab.size, vectors=vocab.vectors)

        self.composition_fn = SimpleTreeLSTM(model_dim / 2,
                                             composition_ln=False)
        self.selection_fn = Linear(initializer=HeKaimingInitializer)(model_dim, 1)

        # TODO: Set up layer norm.

        mlp_input_dim = model_dim * 2 if use_sentence_pair else model_dim

        self.mlp = MLP(mlp_input_dim, mlp_dim, num_classes,
                       num_mlp_layers, mlp_ln, classifier_dropout_rate)

        if self.trainable_temperature:
            self.temperature = nn.Parameter(torch.ones(1, 1), requires_grad=True)

        self.encode = context_args.encoder
        self.reshape_input = context_args.reshape_input
        self.reshape_context = context_args.reshape_context

    def run_hard_pyramid(self, x):
        batch_size, seq_len, model_dim = x.data.size()
        # If we "merge left", that means words before the merge point stay put.
        # Example:
        # 012345678
        #     ++ : 4,5 -> X
        # 0123X678
        #   ++   : 2,3 -> Y
        # 01YX678

        # This is supposed to be done for evaluation only, so I'm not
        # going to even bother with variables.
        words = torch.chunk(x, seq_len, 1) # A list of word vectors
        words = [torch.squeeze(word).data for word in words]

        # Used in phase 1.
        composition_buffers = [torch.zeros(batch_size, model_dim) for w in words]

        # Used in phase 2.
        copy_left_buffer = torch.zeros(batch_size, model_dim)
        copy_right_buffer = torch.zeros(batch_size, model_dim)
        merge_buffer = torch.zeros(batch_size, model_dim)

        while len(words) > 1:
            print(len(words))
            composition_results = composition_buffers[:(len(words) - 1)]
            selection_logits_list = [] # Tensors (not Variable), Bx1 each

            # Phase 1 - evaluate all possible mergers
            for position in range(len(words) - 1):
                left = Variable(words[position])
                right = Variable(words[position + 1])
                composed = self.composition_fn(left, right)
                selection_score = self.selection_fn(composed)
                composition_results[position].copy_(composed.data)
                selection_logits_list.append(selection_score.data)

            # Find the highest scoring position to merge.
            selection_logits = torch.cat(selection_logits_list, 1) # B x L
            merge_points = selection_logits.max(1)[1] # B

            # Phase 2 - apply mergers
            new_words = []
            for position in range(len(words) - 1): # destination position
                left = words[position]
                right = words[position + 1]

                # If merge at 4, positions 0-3 is copy left, 5-... is copy right
                is_copy_left = merge_points.gt(position).float()
                is_copy_right = merge_points.lt(position).float()
                is_merged = (1 - is_copy_left) * (1 - is_copy_right)

                # Expression equivalent to:
                #         is_copy_left.expand_as(left) * left + \
                #         is_copy_right.expand_as(right) * right + \
                #         is_merged.expand_as(left) * composition_results[position]
                copy_left_buffer.copy_(left).mul_(is_copy_left.expand_as(left))
                copy_right_buffer.copy_(right).mul_(is_copy_right.expand_as(left))
                merge_buffer.copy_(composition_results[position]).mul_(is_merged.expand_as(left))
                words[position].copy_(copy_left_buffer).add_(copy_right_buffer).add_(merge_buffer)

            words.pop()

{'fraction_non_alphanumeric': 0.07435745060169366, 'fraction_numerical': 0.007874015748031496, 'mean_word_length': 1.9891207815275311, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4476682', 'n_tokens_mistral': 4050, 'n_tokens_neox': 3785, 'n_words': 710}