def original_query_state_part(query: str, tokenize_fn: TokenizeFn,
                              idf_lookup: IdfLookupFn) -> nqutils.ObsFragment:
  """Computes the original query part of the BERT state."""
  original_tokens = tokenize_fn(query)

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