dictlearn/aggregation_schemes.py
class Perplexity(Mean):

    def get_aggregator(self):
        aggregator = super(Perplexity, self).get_aggregator()
        if aggregator.readout_variable.ndim > 0:
            raise ValueError("can only compute perplexity for a scalar")
        aggregator.readout_variable = tensor.exp(aggregator.readout_variable)
        return aggregator
dictlearn/lookup.py
class LSTMReadDefinitions(Initializable):
    """
    Converts definition into embeddings.

    Parameters
    ----------
    num_input_words: int, default: -1
        If non zero will (a bit confusing name) restrict dynamically vocab.
        WARNING: it assumes word ids are monotonical with frequency!

    emb_dim : int
        Dimensionality of word embeddings

    dim : int
        Dimensionality of the def rnn.

    lookup: None or LookupTable

    fork_and_rnn: None or tuple (Linear, RNN)
    """
    def __init__(self, num_input_words, emb_dim, dim, vocab, lookup=None,
                 fork_and_rnn=None, cache=None, **kwargs):

        self._vocab = vocab
        self._cache = cache
        children = []
        if num_input_words > 0:
            logger.info("Restricting def vocab to " + str(num_input_words))
            self._num_input_words = num_input_words
        else:
            self._num_input_words = vocab.size()


        if lookup is None:
            self._def_lookup = LookupTable(self._num_input_words, emb_dim, name='def_lookup')
        else:
            self._def_lookup = lookup

        if fork_and_rnn is None:
            self._def_fork = Linear(emb_dim, 4 * dim, name='def_fork')
            self._def_rnn = LSTM(dim, name='def_rnn')
        else:
            self._def_fork, self._def_rnn = fork_and_rnn

        children.extend([self._def_lookup, self._def_fork, self._def_rnn])

        super(LSTMReadDefinitions, self).__init__(children=children, **kwargs)


    @application
    def apply(self, application_call,
              defs, def_mask):
        """
        Returns vector per each word in sequence using the dictionary based lookup
        """
        # Short listing
        defs_sl_main = (T.lt(defs, self._num_input_words) * defs
                   + T.ge(defs, self._num_input_words) * self._vocab.unk)
        defs_sl_cache = (T.ge(defs, self._num_input_words) * defs
                   + T.lt(defs, self._num_input_words) * self._vocab.unk)

        application_call.add_auxiliary_variable(
            unk_ratio(defs_sl_main, def_mask, self._vocab.unk),
            name='def_unk_ratio')

        embedded_def_words = self._def_lookup.apply(defs_sl_main)
        cached_embeddings = self._cache.apply(defs_sl_cache) 
        final_embeddings = (T.lt(defs, self._num_input_words).dimshuffle(0,1,'x') * embedded_def_words
                + T.ge(defs, self._num_input_words).dimshuffle(0, 1, 'x') * cached_embeddings)

        def_embeddings = self._def_rnn.apply(
            T.transpose(self._def_fork.apply(final_embeddings), (1, 0, 2)),
            mask=def_mask.T)[0][-1]

        return def_embeddings
dictlearn/theano_util.py
def unk_ratio(words, mask, unk):
    num_unk = (tensor.eq(words, unk) * mask).sum()
    return num_unk / mask.sum()
dictlearn/ops.py
class RetrievalOp(theano.Op):
    """Retrieves the definitions from the dictionary."""
    def __init__(self, retrieval):
        self._retrieval = retrieval

    def make_node(self, input_):
        defs_type = tensor.TensorType('int64', [False, False])
        def_mask_type = tensor.TensorType('float32', [False, False])
        # both type happened to be the same, but this is just a coincidence
        def_map_type = defs_type
        return theano.Apply(
            self, [input_], [defs_type(), def_mask_type(), def_map_type()])

    def perform(self, node, inputs, output_storage):
        defs, def_mask, def_map = self._retrieval.retrieve_and_pad(inputs[0])
        output_storage[0][0] = defs
        output_storage[1][0] = def_mask
        output_storage[2][0] = def_map
dictlearn/ops.py
class WordToCountOp(theano.Op):
    """Replaces words with their counts."""
    def __init__(self, vocab):
        self._vocab = vocab

    def make_node(self, input_):
        input_ = tensor.as_tensor_variable(input_)
        output_type = tensor.TensorType(
            input_.dtype, input_.broadcastable[:-1])
        return theano.Apply(self, [input_], [output_type()])

    def perform(self, node, inputs, output_storage):
        words = inputs[0]
        words_flat = words.reshape(-1, words.shape[-1])
        word_counts = numpy.array([self._vocab.word_freq(vec2str(word))
                                for word in words_flat])
        output_storage[0][0] = word_counts.reshape(words.shape[:-1])
dictlearn/ops.py
class WordToIdOp(theano.Op):
    """Replaces words with their ids."""
    def __init__(self, vocab):
        self._vocab = vocab

    def make_node(self, input_):
        input_ = tensor.as_tensor_variable(input_)
        output_type = tensor.TensorType(
            input_.dtype, input_.broadcastable[:-1])
        return theano.Apply(self, [input_], [output_type()])

    def perform(self, node, inputs, output_storage):
        words = inputs[0]
        words_flat = words.reshape(-1, words.shape[-1])
        word_ids = numpy.array([self._vocab.word_to_id(vec2str(word))
                                for word in words_flat])
        output_storage[0][0] = word_ids.reshape(words.shape[:-1])
dictlearn/lookup.py
class MeanPoolReadDefinitions(Initializable):
    """
    Converts definition into embeddings using simple sum + translation

    Parameters
    ----------
    num_input_words: int, default: -1
        If non zero will (a bit confusing name) restrict dynamically vocab.
        WARNING: it assumes word ids are monotonical with frequency!

    dim : int
        Dimensionality of the def rnn.

    emb_dim : int
        Dimensionality of word embeddings

    """
    def __init__(self, num_input_words, emb_dim, dim, vocab,
                 lookup=None, translate=True, normalize=True, **kwargs):

        if num_input_words > 0:
            logger.info("Restricting def vocab to " + str(num_input_words))
            self._num_input_words = num_input_words
        else:
            self._num_input_words = vocab.size()

        self._vocab = vocab
        self._translate = translate
        self._normalize = normalize

        children = []

        if lookup is None:
            logger.info("emb_dim={}".format(emb_dim))
            self._def_lookup = LookupTable(self._num_input_words, emb_dim, name='def_lookup')
        else:
            self._def_lookup = lookup

        # Makes sense for shared lookup. Then we precondition embeddings.
        # Doesn't makes otherwise (WH = W')
        # TODO(kudkudak): Refactor redundant translate parameter
        if self._translate:
            if emb_dim == dim:
                raise Exception("Redundant layer")

            self._def_translate = Linear(emb_dim, dim, name='def_translate')
            children.extend([self._def_translate])
        else:
            if emb_dim != dim:
                raise Exception("Please pass translate=True if emb_dim != dim")

        children.append(self._def_lookup)

        super(MeanPoolReadDefinitions, self).__init__(children=children, **kwargs)


    @application
    def apply(self, application_call,
              defs, def_mask):
        """
        Returns vector per each word in sequence using the dictionary based lookup
        """
        # Short listing
        defs = (T.lt(defs, self._num_input_words) * defs
                + T.ge(defs, self._num_input_words) * self._vocab.unk)
        # Memory bottleneck:
        # For instance (16101,52,300) ~= 32GB.
        # [(16786, 52, 1), (16786, 52, 100)]
        # TODO: Measure memory consumption here and check if it is in sensible range
        # or maybe introduce some control in Retrieval?
        defs_emb = self._def_lookup.apply(defs)
        application_call.add_auxiliary_variable(
            unk_ratio(defs, def_mask, self._vocab.unk),
            name='def_unk_ratio')

        if self._translate:
            logger.info("Translating in MeanPoolReadDefinitions")
            # Translate. Crucial for recovering useful information from embeddings
            defs_emb = self._def_translate.apply(defs_emb)

        def_emb_mask = def_mask[:, :, None]
        defs_emb = (def_emb_mask * defs_emb).sum(axis=1)
        if self._normalize:
            defs_emb = defs_emb / def_emb_mask.sum(axis=1)

        return defs_emb
dictlearn/stuff.py
class DebugLSTM(BaseRecurrent, Initializable):
    u"""Long Short Term Memory.

    Every unit of an LSTM is equipped with input, forget and output gates.
    This implementation is based on code by Mohammad Pezeshki that
    implements the architecture used in [GSS03]_ and [Grav13]_. It aims to
    do as many computations in parallel as possible and expects the last
    dimension of the input to be four times the output dimension.

    Unlike a vanilla LSTM as described in [HS97]_, this model has peephole
    connections from the cells to the gates. The output gates receive
    information about the cells at the current time step, while the other
    gates only receive information about the cells at the previous time
    step. All 'peephole' weight matrices are diagonal.

    .. [GSS03] Gers, Felix A., Nicol N. Schraudolph, and Jürgen
        Schmidhuber, *Learning precise timing with LSTM recurrent
        networks*, Journal of Machine Learning Research 3 (2003),
        pp. 115-143.
    .. [Grav13] Graves, Alex, *Generating sequences with recurrent neural
        networks*, arXiv preprint arXiv:1308.0850 (2013).
    .. [HS97] Sepp Hochreiter, and Jürgen Schmidhuber, *Long Short-Term
        Memory*, Neural Computation 9(8) (1997), pp. 1735-1780.

    Parameters
    ----------
    dim : int
        The dimension of the hidden state.
    activation : :class:`~.bricks.Brick`, optional
        The activation function. The default and by far the most popular
        is :class:`.Tanh`.
    gate_activation : :class:`~.bricks.Brick` or None
        The brick to apply as activation for gates (input/output/forget).
        If ``None`` a :class:`.Logistic` brick is used.

    Notes
    -----
    See :class:`.Initializable` for initialization parameters.

    """
    @lazy(allocation=['dim'])
    def __init__(self, dim, activation=None, gate_activation=None, **kwargs):
        self.dim = dim

        if not activation:
            activation = Tanh()
        if not gate_activation:
            gate_activation = Logistic()
        self.activation = activation
        self.gate_activation = gate_activation

        children = [self.activation, self.gate_activation]
        kwargs.setdefault('children', []).extend(children)
        super(DebugLSTM, self).__init__(**kwargs)

    def get_dim(self, name):
        if name == 'inputs':
            return self.dim * 4
        if name in ['states', 'cells']:
            return self.dim
        if name == 'mask':
            return 0
        return super(DebugLSTM, self).get_dim(name)


    def _allocate(self):
        self.W_state = shared_floatx_nans((self.dim, 4*self.dim),
                                          name='W_state')
        self.W_cell_to_in = shared_floatx_nans((self.dim,),
                                               name='W_cell_to_in')
        self.W_cell_to_forget = shared_floatx_nans((self.dim,),
                                                   name='W_cell_to_forget')
        self.W_cell_to_out = shared_floatx_nans((self.dim,),
                                                name='W_cell_to_out')
        # The underscore is required to prevent collision with
        # the `initial_state` application method
        self.initial_state_ = shared_floatx_zeros((self.dim,),
                                                  name="initial_state")
        self.initial_cells = shared_floatx_zeros((self.dim,),
                                                 name="initial_cells")
        add_role(self.W_state, WEIGHT)
        add_role(self.W_cell_to_in, WEIGHT)
        add_role(self.W_cell_to_forget, WEIGHT)
        add_role(self.W_cell_to_out, WEIGHT)
        add_role(self.initial_state_, INITIAL_STATE)
        add_role(self.initial_cells, INITIAL_STATE)

        self.parameters = [
            self.W_state, self.W_cell_to_in, self.W_cell_to_forget,
            self.W_cell_to_out, self.initial_state_, self.initial_cells]

    def _initialize(self):
        for weights in self.parameters[:4]:
            self.weights_init.initialize(weights, self.rng)

    @recurrent(sequences=['inputs', 'mask'], states=['states', 'cells'],
               contexts=[], outputs=['states', 'cells', 'input_gates', 'forget_gates', 'output_gates'])
    def inner_apply(self, inputs, states, cells, mask=None):
        """Apply the Long Short Term Memory transition.

        Parameters
        ----------
        states : :class:`~tensor.TensorVariable`
            The 2 dimensional matrix of current states in the shape
            (batch_size, features). Required for `one_step` usage.
        cells : :class:`~tensor.TensorVariable`
            The 2 dimensional matrix of current cells in the shape
            (batch_size, features). Required for `one_step` usage.
        inputs : :class:`~tensor.TensorVariable`
            The 2 dimensional matrix of inputs in the shape (batch_size,
            features * 4). The `inputs` needs to be four times the
            dimension of the LSTM brick to insure each four gates receive
            different transformations of the input. See [Grav13]_
            equations 7 to 10 for more details. The `inputs` are then split
            in this order: Input gates, forget gates, cells and output
            gates.
        mask : :class:`~tensor.TensorVariable`
            A 1D binary array in the shape (batch,) which is 1 if there is
            data available, 0 if not. Assumed to be 1-s only if not given.

        .. [Grav13] Graves, Alex, *Generating sequences with recurrent*
            *neural networks*, arXiv preprint arXiv:1308.0850 (2013).

        Returns
        -------
        states : :class:`~tensor.TensorVariable`
            Next states of the network.
        cells : :class:`~tensor.TensorVariable`
            Next cell activations of the network.

        """
        def slice_last(x, no):
            return x[:, no*self.dim: (no+1)*self.dim]

        activation = tensor.dot(states, self.W_state) + inputs
        in_gate = self.gate_activation.apply(
            slice_last(activation, 0) + cells * self.W_cell_to_in)
        forget_gate = self.gate_activation.apply(
            slice_last(activation, 1) + cells * self.W_cell_to_forget)
        next_cells = (
            forget_gate * cells +
            in_gate * self.activation.apply(slice_last(activation, 2)))
        out_gate = self.gate_activation.apply(
            slice_last(activation, 3) + next_cells * self.W_cell_to_out)
        next_states = out_gate * self.activation.apply(next_cells)

        if mask:
            next_states = (mask[:, None] * next_states +
                           (1 - mask[:, None]) * states)
            next_cells = (mask[:, None] * next_cells +
                          (1 - mask[:, None]) * cells)

        return next_states, next_cells, in_gate, forget_gate, out_gate

    @application(outputs=['states', 'cells'])
    def apply(self, application_call, *args, **kwargs):
        results = self.inner_apply(*args, **kwargs)
        application_call.add_auxiliary_variable(results[2].copy(), name='input_gates')
        application_call.add_auxiliary_variable(results[3].copy(), name='forget_gates')
        application_call.add_auxiliary_variable(results[4].copy(), name='output_gates')
        return results[:2]


    @application(outputs=inner_apply.states)
    def initial_states(self, batch_size, *args, **kwargs):
        return [tensor.repeat(self.initial_state_[None, :], batch_size, 0),
                tensor.repeat(self.initial_cells[None, :], batch_size, 0)]
dictlearn/lookup.py
class MeanPoolCombiner(Initializable):
    """
    Parameters
    ----------
    dim: int
        Dimensionality of def embedding

    dropout_type: str

    dropout: float, defaut: 0.0

    emb_dim: int
        Dimensionality of word embeddings, as well as final output

    compose_type : str
        If 'sum', the definition and word embeddings are averaged
        If 'fully_connected_linear', a learned perceptron compose the 2
        embeddings linearly
        If 'fully_connected_relu', ...
        If 'fully_connected_tanh', ...
    """

    def __init__(self, emb_dim, dim, dropout=0.0,
            def_word_gating="none",
            dropout_type="per_unit", compose_type="sum",
            word_dropout_weighting="no_weighting",
            shortcut_unk_and_excluded=False,
            num_input_words=-1, exclude_top_k=-1, vocab=None,
            **kwargs):

        self._dropout = dropout
        self._num_input_words = num_input_words
        self._exclude_top_K = exclude_top_k
        self._dropout_type = dropout_type
        self._compose_type = compose_type
        self._vocab = vocab
        self._shortcut_unk_and_excluded = shortcut_unk_and_excluded
        self._word_dropout_weighting = word_dropout_weighting
        self._def_word_gating = def_word_gating

        if def_word_gating not in {"none", "self_attention"}:
            raise NotImplementedError()

        if word_dropout_weighting not in {"no_weighting"}:
            raise NotImplementedError("Not implemented " + word_dropout_weighting)

        if dropout_type not in {"per_unit", "per_example", "per_word"}:
            raise NotImplementedError()

        children = []

        if self._def_word_gating=="self_attention":
            self._gate_mlp = Linear(dim, dim)
            self._gate_act = Logistic()
            children.extend([self._gate_mlp, self._gate_act])

        if compose_type == 'fully_connected_linear':
            self._def_state_compose = MLP(activations=[None],
                dims=[emb_dim + dim, emb_dim])
            children.append(self._def_state_compose)
        if compose_type == "gated_sum" or compose_type == "gated_transform_and_sum":
            if dropout_type == "per_word" or dropout_type == "per_example":
                raise RuntimeError("I dont think this combination makes much sense")

            self._compose_gate_mlp = Linear(dim + emb_dim, emb_dim,
                                            name='gate_linear')
            self._compose_gate_act = Logistic()
            children.extend([self._compose_gate_mlp, self._compose_gate_act])
        if compose_type == 'sum':
            if not emb_dim == dim:
                raise ValueError("Embedding has different dim! Cannot use compose_type='sum'")
        if compose_type == 'transform_and_sum' or compose_type == "gated_transform_and_sum":
            self._def_state_transform = Linear(dim, emb_dim, name='state_transform')
            children.append(self._def_state_transform)

        super(MeanPoolCombiner, self).__init__(children=children, **kwargs)

    @application
    def apply(self, application_call,
              word_embs, words_mask,
              def_embeddings, def_map, train_phase=False, word_ids=False, call_name=""):
        """ return a triple:
        * final embeddings
        * unique def embeddings
        * position of the def embeddings in the sentence (so that id can be retrieved)
        """
        batch_shape = word_embs.shape
        # batch_shape[1] = bs?
        flat_indices = def_map[:, 0] * batch_shape[1] + def_map[:, 1] # Index of word in flat
        unique_indices = T.extra_ops.Unique()(flat_indices)

        # def_map is a list of (seq_pos, word_pos, def_index)
        # def_embeddings is (id, emb_dim)

        def_sum = T.zeros((batch_shape[0] * batch_shape[1], def_embeddings.shape[1]))
        def_lens = T.zeros_like(def_sum[:, 0])
        def_lens = T.inc_subtensor(def_lens[flat_indices], 1)

        if self._def_word_gating == "none":
            def_sum = T.inc_subtensor(def_sum[flat_indices],
                def_embeddings[def_map[:, 2]])
            def_mean = def_sum / T.maximum(def_lens[:, None], 1)
        elif self._def_word_gating == "self_attention":
            gates = def_embeddings[def_map[:, 2]]
            gates = self._gate_mlp.apply(gates)[:, 0]
            application_call.add_auxiliary_variable(gates, name='def_gates')

            # Dima: this is numerically unstable. But maybe it can work.
            # If it can work, we can avoid too much coding.
            def_normalization = T.zeros_like(def_lens)
            def_normalization = T.inc_subtensor(
                def_normalization[flat_indices], T.exp(gates))
            gates = T.exp(gates) / def_normalization[flat_indices]

            def_mean = T.inc_subtensor(def_sum[flat_indices],
                gates[:, None] * def_embeddings[def_map[:, 2]])
        else:
            raise NotImplementedError()
        
        # we take the newly computed embeddings.
        # we want to update the lookup but we don't have access to the word ids here
        # we can return the updated embeddings along with their positions in the text
        updated_embeddings = def_mean[unique_indices]

        def_mean = def_mean.reshape((batch_shape[0], batch_shape[1], -1))

        application_call.add_auxiliary_variable(
            masked_root_mean_square(def_mean, words_mask), name=call_name + '_def_mean_rootmean2')

        if train_phase and self._dropout != 0.0:
            if self._dropout_type == "per_unit":
                logger.info("Adding per_unit drop on dict and normal emb")
                word_embs = apply_dropout(word_embs, drop_prob=self._dropout)
                def_mean = apply_dropout(def_mean, drop_prob=self._dropout)
            elif self._dropout_type == "per_example":
                logger.info("Adding per_example drop on dict and normal emb")
                # We dropout mask
                mask_defs = T.ones((batch_shape[0],))
                mask_we = T.ones((batch_shape[0],))

                # Mask dropout
                mask_defs = apply_dropout(mask_defs, drop_prob=self._dropout)
                mask_we = apply_dropout(mask_we, drop_prob=self._dropout)

                # this reduces variance. If both 0 will select both
                where_both_zero = T.eq((mask_defs + mask_we), 0)

                mask_defs = (where_both_zero + mask_defs).dimshuffle(0, "x", "x")
                mask_we = (where_both_zero + mask_we).dimshuffle(0, "x", "x")

                def_mean = mask_defs * def_mean
                word_embs = mask_we * word_embs
            elif self._dropout_type == "per_word_independent":
                # TODO: Maybe we also want to have possibility of including both (like in per_example)
                pass # TODO: implement
            elif self._dropout_type == "per_word":
                # First select if to perform dropout at all
                mask_higher = T.ones((batch_shape[0], batch_shape[1]))
                mask_higher = apply_dropout(mask_higher, drop_prob=self._dropout)
                mask_higher = mask_higher.dimshuffle(0, 1, "x")

                logger.info("Apply per_word dropou on dict and normal emb")
                # And if yes just 50% word vs 50% def
                mask = T.ones((batch_shape[0], batch_shape[1]))
                mask = apply_dropout(mask, drop_prob=0.5)
                mask = mask.dimshuffle(0, 1, "x")

                # Competitive
                def_mean = mask_higher * def_mean + (1 - mask_higher) * mask * def_mean
                word_embs = mask_higher * word_embs + (1 - mask_higher) * (1 - mask) * word_embs

                # TODO: Smarter weighting (at least like divisor in dropout)

                if not self._compose_type == "sum" and not self._compose_type == "transform_and_sum":
                    raise NotImplementedError()

        application_call.add_auxiliary_variable(
            def_mean.copy(),
            name=call_name + '_dict_word_embeddings')

        if self._compose_type == 'sum':
            final_embeddings = word_embs + def_mean
        elif self._compose_type == 'transform_and_sum':
            final_embeddings = (word_embs +
                                self._def_state_transform.apply(def_mean))
        elif self._compose_type == 'gated_sum' or self._compose_type == 'gated_transform_and_sum':
            concat = T.concatenate([word_embs, def_mean], axis=2)
            gates = concat.reshape((batch_shape[0] * batch_shape[1], -1))
            gates = self._compose_gate_mlp.apply(gates)
            gates = self._compose_gate_act.apply(gates)
            gates = gates.reshape((batch_shape[0], batch_shape[1], -1))

            if self._compose_type == 'gated_sum':
                final_embeddings = gates * word_embs + (1 - gates) * def_mean
            else:
                final_embeddings = gates * word_embs + (1 - gates) * self._def_state_transform.apply(def_mean)

            application_call.add_auxiliary_variable(
                masked_root_mean_square(gates.reshape((batch_shape[0], batch_shape[1], -1)), words_mask),
                name=call_name + '_compose_gate_rootmean2')
        elif self._compose_type.startswith('fully_connected'):
            concat = T.concatenate([word_embs, def_mean], axis=2)
            final_embeddings = self._def_state_compose.apply(concat)
        else:
            raise NotImplementedError()

        if self._shortcut_unk_and_excluded:
            # NOTE: It might be better to move it out of Lookup, because it breaks API a bit
            # but at the same time it makes sense to share this code

            # 1. If no def, just go with word emb
            final_embeddings = word_embs * T.lt(word_ids, self._exclude_top_K).dimshuffle(0, 1, "x") + \
                               final_embeddings * T.ge(word_ids, self._exclude_top_K).dimshuffle(0, 1, "x")

            # 2. UNKs always get def embeddings (UNK can happen for dev/test set of course)
            final_embeddings = final_embeddings * T.neq(word_ids, self._vocab.unk).dimshuffle(0, 1, "x") + \
                               def_mean * T.eq(word_ids, self._vocab.unk).dimshuffle(0, 1, "x")


        application_call.add_auxiliary_variable(
            masked_root_mean_square(final_embeddings, words_mask),
            name=call_name + '_merged_input_rootmean2')

        return final_embeddings, updated_embeddings, unique_indices
dictlearn/util.py
def masked_root_mean_square(x, mask):
    """Masked root mean square for a 3D tensor"""
    return (smart_sum((x * mask[:, :, None]) ** 2) / x.shape[2] / mask.sum()) ** 0.5


Question: import theano
from theano import tensor
from blocks.bricks import (Initializable, Linear, NDimensionalSoftmax, MLP,
                           Tanh, Rectifier)
from blocks.bricks.base import application
from blocks.bricks.recurrent import LSTM
from blocks.bricks.lookup import LookupTable
from blocks.initialization import Constant
from dictlearn.theano_util import unk_ratio
from dictlearn.ops import WordToIdOp, RetrievalOp, WordToCountOp
from dictlearn.aggregation_schemes import Perplexity
from dictlearn.stuff import DebugLSTM
from dictlearn.util import masked_root_mean_square
from dictlearn.lookup import (LSTMReadDefinitions, MeanPoolReadDefinitions,
                              MeanPoolCombiner)
"""A dictionary-equipped language model."""





class LanguageModel(Initializable):
    """The dictionary-equipped language model.

    Parameters
    ----------
    emb_dim: int
        The dimension of word embeddings (including for def model if standalone)
    dim : int
        The dimension of the RNNs states (including for def model if standalone)
    num_input_words : int
        The size of the LM's input vocabulary.
    num_output_words : int
        The size of the LM's output vocabulary.
    vocab
        The vocabulary object.
    retrieval
        The dictionary retrieval algorithm. If `None`, the language model
        does not use any dictionary.
    def_reader: either 'LSTM' or 'mean'
    standalone_def_rnn : bool
        If `True`, a standalone RNN with separate word embeddings is used
        to embed definition. If `False` the language model is reused.
    disregard_word_embeddings : bool
        If `True`, the word embeddings are not used, only the information
        from the definitions is used.
    compose_type : str
        If 'sum', the definition and word embeddings are averaged
        If 'fully_connected_linear', a learned perceptron compose the 2
        embeddings linearly
        If 'fully_connected_relu', ...
        If 'fully_connected_tanh', ...

    """
    def __init__(self, emb_dim, emb_def_dim, dim, num_input_words, def_num_input_words,
                 num_output_words,
                 vocab, retrieval=None,
                 def_reader='LSTM',
                 standalone_def_lookup=True,
                 standalone_def_rnn=True,
                 disregard_word_embeddings=False,
                 compose_type='sum',
                 very_rare_threshold=[10],
                 cache_size=0,
                 **kwargs):
        # TODO(tombosc): document
        if emb_dim == 0:
            emb_dim = dim
        if emb_def_dim == 0:
            emb_def_dim = emb_dim
        if num_input_words == 0:
            num_input_words = vocab.size()
        if def_num_input_words == 0:
            def_num_input_words = num_input_words

        if (num_input_words != def_num_input_words) and (not standalone_def_lookup):
            raise NotImplementedError()

        self._very_rare_threshold = very_rare_threshold
        self._num_input_words = num_input_words
        self._num_output_words = num_output_words
        self._vocab = vocab
        self._retrieval = retrieval
        self._disregard_word_embeddings = disregard_word_embeddings
        self._compose_type = compose_type

        self._word_to_id = WordToIdOp(self._vocab)
        self._word_to_count = WordToCountOp(self._vocab)

        children = []
        self._cache = None
        if cache_size > 0:
            #TODO(tombosc) do we implement cache as LookupTable or theano matrix?
            #self._cache = theano.shared(np.zeros((def_num_input_words, emb_dim)))
            self._cache = LookupTable(cache_size, emb_dim,
                                      name='cache_def_embeddings')
            children.append(self._cache)

        if self._retrieval:
            self._retrieve = RetrievalOp(retrieval)

        self._main_lookup = LookupTable(self._num_input_words, emb_dim, name='main_lookup')
        self._main_fork = Linear(emb_dim, 4 * dim, name='main_fork')
        self._main_rnn = DebugLSTM(dim, name='main_rnn') # TODO(tombosc): use regular LSTM?
        children.extend([self._main_lookup, self._main_fork, self._main_rnn])
        if self._retrieval:
            if standalone_def_lookup:
                lookup = None
            else:
                if emb_dim != emb_def_dim:
                    raise ValueError("emb_dim != emb_def_dim: cannot share lookup")
                lookup = self._main_lookup

            if def_reader == 'LSTM':
                if standalone_def_rnn:
                    fork_and_rnn = None
                else:
                    fork_and_rnn = (self._main_fork, self._main_rnn)
                self._def_reader = LSTMReadDefinitions(def_num_input_words, emb_def_dim,
                                                       dim, vocab, lookup,
                                                       fork_and_rnn, cache=self._cache)
            
            elif def_reader == 'mean':

Answer:                 self._def_reader = MeanPoolReadDefinitions(def_num_input_words, emb_def_dim,