<commit_before># ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2021, Numenta, Inc.  Unless you have an agreement
# with Numenta, Inc., for a separate license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------

import os

import torch
from torchvision import transforms
from torchvision.datasets import MNIST


class PermutedMNIST(MNIST):
    """
    The permutedMNIST dataset contains MNIST images where the same random permutation
    of pixel values is applied to each image. More specifically, the dataset can be
    broken down into 'tasks', where each such task is the set of all MNIST images, but
    with a unique pixel-wise permutation applied to all images. `num_tasks` gives the
    number of 10-way classification tasks (each corresponding to a unique pixel-wise
    permutation) that a continual learner will try to learn in sequence.
    """

    def __init__(self, num_tasks, seed, train, root=".", target_transform=None,
                 download=False, normalize=True):

        t = [transforms.ToTensor()]
        if normalize:
            t.append(transforms.Normalize((0.13062755,), (0.30810780,)))
        data_transform = transforms.Compose(t)
        super().__init__(root=root, train=train, transform=data_transform,
                         target_transform=target_transform, download=download)

        self.num_tasks = num_tasks

        # Use a generator object to manually set the seed and generate the same
        # num_tasks random permutations for both training and validation datasets; the
        # first one is the identity permutation (i.e., regular MNIST), represented
        # below as `None`
        g = torch.manual_seed(seed)

        self.permutations = [
            torch.randperm(784, generator=g) for task_id in range(1, num_tasks)
        ]
        self.permutations.insert(0, None)

    def __getitem__(self, index):
        """
        Returns an (image, target) pair.

        In particular, this method retrieves an MNIST image, and based on the value of
        `index`, determines which pixel-wise permutation to apply. Target values are
        also scaled to be unique to each permutation.
        """
        img, target = super().__getitem__(index % len(self.data))

        # Determine which task `index` corresponds to
        task_id = self.get_task_id(index)

        # Apply permutation to `img`
        img = permute(img, self.permutations[task_id])

        # Since target values are not shared between tasks, `target` should be in the
        # range [0 + 10 * task_id, 9 + 10 * task_id]
        target += 10 * task_id
        return img, target

    def __len__(self):
        return self.num_tasks * len(self.data)

    @property
    def processed_folder(self):
        return os.path.join(self.root, "MNIST", "processed")

    def get_task_id(self, index):
        return index // len(self.data)


class ContextDependentPermutedMNIST(PermutedMNIST):
    """
    A variant of permutedMNIST where each permutation (i.e., 'task') is associated with
    a context: a binary sparse vector. The `__getitem__` method returns the context
    vector along with the data sample and target.
    """

    def __init__(self, num_tasks, dim_context, seed, train, root=".",
                 target_transform=None, download=False):

        super().__init__(num_tasks, seed, train, root, target_transform, download)
        self.dim_context = dim_context

        # Initialize random binary sparse context vectors for each permutation
        self.init_contexts(seed)

    def __getitem__(self, index):
        """
        Returns an ((image, context), target) tuple.
        """
        img, target = super().__getitem__(index)
        task_id = self.get_task_id(index)
        return (img, self.contexts[task_id, :]), target

    def init_contexts(self, seed):
        percent_on = 0.05
        num_contexts = self.num_tasks

        num_ones = int(percent_on * self.dim_context)
        num_zeros = self.dim_context - num_ones

        self.contexts = torch.cat((
            torch.zeros((num_contexts, num_zeros)),
            torch.ones((num_contexts, num_ones))
        ), dim=1)

        # Shuffle each context vector i.i.d. to randomize it; use a fixed seed during
        # shuffling so that train & validation contexts are identical
        g = torch.manual_seed(seed)

        for i in range(num_contexts):
            self.contexts[i, :] = self.contexts[i, torch.randperm(self.dim_context,
                                                                  generator=g)]


class OneHotContextPermutedMNIST(PermutedMNIST):
    """
    Dataloader for permutedMNIST that permutes, then flattens, then concatenates
    one hot context vector.
    """
    def __init__(self, num_tasks, seed, train, root=".",
                 target_transform=None, download=False):

        super().__init__(
            num_tasks, seed, train, root, target_transform, download)

    def __getitem__(self, index):

        img, target = super().__getitem__(index)
        task_id = self.get_task_id(index)
        context = self.task_id_to_one_hot(task_id)
        img = permute(img, self.permutations[task_id])
        img = img.flatten()
        img = torch.cat((img, context))

        return img, target

    def task_id_to_one_hot(self, task_id):
        one_hot = torch.zeros(self.num_tasks)
        one_hot[task_id] = 1.
        return one_hot


def permute(x, permutation):
    """
    Applies the permutation specified by `permutation` to `x` and returns the resulting
    tensor.
    """

    # Assume x has shape (1, height, width)
    if permutation is None:
        return x

    _, height, width = x.size()
    x = x.view(-1, 1)
    x = x[permutation, :]
    x = x.view(1, height, width)
    return x
<commit_msg>Refactor ContextDependentPermutedMNIST so it can return multiple types of context<commit_after># ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2021, Numenta, Inc.  Unless you have an agreement
# with Numenta, Inc., for a separate license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------

import os

import torch
from torchvision import transforms
from torchvision.datasets import MNIST


class PermutedMNIST(MNIST):
    """
    The permutedMNIST dataset contains MNIST images where the same random permutation
    of pixel values is applied to each image. More specifically, the dataset can be
    broken down into 'tasks', where each such task is the set of all MNIST images, but
    with a unique pixel-wise permutation applied to all images. `num_tasks` gives the
    number of 10-way classification tasks (each corresponding to a unique pixel-wise
    permutation) that a continual learner will try to learn in sequence.
    """

    def __init__(self, num_tasks, seed, train, root=".", target_transform=None,
                 download=False, normalize=True):

        t = [transforms.ToTensor()]
        if normalize:
            t.append(transforms.Normalize((0.13062755,), (0.30810780,)))
        data_transform = transforms.Compose(t)
        super().__init__(root=root, train=train, transform=data_transform,
                         target_transform=target_transform, download=download)

        self.num_tasks = num_tasks

        # Use a generator object to manually set the seed and generate the same
        # num_tasks random permutations for both training and validation datasets; the
        # first one is the identity permutation (i.e., regular MNIST), represented
        # below as `None`
        g = torch.manual_seed(seed)

        self.permutations = [
            torch.randperm(784, generator=g) for task_id in range(1, num_tasks)
        ]
        self.permutations.insert(0, None)

    def __getitem__(self, index):
        """
        Returns an (image, target) pair.

        In particular, this method retrieves an MNIST image, and based on the value of
        `index`, determines which pixel-wise permutation to apply. Target values are
        also scaled to be unique to each permutation.
        """
        img, target = super().__getitem__(index % len(self.data))

        # Determine which task `index` corresponds to
        task_id = self.get_task_id(index)

        # Apply permutation to `img`
        img = permute(img, self.permutations[task_id])

        # Since target values are not shared between tasks, `target` should be in the
        # range [0 + 10 * task_id, 9 + 10 * task_id]
        target += 10 * task_id
        return img, target

    def __len__(self):
        return self.num_tasks * len(self.data)

    @property
    def processed_folder(self):
        return os.path.join(self.root, "MNIST", "processed")

    def get_task_id(self, index):
        return index // len(self.data)


class ContextDependentPermutedMNIST(PermutedMNIST):
    """
    A variant of permutedMNIST where each permutation (i.e., 'task') is associated with
    a context: a binary sparse vector. The `__getitem__` method returns the context
    vector along with the data sample and target.
    """

    def __init__(self, num_tasks, seed, train, context_type="sparse_binary",
                 combine_context_as="tuple", dim_context=None, root=".",
                 target_transform=None, download=False):

        super().__init__(num_tasks, seed, train, root, target_transform, download)
        self.dim_context = dim_context

        # options for type of context and way of combining with input x
        context_type_choices = ["sparse_binary", "one_hot", "centroid"]
        combine_context_as_choices = ["tuple", "concatenate"]

        # Parse type of context
        if context_type == "sparse_binary":
            # Initialize random binary sparse context vectors for each permutation
            self.init_sparse_binary_contexts(seed)
        elif context_type == "one_hot":
            self.init_one_hot_context(seed)
        elif context_type == "centroid":
            self.init_centroids(seed)
        else:
            error_msg = f"context_type must be one of {context_type_choices}"
            raise ValueError(error_msg)

        # Parse how to combine image and context
        if combine_context_as == "tuple":
            self.combine_context = tuple_context
        elif combine_context_as == "concatenate":
            self.combine_context = concat_context
        else:
            error_msg = f"combine_context_as must be one of {combine_context_as_choices}"
            raise ValueError(error_msg)

    def __getitem__(self, index):
        """
        Returns image, context, and target.
        """
        img, target = super().__getitem__(index)
        task_id = self.get_task_id(index)
        context = self.contexts[task_id,:]
        return self.combine_context(img, context), target

    def init_sparse_binary_contexts(self, seed):
        percent_on = 0.05
        num_contexts = self.num_tasks

        num_ones = int(percent_on * self.dim_context)
        num_zeros = self.dim_context - num_ones

        self.contexts = torch.cat((
            torch.zeros((num_contexts, num_zeros)),
            torch.ones((num_contexts, num_ones))
        ), dim=1)

        # Shuffle each context vector i.i.d. to randomize it; use a fixed seed during
        # shuffling so that train & validation contexts are identical
        g = torch.manual_seed(seed)

        for i in range(num_contexts):
            self.contexts[i, :] = self.contexts[i, torch.randperm(self.dim_context,
                                                                  generator=g)]

    def init_centroids(self):
        """
        Code to compute the mean image from each permutation. Note that you only need
        to compute the mean image for the base dataset. After that you can just apply
        each permutation to the mean vector.
        """
        self.centroids = torch.zeros((self.num_tasks, 28, 28))
        for index in range(len(self.data)):
            img, _ = self[index]
            self.centroids[0] += img

        # This first row has the pixelwise sum for MNIST, divide to get mean
        self.centroids[0] /= len(self.data)

        # Now just apply permutations to the mean vector, one for each remaining task
        for task in range(1, self.num_tasks):
            self.centroids[task] = permute(self.centroids[0], task)

        # 28 x 28 -> 784
        self.contexts = self.centroids.flatten(start_dim=1)

    def init_one_hot_context(self):

        self.contexts = torch.eye(self.num_tasks, self.num_tasks)


def permute(x, permutation):
    """
    Applies the permutation specified by `permutation` to `x` and returns the resulting
    tensor.
    """

    # Assume x has shape (1, height, width)
    if permutation is None:
        return x

    _, height, width = x.size()
    x = x.view(-1, 1)
    x = x[permutation, :]
    x = x.view(1, height, width)
    return x

def tuple_context(x, context):
    return (x, context)

def concat_context(x, context):
    img = x.flatten()
    return torch.cat(img, context)

{'fraction_non_alphanumeric': 0.10332966428178315, 'fraction_numerical': 0.0074986241056686845, 'mean_word_length': 2.4835849508746706, 'pattern_counts': {'":': 5, '<': 3, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '676195', 'n_tokens_mistral': 4205, 'n_tokens_neox': 3946, 'n_words': 1373}