<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Bringing contextual word representations into your models
Step1:  Contents
Overview
General set-up
Hugging Face BERT models and tokenizers
BERT featurization with Hugging Face
Simple feed-forward experiment
A feed-forward experiment with the sst module
An RNN experiment with the sst module
BERT fine-tuning with Hugging Face
HfBertClassifier
HfBertClassifier experiment
Overview
This notebook provides a basic introduction to using pre-trained BERT representations with the Hugging Face library. It is meant as a practical companion to our lecture on contextual word representations. The goal of this notebook is just to help you use these representations in your own work.
If you haven't already, I encourage you to review the notebook vsm_04_contextualreps.ipynb before working with this one. That notebook covers the fundamentals of these models; this one dives into the details more quickly.
A number of the experiments in this notebook are resource-intensive. I've included timing information for the expensive steps, to give you a sense for how long things are likely to take. I ran this notebook on a laptop with a single NVIDIA RTX 2080 GPU. 
General set-up
The following are requirements that you'll already have met if you've been working in this repository. As you can see, we'll use the Stanford Sentiment Treebank for illustrations, and we'll try out a few different deep learning models.
Step2:  The transformers library does a lot of logging. To avoid ending up with a cluttered notebook, I am changing the logging level. You might want to skip this as you scale up to building production systems, since the logging is very good – it gives you a lot of insights into what the models and code are doing.
Step3:  Hugging Face BERT models and tokenizers
We'll illustrate with the BERT-base cased model
Step4:  There are lots other options for pretrained weights. See this Hugging Face directory.
Next, we specify a tokenizer and a model that match both each other and our choice of pretrained weights
Step5:  For modeling (as opposed to creating static representations), we will mostly process examples in batches – generally very small ones, as these models consume a lot of memory. Here's a small batch of texts to use as the starting point for illustrations
Step6:  We will often need to pad (and perhaps truncate) token lists so that we can work with fixed-dimensional tensors
Step7:  The token_type_ids is used for multi-text inputs like NLI. The 'input_ids' field gives the indices for each of the two examples
Step8:  Notice that the final two tokens of the second example are pad tokens.
For fine-tuning, we want to avoid attending to padded tokens. The 'attention_mask' captures the needed mask, which we'll be able to feed directly to the pretrained BERT model
Step9:  Finally, we can run these indices and masks through the pretrained model
Step10:  Hugging Face BERT models create a special pooler_output representation that is the final representation above the [CLS] extended with a single layer of parameters
Step11:  We have two examples, each representented by a single vector of dimension 768, which is $d_{model}$ for BERT base using the notation from the original Transformers paper. This is an easy basis for fine-tuning, as we will see.
We can also access the final output for each state
Step12:  Here, we have 2 examples, each padded to the length of the longer one (12), and each of those representations has dimension 768. These representations can be used for sequence modeling, or pooled somehow for simple classifiers.
Those are all the essential ingredients for working with these parameters in Hugging Face. Of course, the library has a lot of other functionality, but the above suffices to featurize and to fine-tune.
BERT featurization with Hugging Face
To start, we'll use the Hugging Face interfaces just to featurize examples to create inputs to a separate model. In this setting, the BERT parameters are frozen.
Step13:  Simple feed-forward experiment
For a simple feed-forward experiment, we can get the representation of the [CLS] tokens and use them as the inputs to a shallow neural network
Step14:  Next we read in the SST train and dev splits
Step15:  Split the input/output pairs out into separate lists
Step16:  In the next step, we featurize all of the examples. These steps are likely to be the slowest in these experiments
Step17:  Now that all the examples are featurized, we can fit a model and evaluate it
Step18:  A feed-forward experiment with the sst module
It is straightforward to conduct experiments like the above using sst.experiment, which will enable you to do a wider range of experiments without writing or copy-pasting a lot of code.
Step19:  An RNN experiment with the sst module
We can also use BERT representations as the input to an RNN. There is just one key change from how we used these models before
Step20:  BERT fine-tuning with Hugging Face
The above experiments are quite successful – BERT gives us a reliable boost compared to other methods we've explored for the SST task. However, we might expect to do even better if we fine-tune the BERT parameters as part of fitting our SST classifier. To do that, we need to incorporate the Hugging Face BERT model into our classifier. This too is quite straightforward.
HfBertClassifier
The most important step is to create an nn.Module subclass that has, for its parameters, both the BERT model and parameters for our own classifier. Here we define a very simple fine-tuning set-up in which some layers built on top of the output corresponding to [CLS] are used as the basis for the SST classifier
Step21:  As you can see, self.bert does the heavy-lifting
Step22:  HfBertClassifier experiment
That's it! Let's see how we do on the SST binary, root-only problem. Because fine-tuning is expensive, we'll conduct a modest hyperparameter search and run the model for just one epoch per setting evaluation, as we did when assessing NLI models.
Step23:  And now on to the final test-set evaluation, using the best model from above
<ASSISTANT_TASK:>
Python Code:
__author__ = "Christopher Potts"
__version__ = "CS224u, Stanford, Spring 2022"
Explanation: Bringing contextual word representations into your models
End of explanation
import os
from sklearn.metrics import classification_report
import torch
import torch.nn as nn
import transformers
from transformers import BertModel, BertTokenizer
from torch_shallow_neural_classifier import TorchShallowNeuralClassifier
from torch_rnn_classifier import TorchRNNModel
from torch_rnn_classifier import TorchRNNClassifier
from torch_rnn_classifier import TorchRNNClassifierModel
from torch_rnn_classifier import TorchRNNClassifier
import sst
import utils
utils.fix_random_seeds()
SST_HOME = os.path.join("data", "sentiment")
Explanation: Contents
Overview
General set-up
Hugging Face BERT models and tokenizers
BERT featurization with Hugging Face
Simple feed-forward experiment
A feed-forward experiment with the sst module
An RNN experiment with the sst module
BERT fine-tuning with Hugging Face
HfBertClassifier
HfBertClassifier experiment
Overview
This notebook provides a basic introduction to using pre-trained BERT representations with the Hugging Face library. It is meant as a practical companion to our lecture on contextual word representations. The goal of this notebook is just to help you use these representations in your own work.
If you haven't already, I encourage you to review the notebook vsm_04_contextualreps.ipynb before working with this one. That notebook covers the fundamentals of these models; this one dives into the details more quickly.
A number of the experiments in this notebook are resource-intensive. I've included timing information for the expensive steps, to give you a sense for how long things are likely to take. I ran this notebook on a laptop with a single NVIDIA RTX 2080 GPU. 
General set-up
The following are requirements that you'll already have met if you've been working in this repository. As you can see, we'll use the Stanford Sentiment Treebank for illustrations, and we'll try out a few different deep learning models.
End of explanation
transformers.logging.set_verbosity_error()
Explanation: The transformers library does a lot of logging. To avoid ending up with a cluttered notebook, I am changing the logging level. You might want to skip this as you scale up to building production systems, since the logging is very good – it gives you a lot of insights into what the models and code are doing.
End of explanation
weights_name = 'bert-base-cased'
Explanation: Hugging Face BERT models and tokenizers
We'll illustrate with the BERT-base cased model:
End of explanation
bert_tokenizer = BertTokenizer.from_pretrained(weights_name)
bert_model = BertModel.from_pretrained(weights_name)
Explanation: There are lots other options for pretrained weights. See this Hugging Face directory.
Next, we specify a tokenizer and a model that match both each other and our choice of pretrained weights:
End of explanation
example_texts = [
    "Encode sentence 1. [SEP] And sentence 2!",
    "Bert knows Snuffleupagus"]
Explanation: For modeling (as opposed to creating static representations), we will mostly process examples in batches – generally very small ones, as these models consume a lot of memory. Here's a small batch of texts to use as the starting point for illustrations:
End of explanation
example_ids = bert_tokenizer.batch_encode_plus(
    example_texts,
    add_special_tokens=True,
    return_attention_mask=True,
    padding='longest')
example_ids.keys()
Explanation: We will often need to pad (and perhaps truncate) token lists so that we can work with fixed-dimensional tensors: The batch_encode_plus has a lot of options for doing this:
End of explanation
example_ids['input_ids']
Explanation: The token_type_ids is used for multi-text inputs like NLI. The 'input_ids' field gives the indices for each of the two examples:
End of explanation
example_ids['attention_mask']
Explanation: Notice that the final two tokens of the second example are pad tokens.
For fine-tuning, we want to avoid attending to padded tokens. The 'attention_mask' captures the needed mask, which we'll be able to feed directly to the pretrained BERT model:
End of explanation
X_example = torch.tensor(example_ids['input_ids'])
X_example_mask = torch.tensor(example_ids['attention_mask'])
with torch.no_grad():
    reps = bert_model(X_example, attention_mask=X_example_mask)
Explanation: Finally, we can run these indices and masks through the pretrained model:
End of explanation
reps.pooler_output.shape
Explanation: Hugging Face BERT models create a special pooler_output representation that is the final representation above the [CLS] extended with a single layer of parameters:
End of explanation
reps.last_hidden_state.shape
Explanation: We have two examples, each representented by a single vector of dimension 768, which is $d_{model}$ for BERT base using the notation from the original Transformers paper. This is an easy basis for fine-tuning, as we will see.
We can also access the final output for each state:
End of explanation
def bert_phi(text):
    input_ids = bert_tokenizer.encode(text, add_special_tokens=True)
    X = torch.tensor([input_ids])
    with torch.no_grad():
        reps = bert_model(X)
        return reps.last_hidden_state.squeeze(0).numpy()
Explanation: Here, we have 2 examples, each padded to the length of the longer one (12), and each of those representations has dimension 768. These representations can be used for sequence modeling, or pooled somehow for simple classifiers.
Those are all the essential ingredients for working with these parameters in Hugging Face. Of course, the library has a lot of other functionality, but the above suffices to featurize and to fine-tune.
BERT featurization with Hugging Face
To start, we'll use the Hugging Face interfaces just to featurize examples to create inputs to a separate model. In this setting, the BERT parameters are frozen.
End of explanation
def bert_classifier_phi(text):
    reps = bert_phi(text)
    #return reps.mean(axis=0)  # Another good, easy option.
    return reps[0]
Explanation: Simple feed-forward experiment
For a simple feed-forward experiment, we can get the representation of the [CLS] tokens and use them as the inputs to a shallow neural network:
End of explanation
train = sst.train_reader(SST_HOME)
dev = sst.dev_reader(SST_HOME)
Explanation: Next we read in the SST train and dev splits:
End of explanation
X_str_train = train.sentence.values
y_train = train.label.values
X_str_dev = dev.sentence.values
y_dev = dev.label.values
Explanation: Split the input/output pairs out into separate lists:
End of explanation
%time X_train = [bert_classifier_phi(text) for text in X_str_train]
%time X_dev = [bert_classifier_phi(text) for text in X_str_dev]
Explanation: In the next step, we featurize all of the examples. These steps are likely to be the slowest in these experiments:
End of explanation
model = TorchShallowNeuralClassifier(
    early_stopping=True,
    hidden_dim=300)
%time _ = model.fit(X_train, y_train)
preds = model.predict(X_dev)
print(classification_report(y_dev, preds, digits=3))
Explanation: Now that all the examples are featurized, we can fit a model and evaluate it:
End of explanation
def fit_shallow_network(X, y):
    mod = TorchShallowNeuralClassifier(
        hidden_dim=300,
        early_stopping=True)
    mod.fit(X, y)
    return mod
%%time
_ = sst.experiment(
    sst.train_reader(SST_HOME),
    bert_classifier_phi,
    fit_shallow_network,
    assess_dataframes=sst.dev_reader(SST_HOME),
    vectorize=False)  # Pass in the BERT reps directly!
Explanation: A feed-forward experiment with the sst module
It is straightforward to conduct experiments like the above using sst.experiment, which will enable you to do a wider range of experiments without writing or copy-pasting a lot of code.
End of explanation
def fit_rnn(X, y):
    mod = TorchRNNClassifier(
        vocab=[],
        early_stopping=True,
        use_embedding=False)  # Pass in the BERT hidden states directly!
    mod.fit(X, y)
    return mod
%%time
_ = sst.experiment(
    sst.train_reader(SST_HOME),
    bert_phi,
    fit_rnn,
    assess_dataframes=sst.dev_reader(SST_HOME),
    vectorize=False)  # Pass in the BERT hidden states directly!
Explanation: An RNN experiment with the sst module
We can also use BERT representations as the input to an RNN. There is just one key change from how we used these models before:
Previously, we would feed in lists of tokens, and they would be converted to indices into a fixed embedding space. This presumes that all words have the same representation no matter what their context is. 
With BERT, we skip the embedding entirely and just feed in lists of BERT vectors, which means that the same word can be represented in different ways.
TorchRNNClassifier supports this via use_embedding=False. In turn, you needn't supply a vocabulary:
End of explanation
class HfBertClassifierModel(nn.Module):
    def __init__(self, n_classes, weights_name='bert-base-cased'):
        super().__init__()
        self.n_classes = n_classes
        self.weights_name = weights_name
        self.bert = BertModel.from_pretrained(self.weights_name)
        self.bert.train()
        self.hidden_dim = self.bert.embeddings.word_embeddings.embedding_dim
        # The only new parameters -- the classifier:
        self.classifier_layer = nn.Linear(
            self.hidden_dim, self.n_classes)
    def forward(self, indices, mask):
        reps = self.bert(
            indices, attention_mask=mask)
        return self.classifier_layer(reps.pooler_output)
Explanation: BERT fine-tuning with Hugging Face
The above experiments are quite successful – BERT gives us a reliable boost compared to other methods we've explored for the SST task. However, we might expect to do even better if we fine-tune the BERT parameters as part of fitting our SST classifier. To do that, we need to incorporate the Hugging Face BERT model into our classifier. This too is quite straightforward.
HfBertClassifier
The most important step is to create an nn.Module subclass that has, for its parameters, both the BERT model and parameters for our own classifier. Here we define a very simple fine-tuning set-up in which some layers built on top of the output corresponding to [CLS] are used as the basis for the SST classifier:
End of explanation
class HfBertClassifier(TorchShallowNeuralClassifier):
    def __init__(self, weights_name, *args, **kwargs):
        self.weights_name = weights_name
        self.tokenizer = BertTokenizer.from_pretrained(self.weights_name)
        super().__init__(*args, **kwargs)
        self.params += ['weights_name']
    def build_graph(self):
        return HfBertClassifierModel(self.n_classes_, self.weights_name)
    def build_dataset(self, X, y=None):
        data = self.tokenizer.batch_encode_plus(
            X,
            max_length=None,
            add_special_tokens=True,
            padding='longest',
            return_attention_mask=True)
        indices = torch.tensor(data['input_ids'])
        mask = torch.tensor(data['attention_mask'])
        if y is None:
            dataset = torch.utils.data.TensorDataset(indices, mask)
        else:
            self.classes_ = sorted(set(y))
            self.n_classes_ = len(self.classes_)
            class2index = dict(zip(self.classes_, range(self.n_classes_)))
            y = [class2index[label] for label in y]
            y = torch.tensor(y)
            dataset = torch.utils.data.TensorDataset(indices, mask, y)
        return dataset
Explanation: As you can see, self.bert does the heavy-lifting: it reads in all the pretrained BERT parameters, and I've specified self.bert.train() just to make sure that these parameters can be updated during our training process. 
In forward, self.bert is used to process inputs, and then pooler_output is fed into self.classifier_layer. Hugging Face has already added a layer on top of the actual output for [CLS], so we can specify the model as
$$
\begin{align}
[h_{1}, \ldots, h_{n}] &= \text{BERT}([x_{1}, \ldots, x_{n}]) \
h &= \tanh(h_{1}W_{hh} + b_{h}) \
y &= \textbf{softmax}(hW_{hy} + b_{y})
\end{align}$$
for a tokenized input sequence $[x_{1}, \ldots, x_{n}]$. 
The Hugging Face documentation somewhat amusingly says, of pooler_output,
This output is usually not a good summary of the semantic content of the input, you're often better with averaging or pooling the sequence of hidden-states for the whole input sequence.
which is entirely reasonable, but it will require more resources, so we'll do the simpler thing here.
For the training and prediction interface, we can subclass TorchShallowNeuralClassifier so that we don't have to write any of our own data-handling, training, or prediction code. The central changes are using HfBertClassifierModel in build_graph and processing the data with batch_encode_plus.
End of explanation
def bert_fine_tune_phi(text):
    return text
def fit_hf_bert_classifier_with_hyperparameter_search(X, y):
    basemod = HfBertClassifier(
        weights_name='bert-base-cased',
        batch_size=8,  # Small batches to avoid memory overload.
        max_iter=1,  # We'll search based on 1 iteration for efficiency.
        n_iter_no_change=5,   # Early-stopping params are for the
        early_stopping=True)  # final evaluation.
    param_grid = {
        'gradient_accumulation_steps': [1, 4, 8],
        'eta': [0.00005, 0.0001, 0.001],
        'hidden_dim': [100, 200, 300]}
    bestmod = utils.fit_classifier_with_hyperparameter_search(
        X, y, basemod, cv=3, param_grid=param_grid)
    return bestmod
%%time
bert_classifier_xval = sst.experiment(
    sst.train_reader(SST_HOME),
    bert_fine_tune_phi,
    fit_hf_bert_classifier_with_hyperparameter_search,
    assess_dataframes=sst.dev_reader(SST_HOME),
    vectorize=False)  # Pass in the BERT hidden state directly!
Explanation: HfBertClassifier experiment
That's it! Let's see how we do on the SST binary, root-only problem. Because fine-tuning is expensive, we'll conduct a modest hyperparameter search and run the model for just one epoch per setting evaluation, as we did when assessing NLI models.
End of explanation
optimized_bert_classifier = bert_classifier_xval['model']
# Remove the rest of the experiment results to clear out some memory:
del bert_classifier_xval
def fit_optimized_hf_bert_classifier(X, y):
    optimized_bert_classifier.max_iter = 1000
    optimized_bert_classifier.fit(X, y)
    return optimized_bert_classifier
test_df = sst.sentiment_reader(
    os.path.join(SST_HOME, "sst3-test-labeled.csv"))
%%time
_ = sst.experiment(
    sst.train_reader(SST_HOME),
    bert_fine_tune_phi,
    fit_optimized_hf_bert_classifier,
    assess_dataframes=test_df,
    vectorize=False)  # Pass in the BERT hidden state directly!
Explanation: And now on to the final test-set evaluation, using the best model from above:
End of explanation
<END_TASK>