Merge with diff history. Code not connected yet.

changelog

+6/12
+
+Random initialization of generator hiddens, and rework all hidden generation
+Separation of model and loss. Rewrite the training function.
+Generator PG loss runs the generator again
+Removal of oracle loss for surpervised loss. Any oracle call removed
+batchPGLoss was written by iterating over each individual tensor element with no vectorized operation. Fixed
+Dataloader and packed/padded variable length tensor
+Console output cleanup
+Loss implemented as a model
+In discriminator, GRU output used, not hidden.
\ No newline at end of file

dataset.py

+import numpy as np
+import torch
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset
+
+from parameters import *
+
+
+class SupervisedData:
+    def __init__(self, samples, targets, valid_fold=10):
+        # TODO easy discriminator: ends with 1
+        super(SupervisedData, self).__init__()
+        self.samples = samples
+        self.targets = targets
+
+        assert (samples.shape[0] == SUPERVISED_SIZE)
+        assert (targets.shape[0] == SUPERVISED_SIZE)
+
+        self.valid_size = samples.shape[0] // valid_fold
+        self.train_size = SUPERVISED_SIZE - self.valid_size
+
+        self.training_set = None
+        self.validation_set = None
+
+        self.split()
+
+    def split(self):
+        indices = list(range(0, SUPERVISED_SIZE - 1))
+        # random.shuffle(indices)
+        train_indices = indices[0:self.train_size]
+        valid_indices = indices[self.train_size:]
+        self.training_set = SupervisedDataset(self, train_indices)
+        self.validation_set = SupervisedDataset(self, valid_indices)
+
+
+class SupervisedDataset(Dataset):
+    def __init__(self, supervised_data, indices):
+        super(SupervisedDataset, self).__init__()
+        self.indices = indices
+        self.supervised_data = supervised_data
+
+    def __len__(self):
+        return len(self.indices)
+
+    def __getitem__(self, idx):
+        data_idx = self.indices[idx]
+        x = self.supervised_data.samples[data_idx]
+        y = self.supervised_data.targets[data_idx]
+
+        # dumb and slow consistency check to ensure that padding is only at the end
+        first_one = None
+        for idx in range(x.size()[0]):
+            if x[idx] == 1 and first_one is None:
+                first_one = idx
+            elif first_one is not None:
+                assert x[idx] == 1
+
+        if first_one is not None:
+            xx = x[:first_one]
+        else:
+            xx = x
+        return xx, y
+
+
+class GeneratedDataset(Dataset):
+    """
+    Receives the sequences generated from the generator
+    Calls the zork and gcov to run the sequences generated.
+    """
+
+    def __init__(self, generated_sequences, random_truncate=False):
+        self.generated_sequences = generated_sequences.cpu()
+        self.random_truncate = random_truncate
+        self.truncate_starts = self.generated_sequences.shape[1] - \
+                               np.random.choice(range(0, 7), self.generated_sequences.shape[0])
+
+    def measure_validity_coverage(self, sequence):
+        # TODO magic, maybe run in __init__
+        label = torch.zeros((3,))
+        # label[1]=validity(sequence)
+        # label[2]=coverage(sequence)
+        return label
+
+    def get_sequence(self, idx):
+        if not self.random_truncate:
+            return self.generated_sequences[idx, :]
+        else:
+            return self.generated_sequences[idx, : self.truncate_starts[idx]]
+
+    def __getitem__(self, idx):
+        sequence = self.get_sequence(idx),
+        return self.get_sequence(idx), self.measure_validity_coverage(sequence)
+
+    def __len__(self):
+        return self.generated_sequences.shape[0]
+
+
+def auto_batch_target_length(sequences, random_truncate=False):
+    sequences_list = []
+    for i in range(sequences.shape[0]):
+        seq = sequences[i, :]
+        if random_truncate:
+            seq = seq[:seq.shape[0] - np.random.choice(range(0, 7))]
+        sequences_list.append(seq)
+
+    x_lens = [len(x) for x in sequences_list]
+    xx_pad = pad_sequence(sequences_list, batch_first=True, padding_value=1)
+
+    return xx_pad, x_lens
+
+
+class ConcatDataset(Dataset):
+    """
+    Not shuffled. Dataloader does the shuffling.
+    """
+
+    def __init__(self, *datasets):
+        self.datasets = datasets
+
+    def __getitem__(self, i):
+        for dataset in self.datasets:
+            if i < len(dataset):
+                return dataset[i]
+            else:
+                i = i - len(dataset)
+
+    def __len__(self):
+        sum_len = 0
+        for ds in self.datasets:
+            sum_len += len(ds)
+        return sum_len
+
+
+def pad_collate(batch):
+    (xx, yy) = zip(*batch)
+    x_lens = [len(x) for x in xx]
+
+    xx_pad = pad_sequence(xx, batch_first=True, padding_value=1)
+    yy_stack = torch.stack(yy)
+
+    return xx_pad, yy_stack, x_lens
+
+
+def test_data_loader():
+    supervised_samples = torch.load(supervised_sequences_path)
+    supervised_targets = torch.load(supervised_val_cov_path)
+    zork_data = SupervisedData(supervised_samples, supervised_targets)
+    train_set = DataLoader(dataset=zork_data.training_set, batch_size=BATCH_SIZE,
+                           shuffle=True, collate_fn=pad_collate)
+    for d in train_set:
+        print(d)

methods_665.py

-import torch
-from torch.autograd import Variable
 from math import ceil
-import sys
 
-def prepare_generator_batch(samples, sample_lengths, start_letter=0, gpu=False):
+import torch
+
+def tensor_to_array(tensor):
+    return tensor.cpu().numpy()
+
+def one_token_shift_sequences(sequences, start_letter=0, gpu=False):
     """
-    Takes samples (a batch) and returns
+    Take a batch of sequences, one shift the sequences so that the model uses token i to predict token i+1.
 
     Inputs: samples, start_letter, cuda
-        - samples: batch_size x seq_len (Tensor with a sample in each row)
+        - samples: batch_size x seq_len
 
     Returns: inp, target
-        - inp: batch_size x seq_len (same as target, but with start_letter prepended)
-        - target: batch_size x seq_len (Variable same as samples)
+        - inp: previous tokens, prepend with start_letter
+        batch_size x seq_len (same as target, but with start_letter prepended)
+        - target: tokens to be predicted, sequences themselves
+        batch_size x seq_len
     """
 
-    batch_size, seq_len = samples.size()
+    batch_size, seq_len = sequences.size()
 
-    inp = torch.zeros(batch_size, seq_len)
-    target = samples
+    inp = torch.zeros(batch_size, seq_len, device=sequences.device)
+    target = sequences
     inp[:, 0] = start_letter
-    inp[:, 1:] = target[:, :seq_len-1]
-
-    inp = Variable(inp).type(torch.LongTensor)
-    target = Variable(target).type(torch.LongTensor)
-
-    if gpu:
-        inp = inp.cuda()
-        target = target.cuda()
-     #   samples_lengths = sample_lengths.cuda() #?
-
-    return inp, sample_lengths, target
-
+    inp[:, 1:] = target[:, :seq_len - 1]
 
-def prepare_discriminator_data(pos_samples, neg_samples, gpu=False, real_target=None):
-    """
-    Takes positive (target) samples, negative (generator) samples and prepares inp and target data for discriminator.
-
-    Inputs: pos_samples, neg_samples
-        - pos_samples: pos_size x seq_len
-        - neg_samples: neg_size x seq_len
-
-    Returns: inp, target
-        - inp: (pos_size + neg_size) x seq_len
-        - target: pos_size + neg_size (boolean 1/0)
-    """
-
-    #print(pos_samples.size())
-  #  print(neg_samples.size())
-  #  pos_samples = pos_samples.narrow(1,0,20)
-    pos_samples = pos_samples.type(torch.LongTensor)
-    neg_samples = neg_samples.type(torch.LongTensor)
-    inp = torch.cat((pos_samples, neg_samples), 0).type(torch.LongTensor)
-    target = torch.ones([pos_samples.size()[0] + neg_samples.size()[0], 3])
-    target[pos_samples.size()[0]:] = 0
-    
-    if real_target != None:
-        target[0:pos_samples.size()[0]] = real_target
-
-    # shuffle
-    perm = torch.randperm(target.size()[0])
-    target = target[perm]
-    inp = inp[perm]
-
-    inp = Variable(inp)
-    target = Variable(target)
+    inp = inp.long()
+    target = target.long()
 
     if gpu:
         inp = inp.cuda()
         target = target.cuda()
+    #   samples_lengths = sample_lengths.cuda() #?
 
     return inp, target
 
 
-def batchwise_sample(gen, num_samples, batch_size):
-    """
-    Sample num_samples samples batch_size samples at a time from gen.
-    Does not require gpu since gen.sample() takes care of that.
-    """
-
-    samples = []
-    for i in range(int(ceil(num_samples/float(batch_size)))):
-        samples.append(gen.sample(batch_size))
-
-    return torch.cat(samples, 0)[:num_samples]
-
+# def prepare_discriminator_data(pos_sequences, neg_sequences, gpu=False, real_target=None):
+#     """
+#     Takes positive (target) samples, negative (generator) samples and prepares inp and target data for discriminator.
+#
+#     Inputs: pos_samples, neg_samples
+#         - pos_samples: pos_size x seq_len
+#         - neg_samples: neg_size x seq_len
+#
+#     Returns: inp, target
+#         - inp: (pos_size + neg_size) x seq_len
+#         - target: pos_size + neg_size (boolean 1/0)
+#     """
+#
+#     # print(pos_samples.size())
+#     #  print(neg_samples.size())
+#     #  pos_samples = pos_samples.narrow(1,0,20)
+#     pos_sequences = pos_sequences.long()
+#     neg_sequences = neg_sequences.long()
+#     inp = torch.cat((pos_sequences, neg_sequences), 0).long()
+#     # targets are ones
+#     target = torch.ones([pos_sequences.size()[0] + neg_sequences.size()[0], 3])
+#     target[pos_sequences.size()[0]:] = 0
+#
+#     if real_target is not None:
+#         target[0:pos_sequences.size()[0]] = real_target
+#
+#     # shuffle
+#     perm = torch.randperm(target.size()[0])
+#     target = target[perm]
+#     inp = inp[perm]
+#
+#     if gpu:
+#         inp = inp.cuda()
+#         target = target.cuda()
+#
+#     return inp, target
+
+
+# def multi_batch_sample(gen, num_batches, batch_size):
+#     samples = []
+#     for i in range(num_batches):
+#         samples.append(gen.generate_sequences(batch_size))
+#
+#     return torch.cat(samples, 0)
+#
+
+# def supervised_nll(gen, num_batches, batch_size, max_seq_len, start_letter=0, gpu=False):
+#     # after concatenating this guy slice the tensors later, lmao
+#     generated_batches = []
+#     for i in range(int(ceil(num_batches / float(batch_size)))):
+#         generated_batches.append(gen.generate_sequences(batch_size))
+#     # samples = torch.cat(samples, 0)
+#
+#     oracle_nll = 0
+#     for i in range(0, num_batches):
+#         inp, target = one_shift_sequences(generated_batches[i], start_letter, gpu)
+#         oracle_loss = gen.batchNLLLoss(inp, target) / max_seq_len
+#         oracle_nll += oracle_loss.data.item()
+#
+#     return oracle_nll / num_batches
 
-def batchwise_oracle_nll(gen, oracle, num_samples, batch_size, max_seq_len, start_letter=0, gpu=False):
-    s = batchwise_sample(gen, num_samples, batch_size)
-    oracle_nll = 0
-    for i in range(0, num_samples, batch_size):
-        inp, inp_lengths, target = prepare_generator_batch(s[i:i+batch_size], max_seq_len, start_letter, gpu)
-      #  if(gpu):
-      #      print("HERE")
-      #      inp = inp.cuda()
-       #     target = target.cuda()
-        oracle_loss = oracle.batchNLLLoss(inp, inp_lengths, target) / max_seq_len
-        oracle_nll += oracle_loss.data.item()
 
-    return oracle_nll/(num_samples/batch_size)
-    

models_665.py

 import torch
-import torch.autograd as autograd
 import torch.nn as nn
-import pdb
 import torch.nn.functional as F
-import numpy as np
-import math
 import torch.nn.init as init
+from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
+
+
+# TODO check parameter initialization
 
 class Discriminator(nn.Module):
 
@@ -18,59 +18,93 @@ class Discriminator(nn.Module):
 
         self.embeddings = nn.Embedding(vocab_size, embedding_dim)
         self.gru = nn.GRU(embedding_dim, hidden_dim, num_layers=2, bidirectional=True, dropout=dropout)
-        self.gru2hidden = nn.Linear(2*2*hidden_dim, hidden_dim)
+        self.gru2hidden = nn.Linear(2 * 2* hidden_dim, hidden_dim)
+        self.gru_out_linear = nn.Linear(2 * hidden_dim, hidden_dim)
         self.dropout_linear = nn.Dropout(p=dropout)
         self.hidden2out = nn.Linear(hidden_dim, 3)
 
-    def init_hidden(self, batch_size):
-        h = autograd.Variable(torch.zeros(2*2*1, batch_size, self.hidden_dim))
-
-        if self.gpu:
-            return h.cuda()
-        else:
-            return h
-
-    def forward(self, input, lengths, hidden):
-        # input dim                                                # batch_size x seq_len
-        emb = self.embeddings(input)                               # batch_size x seq_len x embedding_dim
-        emb = emb.permute(1, 0, 2)                                 # seq_len x batch_size x embedding_dim
-        _, hidden = self.gru(emb, hidden)                          # 4 x batch_size x hidden_dim
-        hidden = hidden.permute(1, 0, 2).contiguous()              # batch_size x 4 x hidden_dim
-        out = self.gru2hidden(hidden.view(-1, 4*self.hidden_dim))  # batch_size x 4*hidden_dim
-        out = torch.tanh(out)
-        out = self.dropout_linear(out)
-        out = self.hidden2out(out)                                 # batch_size x 1
-        out = torch.sigmoid(out)
-        return out
-
-    def batchClassify(self, inp):
+    # def hidden_forward(self, input, length, hidden=None):
+    #     """
+    #
+    #     :param input: whole sequence, batch_size x seq_len
+    #     :param hidden: default to zeros
+    #     :return:
+    #     """
+    #     emb = self.embeddings(input)  # batch_size x seq_len x embedding_dim
+    #     # emb = emb.permute(1, 0, 2)  # seq_len x batch_size x embedding_dim
+    #     # TODO pack tensors here
+    #     emb_packed = pack_padded_sequence(emb, length, batch_first=True, enforce_sorted=False)
+    #     # Because self.gru automatically initializes hidden
+    #     _, hidden = self.gru(emb_packed, hidden)  # 4 x batch_size x hidden_dim
+    #     hidden = hidden.permute(1, 0, 2).contiguous()  # batch_size x 4 x hidden_dim
+    #     out = self.gru2hidden(hidden.view(-1, 4 * self.hidden_dim))  # batch_size x 4*hidden_dim
+    #     out = torch.tanh(out)
+    #     out = self.dropout_linear(out)
+    #     out = self.hidden2out(out)  # batch_size x 3
+    #     out = torch.sigmoid(out)
+    #     return out
+
+    def forward(self, input, length, hidden=None):
         """
-        Classifies a batch of sequences.
 
-        Inputs: inp
-            - inp: batch_size x seq_len
-
-        Returns: out
-            - out: batch_size ([0,1] score)
+        :param input: whole sequence, batch_size x seq_len
+        :param hidden: default to zeros
+        :return:
         """
-
-        h = self.init_hidden(inp.size()[0])
-        out = self.forward(inp, 0, h)
+        emb = self.embeddings(input)  # batch_size x seq_len x embedding_dim
+        # emb = emb.permute(1, 0, 2)  # seq_len x batch_size x embedding_dim
+        # TODO pack tensors here
+        emb_packed = pack_padded_sequence(emb, length, batch_first=True, enforce_sorted=False)
+        # Because self.gru automatically initializes hidden
+        out, _ = self.gru(emb_packed, hidden)
+        out, lengths = pad_packed_sequence(out, batch_first=True)
+        outt = out[torch.arange(0, lengths.shape[0]), lengths-1, :] # (batch_size, 2*hidden_dim)
+        out = self.gru_out_linear(outt)  # batch_size x 4*hidden_dim
+        out = torch.tanh(out)
+        out = self.dropout_linear(out)
+        out = self.hidden2out(out)  # batch_size x 3
+        out = torch.sigmoid(out)
         return out
 
-    def batchBCELoss(self, inp, target):
-        """
-        Returns Binary Cross Entropy Loss for discriminator.
-
-         Inputs: inp, target
-            - inp: batch_size x seq_len
-            - target: batch_size (binary 1/0)
-        """
 
-        loss_fn = nn.BCELoss()
-        h = self.init_hidden(inp.size()[0])
-        out = self.forward(inp, h)
-        return loss_fn(out, target)
+# model and loss are separated
+
+# removed this function.
+# what does discriminator do if not batch classify? hilarious
+# def batchClassify(self, inp):
+#     """
+#     Classifies a batch of sequences.
+#
+#     Author did not read GRU implementation to know that hidden can be defaulted to None.
+#     Moreover, all tensors are batched by default in PyTorch.
+#     Original code also uses self.forward(),
+#     https://discuss.pytorch.org/t/any-different-between-model-input-and-model-forward-input/3690
+#
+#     Inputs: inp
+#         - inp: batch_size x seq_len
+#
+#     Returns: out
+#         - out: batch_size ([0,1] score)
+#     """
+#
+#     h = self.init_hidden(inp.size()[0])
+#     out = self(inp, h)
+#     return out
+#
+# def batchBCELoss(self, inp, target):
+#     """
+#     Returns Binary Cross Entropy Loss for discriminator.
+#     Model and loss need to be separate.
+#
+#      Inputs: inp, target
+#         - inp: batch_size x seq_len
+#         - target: batch_size (binary 1/0)
+#     """
+#
+#     loss_fn = nn.BCELoss()
+#     h = self.init_hidden(inp.size()[0])
+#     out = self(inp, h)
+#     return loss_fn(out, target)
 
 
 class Generator(nn.Module):
@@ -93,142 +127,149 @@ class Generator(nn.Module):
             for p in self.parameters():
                 init.normal_(p, 0, 1)
 
-    def init_hidden(self, batch_size=1):
-        h = autograd.Variable(torch.zeros(1, batch_size, self.hidden_dim))
-
-        if self.gpu:
-            return h.cuda()
-        else:
-            return h
-
-    def forward(self, inp, inp_length, hidden):
+    def forward(self, inp, hidden=None):
         """
         Embeds input and applies GRU one token at a time (seq_len = 1)
+
+        :param inp: a batch of tokens, (batch_size,)
+        :param hidden:
+        :return: out: log probability for each vocab class, (batch_size, vocab_size)
+                 hidden: hidden vector to be passed to the next time step, (1, batch_size, hidden_dim)
         """
-      #  print("in forward <<< ")
-       # print(inp)
-     #   print(inp_length)
-        
-        # input dim                                             # batch_size
-        emb = self.embeddings(inp)                              # batch_size x embedding_dim
-        
-        
-      #  print(emb)
-        
-       # inp_length, perm_idx = inp_length.sort(0, descending=True)
-       # emb = emb[perm_idx]
-        
-     #   print(emb)
-        
-       # packed = rnn_utils.pack_padded_sequence(emb, inp_length.data.cpu().numpy(), batch_first=True)
-        
-        emb = emb.view(1, -1, self.embedding_dim)               # 1 x batch_size x embedding_dim
-        
-     #   print(emb)
-        
-        out, hidden = self.gru(emb, hidden)                     # 1 x batch_size x hidden_dim (out)
-            
-        out = self.gru2out(out.view(-1, self.hidden_dim))       # batch_size x vocab_size
+
+        if hidden is None:
+            # if without xavier_uniform, all sequences generated will be the same
+            hidden = torch.empty((1, inp.size(0), self.hidden_dim), device=inp.device)
+            init.xavier_uniform_(hidden)
+
+        emb = self.embeddings(inp)  # batch_size x embedding_dim
+        emb = emb.unsqueeze(0)  # 1 x batch_size x embedding_dim
+        out, hidden = self.gru(emb, hidden)  # 1 x batch_size x hidden_dim (out)
+        out = self.gru2out(out.view(-1, self.hidden_dim))  # batch_size x vocab_size
         out = F.log_softmax(out, dim=1)
         return out, hidden
 
-    def sample(self, num_samples, start_letter=0):
+    def generate_sequences(self, batch_size, start_letter=0):
         """
-        Samples the network and returns num_samples samples of length max_seq_len.
+        Samples the network and returns batch_size samples of self.length max_seq_len.
 
-        Outputs: samples, hidden
-            - samples: num_samples x max_seq_length (a sampled sequence in each row)
+        :param batch_size:
+        :param start_letter:
+        :return: sequences: batch_size x self.max_seq_length (a sampled sequence in each row)
+                 logits: the log probability of the sampled token, used for PG loss
         """
 
-        samples = torch.zeros(num_samples, self.max_seq_len).type(torch.LongTensor)
-
-        h = self.init_hidden(num_samples)
-        inp = autograd.Variable(torch.LongTensor([start_letter]*num_samples))
-        #print("IN SAMPLE: ", inp)
-        inp_lengths = [self.max_seq_len] * len(inp)
-        
-        if self.gpu:
-            samples = samples.cuda()
+        inp = torch.LongTensor([start_letter] * batch_size)
+        sequences = torch.zeros(batch_size, self.max_seq_len).long()
+        logits = torch.zeros(batch_size, self.max_seq_len)
+        if self.cuda:
+            sequences = sequences.cuda()
             inp = inp.cuda()
+            logits = logits.cuda()
 
+        h = None
         for i in range(self.max_seq_len):
-            out, h = self.forward(inp, inp_lengths, h)               # out: num_samples x vocab_size
-            out = torch.multinomial(torch.exp(out), 1)  # num_samples x 1 (sampling from each row)
-            samples[:, i] = out.view(-1).data
-
-            inp = out.view(-1)
-
-        return samples
-
-    def batchNLLLoss(self, inp, inp_lengths, target):
-        """
-        Returns the NLL Loss for predicting target sequence.
-
-        Inputs: inp, target
-            - inp: batch_size x seq_len
-            - target: batch_size x seq_len
-
-            inp should be target with <s> (start letter) prepended
-        """
-        
-       # inp = torch.transpose(inp,0,1)
-       # target = torch.transpose(target,0,1)
-
-        loss_fn = nn.NLLLoss()
-        batch_size, seq_len = inp.size()
-        inp = inp.permute(1, 0)           # seq_len x batch_size
-        target = target.permute(1, 0)     # seq_len x batch_size
-        h = self.init_hidden(batch_size)
-
-        loss = 0
-        
-       # print("In batchNLLLoss <<<")
-        
-       # print(inp[0], inp_lengths[0], h)
-        
-      #  print("after batchNLLLoss prints")
-        
-        for i in range(seq_len):
-            out, h = self.forward(inp[i], inp_lengths, h)
-            loss += loss_fn(out, target[i])
-
-        return loss     # per batch
-
-    def batchPGLoss(self, inp, target, reward):
-        """
-        Returns a pseudo-loss that gives corresponding policy gradients (on calling .backward()).
-        Inspired by the example in http://karpathy.github.io/2016/05/31/rl/
-
-        Inputs: inp, target
-            - inp: batch_size x seq_len
-            - target: batch_size x seq_len
-            - reward: batch_size (discriminator reward for each sentence, applied to each token of the corresponding
-                      sentence)
-
-            inp should be target with <s> (start letter) prepended
-        """
-        weights =[0.4,0.10,0.50]
-        
-        batch_size, seq_len = inp.size()
-        inp = inp.permute(1, 0)          # seq_len x batch_size
-        target = target.permute(1, 0)    # seq_len x batch_size
-        h = self.init_hidden(batch_size)
-        inp_lengths = [self.max_seq_len] * len(inp)
-
-                                                                    # punish repeat targets?
-                                                                    # how to define repeat target?
-        
-        loss = 0
-        for i in range(seq_len):
-            out, h = self.forward(inp[i], inp_lengths, h)
-            # TODO: should h be detached from graph (.detach())?
-            for j in range(batch_size):
-                loss += -out[j][target.data[i][j]]*reward[j]     # log(P(y_t|Y_1:Y_{t-1})) * Q
-                                                                                        # replace with monte carlo?
-      #  print(loss)
-        #loss = obj_param*loss[0] + (1-obj_param)*(loss[1] + loss[2]) 
-        loss = weights[0]*loss[0] + weights[1]*loss[1] + weights[2]*loss[2] 
-            
-        return loss/batch_size
-
-
+            out, h = self(inp, h)  # out: num_samples x vocab_size
+            # is autograd successfully passing the sampling?
+            tokens = torch.multinomial(torch.exp(out), 1).squeeze(1)  # num_samples x 1 (sampling from each row)
+            sequences[:, i] = tokens
+            logits[:, i] = out[torch.arange(0, batch_size), tokens]
+            inp = tokens
+
+        return sequences, logits.contiguous()
+
+    # def batchNLLLoss(self, inp, target):
+    #     # TODO fix all of this. The training logic need to be exposed in parallel in the calling function.
+    #     """
+    #     Returns the NLL Loss for predicting target sequence.
+    #
+    #     Inputs: inp, target
+    #         - inp: batch_size x seq_len
+    #         - target: batch_size x seq_len
+    #
+    #         inp should be target with <s> (start letter) prepended
+    #     """
+    #
+    #     # inp = torch.transpose(inp,0,1)
+    #     # target = torch.transpose(target,0,1)
+    #
+    #     loss_fn = nn.NLLLoss()
+    #     batch_size, seq_len = inp.size()
+    #     # the hiddens are correct?
+    #     inp = inp.permute(1, 0)  # seq_len x batch_size
+    #     target = target.permute(1, 0)  # seq_len x batch_size
+    #     h = self.init_hidden(batch_size)
+    #
+    #     loss = 0
+    #
+    #     # print("In batchNLLLoss <<<")
+    #
+    #     # print(inp[0], inp_lengths[0], h)
+    #
+    #     #  print("after batchNLLLoss prints")
+    #
+    #     for i in range(seq_len):
+    #         out, h = self(inp[i], h)
+    #         loss += loss_fn(out, target[i])
+    #
+    #     return loss  # per batch
+    #
+    # def batchPGLoss(self, inp, target, reward):
+    #     """
+    #
+    #     Inputs: inp, target
+    #         - inp: batch_size x seq_len
+    #         - target: batch_size x seq_len
+    #         - reward: batch_size (discriminator reward for each sentence, applied to each token of the corresponding
+    #                   sentence)
+    #
+    #         inp should be target with <s> (start letter) prepended
+    #     """
+    #     weights = [0.4, 0.10, 0.50]
+    #
+    #     batch_size, seq_len = inp.size()
+    #     # inp = inp.permute(1, 0)  # seq_len x batch_size
+    #     # target = target.permute(1, 0)  # seq_len x batch_size
+    #     # h = self.init_hidden(batch_size)
+    #     # inp_lengths = [self.max_seq_len] * len(inp)
+    #
+    #     # punish repeat targets?
+    #     # how to define repeat target?
+    #
+    #     loss = 0
+    #     for i in range(seq_len):
+    #         tokens = inp[:, i]
+    #         out, h = self(tokens)
+    #         # TODO: should h be detached from graph (.detach())?
+    #         for j in range(batch_size):
+    #             # TODO what is this? Why are we indexing?
+    #             # TODO what is this? Look at this ridiculous unvectorized indexing fiesta
+    #             loss += -out[j][target.data[j][i]] * reward[j]  # -log(P(y_t|Y_1:Y_{t-1})) * Q
+    #             # replace with monte carlo?
+    #     loss = weights[0] * loss[0] + weights[1] * loss[1] + weights[2] * loss[2]
+    #
+    #     return loss / batch_size
+    #
+
+
+class PolicyGradientLoss(nn.Module):
+    def __init__(self, size_average=True, reduce=True):
+        super(PolicyGradientLoss, self).__init__()
+        self.weights = [0.4, 0.10, 0.50]
+        # TODO lacks documentation as to which is which
+        # self.fake_weight=0.4
+        # self.val_weight=0.1
+        # self.cov_weight=0.5
+        self.size_average = size_average
+        self.reduce = reduce
+
+    def forward(self, logits, dis_preds):
+        # TODO punish repeat targets?
+        # how to define repeat target?
+
+        sum_logits = logits.sum(1)
+        loss = torch.matmul(sum_logits.cuda(), dis_preds)
+        weighted_sum = self.weights[0] * loss[0] + self.weights[1] * loss[1] + self.weights[2] * loss[2]
+        weighted_sum = - weighted_sum
+
+        return weighted_sum / logits.shape[0] / logits.shape[1]

parameters.py

+CUDA = True
+VOCAB_SIZE = 400  ## TODO what is the correct vocab size?
+MAX_SEQ_LEN = 500
+START_LETTER = 0
+BATCH_SIZE = 256 # 32
+SUPERVISED_TRAIN_EPOCHS = 3
+ADV_TRAIN_EPOCHS = 20
+SUPERVISED_SIZE = 704
+# NUM_BATCHES = 22
+VALIDATION_SIZE=64
+
+GEN_EMBEDDING_DIM = 32
+GEN_HIDDEN_DIM = 32
+DIS_EMBEDDING_DIM = 64
+DIS_HIDDEN_DIM = 32
+
+DATALOADER_NUM_WORKERS=4
+
+# oracle_samples_path = './oracle_samples.trc'
+supervised_sequences_path = './zorkdata_long_500.pt'  ## INPUT DATA
+supervised_val_cov_path = './zorkdata_long_500_target.pt'  ## LABEL DATA
+supervised_lengths_path = './zorkdata_lengths_500.pt'  ## INPUT LENGTHS
+supervised_state_dict_path = './oracle_ZORK20_60_v2.trc'  ## can ignore
+pretrained_gen_path = './gen_ZORK500_MLEtrain_v2.trc'  ## GENERATOR SAVE FILE
+pretrained_dis_path = './dis_ZORK500_pretrain_v2.trc'  ## DISCRIMINATOR SAVE FILE
+vocab_path = 'zork_vocab.txt'  ## VOCAB LIST
+output_path = 'generated_inputs/gen_input'  ## output generatred samples
\ No newline at end of file