def next_tokens(self) -> Set[int]:
        """Returns the list of tokens that could come next.
        These are (a) all tokens extending the root state and, for
        non-root states, additionally all tokens extending the current
        state."""

        if self.node != self.root:
            return self.root.next_tokens().union(self.node.next_tokens())
        else:
            return self.root.next_tokens()
Returns the list of tokens that could come next.
        These are (a) all tokens extending the root state and, for
        non-root states, additionally all tokens extending the current
        state.
understand_icl/fairseq/fairseq/token_generation_constraints.py