<jupyter_start><jupyter_text>3. An LSTM language model(consult *Recurrent Neural Network Regularization*, https://arxiv.org/pdf/1409.2329.pdf)<jupyter_code>!pip install -q torchtext 

import torch
import torchtext
from torchtext.vocab import Vectors
from torchtext.data.iterator import BPTTIterator
from torchtext.data import Batch, Dataset
import math


import numpy as np

from tqdm import tqdm_notebook
from collections import Counter, defaultdict
import matplotlib.pyplot as plt

import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torch.optim as optim

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Our input $x$
TEXT = torchtext.data.Field()

!curl -qO https://raw.githubusercontent.com/harvard-ml-courses/cs287-s18/master/HW2/input.txt
!curl -qO https://raw.githubusercontent.com/harvard-ml-courses/cs287-s18/master/HW2/train.5k.txt
!curl -qO https://raw.githubusercontent.com/harvard-ml-courses/cs287-s18/master/HW2/train.txt
!curl -qO https://raw.githubusercontent.com/harvard-ml-courses/cs287-s18/master/HW2/valid.txt
!curl -qO https://raw.githubusercontent.com/wojzaremba/lstm/master/data/ptb.test.txt

# Data distributed with the assignment
train, val, test = torchtext.datasets.LanguageModelingDataset.splits(
    path=".", 
    train="train.txt", validation="valid.txt", test="ptb.test.txt", text_field=TEXT)

train_iter, val_iter, test_iter = BPTTIterator.splits(
    (train, val, test), batch_size=10, device=torch.device("cuda"), bptt_len=32, repeat=False)
TEXT.build_vocab(train) #for whole vocab
# TEXT.build_vocab(train, max_size=1000) #for sampling
class LSTM(nn.Module):
  def __init__(self, embed_size, V, batch_size, hidden_size, dropout, weight_bound):
    super(LSTM, self).__init__()
    self.batch_size = batch_size
    self.hidden_size = hidden_size

    self.embeddings = nn.Embedding(V, embed_size)
    self.lstm = nn.LSTM(input_size=embed_size, hidden_size=hidden_size, dropout=dropout)
    self.hidden = nn.Linear(in_features=hidden_size, out_features=V)
    self.dropout = nn.Dropout(p=dropout)

    # Init all parameters to a uniform distribution between the given weight bounds
    for params in self.lstm._all_weights:
      for param in params:
        if 'weight' in param:
          nn.init.uniform_(self.lstm.__getattr__(param), -weight_bound, weight_bound)
  
  def create_hidden_state(self):
    if torch.cuda.is_available():
      layer = Variable(torch.zeros(self.lstm.num_layers, self.batch_size, self.hidden_size)).cuda()
    else:
      layer = Variable(torch.zeros(self.lstm.num_layers, self.batch_size, self.hidden_size))
    return (layer, layer)
  
  def detach(self, hidden_state):
    return tuple([x.detach() for x in hidden_state])

  def forward(self, x, last_hidden, train=True):
    # Takes in hidden state, produces next hidden state
    embeds = self.embeddings(x)
    if train:
      embeds = self.dropout(embeds)
    out, next_hidden = self.lstm(embeds, last_hidden)
    out = out.view(-1, out.size(2))
    if train:
      out = self.dropout(out)
    out = self.hidden(out)
    logits = F.log_softmax(out, dim=1)
    return logits, self.detach(next_hidden)

class LSTMPretrained(nn.Module):
  def __init__(self, embed_size, V, batch_size, hidden_size, dropout, weight_bound, TEXT):
    super(LSTMPretrained, self).__init__()
    self.batch_size = batch_size
    self.hidden_size = hidden_size

    V, embed_size = TEXT.vocab.vectors.shape
    self.embeddings = nn.Embedding(V, embed_size)
    self.embeddings.weight.data.copy_(TEXT.vocab.vectors)

    self.lstm = nn.LSTM(input_size=embed_size, hidden_size=hidden_size, dropout=dropout)
    self.hidden = nn.Linear(in_features=hidden_size, out_features=V)
    self.dropout = nn.Dropout(p=dropout)

    # Init all parameters to a uniform distribution between the given weight bounds
    for params in self.lstm._all_weights:
      for param in params:
        if 'weight' in param:
          nn.init.uniform_(self.lstm.__getattr__(param), -weight_bound, weight_bound)
  
  def create_hidden_state(self):
    if torch.cuda.is_available():
      layer = Variable(torch.zeros(self.lstm.num_layers, self.batch_size, self.hidden_size)).cuda()
    else:
      layer = Variable(torch.zeros(self.lstm.num_layers, self.batch_size, self.hidden_size))
    return (layer, layer)
  
  def detach(self, hidden_state):
    return tuple([x.detach() for x in hidden_state])

  def forward(self, x, last_hidden, train=True):
    # Takes in hidden state, produces next hidden state
    embeds = self.embeddings(x)
    if train:
      embeds = self.dropout(embeds)
    out, next_hidden = self.lstm(embeds, last_hidden)
    out = out.view(-1, out.size(2))
    if train:
      out = self.dropout(out)
    out = self.hidden(out)
    logits = F.log_softmax(out, dim=1)
    return logits, self.detach(next_hidden)

class LSTMTrainer:
  def __init__(self, train_iter, val_iter):
    self.train_iter = train_iter
    self.val_iter = val_iter
  
  def train(self, model, epochs, lr):
    loss_function = nn.NLLLoss(reduction="mean") #use negative log likelihood as the loss function, get ppl for free
    params = filter(lambda x: x.requires_grad, model.parameters())
    optimizer = optim.Adam(params=params, lr=lr)

    train_ppls = list()
    val_ppls = list()
    for epoch in tqdm_notebook(range(epochs)):
      hidden = model.create_hidden_state()
      epoch_losses = list()
      model.train()
      for batch in tqdm_notebook(self.train_iter):
        if torch.cuda.is_available():
          x, y = batch.text.cuda(), batch.target.view(-1).cuda()
        else:
          x, y = batch.text, batch.target.view(-1)
        optimizer.zero_grad()
        try:
          yhat, hidden = model.forward(x, hidden, train=True)
        except:
          print(x)
          print(hidden)
          raise
        loss = loss_function(yhat, y)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.lstm.parameters(), max_norm=5) #gradient clipping per Zaremba
        optimizer.step()
        epoch_losses.append(loss.item())

      # Get val acc at every epoch
      model.eval()
      train_ppl = np.exp(np.mean(epoch_losses))
      val_ppl = self.validate(model)
      print('Epoch={}\ttrain_ppl={}\tval_ppl={}\tloss={}'.format(epoch, train_ppl, val_ppl, np.mean(epoch_losses)))
      train_ppls.append(train_ppl)
      val_ppls.append(val_ppl)
    
    return train_ppls, val_ppls
  
  def validate(self, model):
    loss_function = nn.NLLLoss()
    hidden = model.create_hidden_state()
    losses = list()
    for batch in self.val_iter:
      x, y = batch.text, batch.target.view(-1)
      yhat, _ = model.forward(x, hidden, train=False)
      loss = loss_function(yhat, y)
      losses.append(loss.item())
    pp_val = np.exp(np.mean(losses))
    return pp_val
V = len(TEXT.vocab)
embed_size = 300
hidden_size = 650
dropout = 0.5
batch_size = 10
weight_bound = 0.05
print("vocab of len={}".format(V))
epochs=5
lr=1e-3
lstm = LSTM(embed_size=embed_size, V=V, batch_size=batch_size, hidden_size=hidden_size, dropout=dropout, weight_bound=weight_bound)
lstm.to(device)
trainer = LSTMTrainer(train_iter, val_iter)
train_ppls, val_ppls = trainer.train(lstm, epochs, lr)
plt.plot(train_ppls, label="PP(train)")
plt.plot(val_ppls, label="PP(val)")
plt.legend()
plt.xlabel('Epochs')
plt.ylabel('Perplexity')
plt.title('Perplexity over Epochs for Naive LSTM')
plt.show()
def get_test_pp(model, test_iter, trainer):
  with torch.no_grad():
    loss_function = nn.NLLLoss()
    losses = list()
    hidden = model.create_hidden_state()
    for batch in tqdm_notebook(test_iter, position=0, leave=True):
      x, y = batch.text, batch.target.view(-1)
      yhat, _ = model.forward(x, hidden, train=False)
      loss = loss_function(yhat, y)
      losses.append(loss.item())
    pp_test = np.exp(np.mean(losses))
    return pp_test
print('test pp: {}'.format(get_test_pp(lstm, test_iter, trainer)))
url = 'https://dl.fbaipublicfiles.com/fasttext/vectors-wiki/wiki.simple.vec'
TEXT.vocab.load_vectors(vectors=Vectors('wiki.simple.vec', url=url))
# TEXT.vocab.load_vectors(vectors=Vectors('wiki.simple.vec', url=url))
pretrained_lstm = LSTMPretrained(embed_size=embed_size, V=V, batch_size=batch_size, hidden_size=hidden_size, dropout=dropout, TEXT=TEXT, weight_bound=weight_bound)
pretrained_lstm.to(device)
pretrained_trainer = LSTMTrainer(train_iter, val_iter)
pretrained_train_ppls, pretrained_val_ppls = pretrained_trainer.train(pretrained_lstm, epochs, lr)
plt.plot(pretrained_train_ppls, label="PP(train)")
plt.plot(pretrained_val_ppls, label="PP(val)")
plt.legend()
plt.xlabel('Epochs')
plt.ylabel('Perplexity')
plt.title('Perplexity over Epochs for Pretrained LSTM')
plt.show()
print('pretrained test pp: {}'.format(get_test_pp(pretrained_lstm, test_iter, pretrained_trainer)))
def get_H(x, y, model, trainer):
  with torch.no_grad():
    loss_function = nn.NLLLoss()
    hidden = model.create_hidden_state()
    yhat, _ = model.forward(x, hidden, train=False)
    H = loss_function(yhat, y)
    return H
examples = list()
for i, batch in enumerate(train_iter):
  if i in [1, 7]:
    examples.append((batch.text.data, batch.target.data))
for i, batch in enumerate(test_iter):
  if i in [117, 231]:
    examples.append((batch.text.data, batch.target.data))
entropies = list()
for context in examples:
  H = get_H(context[0], context[1].view(-1), pretrained_lstm, pretrained_trainer)
  entropies.append(H.item())
entropies<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12496058854440358, 'fraction_numerical': 0.008512874408828166, 'mean_word_length': 4.1410048622366284, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 7, '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}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '272790', 'license': 'MIT', 'n_tokens_mistral': 3480, 'n_tokens_neox': 3374, 'n_words': 690, 'path': 'hw2/submission/HW2_LSTM.ipynb', 'repo_name': 'emtseng/cs6741', 'size': '9515'}