Nucleus Sampling Candidates

Problem

Given a normalized vocabulary distribution, sort token IDs by descending probability, breaking ties by smaller ID. Retain the shortest nonempty prefix whose cumulative probability is at least p, then renormalize that prefix. Return {"ids": ranked retained IDs, "probabilities": their normalized probabilities}.

At p=0 retain one token. Compare cumulative mass to p using a 1e-12 absolute boundary allowance to avoid including an extra token because of decimal summation error. Zero-probability tails are excluded once the threshold is reached.

def nucleus_candidates(probabilities, p):

Constraints

• 1..32 finite nonnegative probabilities summing to one within 1e-12.
• 0 ≤ p ≤ 1; output IDs remain in probability-ranked order.
• Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.

Examples

Arguments: [
  [
    0.5,
    0.3,
    0.2
  ],
  0.6
]
Expected result: {
  "ids": [
    0,
    1
  ],
  "probabilities": [
    0.625,
    0.375
  ]
}

Arguments: [
  [
    0.5,
    0.3,
    0.2
  ],
  0.5
]
Expected result: {
  "ids": [
    0
  ],
  "probabilities": [
    1
  ]
}