code: def masked_token_mean(vectors, masks):
      """Mean over tokens.
  
      Args:
        vectors: <tf.float32>[batch_size, num_tokens, emb_dim]
        masks: <tf.int32>[batch_size, num_tokens]
  
      Returns:
        <tf.float32>[batch_size, emb_dim]
      """
      masks = tf.cast(masks, tf.float32)
      weights = masks / tf.reduce_sum(masks, axis=1, keepdims=True)
      return tf.reduce_sum(
          vectors * tf.expand_dims(weights, axis=-1), axis=1
      )
  
imports: import re
  import attr
  from lit_nlp.api import model
  from lit_nlp.api import types
  from lit_nlp.examples.models import model_utils
  from lit_nlp.lib import file_cache
  from lit_nlp.lib import utils
  import tensorflow
  import tensorflow_text
  import transformers
  from rouge_score import rouge_scorer
label_desc: Use image node_id 1 for calling a global function with example usage: masked_token_mean(vectors, masks) and returns: tf