How about going one step further and constrain transformer output with a context-free grammar? That way you can generate more conformant code such as Python or C.
Wouldn't even need to beam search if you restrict it to deterministic context free grammars, which would satisfy > 95% of these "generate some JSON schema" use-cases. For DCFGs you can just zero-out the probability for any token that is invalid in the context, no lookahead or search needed. Wouldn't work for truly context free things like most programming languages, though.
