Mask Logits with a Grammar State
Easy · Solved in PyTorch · Machine learning coding practice
Problem
Apply a grammar mask to a batch of vocabulary logits, so that only the tokens the grammar permits from a request's current state can be sampled.
logits x has shape (B, V) and is finite float32 or float64. transitions T is int64 of shape (S, V), where T_{k,v} is the state reached by emitting token v from state k, or -1 if that token is forbidden. states s is int64 of shape (B,) and holds each request's current state. Return
Topics: inference, vllm, pytorch
The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Mask Logits with a Grammar State in the browser and run it against the tests.