code: def construct_prefix_tree(targets, target_mask=None):
      """
      Takes a list of valid target action sequences and creates a mapping from all possible
      (valid) action prefixes to allowed actions given that prefix.  While the method is called
      ``construct_prefix_tree``, we're actually returning a map that has as keys the paths to
      `all internal nodes of the trie`, and as values all of the outgoing edges from that node.
  
      ``targets`` is assumed to be a tensor of shape ``(batch_size, num_valid_sequences,
      sequence_length)``.  If the mask is not ``None``, it is assumed to have the same shape, and
      we will ignore any value in ``targets`` that has a value of ``0`` in the corresponding
      position in the mask.  We assume that the mask has the format 1*0* for each item in
      ``targets`` - that is, once we see our first zero, we stop processing that target.
  
      For example, if ``targets`` is the following tensor: ``[[1, 2, 3], [1, 4, 5]]``, the return
      value will be: ``{(): set([1]), (1,): set([2, 4]), (1, 2): set([3]), (1, 4): set([5])}``.
  
      This could be used, e.g., to do an efficient constrained beam search, or to efficiently
      evaluate the probability of all of the target sequences.
      """
      batched_allowed_transitions = []
  
      if not isinstance(targets, list):
          assert (
              targets.dim() == 3
          ), "targets tensor needs to be batched!"
          targets = targets.detach().cpu().numpy().tolist()
      if target_mask is not None:
          target_mask = target_mask.detach().cpu().numpy().tolist()
      else:
          target_mask = [None for _ in targets]
  
      for instance_targets, instance_mask in izip(targets, target_mask):
          allowed_transitions = defaultdict(set)
          for i, target_sequence in enumerate(instance_targets):
              history = ()
              for j, action in enumerate(target_sequence):
                  if instance_mask and instance_mask[i][j] == 0:
                      break
                  allowed_transitions[history].add(action)
                  history = history + (action,)
          batched_allowed_transitions.append(allowed_transitions)
      return batched_allowed_transitions
  
imports: from __future__ import absolute_import
  from collections import defaultdict
  import torch
label_desc: Use image node_id 1 for calling a global function with example usage: construct_prefix_tree(targets, target_mask) and returns: batched_allowed_transitions