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 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 =[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?