def compute_bert_state(query: str,
                       documents: Sequence[environment_pb2.Document],
                       idf_lookup: IdfLookupFn, context_size: int,
                       title_size: int, seq_length: int,
                       tokenize_fn: TokenizeFn,
                       tokens_to_ids_fn: TokensToIdsFn) -> Dict[str, Any]:
  """Computes the BERT state as in the agent."""

  obs_fragments = make_bert_state_impl(
      query=query,
      tree=state_tree.NQStateTree(nltk.CFG.fromstring("Q -> 'unused'")),
      documents=documents,
      idf_lookup=idf_lookup,
      tokenize_fn=tokenize_fn,
      context_size=context_size,
      max_length=seq_length,
      max_title_length=title_size)

  token_ids, type_ids, float_values = nqutils.ObsFragment.combine_and_expand(
      fragments=obs_fragments,
      length=seq_length,
      type_vocabs=TYPE_VOCABS,
      float_names=FLOAT_NAMES,
      tokens_to_id_fn=tokens_to_ids_fn,
  )
  token_ids = np.array(token_ids, np.int32)
  type_ids = np.array(type_ids, np.int32).T
  float_values = np.array(float_values, np.float32).T

  return {
      'obs_fragments': obs_fragments,
      'token_ids': token_ids,
      'type_ids': type_ids,
      'float_values': float_values,
  }
Computes the BERT state as in the agent.
language/search_agents/muzero/bert_state_lib.py