code: def predict_duration(tokens, ckpt_file):
      def fwd_(x):
          return DurationModel(is_training=False)(x)
  
      forward_fn = jax.jit(hk.transform_with_state(fwd_).apply)
      with open(ckpt_file, "rb") as f:
          dic = pickle.load(f)
      x = DurationInput(
          np.array(tokens, dtype=np.int32)[None, :],
          np.array([len(tokens)], dtype=np.int32),
          None,
      )
      return forward_fn(dic["params"], dic["aux"], dic["rng"], x)[0]
  
file_path: underthesea/underthesea/pipeline/say/viettts_/nat/text2mel.py
imports: import pickle
  from argparse import ArgumentParser
  from pathlib import Path
  import haiku
  import jax
  import jax.numpy
  import matplotlib.pyplot
  import numpy
  from .config import FLAGS, DurationInput
  from .data_loader import load_phonemes_set
  from .model import AcousticModel, DurationModel
label_desc: Use image node_id 2 for calling a global function with example usage: predict_duration(tokens, ckpt_file) and returns: forward_fn, DurationModel