def state_tree_state_part(tree: state_tree.NQStateTree,
                          idf_lookup: IdfLookupFn) -> nqutils.ObsFragment:
  """Computes the state tree part of the BERT state."""
  leave_tokens = tree.root.leaves()

  return nqutils.ObsFragment(
      text=nqutils.Text(tokens=[
          *leave_tokens,
          '[SEP]',
      ]),
      type_values={
          'state_part': ['tree_leaves'] * len(leave_tokens) + ['[SEP]'],
      },
      float_values={
          'mr_score': [0.0] * (len(leave_tokens) + 1),
          'idf_score': [
              idf_lookup[token]
              for token in nqutils.bert_tokens_to_words(leave_tokens)
          ] + [0.]
      },
  )
Computes the state tree part of the BERT state.
language/search_agents/muzero/bert_state_lib.py