def _expand_vocab_type_grammar(grammar_productions: List[str]) -> List[str]:
  """Expands the `Word` non-terminal to question|answer|document subtypes.

  Args:
    grammar_productions: Current grammar with the basic `Word` non-terminal.

  Returns:
    Rules where each `Word` non-terminal (_W_) is expanded into question (_Wq_),
    answer (_Wa_), document (_Wd_), and document title (_Wt_) subtypes.
  """
  productions = []
  for production in grammar_productions:
    if '_W_' in production:
      for term_type in get_term_types():
        productions.append(production.replace('_W_', term_type))
    else:
      productions.append(production)
  return productions
Expands the `Word` non-terminal to question|answer|document subtypes.

  Args:
    grammar_productions: Current grammar with the basic `Word` non-terminal.

  Returns:
    Rules where each `Word` non-terminal (_W_) is expanded into question (_Wq_),
    answer (_Wa_), document (_Wd_), and document title (_Wt_) subtypes.
  
language/search_agents/muzero/grammar_lib.py