text: import shutil
  from pathlib import Path
  
  import fairseq
  import numpy as np
  import sacrebleu
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import tqdm.auto as tqdm
  import wandb
  from config import config, task
  from fairseq.data import iterators
  from models.RNN import RNNDecoder, RNNEncoder, Seq2Seq
  from torch.cuda.amp import GradScaler, autocast
  
  from utils.data import load_data_iterator
  
  
  def build_model(args, task):
      """ build a model instance based on hyperparameters """
      src_dict, tgt_dict = task.source_dictionary, task.target_dictionary
  
      # token embeddings
      encoder_embed_tokens = nn.Embedding(len(src_dict), args.encoder_embed_dim,
                                          src_dict.pad())
      decoder_embed_tokens = nn.Embedding(len(tgt_dict), args.decoder_embed_dim,
                                          tgt_dict.pad())
  
      # encoder decoder
      # HINT: TODO: switch to TransformerEncoder & TransformerDecoder
      encoder = RNNEncoder(args, src_dict, encoder_embed_tokens)
      decoder = RNNDecoder(args, tgt_dict, decoder_embed_tokens)
  
      # sequence to sequence model
      model = Seq2Seq(args, encoder, decoder)
  
      # initialization for seq2seq model is important, requires extra handling
      def init_params(module):
          from fairseq.modules import MultiheadAttention
          if isinstance(module, nn.Linear):
              module.weight.data.normal_(mean=0.0, std=0.02)
              if module.bias is not None:
                  module.bias.data.zero_()
          if isinstance(module, nn.Embedding):
              module.weight.data.normal_(mean=0.0, std=0.02)
              if module.padding_idx is not None:
                  module.weight.data[module.padding_idx].zero_()
          if isinstance(module, MultiheadAttention):
              module.q_proj.weight.data.normal_(mean=0.0, std=0.02)
              module.k_proj.weight.data.normal_(mean=0.0, std=0.02)
              module.v_proj.weight.data.normal_(mean=0.0, std=0.02)
          if isinstance(module, nn.RNNBase):
              for name, param in module.named_parameters():
                  if "weight" in name or "bias" in name:
                      param.data.uniform_(-0.1, 0.1)
  
      # weight initialization
      model.apply(init_params)
      return model
  
  
  def train_one_epoch(logger,
                      epoch_itr,
                      model,
                      task,
                      criterion,
                      optimizer,
                      device,
                      accum_steps=1):
      itr = epoch_itr.next_epoch_itr(shuffle=True)
      itr = iterators.GroupedIterator(
          itr, accum_steps)  # 梯度累積: 每 accum_steps 個 sample 更新一次
  
      stats = {"loss": []}
      scaler = GradScaler()  # 混和精度訓練 automatic mixed precision (amp)
  
      model.train()
      progress = tqdm.tqdm(itr,
                           desc=f"train epoch {epoch_itr.epoch}",
                           leave=False)
      for samples in progress:
          model.zero_grad()
          accum_loss = 0
          sample_size = 0
          # 梯度累積: 每 accum_steps 個 sample 更新一次
          for i, sample in enumerate(samples):
              if i == 1:
                  # emptying the CUDA cache after the first step can reduce the chance of OOM
                  torch.cuda.empty_cache()
  
              sample = fairseq.utils.move_to_cuda(sample, device=device)
              target = sample["target"]
              sample_size_i = sample["ntokens"]
              sample_size += sample_size_i
  
              # 混和精度訓練
              with autocast():
                  net_output = model.forward(**sample["net_input"])
                  lprobs = F.log_softmax(net_output[0], -1)
                  loss = criterion(lprobs.view(-1, lprobs.size(-1)),
                                   target.view(-1))
  
                  # logging
                  accum_loss += loss.item()
                  # back-prop
                  scaler.scale(loss).backward()
  
          scaler.unscale_(optimizer)
          # (sample_size or 1.0) handles the case of a zero gradient
          optimizer.multiply_grads(1 / (sample_size or 1.0))
          gnorm = nn.utils.clip_grad_norm_(model.parameters(),
                                           config.clip_norm)  # 梯度裁剪 防止梯度爆炸
  
          scaler.step(optimizer)
          scaler.update()
  
          # logging
          loss_print = accum_loss / sample_size
          stats["loss"].append(loss_print)
          progress.set_postfix(loss=loss_print)
          if config.use_wandb:
              wandb.log({
                  "train/loss": loss_print,
                  "train/grad_norm": gnorm.item(),
                  "train/lr": optimizer.rate(),
                  "train/sample_size": sample_size,
              })
  
      loss_print = np.mean(stats["loss"])
      logger.info(f"training loss: {loss_print:.4f}")
      return stats
  
  
  def decode(toks, dictionary):
      # 從 Tensor 轉成人看得懂的句子
      s = dictionary.string(
          toks.int().cpu(),
          config.post_process,
      )
      return s if s else "<unk>"
  
  
  def inference_step(sample, model):
      sequence_generator = task.build_generator([model], config)
      gen_out = sequence_generator.generate([model], sample)
      srcs = []
      hyps = []
      refs = []
      for i in range(len(gen_out)):
          # 對於每個 sample, 收集輸入，輸出和參考答案，稍後計算 BLEU
          srcs.append(
              decode(
                  fairseq.utils.strip_pad(sample["net_input"]["src_tokens"][i],
                                          task.source_dictionary.pad()),
                  task.source_dictionary,
              ))
          hyps.append(
              decode(
                  gen_out[i][0]["tokens"],  # 0 代表取出 beam 內分數第一的輸出結果
                  task.target_dictionary,
              ))
          refs.append(
              decode(
                  fairseq.utils.strip_pad(sample["target"][i],
                                          task.target_dictionary.pad()),
                  task.target_dictionary,
              ))
      return srcs, hyps, refs
  
  
  def validate(logger, model, task, criterion, device, log_to_wandb=True):
      logger.info('begin validation')
      itr = load_data_iterator(task, "valid", 1, config.max_tokens,
                               config.num_workers).next_epoch_itr(shuffle=False)
  
      stats = {"loss": [], "bleu": 0, "srcs": [], "hyps": [], "refs": []}
      srcs = []
      hyps = []
      refs = []
  
      model.eval()
      progress = tqdm.tqdm(itr, desc=f"validation", leave=False)
      with torch.no_grad():
          for i, sample in enumerate(progress):
              # validation loss
              sample = fairseq.utils.move_to_cuda(sample, device=device)
              net_output = model.forward(**sample["net_input"])
  
              lprobs = F.log_softmax(net_output[0], -1)
              target = sample["target"]
              sample_size = sample["ntokens"]
              loss = criterion(lprobs.view(-1, lprobs.size(-1)),
                               target.view(-1)) / sample_size
              progress.set_postfix(valid_loss=loss.item())
              stats["loss"].append(loss)
  
              # 進行推論
              s, h, r = inference_step(sample, model)
              srcs.extend(s)
              hyps.extend(h)
              refs.extend(r)
  
      tok = 'zh' if task.cfg.target_lang == 'zh' else '13a'
      stats["loss"] = torch.stack(stats["loss"]).mean().item()
      stats["bleu"] = sacrebleu.corpus_bleu(hyps, [refs],
                                            tokenize=tok)  # 計算BLEU score
      stats["srcs"] = srcs
      stats["hyps"] = hyps
      stats["refs"] = refs
  
      if config.use_wandb and log_to_wandb:
          wandb.log(
              {
                  "valid/loss": stats["loss"],
                  "valid/bleu": stats["bleu"].score,
              },
              commit=False)
  
      showid = np.random.randint(len(hyps))
      logger.info("example source: " + srcs[showid])
      logger.info("example hypothesis: " + hyps[showid])
      logger.info("example reference: " + refs[showid])
  
      # show bleu results
      logger.info(f"validation loss:\t{stats['loss']:.4f}")
      logger.info(stats["bleu"].format())
      return stats
  
  
  def validate_and_save(logger,
                        model,
                        task,
                        criterion,
                        optimizer,
                        device,
                        epoch,
                        save=True):
      stats = validate(logger, model, task, criterion, device)
      bleu = stats['bleu']
      loss = stats['loss']
      if save:
          # save epoch checkpoints
          savedir = Path(config.savedir).absolute()
          savedir.mkdir(parents=True, exist_ok=True)
  
          check = {
              "model": model.state_dict(),
              "stats": {
                  "bleu": bleu.score,
                  "loss": loss
              },
              "optim": {
                  "step": optimizer._step
              }
          }
          torch.save(check, savedir / f"checkpoint{epoch}.pt")
          shutil.copy(savedir / f"checkpoint{epoch}.pt",
                      savedir / f"checkpoint_last.pt")
          logger.info(f"saved epoch checkpoint: {savedir}/checkpoint{epoch}.pt")
  
          # save epoch samples
          with open(
                  savedir /
                  f"samples{epoch}.{config.source_lang}-{config.target_lang}.txt",
                  "w") as f:
              for s, h in zip(stats["srcs"], stats["hyps"]):
                  f.write(f"{s}\t{h}\n")
  
          # get best valid bleu
          if getattr(validate_and_save, "best_bleu", 0) < bleu.score:
              validate_and_save.best_bleu = bleu.score
              torch.save(check, savedir / f"checkpoint_best.pt")
  
          del_file = savedir / f"checkpoint{epoch - config.keep_last_epochs}.pt"
          if del_file.exists():
              del_file.unlink()
      return stats
  
  
  def try_load_checkpoint(logger, model, optimizer=None, name=None):
      name = name if name else "checkpoint_last.pt"
      checkpath = Path(config.savedir) / name
      if checkpath.exists():
          check = torch.load(checkpath)
          model.load_state_dict(check["model"])
          stats = check["stats"]
          step = "unknown"
          if optimizer != None:
              optimizer._step = step = check["optim"]["step"]
          logger.info(
              f"loaded checkpoint {checkpath}: step={step} loss={stats['loss']} bleu={stats['bleu']}"
          )
      else:
          logger.info(f"no checkpoints found at {checkpath}!")
  