MediumNLPnlp-312

Logprob of a Token Given Top-k

Problem

Estimate the log probability of a token when only the top-k log probabilities are known. topk_logprobs maps the k most likely tokens to their log probabilities, and vocab_size is the full vocabulary size.

When the token is in the top-k, return its stored log probability. Otherwise distribute the remaining probability mass evenly over the other tokens: log((1 - sum of top-k probabilities) / (vocab_size - k)). Return float('-inf') when no mass or no remaining tokens exist.

Hint

The unseen-token mass is shared uniformly among the vocabulary entries outside the top-k.

Test cases

case 1

in: ["c",{"a":-0.6931471805599453,"b":-2.302585092994046},10]

expected: -2.995732273553991

case 2

in: ["a",{"a":-0.6931471805599453,"b":-2.302585092994046},10]

expected: -0.6931471805599453

case 3

in: ["x",{},4]

expected: -1.3862943611198906

case 4

in: ["b",{"a":-0.22314355131420976},3]

expected: -2.302585092994046

case 5

in: ["z",{"a":-1.6094379124341003,"b":-1.6094379124341003},5]

expected: -1.6094379124341005

Starter codePython

import math
def logprob_given_topk(token, topk_logprobs, vocab_size):
    # Your code here
    pass

5 test cases · run your code against them in the editor