Finite-State Constrained Decoding

Problem

Perform one greedy decoding step constrained by a finite-state machine. transitions maps string state IDs to dictionaries mapping string token IDs to integer next states. Only transitions from the current state are legal.

Choose the legal token with greatest logit, breaking ties by smaller numeric token ID. Return {"token": chosen ID, "state": next state}. If no legal token exists, return token −1 and preserve the current state. Missing state entries act like empty transition maps.

def constrained_step(logits, state, transitions):

Constraints

• logits has 1..32 finite entries; all transition token IDs are valid vocabulary IDs.
• state and next states are nonnegative integers; JSON object keys are their decimal strings.
• Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.

Examples

Arguments: [
  [
    9,
    2,
    3
  ],
  0,
  {
    "0": {
      "1": 1,
      "2": 2
    }
  }
]
Expected result: {
  "token": 2,
  "state": 2
}

Arguments: [
  [
    1,
    2
  ],
  3,
  {}
]
Expected result: {
  "token": -1,
  "state": 3
}