code: def continuation(
      args, dataset, model, criterion, tgt_dict, rank, world_size
  ):
      is_discrete_duration = dataset.discrete_dur
      is_discrete_f0 = dataset.discrete_f0
  
      f0_decoder = None
      if args.dequantize_prosody:
          assert dataset.discrete_f0
          print("Reporting MAE F0 for a discrete model")
          f0_decoder = Naive_F0_Decoder(
              args.f0_discretization_bounds,
              dataset.config.f0_vq_n_units,
          ).cuda()
  
      dataset = InferenceDataset(
          dataset,
          args.prefix_length,
          filter_short=True,
          presort_by_length=True,
      )
      sampler = (
          None
          if world_size == 1
          else DistributedSampler(
              dataset, num_replicas=world_size, rank=rank, shuffle=False
          )
      )
      dataloader = DataLoader(
          dataset,
          batch_size=1,
          shuffle=False,
          collate_fn=dataset.collater,
          sampler=sampler,
      )
  
      Ts = args.T_token, args.T_duration, args.T_f0
      decoder = TemperatureDecoder(
          Ts,
          discrete_dur=is_discrete_duration,
          discrete_f0=is_discrete_f0,
      )
  
      running_stats = SimpleNamespace(
          token_bleu=0.0,
          duration_nll=0.0,
          duration_mae=0.0,
          f0_nll=0.0,
          f0_mae=0.0,
          n_tokens=0.0,
          n_sentences=0.0,
          f0_sum=0.0,
          f0_sum_sq=0.0,
          dur_sum=0.0,
          dur_sum_sq=0.0,
      )
  
      for i, batch in enumerate(dataloader):
          batch = explode_batch(batch, args.batch_explosion_rate)
          bsz = batch["target"].size(0)
  
          batch = move_to_cuda(batch)
          prefix = batch["prefix"][0]
  
          max_length_to_unroll = batch["target"].size(1)
          prefix_length = batch["net_input"]["src_tokens"].size(1)
          steps = max_length_to_unroll - prefix_length + 1
  
          assert steps > 0
          (
              produced_tokens,
              produced_durations,
              produced_f0,
              outputs,
          ) = do_sampling(
              model,
              batch,
              tgt_dict.eos(),
              decoder,
              autoregressive_steps=steps,
              teacher_force_tokens=args.teacher_force_tokens,
              teacher_force_duration=args.teacher_force_duration,
              teacher_force_f0=args.teacher_force_f0,
          )
  
          if args.teacher_force_tokens:
              assert (produced_tokens[:, 1:] == batch["target"]).all()
          if args.teacher_force_duration:
              assert (
                  produced_durations[:, 1:] == batch["dur_target"]
              ).all()
          if args.teacher_force_f0:
              assert (produced_f0[:, 1:] == batch["f0_target"]).all()
  
          dur_target = batch["dur_target"][:, prefix - 1 :].contiguous()
          f0_target = batch["f0_target"][:, prefix - 1 :].contiguous()
  
          f0_mask = batch["f0_mask"][:, prefix - 1 :].contiguous()
          dur_mask = batch["dur_mask"][:, prefix - 1 :].contiguous()
  
          duration_mae = mae_loss(
              produced_durations[:, prefix:].float(),
              dur_target.float(),
              dur_mask,
              reduce=False,
          )
          min_duration_mae = (
              duration_mae.view(bsz, -1).sum(dim=-1).min(dim=0)[0]
          )
          running_stats.duration_mae += min_duration_mae
  
          running_stats.dur_sum += (
              produced_durations[:, prefix:].float() * (~dur_mask)
          ).sum() / args.batch_explosion_rate
          running_stats.dur_sum_sq += (
              produced_durations[:, prefix:].float() * (~dur_mask)
          ).pow(2.0).sum() / args.batch_explosion_rate
  
          if is_discrete_duration:
              duration_loss = criterion.dur_loss_fn(
                  torch.stack([x[1] for x in outputs], dim=1),
                  dur_target,
                  dur_mask,
                  reduce=False,
              )
              min_duration_loss = (
                  duration_loss.view(bsz, -1).sum(dim=-1).min(dim=0)[0]
              )
              running_stats.duration_nll += min_duration_loss
  
          if f0_decoder:  # can only exist for discrete F0 models
              decoded_produced_f0 = f0_decoder(produced_f0[:, prefix:])
              decoded_f0_target = batch["raw_f0"][
                  :, prefix - 1 :
              ].contiguous()
  
              if produced_f0.ndim == 3:
                  decoded_produced_f0 = decoded_produced_f0.squeeze(2)
                  decoded_f0_target = decoded_f0_target.squeeze(2)
  
              f0_mae = mae_loss(
                  decoded_produced_f0,
                  decoded_f0_target,
                  f0_mask,
                  reduce=False,
              )
              f0_mae = f0_mae.view(bsz, -1).sum(dim=-1).min(dim=0)[0]
              running_stats.f0_mae += f0_mae
  
              f0_loss = criterion.f0_loss_fn(
                  torch.stack([x[2] for x in outputs], dim=1),
                  f0_target.long(),
                  f0_mask,
                  reduce=False,
              )
              f0_loss = f0_loss.view(bsz, -1).sum(dim=-1).min(dim=0)[0]
              running_stats.f0_nll += f0_loss
  
              running_stats.f0_sum += (
                  decoded_produced_f0 * (~f0_mask)
              ).sum() / args.batch_explosion_rate
              running_stats.f0_sum_sq += (
                  decoded_produced_f0 * (~f0_mask)
              ).pow(2.0).sum() / args.batch_explosion_rate
  
          else:
              assert not is_discrete_duration
  
              f0_loss = mae_loss(
                  produced_f0[:, prefix:],
                  f0_target,
                  f0_mask,
                  reduce=False,
              )
              f0_loss = f0_loss.view(bsz, -1).sum(dim=-1).min(dim=0)[0]
              running_stats.f0_mae += f0_loss
  
              running_stats.f0_sum += (
                  produced_f0[:, prefix:].sum()
                  / args.batch_explosion_rate
              )
              running_stats.f0_sum_sq += (
                  produced_f0[:, prefix:].pow(2.0).sum()
                  / args.batch_explosion_rate
              )
  
          running_stats.n_tokens += (~dur_mask)[0, ...].sum()
  
          token_loss = get_bleu(
              produced_tokens[:, prefix:],
              batch["target"][0, prefix - 1 :],
              tgt_dict,
          )
          running_stats.token_bleu += token_loss
          running_stats.n_sentences += 1
  
          if args.debug:
              break
  
      values = torch.tensor(
          [
              running_stats.token_bleu,
              running_stats.duration_nll,
              running_stats.duration_mae,
              running_stats.f0_nll,
              running_stats.f0_mae,
              running_stats.f0_sum,
              running_stats.f0_sum_sq,
              running_stats.dur_sum,
              running_stats.dur_sum_sq,
          ]
      )
      normalizers = torch.tensor(
          [running_stats.n_sentences] + [running_stats.n_tokens] * 8
      )
  
      return values, normalizers
  
imports: import os
  import numpy
  import scipy
  import torch
  import torch.multiprocessing
  from fairseq import checkpoint_utils, options
  from fairseq.data.codedataset import CodeDataset, ExpressiveCodeDataConfig
  from fairseq.dataclass.utils import convert_namespace_to_omegaconf
  from torch.utils.data import DataLoader, DistributedSampler
  from fairseq.utils import move_to_cuda
  from fairseq import utils
  from fairseq.criterions.speech_ulm_criterion import nll_loss, mae_loss
  import time
  from types import SimpleNamespace
  import sys, pathlib
  from naive_decoder import Naive_F0_Decoder
  from inference_dataset import InferenceDataset, explode_batch
  from sample.sample import do_sampling, TemperatureDecoder, FilterNamesDataset
label_desc: Use image node_id 3 for calling a global function with example usage: continuation(args, dataset, model, criterion, tgt_dict, rank, world_size) and returns: values,  normalizers