When using the transformers library for text generation with the prefix_allowed_tokens_fn parameter to constrain model outputs, a specific error may appear intermittently:

ValueError: prefix_allowed_tokens_fn returned an empty list for batch ID 0. This means that the constraint is unsatisfiable. Please check your implementation of prefix_allowed_tokens_fn

This error indicates that at some point during generation, the constraint function returned no valid tokens, making it impossible to continue generating. The root cause typically stems from version incompatibility rather than implementation bugs.

Implementation Example

The following demonstrates a working implementation using a Trie structure to constrain generation to a predefined set of allowed token sequences:

from typing import Dict, List, Optional

class TokenTrie:
    """Trie data structure for storing and retrieving allowed token sequences."""
    
    def __init__(self, sequences: Optional[List[List[int]]] = None):
        self.tree: Dict = {}
        self.count = 0
        if sequences:
            for seq in sequences:
                self._insert_sequence(seq)
                self.count += 1
        
        self.extension: Optional[TokenTrie] = None
        self.start_token: Optional[int] = None

    def _insert_sequence(self, sequence: List[int]) -> None:
        """Recursively insert a sequence into the trie."""
        if sequence:
            if sequence[0] not in self.tree:
                self.tree[sequence[0]] = {}
            self._insert_sequence(sequence[1:], self.tree[sequence[0]])

    def _insert_sequence(self, sequence: List[int], node: Dict) -> None:
        if not sequence:
            return
        if sequence[0] not in node:
            node[sequence[0]] = {}
        self._insert_sequence(sequence[1:], node[sequence[0]])

    def append(self, trie: 'TokenTrie', bos_token: int) -> None:
        """Append another trie for continued valid paths."""
        self.extension = trie
        self.start_token = bos_token

    def add(self, sequence: List[int]) -> None:
        """Add a single sequence to the trie."""
        self._insert_sequence(sequence)
        self.count += 1

    def query(self, prefix: List[int]) -> List[int]:
        """Get all valid next tokens for the given prefix."""
        return self._find_in_tree(prefix, self.tree)

    def _find_in_tree(
        self,
        prefix: List[int],
        node: Dict,
        append_trie: Optional['TokenTrie'] = None,
        bos_token: Optional[int] = None
    ) -> List[int]:
        """Recursively find valid tokens in the trie."""
        if not prefix:
            valid_tokens = list(node.keys())
            if append_trie and bos_token in valid_tokens:
                valid_tokens.remove(bos_token)
                valid_tokens.extend(list(append_trie.tree.keys()))
            return valid_tokens
        
        if prefix[0] in node:
            return self._find_in_tree(
                prefix[1:],
                node[prefix[0]],
                append_trie,
                bos_token
            )
        
        if append_trie:
            return append_trie.query(prefix)
        return []

    def __iter__(self):
        """Iterate over all sequences in the trie."""
        def traverse(current_prefix: List[int], current_node: Dict):
            if current_node:
                for token in current_node:
                    yield from traverse(
                        current_prefix + [token],
                        current_node[token]
                    )
            else:
                yield current_prefix
        
        return traverse([], self.tree)

    def __len__(self) -> int:
        return self.count

    def __getitem__(self, key: List[int]) -> List[int]:
        return self.query(key)


def create_token_constraint(trie: TokenTrie):
    """Factory function to create the constraint callback for generation."""
    def constraint_validator(batch_id: int, token_sequence):
        tokens = token_sequence.tolist()
        valid_next = trie.query(tokens)
        return valid_next
    
    return constraint_validator

Usage with Model Generation

allowed_trie = TokenTrie(
    [ [0] + tokenizer.encode(candidate) for candidate in allowed_phrases ]
)

constraint_fn = create_token_constraint(allowed_trie)

outputs = model.generate(
    input_ids=input_ids,
    attention_mask=attention_mask,
    max_length=30,
    prefix_allowed_tokens_fn=constraint_fn,
    num_beams=beam_count,
    num_return_sequences=beam_count,
    output_scores=True,
    return_dict_in_generate=True,
)

Version Compatibility Issue

After extensive debugging, the actual issue was identified as a transformers library version incompatibility. Testing confirmed that:

• transformers==4.26.0 runs without errors
• transformers==4.40.0 and above triggers the ValueError

The behavior change appears to be related to how newer versions handle empty token lists returned by the constraint function during beam search expansion. The breaking change in version 4.40 results in stricter validation that causes generation to fail when any valid token path temporarily becomes unavailable during the search process.

To resolve this issue, either downgrade to version 4.26.0 or adjust the constraint logic to ensure valid tokens are always available, including handling intermediate states during beam search.