def expand(
      self, type_vocabs: collections.OrderedDict, float_names: List[str],
      tokens_to_id_fn: Callable[[List[str]], List[int]]
  ) -> Tuple[List[int], List[List[int]], List[List[float]]]:
    """Turns this fragment into encoder-compatible inputs.

    - Maps the `text` tokens to integer ids, also expanding contextual
    pre-terminals such as `Vap0` by way of `type_vocabs`.
    - Maps the `type_values` to integer ids and assembles the different kinds
      into a list.
    - Assembles the `float_values` into a list.

    Args:
      type_vocabs:  Segment types and their valid values.
      float_names:  Name of the float features to expand.
      tokens_to_id_fn:  Performs the mapping from text to integer ids.

    Returns:
      A list with three elements, [text_ids, type_ids, float_values].
      Note that `type_ids` and `float_values` themselves are lists of lists,
      whereas `text_ids` is simply a list of ints.
    """

    text_ids = self.text.expand(tokens_to_id_fn)

    all_type_values = []
    for type_name, type_vocab in type_vocabs.items():
      type_values = self.type_values[type_name]
      if len(type_values) != len(text_ids):
        raise ValueError(
            'Type values for {} ({}) must match length of {}'.format(
                type_name, type_values, str(self.text)))
      all_type_values.append([type_vocab.index(t) for t in type_values])

    all_float_values = []
    for fname in float_names:
      fvalues = self.float_values[fname]
      if len(text_ids) != len(fvalues):
        raise ValueError(
            'Float value for {} ({}) must match length of {}'.format(
                fname, fvalues, str(self.text)))
      all_float_values.append(fvalues)

    return (text_ids, all_type_values, all_float_values)
Turns this fragment into encoder-compatible inputs.

    - Maps the `text` tokens to integer ids, also expanding contextual
    pre-terminals such as `Vap0` by way of `type_vocabs`.
    - Maps the `type_values` to integer ids and assembles the different kinds
      into a list.
    - Assembles the `float_values` into a list.

    Args:
      type_vocabs:  Segment types and their valid values.
      float_names:  Name of the float features to expand.
      tokens_to_id_fn:  Performs the mapping from text to integer ids.

    Returns:
      A list with three elements, [text_ids, type_ids, float_values].
      Note that `type_ids` and `float_values` themselves are lists of lists,
      whereas `text_ids` is simply a list of ints.
    
language/search_agents/muzero/utils.py