generate_beam_viz.py 文件源码

def create_graph(predicted_ids, parent_ids, scores, vocab=None):
  def get_node_name(pred):
    return vocab[pred] if vocab else str(pred)

  seq_length = predicted_ids.shape[0]
  graph = nx.DiGraph()
  for level in range(seq_length):
    names = [get_node_name(pred) for pred in predicted_ids[level]]
    _add_graph_level(graph, level + 1, parent_ids[level], names, scores[level])
  graph.node[(0, 0)]["name"] = "START"
  return graph