import numpy as np import random import torch from utils import randvec __author__ = "Atticus Geiger" __version__ = "CS224u, Stanford, Spring 2023" def get_IIT_equality_dataset_both(embed_dim, size): train_dataset = IIT_PremackDatasetBoth( embed_dim=embed_dim, size=size) X_base_train, X_sources_train, y_base_train, y_IIT_train, interventions = train_dataset.create() X_base_train = torch.tensor(X_base_train) X_sources_train = [torch.tensor(X_source_train) for X_source_train in X_sources_train] y_base_train = torch.tensor(y_base_train) y_IIT_train = torch.tensor(y_IIT_train) interventions = torch.tensor(interventions) return X_base_train, X_sources_train, y_base_train, y_IIT_train, interventions def get_IIT_equality_dataset(variable, embed_dim, size): class_size = size/2 train_dataset = IIT_PremackDataset( variable, embed_dim=embed_dim, n_pos=class_size, n_neg=class_size) X_base_train, X_sources_train, y_base_train, y_IIT_train, interventions = train_dataset.create() X_base_train = torch.tensor(X_base_train) X_sources_train = [torch.tensor(X_source_train) for X_source_train in X_sources_train] y_base_train = torch.tensor(y_base_train) y_IIT_train = torch.tensor(y_IIT_train) interventions = torch.tensor(interventions) return X_base_train, X_sources_train, y_base_train, y_IIT_train, interventions def get_equality_dataset(embed_dim, size): class_size = size/2 train_dataset = PremackDataset( embed_dim=embed_dim, n_pos=class_size, n_neg=class_size) X_train, y_train = train_dataset.create() test_dataset = PremackDataset( embed_dim=embed_dim, n_pos=class_size, n_neg=class_size) X_test, y_test = test_dataset.create() train_dataset.test_disjoint(test_dataset) X_train = torch.tensor(X_train) X_test = torch.tensor(X_test) return X_train, X_test, y_train, y_test, test_dataset class EqualityDataset: POS_LABEL = 1 NEG_LABEL = 0 def __init__(self, embed_dim=50, n_pos=500, n_neg=500, flatten=True): """Creates simple equality datasets, which are basically lists of `((vec1, vec2), label)` instances, where `label == POS_LABEL` if `vec1 == vec2`, else `label == NEG_LABEL`. With `flatten=True`, the instances become `(vec1;vec2, label)`. Parameters ---------- embed_dim : int Sets the dimensionality of the individual component vectors. n_pos : int n_neg : int flatten : bool If False, instances are of the form ((vec1, vec2), label). If True, vec1 and vec2 are concatenated, creating instances (x, label) where len(x) == embed_dim*2. Usage ----- dataset = EqualityDataset() X, y = dataset.create() Attributes ---------- embed_dim : int n_pos : int n_neg : int flatten : bool """ self.embed_dim = embed_dim self.n_pos = n_pos self.n_neg = n_neg self.flatten = flatten def create(self): """Main interface Attributes ---------- data : list Shuffled version of the raw instances, ignoring `self.flatten`. Thus, these are all of the form `((vec1, vec2), label)` X : np.array The dimensionality depends on `self.flatten`. If it is False, then `X.shape == (n_pos+n_neg, 2, embed_dim)`. If it is True, then `X.shape == (n_pos+n_neg, embed_dim*2)`. y : list Containing `POS_LABEL` and `NEG_LABEL`. Length: n_pos+n_neg Returns ------- self.X, self.y """ self.data = [] self.data += self._create_pos() self.data += self._create_neg() random.shuffle(self.data) data = self.data.copy() if self.flatten: data = [(np.concatenate(x), label) for x, label in data] X, y = zip(*data) self.X = np.array(X) self.y = y return self.X, self.y def test_disjoint(self, other_dataset): these_vecs = {tuple(x) for pair, label in self.data for x in pair} other_vecs = {tuple(x) for pair, label in other_dataset.data for x in pair} shared = these_vecs & other_vecs assert len(shared) == 0, \ f"This dataset and the other dataset shared {len(shared)} word-level reps." def _create_pos(self): data = [] for _ in range(self.n_pos): vec = randvec(self.embed_dim) rep = (vec, vec) data.append((rep, self.POS_LABEL)) return data def _create_neg(self): data = [] for _ in range(self.n_neg): vec1 = randvec(self.embed_dim) vec2 = vec1.copy() while np.array_equal(vec1, vec2): vec2 = randvec(self.embed_dim) rep = (vec1, vec2) data.append((rep, self.NEG_LABEL)) return data class PremackDataset: POS_LABEL = 1 NEG_LABEL = 0 def __init__(self, embed_dim=50, n_pos=500, n_neg=500, flatten_root=True, flatten_leaves=True, intermediate=False): """Creates Premack datasets. Conceptually, the instances are (((a, b), (c, d)), label) where `label == POS_LABEL` if (a == b) == (c == d), else `label == NEG_LABEL`. With `flatten_leaves=True`, these become ((a;b, c;d), label) and with `flatten_root=True`, these become (a;b;c;d, label) and `flatten_root=True` means that `flatten_leaves=True`, since we can't flatten the roof without flattening the leaves. Parameters ---------- embed_dim : int Sets the dimensionality of the individual component vectors. n_pos : int n_neg : int flatten_root : bool flatten_leaves : bool Usage ----- dataset = EqualityDataset() X, y = dataset.create() Attributes ---------- embed_dim : int n_pos : int n_neg : int flatten_root : bool flatten_leaves : bool n_same_same : n_pos / 2 n_diff_diff : n_pos / 2 n_same_diff : n_neg / 2 n_diff_same : n_neg / 2 Raises ------ ValueError If `n_pos` or `n_neg` is not even, since this means we can't get an even distribtion of the two sub-types of each of those classes while also staying faithful to user's expected number of examples for each class. """ self.embed_dim = embed_dim self.n_pos = n_pos self.n_neg = n_neg for n, v in ((n_pos, 'n_pos'), (n_neg, 'n_neg')): if n % 2 != 0: raise ValueError( f"The value of {v} must be even to ensure a balanced " f"split across its two sub-types of the {v} class.") self.n_same_same = int(n_pos / 2) self.n_diff_diff = int(n_pos / 2) self.n_same_diff = int(n_neg / 2) self.n_diff_same = int(n_neg / 2) self.flatten_root = flatten_root self.flatten_leaves = flatten_leaves self.intermediate = intermediate def create(self): """Main interface Attributes ---------- data : list Shuffled version of the raw instances, ignoring `self.flatten_root` and `self.flatten_leaves`. Thus, these are all of the form `(((a, b), (c, d)), label)` X : np.array The dimensionality depends on `self.flatten_root` and `self.flatten_leaves`. If both are False, then `X.shape == (n_pos+n_neg, 2, 2, embed_dim)` If `self.flatten_root`, then `X.shape == (n_pos+n_neg, embed_dim*4)` If only `self.flatten_leaves`, then `X.shape == (n_pos+n_neg, 2, embed_dim*2)` y : list Containing `POS_LABEL` and `NEG_LABEL`. Length: n_pos+n_neg Returns ------- self.X, self.y """ self.data = [] self.data += self._create_same_same() self.data += self._create_diff_diff() self.data += self._create_same_diff() self.data += self._create_diff_same() random.shuffle(self.data) data = self.data.copy() if self.flatten_root or self.flatten_leaves: data = [((np.concatenate(x1), np.concatenate(x2)), label) for (x1, x2), label in data] if self.flatten_root: data = [(np.concatenate(x), label) for x, label in data] X, y = zip(*data) self.X = np.array(X) self.y = y return self.X, self.y def test_disjoint(self, other_dataset): these_vecs = {tuple(x) for root_pair, label in self.data for pair in root_pair for x in pair} other_vecs = {tuple(x) for root_pair, label in other_dataset.data for pair in root_pair for x in pair} shared = these_vecs & other_vecs assert len(shared) == 0, \ f"This dataset and the other dataset shared {len(shared)} word-level reps." def _create_same_same(self): data = [] for _ in range(self.n_same_same): left = self._create_same_pair() right = self._create_same_pair() rep = (left, right) data.append((rep, self.POS_LABEL)) return data def _create_diff_diff(self): data = [] for _ in range(self.n_diff_diff): left = self._create_diff_pair() right = self._create_diff_pair() rep = (left, right) data.append((rep, self.POS_LABEL)) return data def _create_same_diff(self): data = [] for _ in range(self.n_same_diff): left = self._create_same_pair() right = self._create_diff_pair() rep = (left, right) data.append((rep, self.NEG_LABEL)) return data def _create_diff_same(self): data = [] for _ in range(self.n_diff_same): left = self._create_diff_pair() right = self._create_same_pair() rep = (left, right) data.append((rep, self.NEG_LABEL)) return data def _create_same_pair(self): vec = randvec(self.embed_dim) return (vec, vec) def _create_diff_pair(self): vec1 = randvec(self.embed_dim) vec2 = randvec(self.embed_dim) assert not np.array_equal(vec1, vec2) return (vec1, vec2) class PremackDatasetLeafFlattened(PremackDataset): def __init__(self, embed_dim=50, n_pos=500, n_neg=500): super().__init__( embed_dim=embed_dim, n_pos=n_pos, n_neg=n_neg, flatten_leaves=True, flatten_root=False, intermediate=False) class IIT_PremackDataset: V1 = 0 V2 = 1 POS_LABEL = 1 NEG_LABEL = 0 def __init__(self, variable, embed_dim=50, n_pos=500, n_neg=500, flatten_root=True, flatten_leaves=True, intermediate=False): self.variable = variable self.embed_dim = embed_dim self.n_pos = n_pos self.n_neg = n_neg for n, v in ((n_pos, 'n_pos'), (n_neg, 'n_neg')): if n % 2 != 0: raise ValueError( f"The value of {v} must be even to ensure a balanced " f"split across its two sub-types of the {v} class.") self.n_same_same_to_same = int(n_pos / 4) self.n_diff_diff_to_same = int(n_neg / 4) self.n_same_diff_to_same = int(n_neg / 4) self.n_diff_same_to_same = int(n_neg / 4) self.n_same_same_to_diff = int(n_neg / 4) self.n_diff_diff_to_diff = int(n_neg / 4) self.n_same_diff_to_diff = int(n_neg / 4) self.n_diff_same_to_diff = int(n_neg / 4) self.flatten_root = flatten_root self.flatten_leaves = flatten_leaves self.intermediate = intermediate def create(self): self.data = [] self.data += self._create_same_same_to_same() self.data += self._create_diff_diff_to_same() self.data += self._create_same_diff_to_same() self.data += self._create_diff_same_to_same() self.data += self._create_same_same_to_diff() self.data += self._create_diff_diff_to_diff() self.data += self._create_same_diff_to_diff() self.data += self._create_diff_same_to_diff() random.shuffle(self.data) data = self.data.copy() if self.flatten_root or self.flatten_leaves: data = [((np.concatenate(x1), np.concatenate(x2)),(np.concatenate(x3), np.concatenate(x4)), base_label, IIT_label, intervention) for (x1, x2,x3,x4), base_label, IIT_label, intervention in data] if self.flatten_root: data = [(np.concatenate(base), np.concatenate(source), label, IIT_label, intervention) for base, source, label, IIT_label, intervention in data] base, source, y, IIT_y, interventions = zip(*data) self.base = np.array(base) self.source = np.array(source) self.y = np.array(y) self.IIT_y = np.array(IIT_y) self.interventions = np.array(interventions) self.sources = list() self.sources.append(self.source) return self.base, self.sources, self.y, self.IIT_y, self.interventions def _create_same_same_to_same(self): data = [] for _ in range(self.n_same_same_to_same): base_left = self._create_same_pair() base_right = self._create_same_pair() base_label = self.POS_LABEL if self.variable == "V1": source_left = self._create_same_pair() source_right = self._create_random_pair() IIT_label = self.POS_LABEL intervention = self.V1 if self.variable == "V2": source_left = self._create_random_pair() source_right = self._create_same_pair() intervention = self.V2 IIT_label = self.POS_LABEL rep = (base_left, base_right, source_left,source_right) data.append((rep, base_label, IIT_label, intervention)) return data def _create_diff_diff_to_same(self): data = [] for _ in range(self.n_diff_diff_to_same): base_left = self._create_diff_pair() base_right = self._create_diff_pair() base_label = self.POS_LABEL if self.variable == "V1": source_left = self._create_same_pair() source_right = self._create_random_pair() IIT_label = self.NEG_LABEL intervention = self.V1 if self.variable == "V2": source_left = self._create_random_pair() source_right = self._create_same_pair() IIT_label = self.NEG_LABEL intervention = self.V2 rep = (base_left, base_right, source_left,source_right) data.append((rep, base_label, IIT_label, intervention)) return data def _create_same_diff_to_same(self): data = [] for _ in range(self.n_same_diff_to_same): base_left = self._create_same_pair() base_right = self._create_diff_pair() base_label = self.NEG_LABEL if self.variable == "V1": source_left = self._create_same_pair() source_right = self._create_random_pair() IIT_label = self.NEG_LABEL intervention = self.V1 if self.variable == "V2": source_left = self._create_random_pair() source_right = self._create_same_pair() IIT_label = self.POS_LABEL intervention = self.V2 rep = (base_left, base_right, source_left,source_right) data.append((rep, base_label, IIT_label, intervention)) return data def _create_diff_same_to_same(self): data = [] for _ in range(self.n_diff_same_to_same): base_left = self._create_diff_pair() base_right = self._create_same_pair() base_label = self.NEG_LABEL if self.variable == "V1": source_left = self._create_same_pair() source_right = self._create_random_pair() IIT_label = self.POS_LABEL intervention = self.V1 if self.variable == "V2": source_left = self._create_random_pair() source_right = self._create_same_pair() IIT_label = self.NEG_LABEL intervention = self.V2 rep = (base_left, base_right, source_left,source_right) data.append((rep, base_label, IIT_label, intervention)) return data def _create_same_same_to_diff(self): data = [] for _ in range(self.n_same_same_to_diff): base_left = self._create_same_pair() base_right = self._create_same_pair() base_label = self.POS_LABEL if self.variable == "V1": source_left = self._create_diff_pair() source_right = self._create_random_pair() IIT_label = self.NEG_LABEL intervention = self.V1 if self.variable == "V2": source_left = self._create_random_pair() source_right = self._create_diff_pair() IIT_label = self.NEG_LABEL intervention = self.V2 rep = (base_left, base_right, source_left,source_right) data.append((rep, base_label, IIT_label, intervention)) return data def _create_diff_diff_to_diff(self): data = [] for _ in range(self.n_diff_diff_to_diff): base_left = self._create_diff_pair() base_right = self._create_diff_pair() base_label = self.POS_LABEL if self.variable == "V1": source_left = self._create_diff_pair() source_right = self._create_random_pair() IIT_label = self.POS_LABEL intervention = self.V1 if self.variable == "V2": source_left = self._create_random_pair() source_right = self._create_diff_pair() IIT_label = self.POS_LABEL intervention = self.V2 rep = (base_left, base_right, source_left,source_right) data.append((rep, base_label, IIT_label, intervention)) return data def _create_same_diff_to_diff(self): data = [] for _ in range(self.n_same_diff_to_diff): base_left = self._create_same_pair() base_right = self._create_diff_pair() base_label = self.NEG_LABEL if self.variable == "V1": source_left = self._create_diff_pair() source_right = self._create_random_pair() IIT_label = self.POS_LABEL intervention = self.V1 if self.variable == "V2": source_left = self._create_random_pair() source_right = self._create_diff_pair() IIT_label = self.NEG_LABEL intervention = self.V2 rep = (base_left, base_right, source_left,source_right) data.append((rep, base_label, IIT_label, intervention)) return data def _create_diff_same_to_diff(self): data = [] for _ in range(self.n_diff_same_to_diff): base_left = self._create_diff_pair() base_right = self._create_same_pair() base_label = self.NEG_LABEL if self.variable == "V1": source_left = self._create_diff_pair() source_right = self._create_random_pair() IIT_label = self.NEG_LABEL intervention = self.V1 if self.variable == "V2": source_left = self._create_random_pair() source_right = self._create_diff_pair() IIT_label = self.POS_LABEL intervention = self.V2 rep = (base_left, base_right, source_left,source_right) data.append((rep, base_label, IIT_label, intervention)) return data def _create_random_pair(self): if random.choice([True,False]): return self._create_same_pair() else: return self._create_diff_pair() def _create_same_pair(self): vec = randvec(self.embed_dim) return (vec, vec) def _create_diff_pair(self): vec1 = randvec(self.embed_dim) vec2 = randvec(self.embed_dim) assert not np.array_equal(vec1, vec2) return (vec1, vec2) class IIT_PremackDatasetBoth: V1 = 0 V2 = 1 POS_LABEL = 1 NEG_LABEL = 0 both_coord_id = 2 def __init__(self, size= 1000, embed_dim=50, flatten_root=True, flatten_leaves=True, intermediate=False): self.embed_dim = embed_dim self.size= size self.flatten_root = flatten_root self.flatten_leaves = flatten_leaves self.intermediate = intermediate def create(self): data = [] for _ in range(self.size): rep = [self._create_random_pair() for _ in range(6)] if (rep[0][0] == rep[0][1]).all() == (rep[1][0] == rep[1][1]).all(): base_label = self.POS_LABEL else: base_label = self.NEG_LABEL if (rep[2][0] == rep[2][1]).all() == (rep[5][0] == rep[5][1]).all(): IIT_label = self.POS_LABEL else: IIT_label = self.NEG_LABEL data.append((rep,base_label, IIT_label, self.both_coord_id)) random.shuffle(data) data = data.copy() if self.flatten_root or self.flatten_leaves: data = [ ( ( (np.concatenate(x1), np.concatenate(x2)), (np.concatenate(x3), np.concatenate(x4)), (np.concatenate(x5), np.concatenate(x6)) ), base_label, IIT_label, intervention ) for (x1, x2,x3,x4,x5,x6), base_label, IIT_label, intervention in data ] if self.flatten_root: data = [(np.concatenate(base), np.concatenate(source),np.concatenate(source2), label, IIT_label, intervention) for (base, source, source2), label, IIT_label, intervention in data] base, source, source2, y, IIT_y, interventions = zip(*data) self.base = np.array(base) self.source = np.array(source) self.source2 = np.array(source2) self.y = np.array(y) self.IIT_y = np.array(IIT_y) self.interventions = np.array(interventions) return self.base, [self.source, self.source2], self.y, self.IIT_y, self.interventions def _create_random_pair(self): if random.choice([True,False]): return self._create_same_pair() else: return self._create_diff_pair() def _create_same_pair(self): vec = randvec(self.embed_dim) return (vec, vec) def _create_diff_pair(self): vec1 = randvec(self.embed_dim) vec2 = randvec(self.embed_dim) assert not np.array_equal(vec1, vec2) return (vec1, vec2)