Estimate token probability/logits given a sentence without computing the entire sentence

user_7872025 0
Viewed 1813

I have a sentence like: "I like sitting in my new chair and _____ about life".

And I have a SPECIFIC set of tokens like ["watch", "run", "think", "apple", "light"]

I would like to calculate the probability of each of those tokens to appear as the next word in that incomplete sentence. Hopefully I should get that the probability of "think" is higher that "apple" for instance.

I am working with pytorch-transformers (GPT2LMHeadModel specifically), and a possible solution is to evaluate the score of the full sentence with each of the tokens, but when number of tokens to evaluate is on the order of 100 or 1000 then the computation time starts to be too long.

It must be possible to process the sentence only once and somehow use the hidden states to calculate the probabilities of the set of tokens, but I don't know how to do it.

Any ideas? Thanks in advance

EDIT:

The actual code looks like the one below (estimating the probability for the full sentence every time). For every sentence it takes about 0.1 seconds to run the score() method, which turns into hours if I want to evaluate some thousands of words.

from pytorch_transformers import GPT2Tokenizer, GPT2LMHeadModel
import pandas as pd

model = GPT2LMHeadModel.from_pretrained("gpt2")
model.eval()
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")


def score(sentence):
    tokenize_input = tokenizer.tokenize(sentence)
    tensor_input = torch.tensor([tokenizer.convert_tokens_to_ids(tokenize_input)])
    loss = model(tensor_input, labels=tensor_input)
    return -loss[0].item()


candidates = ["watch", "run", "think", "apple", "light"]
sent_template = "I like sitting in my new chair and {} about life"
print({candidate: score(sent_template.format(candidate)) for candidate in candidates})
huggingface-transformers nlp python
1 Answers

Your example produced the following output and took around 48.5 seconds with 282 candiates to finish in my environment (I only conducted a 3 runs):

{'watch': -5.406847953796387
, 'run': -5.533411502838135
, 'think': -4.525279521942139
, 'apple': -6.158637046813965
, 'light': -5.835141658782959}

As mentioned in the comments I think you can spare some calculation with the past parameter and the fast tokenizer as shown in the commented example below:

import torch

from  transformers import GPT2TokenizerFast, GPT2LMHeadModel
from torch.nn import CrossEntropyLoss

model = GPT2LMHeadModel.from_pretrained("gpt2")
model.eval()
tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")

###We calculate the hidden_states and the past of the common left part of the sentence
past = "I like sitting in my new chair and"
past_tokenize_input = tokenizer.tokenize(past)
past_tensor_input = torch.tensor([tokenizer.convert_tokens_to_ids(past_tokenize_input)])

past_last_hidden_state, past = model.transformer(past_tensor_input)

def score(sentence, past, past_last_hidden_state, past_tensor_input):
    tokenize_input = tokenizer.tokenize(sentence, )
    tensor_input = torch.tensor([tokenizer.convert_tokens_to_ids(tokenize_input)])

    ###the following code is slightly modified from https://github.com/huggingface/transformers/blob/09a2f40684f77e62d0fd8485fe9d2d610390453f/src/transformers/modeling_gpt2.py#L604
    ###now we calculate the right part of the sentence with the already calculated past
    transformer_outputs = model.transformer(
            tensor_input,
            past=past,
            attention_mask=None,
            token_type_ids=None,
            position_ids=None,
            head_mask=None,
            inputs_embeds=None,
            use_cache=None,
            output_attentions=None,
            output_hidden_states=None,
        )
    ###and concatenate the output of with the hidden_state of the left part of the sentence
    hidden_states = torch.cat((past_last_hidden_state, transformer_outputs[0]), dim=1)
    
    ###the following part is exactly the same as https://github.com/huggingface/transformers/blob/09a2f40684f77e62d0fd8485fe9d2d610390453f/src/transformers/modeling_gpt2.py#L604
    lm_logits = model.lm_head(hidden_states)

    labels_input = torch.cat((past_tensor_input, tensor_input), dim=1)

    # Shift so that tokens < n predict n
    shift_logits = lm_logits[..., :-1, :].contiguous()
    shift_labels = labels_input[..., 1:].contiguous()
    # Flatten the tokens
    loss_fct = CrossEntropyLoss()
    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
    return -loss.item()

candidates = ["watch", "run", "think", "apple", "light"]

sent_template = " {} about life"

print({candidate: score(sent_template.format(candidate), past, past_last_hidden_state, past_tensor_input) for candidate in candidates})

Output:

{'watch': -5.406846046447754
, 'run': -5.533413887023926
, 'think': -4.525280952453613
, 'apple': -6.158637046813965
, 'light': -5.835141181945801}

The runtime here was 40.5 seconds with 282 candidates (3 cycles again). You also see that I lost some precision.

Many thanks to patrickvonplaten who gave me a good explanation about the past implementation.

user_66648720
Related