<reponame>DarkMatterAI/mrl
# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/10_generative_models.lstm_lm.ipynb (unless otherwise specified).

__all__ = ['LSTM_LM', 'Conditional_LSTM_LM']

# Cell

from ..imports import *
from ..torch_imports import *
from ..torch_core import *
from ..layers import *
from .generative_base import *

# Cell

class LSTM_LM(GenerativeModel):
    '''
    LSTM_LM - LSTM language model

    Inputs:

    - `d_vocab int`: vocab size

    - `d_embedding int`: embedding size

    - `d_hidden int`: hidden dimension

    - `n_layers int`: number of LSTM layers

    - `input_dropout float`: dropout on the embedding layer

    - `lstm_dropout float`: dropout on the LSTM layers

    - `bos_idx int`: beginning of sentence token

    - `bidir bool`: if the LSTMs are bidirectional.
    Must be False for use in generative models

    - `tie_weights bool`: if True, tie the weights of
    the embedding and the output layer
    '''
    def __init__(self, d_vocab, d_embedding, d_hidden, n_layers,
                input_dropout=0., lstm_dropout=0., bos_idx=0,
                bidir=False, tie_weights=False):
        super().__init__()

        self.embedding = nn.Embedding(d_vocab, d_embedding)

        self.lstm = LSTM(d_embedding, d_hidden, d_embedding, n_layers,
                                     bidir=bidir, input_dropout=input_dropout,
                                     lstm_dropout=lstm_dropout)

        self.head = nn.Linear(d_embedding, d_vocab)

        if tie_weights:
            self.embedding.weight = self.head.weight

        self.bos_idx = bos_idx

    def _forward(self, x, hiddens=None):
        x = self.embedding(x)
        encoded, hiddens = self.lstm(x, hiddens)
        output = self.head(encoded)
        return output, hiddens, encoded

    def forward(self, x, hiddens=None):
        output, hiddens, encoded = self._forward(x, hiddens)

        return output

    def x_to_latent(self, x):
        return None

    def sample(self, bs, sl, z=None, temperature=1., multinomial=True):
        '''
        sample - sample from the model

        Inputs:

        - `bs int`: batch size

        - `sl int`: maximum sequence length

        - `z Optional[torch.Tensor[bs, d_latent]]`: latent vector

        - `temperature float`: sample temperature

        - `multinomial bool`: if True, use multinomial sampling.
        If False, use argmax greedy sampling

        Returns:

        - `preds torch.LongTensor[bs, sl]`: predicted sequence tokens

        - `lps torch.FloatTensor[bs, sl, d_vocab]`: prediction log probabilities
        '''

        current_device = next(self.parameters()).device

        start_idx = torch.tensor([self.bos_idx]*bs).long().unsqueeze(-1)

        preds = idxs = to_device(start_idx, device=current_device)
        lps = []

        hiddens = None

        for i in range(sl):
            x, hiddens, encoded = self._forward(idxs, hiddens)
            x.div_(temperature)

            idxs, lp = x_to_preds(x, multinomial=multinomial)

            lps.append(lp)
            preds = torch.cat([preds, idxs], -1)

        return preds[:, 1:], torch.cat(lps,-1)

    def sample_no_grad(self, bs, sl, z=None, temperature=1., multinomial=True):
        '''
        sample_no_grad - sample from the model without saving values for
        gradient calculation

        Inputs:

        - `bs int`: batch size

        - `sl int`: maximum sequence length

        - `z Optional[torch.Tensor[bs, d_latent]]`: latent vector

        - `temperature float`: sample temperature

        - `multinomial bool`: if True, use multinomial sampling.
        If False, use argmax greedy sampling

        Returns:

        - `preds torch.LongTensor[bs, sl]`: predicted sequence tokens

        - `lps torch.FloatTensor[bs, sl, d_vocab]`: prediction log probabilities
        '''
        with torch.no_grad():
            return self.sample(bs, sl, temperature=temperature, multinomial=multinomial)

    def get_rl_tensors(self, x, y, temperature=1., latent=None):
        '''
        get_rl_tensors - generate values needed for RL training

        Inputs:

        - `x torch.LongTensor[bs, sl]`: x value

        - `y torch.LongTensor[bs, sl]`: y value

        - `temperature float`: sample temperature

        - `latent None`: latent vector. LSTM_LM does not use latent
        vectors, this keyword is included for compatibility

        Returns:

        - `output torch.FloatTensor[bs, sl, d_vocab]`: output of the model

        - `lps torch.FloatTensor[bs, sl, d_vocab]`: log probabilities.
        Log softmax of `output` values

        - `lps_gathered torch.FloatTensor[bs, sl]`: log probabilities
        gathered by the values in `y`

        - `encoded torch.FloatTensor[bs, sl, d_embedding]`: output from
        final LSTM layer
        '''
        output, hiddens, encoded = self._forward(x)
        output.div_(temperature)
        lps = F.log_softmax(output, -1)
        lps_gathered = gather_lps(lps, y)
        return output, lps, lps_gathered, encoded

# Cell

class Conditional_LSTM_LM(GenerativeModel):
    '''
    LSTM_LM - LSTM language model

    Inputs:

    - `encoder nn.Module`: encoder model

    - `d_vocab int`: vocab size

    - `d_embedding int`: embedding dimension

    - `d_hidden int`: hidden dimension

    - `d_latent int`: latent vector dimension

    - `n_layers int`: number of LSTM layers

    - `input_dropout float`: dropout percentage on inputs

    - `lstm_dropout float`: dropout on LSTM layers

    - `norm_latent bool`: if True, latent vectors are scaled to a norm of 1

    - `condition_hidden bool`: if True, latent vector is used to initialize the
    hidden state

    - `condition_output bool`: if True, latent vector is concatenated to inputs

    - `bos_idx int`: beginning of sentence token

    - `prior Optional[nn.Module]`: optional prior distribution to sample from.
    see `Prior`

    - `forward_rollout bool`: if True, run supervised training using
    rollout with teacher forcing. This is a technique used in some
    seq2seq models and should not be used for pure generative models

    - `p_force float`: teacher forcing probabiliy

    - `force_decay float`: rate of decay of p_force
    '''
    def __init__(self, encoder, d_vocab, d_embedding, d_hidden, d_latent, n_layers,
                 input_dropout=0., lstm_dropout=0., norm_latent=True,
                 condition_hidden=True, condition_output=False, bos_idx=0,
                 prior=None, forward_rollout=False, p_force=0., force_decay=0.99):
        super().__init__()

        self.encoder = encoder

        self.decoder = Conditional_LSTM_Block(d_vocab, d_embedding, d_hidden, d_embedding,
                                d_latent, n_layers, input_dropout=input_dropout,
                                lstm_dropout=lstm_dropout,
                                condition_hidden=condition_hidden, condition_output=condition_output)

        self.norm_latent = norm_latent
        self.bos_idx = bos_idx

        if prior is None:
            prior = SphericalPrior(torch.zeros((encoder.d_latent)), torch.zeros((encoder.d_latent)),
                                trainable=False)

        self.prior = prior
        self.forward_rollout = forward_rollout
        self.p_force = p_force
        self.force_decay = force_decay

    def forward(self, x, condition=None, hiddens=None):
        if condition is None:
            condition = x

        z = self.encoder(condition)

        z = self.maybe_norm(z)

        x, hiddens, encoded = self.decoder_forward(x, z, hiddens)
        return x

    def decoder_forward(self, x, z, hiddens=None):
        '''
        decoder_forward

        If using forward rollout, the model will be sampled
        from at each time step. The model has a `p_force` chance
        to recieve the correct next input, and a `1-p_force`
        chance to recieve the model's own prediction.

        If forward rollout is not used, `x` and `z` are
        simply passed to the decoder
        '''
        if self.forward_rollout:
#             current_device = next(self.parameters()).device
            sl = x.shape[1]
            bs = x.shape[0]

            idxs = x[:,0].unsqueeze(-1)

#             start_idx = torch.tensor([self.bos_idx]*bs).long().unsqueeze(-1)

#             idxs = to_device(start_idx, device=current_device)

            output = []
            encoded = []

            for i in range(sl):
                output_iter, hiddens, encoded_iter = self.decoder(idxs,z,hiddens)
                output.append(output_iter)
                encoded.append(encoded_iter)

                if np.random.random()<self.p_force:
                    idxs = x[:,i].unsqueeze(-1)

                else:
                    with torch.no_grad():
                        idxs = F.softmax(output_iter,-1).argmax(-1)

            output = torch.cat(output, 1)
            encoded = torch.cat(encoded, 1)
            self.p_force = self.p_force * self.force_decay

        else:
            output, hiddens, encoded = self.decoder(x, z, hiddens)

        return output, hiddens, encoded

    def maybe_norm(self, z):
        '''
        normalize `z` if applicable
        '''
        if self.norm_latent:
            z = F.normalize(z, p=2, dim=-1)

        return z

    def x_to_latent(self, x):

        x, condition = x
        z = self.encoder(condition)
        z = self.maybe_norm(z)

        return z

    def sample(self, bs, sl, z=None, temperature=1., multinomial=True):

        '''
        sample - sample from the model

        Inputs:

        - `bs int`: batch size

        - `sl int`: maximum sequence length

        - `z Optional[torch.Tensor[bs, d_latent]]`: latent vector

        - `temperature float`: sample temperature

        - `multinomial bool`: if True, use multinomial sampling.
        If False, use argmax greedy sampling

        Returns:

        - `preds torch.LongTensor[bs, sl]`: predicted sequence tokens

        - `lps torch.FloatTensor[bs, sl, d_vocab]`: prediction log probabilities
        '''

        current_device = next(self.parameters()).device

        if z is None:
            if self.prior is not None:
                z = to_device(self.prior.rsample([bs]), device=current_device)
            else:
                z = to_device(torch.randn((bs, self.encoder.d_latent)), device=current_device)
        else:
            bs = z.shape[0]

        z = self.maybe_norm(z)

        start_idx = torch.tensor([self.bos_idx]*bs).long().unsqueeze(-1)

        preds = idxs = to_device(start_idx, device=current_device)

        lps = []

        hiddens = None

        for i in range(sl):

            x, hiddens, encoded = self.decoder(idxs,z,hiddens)
            x.div_(temperature)

            idxs, lp = x_to_preds(x, multinomial=multinomial)

            lps.append(lp)
            preds = torch.cat([preds, idxs], -1)

        return preds[:, 1:], torch.cat(lps,-1)

    def sample_no_grad(self, bs, sl, z=None, temperature=1., multinomial=True):
        '''
        sample_no_grad - sample from the model without saving values for
        gradient calculation

        Inputs:

        - `bs int`: batch size

        - `sl int`: maximum sequence length

        - `z Optional[torch.Tensor[bs, d_latent]]`: latent vector

        - `temperature float`: sample temperature

        - `multinomial bool`: if True, use multinomial sampling.
        If False, use argmax greedy sampling

        Returns:

        - `preds torch.LongTensor[bs, sl]`: predicted sequence tokens

        - `lps torch.FloatTensor[bs, sl, d_vocab]`: prediction log probabilities
        '''
        with torch.no_grad():
            return self.sample(bs, sl, z=z, temperature=temperature, multinomial=multinomial)

    def get_rl_tensors(self, x, y, temperature=1., latent=None):
        '''
        get_rl_tensors - generate values needed for RL training

        Inputs:

        - `x torch.LongTensor[bs, sl]`: x value

        - `y torch.LongTensor[bs, sl]`: y value

        - `temperature float`: sample temperature

        - `latent Optonal[torch.FloatTensor[bs, d_latent]]`: latent vector

        Returns:

        - `output torch.FloatTensor[bs, sl, d_vocab]`: output of the model

        - `lps torch.FloatTensor[bs, sl, d_vocab]`: log probabilities.
        Log softmax of `output` values

        - `lps_gathered torch.FloatTensor[bs, sl]`: log probabilities
        gathered by the values in `y`

        - `encoded torch.FloatTensor[bs, sl, d_embedding]`: output from
        final LSTM layer
        '''

        x,c = x
        if latent is None:
            latent = self.encoder(c)

        latent = self.maybe_norm(latent)

        output, hiddens, encoded = self.decoder(x,latent)
        output.div_(temperature)
        lps = F.log_softmax(output, -1)

        if self.prior.trainable:
            prior_lps = self.prior.log_prob(latent)
            prior_lps = prior_lps.mean(-1).unsqueeze(-1).unsqueeze(-1)
            pass_through = torch.zeros(prior_lps.shape).float().to(prior_lps.device)
            pass_through = pass_through + prior_lps - prior_lps.detach() # add to gradient chain
            lps = lps + pass_through

        lps_gathered = gather_lps(lps, y)
        return output, lps, lps_gathered, encoded

    def set_prior_from_latent(self, z, logvar, trainable=False):
        '''
        set_prior_from_latent - set prior distribution
        from `z` and `logvar`

        Inputs:

        - `z torch.FloatTensor[bs, d_latent]`: latent vector

        - `logvar torch.FloatTensor[bs, d_laten]`: log variance vector

        - `trainable bool`: if True, prior will be updated by gradient descent
        '''
        z = z.detach()
        logvar = logvar.detach()
        self.prior = SphericalPrior(z, logvar, trainable)

    def set_prior_from_encoder(self, condition, logvar, trainable=False):
        '''
        set_prior_from_encoder - computes a latent vector
        from `condition` and uses it to set the model's prior

        Inputs:

        - `condition`: input condition (depends on encoder)

        - `logvar torch.FloatTensor[bs, d_latent]`: log variance vector

        - `trainable bool`: if True, prior will be updated by gradient descent
        '''
        assert condition.shape[0]==1
        z = self.encoder(condition)
        z = self.maybe_norm(z)
        z = z.squeeze(0)
        self.set_prior_from_latent(z, logvar, trainable)