def bert_tokens_to_words(bert_tokens: List[str]) -> List[str]:
  """Maps word-pieces to the full word they are a part of.

  Given a word-piece tokenized sequence, returns a sequence of the same length
  which, at position i, contains the full word to which the word piece at
  position i belongs.
  For example:
    input:   wh   ##y  do   ##es  it  do  th   ##is
    output:  why  why  does does  it  do  this this

  Args:
    bert_tokens:  A sequence of word pieces.

  Returns:
    A sequence of equal length as `bert_tokens` with the full word to which each
    word-piece in `bert_tokens` belongs.
  """

  words = []
  i = 0
  while i < len(bert_tokens):
    full_word = bert_tokens[i]
    num_subwords = 1
    while i < (len(bert_tokens) - 1) and bert_tokens[i + 1].startswith('##'):
      full_word += bert_tokens[i + 1].replace('##', '')
      i += 1
      num_subwords += 1
    i += 1
    words += [full_word] * num_subwords
  return words
Maps word-pieces to the full word they are a part of.

  Given a word-piece tokenized sequence, returns a sequence of the same length
  which, at position i, contains the full word to which the word piece at
  position i belongs.
  For example:
    input:   wh   ##y  do   ##es  it  do  th   ##is
    output:  why  why  does does  it  do  this this

  Args:
    bert_tokens:  A sequence of word pieces.

  Returns:
    A sequence of equal length as `bert_tokens` with the full word to which each
    word-piece in `bert_tokens` belongs.
  
language/search_agents/muzero/utils.py