Getting vectors from a language model

I’m trying to adapt the language model created in IMBD example from lesson 10.

I’d like to create sentence vectors for a sentence/document similarity task.

It’s not clear to me how to take the pre-trained LM from the lesson and create sentence vectors from the hidden states.

Would it make sense to just take the PoolingLinearClassifier from lm_rnn.py and have it return the pooled output (without running it through the linear layers) ?

For example:

class PoolingLayer(nn.Module):

    def pool(self, x, bs, is_max):
        f = F.adaptive_max_pool1d if is_max else F.adaptive_avg_pool1d
        return f(x.permute(1,2,0), (1,)).view(bs,-1)

    def forward(self, input):
        outputs = input
        output = outputs[-1]
        sl,bs,_ = output.size()
        avgpool = self.pool(output, bs, False)
        mxpool = self.pool(output, bs, True)
        x = torch.cat([output[-1], mxpool, avgpool], 1)
        return x

Then I would create the new model as in get_rnn_classifer but like this:

def get_rnn_pooler(bptt, max_seq, n_class, n_tok, emb_sz, n_hid, n_layers, pad_token, layers, drops, bidir=False,
                  dropouth=0.3, dropouti=0.5, dropoute=0.1, wdrop=0.5, qrnn=False):
    rnn_enc = MultiBatchRNN(bptt, max_seq, n_tok, emb_sz, n_hid, n_layers, pad_token=pad_token, bidir=bidir,
                  dropouth=dropouth, dropouti=dropouti, dropoute=dropoute, wdrop=wdrop, qrnn=qrnn)
    return SequentialRNN(rnn_enc, PoolingLayer())

Am I on the right path?

Thanks
Brian