def history_state_part(documents: Sequence[environment_pb2.Document],
                       tokenize_fn: TokenizeFn, idf_lookup: IdfLookupFn,
                       context_size: int, max_length: int,
                       max_title_length: int) -> List[nqutils.ObsFragment]:
  """Computes the history part of the BERT state."""
  history_fragments = []

  for doc in documents:
    answer_tokens = tokenize_fn(doc.answer.answer if doc.answer else '')
    answer_tokens_words = nqutils.bert_tokens_to_words(answer_tokens)
    _, context_window = types.HistoryEntry.get_window_around_substr(
        doc.content, doc.answer.answer, context_size)
    context_tokens = tokenize_fn(context_window)
    context_tokens_words = nqutils.bert_tokens_to_words(context_tokens)
    title_tokens = tokenize_fn(doc.title if doc.title else '')
    title_tokens = title_tokens[:max_title_length]
    title_tokens_words = nqutils.bert_tokens_to_words(title_tokens)
    state_part = [
        p for tokens, part in
        zip([answer_tokens, context_tokens, title_tokens],
            ['history_answer', 'history_context', 'history_title'])
        for p in [part] * len(tokens) + ['[SEP]']
    ]

    length = len(context_tokens) + len(answer_tokens) + len(title_tokens) + 3
    history_fragments.append(
        nqutils.ObsFragment(
            text=nqutils.Text(tokens=[
                *answer_tokens,
                '[SEP]',
                *context_tokens,
                '[SEP]',
                *title_tokens,
                '[SEP]',
            ]),
            type_values={
                'state_part': state_part,
            },
            float_values={
                'mr_score': [doc.answer.mr_score] * length,
                'idf_score':
                    [idf_lookup[word] for word in answer_tokens_words] + [0.] +
                    [idf_lookup[word] for word in context_tokens_words] + [0.] +
                    [idf_lookup[word] for word in title_tokens_words] + [0.]
            },
        ))
    # stop early if we already have enough history to fill up the whole state
    if sum([len(fragment.text.tokens) for fragment in history_fragments
           ]) >= max_length:
      logging.info('BERT state reached max_length: %d', max_length)
      break

  return history_fragments
Computes the history part of the BERT state.
language/search_agents/muzero/bert_state_lib.py