<commit_before>import random
import torch
import torch.nn as nn
import torch.utils.data
from torch_model_base import TorchModelBase
from utils import progress_bar

__author__ = "<NAME>"
__version__ = "CS224u, Stanford, Spring 2019"


class TorchTreeNNModel(nn.Module):
    def __init__(self, vocab, embed_dim, embedding, output_dim, hidden_activation):
        super().__init__()
        self.vocab = vocab
        self.vocab_size = len(vocab)
        self.vocab_lookup = dict(zip(self.vocab, range(self.vocab_size)))
        self.embed_dim = embed_dim
        self.hidden_dim = embed_dim * 2
        self.hidden_activation = hidden_activation
        self.output_dim = output_dim
        self.tree_layer = nn.Linear(self.hidden_dim, self.embed_dim)
        self.embedding = self._define_embedding(embedding)
        self.classifier_layer = nn.Linear(self.embed_dim, self.output_dim)

    def _define_embedding(self, embedding):
        if embedding is None:
            return nn.Embedding(self.vocab_size, self.embed_dim)
        else:
            embedding = torch.tensor(embedding, dtype=torch.float)
            return nn.Embedding.from_pretrained(embedding)

    def forward(self, tree):
        """Recursively interprets `tree`, applying a classifier layer
        to the final representation. The label comes from the root
        of the tree itself.

        Parameters
        ----------
        tree : nltk.tree.Tree

        Returns
        -------
        torch.LongTensor, label (str)

        """
        root = self.interpret(tree)
        return self.classifier_layer(root), tree.label()

    def interpret(self, subtree):
        # Terminal nodes are str:
        if isinstance(subtree, str):
            i = self.vocab_lookup.get(subtree, self.vocab_lookup['$UNK'])
            ind = torch.tensor([i], dtype=torch.long)
            return self.embedding(ind)
        # Non-branching nodes:
        elif len(subtree) == 1:
            return self.interpret(subtree[0])
        # Branching nodes:
        else:
            left_subtree, right_subtree = subtree[0], subtree[1]
            left_subtree = self.interpret(left_subtree)
            right_subtree = self.interpret(right_subtree)
            combined = torch.cat((left_subtree, right_subtree), dim=1)
            root_rep = self.hidden_activation(self.tree_layer(combined))
            return root_rep


class TorchTreeNN(TorchModelBase):
    def __init__(self, vocab, embedding=None, embed_dim=50, **kwargs):
        self.vocab = vocab
        self.embedding = embedding
        self.embed_dim = embed_dim
        if self.embedding is not None:
            self.embed_dim = embedding.shape[1]
        super(TorchTreeNN, self).__init__(**kwargs)
        self.params += ['embed_dim', 'embedding']
        self.device = 'cpu'

    def build_graph(self):
        self.model = TorchTreeNNModel(
            vocab=self.vocab,
            embedding=self.embedding,
            embed_dim=self.embed_dim,
            output_dim=self.n_classes_,
            hidden_activation=self.hidden_activation)

    def fit(self, X, **kwargs):
        """Fairly standard `fit` method except that the labels `y` are
        presumed to come from the root nodes of the trees in `X`.

        Parameters
        ----------
        X : list of `nltk.Tree` instances
        kwargs : dict
            For passing other parameters. If 'X_dev' is included,
            then performance is monitored every 10 epochs; use
            `dev_iter` to control this number.

        Returns
        -------
        self

        """
        # Incremental performance:
        X_dev = kwargs.get('X_dev')
        if X_dev is not None:
            dev_iter = kwargs.get('dev_iter', 10)
        # Data prep:
        self.classes_ = self.get_classes(X)
        self.n_classes_ = len(self.classes_)
        self.class2index = dict(zip(self.classes_, range(self.n_classes_)))
        # Model:
        self.build_graph()
        self.model.to(self.device)
        # Optimization:
        loss = nn.CrossEntropyLoss()
        optimizer = self.optimizer(self.model.parameters(), lr=self.eta)
        # Train:
        for iteration in range(1, self.max_iter+1):
            epoch_error = 0.0
            random.shuffle(X)
            for tree in X:
                pred, label = self.model.forward(tree)
                label = self.convert_label(label)
                err = loss(pred, label)
                epoch_error += err.item()
                optimizer.zero_grad()
                err.backward()
                optimizer.step()
            # Incremental predictions where possible:
            if X_dev is not None and iteration > 0 and iteration % dev_iter == 0:
                self.dev_predictions[iteration] = self.predict(X_dev)
            self.errors.append(epoch_error)
            progress_bar(
                "Finished epoch {} of {}; error is {}".format(
                    iteration, self.max_iter, epoch_error))
        return self

    @staticmethod
    def get_classes(X):
        """Classes as given by the root nodes in `X`.

        Parameters
        ----------
        X : list of nltk.tree.Tree

        Returns
        -------
        list

        """
        return sorted({t.label() for t in X})

    def convert_label(self, label):
        """Convert a class label to a format that PyTorch can handle.

        Parameters
        ----------
        label : member of `self.classes_`

        Returns<commit_msg>fix torch_tree_nn loss normalization error<commit_after>import random
import torch
import torch.nn as nn
import torch.utils.data
from torch_model_base import TorchModelBase
from utils import progress_bar

__author__ = "<NAME>"
__version__ = "CS224u, Stanford, Spring 2019"


class TorchTreeNNModel(nn.Module):
    def __init__(self, vocab, embed_dim, embedding, output_dim, hidden_activation):
        super().__init__()
        self.vocab = vocab
        self.vocab_size = len(vocab)
        self.vocab_lookup = dict(zip(self.vocab, range(self.vocab_size)))
        self.embed_dim = embed_dim
        self.hidden_dim = embed_dim * 2
        self.hidden_activation = hidden_activation
        self.output_dim = output_dim
        self.tree_layer = nn.Linear(self.hidden_dim, self.embed_dim)
        self.embedding = self._define_embedding(embedding)
        self.classifier_layer = nn.Linear(self.embed_dim, self.output_dim)

    def _define_embedding(self, embedding):
        if embedding is None:
            return nn.Embedding(self.vocab_size, self.embed_dim)
        else:
            embedding = torch.tensor(embedding, dtype=torch.float)
            return nn.Embedding.from_pretrained(embedding)

    def forward(self, tree):
        """Recursively interprets `tree`, applying a classifier layer
        to the final representation. The label comes from the root
        of the tree itself.

        Parameters
        ----------
        tree : nltk.tree.Tree

        Returns
        -------
        torch.LongTensor, label (str)

        """
        root = self.interpret(tree)
        return self.classifier_layer(root), tree.label()

    def interpret(self, subtree):
        # Terminal nodes are str:
        if isinstance(subtree, str):
            i = self.vocab_lookup.get(subtree, self.vocab_lookup['$UNK'])
            ind = torch.tensor([i], dtype=torch.long)
            return self.embedding(ind)
        # Non-branching nodes:
        elif len(subtree) == 1:
            return self.interpret(subtree[0])
        # Branching nodes:
        else:
            left_subtree, right_subtree = subtree[0], subtree[1]
            left_subtree = self.interpret(left_subtree)
            right_subtree = self.interpret(right_subtree)
            combined = torch.cat((left_subtree, right_subtree), dim=1)
            root_rep = self.hidden_activation(self.tree_layer(combined))
            return root_rep


class TorchTreeNN(TorchModelBase):
    def __init__(self, vocab, embedding=None, embed_dim=50, **kwargs):
        self.vocab = vocab
        self.embedding = embedding
        self.embed_dim = embed_dim
        if self.embedding is not None:
            self.embed_dim = embedding.shape[1]
        super(TorchTreeNN, self).__init__(**kwargs)
        self.params += ['embed_dim', 'embedding']
        self.device = 'cpu'

    def build_graph(self):
        self.model = TorchTreeNNModel(
            vocab=self.vocab,
            embedding=self.embedding,
            embed_dim=self.embed_dim,
            output_dim=self.n_classes_,
            hidden_activation=self.hidden_activation)

    def fit(self, X, **kwargs):
        """Fairly standard `fit` method except that the labels `y` are
        presumed to come from the root nodes of the trees in `X`.

        Parameters
        ----------
        X : list of `nltk.Tree` instances
        kwargs : dict
            For passing other parameters. If 'X_dev' is included,
            then performance is monitored every 10 epochs; use
            `dev_iter` to control this number.

        Returns
        -------
        self

        """
        # Incremental performance:
        X_dev = kwargs.get('X_dev')
        if X_dev is not None:
            dev_iter = kwargs.get('dev_iter', 10)
        # Data prep:
        self.classes_ = self.get_classes(X)
        self.n_classes_ = len(self.classes_)
        self.class2index = dict(zip(self.classes_, range(self.n_classes_)))
        # Model:
        self.build_graph()
        self.model.to(self.device)
        # Optimization:
        loss = nn.CrossEntropyLoss()
        optimizer = self.optimizer(self.model.parameters(), lr=self.eta)
        # Train:
        for iteration in range(1, self.max_iter+1):
            epoch_error = 0.0
            random.shuffle(X)
            for tree in X:
                pred, label = self.model.forward(tree)
                label = self.convert_label(label)
                err = loss(pred, label)
                epoch_error += err.item()
                optimizer.zero_grad()
                err.backward()
                optimizer.step()
            # Incremental predictions where possible:
            if X_dev is not None and iteration > 0 and iteration % dev_iter == 0:
                self.dev_predictions[iteration] = self.predict(X_dev)
            self.errors.append(epoch_error)
            progress_bar(
                "Finished epoch {} of {}; error is {}".format(
                    iteration, self.max_iter, epoch_error/len(X)))
        return self

    @staticmethod
    def get_classes(X):
        """Classes as given by the root nodes in `X`.

        Parameters
        ----------
        X : list of nltk.tree.Tree

        Returns
        -------
        list

        """
        return sorted({t.label() for t in X})

    def convert_label(self, label):
        """Convert a class label to a format that PyTorch can handle.

        Parameters
        ----------
        label : member of `self.classes_`

        Returns
{'fraction_non_alphanumeric': 0.10190909090909091, 'fraction_numerical': 0.004909090909090909, 'mean_word_length': 1.9796858071505958, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, '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': '1498998', 'n_tokens_mistral': 3230, 'n_tokens_neox': 3130, 'n_words': 765}