code: def build_model(
      hparams_set, model_name, data_dir, problem_name, beam_size=1
  ):
      """Build the graph required to fetch the attention weights.
  
      Args:
        hparams_set: HParams set to build the model with.
        model_name: Name of model.
        data_dir: Path to directory containing training data.
        problem_name: Name of problem.
        beam_size: (Optional) Number of beams to use when decoding a translation.
            If set to 1 (default) then greedy decoding is used.
  
      Returns:
        Tuple of (
            inputs: Input placeholder to feed in ids to be translated.
            targets: Targets placeholder to feed to translation when fetching
                attention weights.
            samples: Tensor representing the ids of the translation.
            att_mats: Tensors representing the attention weights.
        )
      """
      hparams = trainer_lib.create_hparams(
          hparams_set, data_dir=data_dir, problem_name=problem_name
      )
      translate_model = registry.model(model_name)(
          hparams, tf_estimator.ModeKeys.EVAL
      )
  
      inputs = tf.placeholder(
          tf.int32, shape=(1, None, 1, 1), name="inputs"
      )
      targets = tf.placeholder(
          tf.int32, shape=(1, None, 1, 1), name="targets"
      )
      translate_model(
          {
              "inputs": inputs,
              "targets": targets,
          }
      )
  
      # Must be called after building the training graph, so that the dict will
      # have been filled with the attention tensors. BUT before creating the
      # inference graph otherwise the dict will be filled with tensors from
      # inside a tf.while_loop from decoding and are marked unfetchable.
      att_mats = get_att_mats(translate_model)
  
      with tf.variable_scope(tf.get_variable_scope(), reuse=True):
          samples = translate_model.infer(
              {
                  "inputs": inputs,
              },
              beam_size=beam_size,
          )["outputs"]
  
      return inputs, targets, samples, att_mats
  
file_path: tensor2tensor/tensor2tensor/visualization/visualization.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  from tensor2tensor import models
  from tensor2tensor import problems
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import trainer_lib
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 1 for calling a global function with example usage: build_model(hparams_set, model_name, data_dir, problem_name, beam_size) and returns: inputs,  targets,  samples,  att_mats