<commit_before>"""
This module provides two RNN Walkers (one is random, the other deterministic) 
and an a simple torch-based RNN model. The RNN Walker and RNN model communicate
through an RNN Adapter, so that several kinds of models can be used with an RNN Walker.
"""
from collections import namedtuple
import numpy as np
import math

import torch
import torch.nn as nn
from torch.autograd import Variable

from .grammar import Walker
from .grammar import DeterministicWalker
from .grammar import as_str

from .types import Int
from .types import Float


class RnnModel(nn.Module):
    """
    A simple LSTM torch based model.
    It implements `next_token` and `next_value` which are necessary
    to communicate with `RnnAdapter`.

    Parameters
    ----------
    
    vocab_size : int
        vocabulary size.
        Practically, it corresponds to the number of possible production
        rules. From the RNN perspective, each production rule is a token.
    
    emb_size : int
        Embedding size

    hidden_size : int
        number of hidden units of the lSTM

    nb_features : int
        number of real-valued features to predict when a value is needed.
        Practically, it corresponds to the number of statistics to predict 
        needed to sample from a `Type`.

    use_cuda : bool
        if True, use cuda.
        
    """
    def __init__(self, vocab_size=10, emb_size=128, hidden_size=128, num_layers=1, nb_features=1, use_cuda=False):
        super().__init__()   
        self.vocab_size = vocab_size
        self.emb_size = emb_size
        self.hidden_size = hidden_size
        self.use_cuda = use_cuda

        self.emb = nn.Embedding(vocab_size, emb_size)
        self.lstm = nn.LSTM(emb_size, hidden_size, batch_first=True, num_layers=num_layers)
        self.out_token  = nn.Linear(hidden_size, vocab_size)
        self.out_value = nn.Linear(hidden_size, nb_features)
    
    def forward(self, inp):
        x = self.emb(inp)
        o, _ = self.lstm(x)
        o = o.contiguous()
        o = o.view(o.size(0) * o.size(1), o.size(2))
        o = self.out_token(o)
        return o

    def next_token(self, inp, state):
        if self.use_cuda:
            inp = inp.cuda()
        x = self.emb(inp)
        _, state = self.lstm(x, state)
        h, c = state
        h = h.view(h.size(0) * h.size(1), h.size(2))
        o = self.out_token(h)
        return o, state
    
    def next_value(self, inp, state):
        if self.use_cuda:
            inp = inp.cuda()
        x = self.emb(inp)
        _, state = self.lstm(x, state)
        h, c = state
        h = h.view(h.size(0) * h.size(1), h.size(2))
        o = self.out_value(h)
        return o, state


class RnnAdapter:
    """
    `RnnAdapter` adapts a Model to communicate with an RNN `Walker` so that
    different Models, adapted with the same `RnnAdapter` can be used in an
    RNN `Walker`.

    Parameters
    ----------

    model : Model to adapt

    tok_to_id : dict of Token->int
        Practically, `Token` corresponds to a token, which in our case is a 
        `parsimonious` `Expression`, which refers to a production rule.
        So this dict convert a production rule to an integer identifying the
        produciton rule.

    begin_tok : Token
        this Token (practically corresponding to a production rule, which is
        a `parsimonious` `Expression`) is given as input in the first step
        the RNN. If not provided, the token with ID 0 (that is, the `tok` for which tok_to_id[tok] ==0)
        is used.

    """
    def __init__(self, model, tok_to_id, begin_tok=None, random_state=None):
        assert model.vocab_size == len(tok_to_id)
        self.model = model
        self.tok_to_id = tok_to_id
        self.id_to_tok = {v: k for k, v in tok_to_id.items()}
        if begin_tok is None:
            begin_tok = self.id_to_tok[0]
        self.begin_tok = begin_tok
        self.rng = np.random.RandomState(random_state)

    def predict_next_token(self, inp, state, temperature=1.0):
        x = self._preprocess_input(inp)
        o, state = self.model.next_token(x, state)
        o = o / temperature
        pr = nn.Softmax()(o)
        return pr, state

    def generate_next_token(self, pr, allowed=None):
        pr = pr[0].data.clone()
        if allowed is not None:
            assert len(allowed)
            full = set(self.tok_to_id.keys())
            forbidden = full - set(allowed)
            for tok in forbidden:
                pr[self.tok_to_id[tok]] = 0.
        pr = pr.tolist()
        # - the nan thing happened when I added a random
        #   exploration (prop of time generate randomly), not sure
        #   why
        # - the second condition is to prevent cases
        #   where the forbidden tokens have high proba
        #   and the rest essiantially zero proba, it means
        #   that pr will be close to zero
        # - in both cases, solve the problem by generating
        #   uniformly from the allowed tokens<commit_msg>select argmax token when temperate is zero in RNN<commit_after>"""
This module provides two RNN Walkers (one is random, the other deterministic) 
and an a simple torch-based RNN model. The RNN Walker and RNN model communicate
through an RNN Adapter, so that several kinds of models can be used with an RNN Walker.
"""
from collections import namedtuple
import numpy as np
import math

import torch
import torch.nn as nn
from torch.autograd import Variable

from .grammar import Walker
from .grammar import DeterministicWalker
from .grammar import as_str

from .types import Int
from .types import Float


class RnnModel(nn.Module):
    """
    A simple LSTM torch based model.
    It implements `next_token` and `next_value` which are necessary
    to communicate with `RnnAdapter`.

    Parameters
    ----------
    
    vocab_size : int
        vocabulary size.
        Practically, it corresponds to the number of possible production
        rules. From the RNN perspective, each production rule is a token.
    
    emb_size : int
        Embedding size

    hidden_size : int
        number of hidden units of the lSTM

    nb_features : int
        number of real-valued features to predict when a value is needed.
        Practically, it corresponds to the number of statistics to predict 
        needed to sample from a `Type`.

    use_cuda : bool
        if True, use cuda.
        
    """
    def __init__(self, vocab_size=10, emb_size=128, hidden_size=128, num_layers=1, nb_features=1, use_cuda=False):
        super().__init__()   
        self.vocab_size = vocab_size
        self.emb_size = emb_size
        self.hidden_size = hidden_size
        self.use_cuda = use_cuda

        self.emb = nn.Embedding(vocab_size, emb_size)
        self.lstm = nn.LSTM(emb_size, hidden_size, batch_first=True, num_layers=num_layers)
        self.out_token  = nn.Linear(hidden_size, vocab_size)
        self.out_value = nn.Linear(hidden_size, nb_features)
    
    def forward(self, inp):
        x = self.emb(inp)
        o, _ = self.lstm(x)
        o = o.contiguous()
        o = o.view(o.size(0) * o.size(1), o.size(2))
        o = self.out_token(o)
        return o

    def next_token(self, inp, state):
        if self.use_cuda:
            inp = inp.cuda()
        x = self.emb(inp)
        _, state = self.lstm(x, state)
        h, c = state
        h = h.view(h.size(0) * h.size(1), h.size(2))
        o = self.out_token(h)
        return o, state
    
    def next_value(self, inp, state):
        if self.use_cuda:
            inp = inp.cuda()
        x = self.emb(inp)
        _, state = self.lstm(x, state)
        h, c = state
        h = h.view(h.size(0) * h.size(1), h.size(2))
        o = self.out_value(h)
        return o, state


class RnnAdapter:
    """
    `RnnAdapter` adapts a Model to communicate with an RNN `Walker` so that
    different Models, adapted with the same `RnnAdapter` can be used in an
    RNN `Walker`.

    Parameters
    ----------

    model : Model to adapt

    tok_to_id : dict of Token->int
        Practically, `Token` corresponds to a token, which in our case is a 
        `parsimonious` `Expression`, which refers to a production rule.
        So this dict convert a production rule to an integer identifying the
        produciton rule.

    begin_tok : Token
        this Token (practically corresponding to a production rule, which is
        a `parsimonious` `Expression`) is given as input in the first step
        the RNN. If not provided, the token with ID 0 (that is, the `tok` for which tok_to_id[tok] ==0)
        is used.

    """
    def __init__(self, model, tok_to_id, begin_tok=None, random_state=None):
        assert model.vocab_size == len(tok_to_id)
        self.model = model
        self.tok_to_id = tok_to_id
        self.id_to_tok = {v: k for k, v in tok_to_id.items()}
        if begin_tok is None:
            begin_tok = self.id_to_tok[0]
        self.begin_tok = begin_tok
        self.rng = np.random.RandomState(random_state)

    def predict_next_token(self, inp, state, temperature=1.0):
        x = self._preprocess_input(inp)
        o, state = self.model.next_token(x, state)
        if temperature > 0:
            o = o / temperature
            pr = nn.Softmax()(o)
        else:
            # when temperature is 0, it is equivalent
            # to take the argmax
            pr = o.clone()
            _, ids = pr.max(1)
            pr.data.fill_(0.)
            pr.data[0, ids.data[0]] = 1.0
        return pr, state

    def generate_next_token(self, pr, allowed=None):
        pr = pr[0].data.clone()
        if allowed is not None:
            assert len(allowed)
            full = set(self.tok_to_id.keys())
            forbidden = full - set(allowed)
            for tok in forbidden:
                pr[self.tok_to_id[tok]] = 0.
        pr = pr.tolist()
        # - the nan thing happened when I added a random
        #   exploration (prop of time generate randomly), not sure
        #   why
        # - the second condition is to prevent cases
        #   where the forbidden tokens have high proba
        #   and the rest essiantially zero proba, it means
        #   that pr will be close to zero
        # - in both cases, solve the problem by generating
        #   uniformly from the allowed tokens
{'fraction_non_alphanumeric': 0.08669960282863509, 'fraction_numerical': 0.005812263876780006, 'mean_word_length': 2.17173579109063, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3990262', 'n_tokens_mistral': 3244, 'n_tokens_neox': 3174, 'n_words': 1080}