import numbers

import numpy as np

import theano

import theano.tensor as T

import warnings

from collections import Counter

from lasagne.layers import InputLayer, DropoutLayer, DenseLayer, EmbeddingLayer, NonlinearityLayer

from lasagne.layers import NINLayer, FeaturePoolLayer, ConcatLayer, SliceLayer, ElemwiseMergeLayer

from lasagne.layers import dimshuffle, reshape, get_output

from lasagne.layers.recurrent import Gate

from lasagne.init import Constant

from lasagne.objectives import categorical_crossentropy

from lasagne.nonlinearities import softmax

from lasagne.updates import rmsprop

from stanza.monitoring import progress

from stanza.research import config, instance, iterators, rng

import color_instances

import speaker

import data_aug

from helpers import ForgetSizeLayer, GaussianScoreLayer, logit_softmax_nd

from neural import NeuralLearner, SimpleLasagneModel

from neural import NONLINEARITIES, OPTIMIZERS, CELLS, sample

from vectorizers import SequenceVectorizer, BucketsVectorizer, SymbolVectorizer

from vectorizers import strip_invalid_tokens, COLOR_REPRS

from tokenizers import TOKENIZERS

random = rng.get_rng()

parser = config.get_options_parser()

parser.add_argument('--listener_cell_size', type=int, default=20,

help='The number of dimensions of all hidden layers and cells in '

'the listener model. If 0 and using the AtomicListenerLearner, '

'remove all hidden layers and only train a linear classifier.')

parser.add_argument('--listener_forget_bias', type=float, default=5.0,

help='The initial value of the forget gate bias in LSTM cells in '

'the listener model. A positive initial forget gate bias '

'encourages the model to remember everything by default.')

parser.add_argument('--listener_nonlinearity', choices=NONLINEARITIES.keys(), default='tanh',

help='The nonlinearity/activation function to use for dense and '

'LSTM layers in the listener model.')

parser.add_argument('--listener_cell', choices=CELLS.keys(), default='LSTM',

help='The recurrent cell to use for the listener model.')

parser.add_argument('--listener_bidi', type=config.boolean, default=False,

help='If True, recurrent listener models will use bidirectional RNNs; '

'otherwise, only the final output of the forward RNN is used.')

parser.add_argument('--listener_dropout', type=float, default=0.2,

help='The dropout rate (probability of setting a value to zero). '

'Dropout will be disabled if nonpositive.')

parser.add_argument('--listener_color_resolution', type=int, nargs='+', default=[4],

help='The number of buckets along each dimension of color space '

'for the output of the listener model.')

parser.add_argument('--listener_hidden_color_layers', type=int, default=0,

help='The number of dense layers after the color representation '

'(ContextListenerLearner only).')

parser.add_argument('--listener_hsv', type=config.boolean, default=False,

help='If True, output color buckets are in HSV space; otherwise, '

'color buckets will be in RGB. Final output instances will be in HSV '

'regardless; this sets the internal representation for training '

'and prediction.')

parser.add_argument('--listener_eval_batch_size', type=int, default=65536,

help='The number of examples per batch for evaluating the listener '

'model. Higher means faster but more memory usage. This should '

'not affect modeling accuracy.')

parser.add_argument('--listener_optimizer', choices=OPTIMIZERS.keys(), default='rmsprop',

help='The optimization (update) algorithm to use for listener training.')

parser.add_argument('--listener_learning_rate', type=float, default=0.1,

help='The learning rate to use for listener training.')

parser.add_argument('--listener_grad_clipping', type=float, default=0.0,

help='The maximum absolute value of the gradient messages for the'

'LSTM component of the listener model.')

parser.add_argument('--listener_color_repr', choices=COLOR_REPRS.keys(), default='buckets',

help='The representation of the color to use in the listener model: a regular '

'grid of `buckets`, overlapping bucket grids at multiple resolutions '

'(`ms`), `raw` RGB/HSV values, or a `fourier` transform-based '

'representation. Only used for ContextListenerLearner.')

parser.add_argument('--listener_tokenizer', choices=TOKENIZERS.keys(), default='whitespace',

help='The tokenization/preprocessing method to use for the listener model.')

parser.add_argument('--listener_unk_threshold', type=int, default=0,

help="The maximum number of occurrences of a token in the training data "

"before it's assigned a non-<unk> token index. 0 means nothing in "

"the training data is to be treated as unknown words; 1 means "

"single-occurrence words (hapax legomena) will be replaced with <unk>.")

class UnigramPrior(object):

for seq in sequences]

class AtomicUniformPrior(object):

return [instance.Instance(c) for c in self.vec.unvectorize_all(indices)]

class UnigramContextPrior(UnigramPrior):

return color_instances.reference_game(insts, color_instances.uniform, listener=True)

class AtomicUniformContextPrior(AtomicUniformPrior):

return color_instances.reference_game(insts, color_instances.uniform, listener=True)

PRIORS = {

}

parser.add_argument('--listener_prior', choices=PRIORS.keys(), default='Unigram',

help='The prior model for the listener (prior over utterances). '

'Only used in RSA learner.')

class ListenerLearner(NeuralLearner):

return self.prior_smooth.sample(num_samples)

class ContextListenerLearner(ListenerLearner):

extra_vars)

class TwoStreamListenerLearner(ContextListenerLearner):

return l_scores, [l_in] + context_inputs

class GaussianContextListenerLearner(ContextListenerLearner):

return mean[0], covar[0]

class ContextVecListenerLearner(ContextListenerLearner):

return l_scores, [l_in] + context_inputs

def broadcast_sub_layer(l_pred, l_targets, feature_dim, id_tag):

l_broadcast = dimshuffle(l_pred, (0, 1, 'x'), name=id_tag + 'sub_broadcast')

l_forget = ForgetSizeLayer(l_broadcast, axis=2, name=id_tag + 'sub_nosize')

return ElemwiseMergeLayer((l_forget, l_targets), T.sub, name=id_tag + 'broadcast_sub')

def broadcast_dot_layer(l_pred, l_targets, feature_dim, id_tag):

return reshape(l_pool, ([0], [2]), name=id_tag + 'broadcast_dot')

class AtomicListenerLearner(ListenerLearner):