text: from PIL import Image
  
  from src.models.mnist_model import LitModelMNIST
  from src.models.nmt_model import NMTLitModel
  from src.transforms import mnist_transforms
  import torch
  import pdb
  import numpy as np
  from tokenizers import Tokenizer
  
  def predict():
      """Example of inference with trained model.
      It loads trained image classification model from checkpoint.
      Then it loads example image and predicts its label.
      """
  
      # ckpt can be also a URL!
      CKPT_PATH = "last.ckpt"
  
      # load model from checkpoint
      # model __init__ parameters will be loaded from ckpt automatically
      # you can also pass some parameter explicitly to override it
      trained_model = LitModelMNIST.load_from_checkpoint(checkpoint_path=CKPT_PATH)
  
      # print model hyperparameters
      print(trained_model.hparams)
  
      # switch to evaluation mode
      trained_model.eval()
      trained_model.freeze()
  
      # load data
      img = Image.open("data/example_img.png").convert("L")  # convert to black and white
      # img = Image.open("data/example_img.png").convert("RGB")  # convert to RGB
  
      # preprocess
      img = mnist_transforms.mnist_test_transforms(img)
      img = img.reshape((1, *img.size()))  # reshape to form batch of size 1
  
      # inference
      output = trained_model(img)
      print(output)
  
  
  CKPT_PATH = "../../logs/runs/2021-04-01/17-37-02/checkpoints/last.ckpt"
  litmodel = NMTLitModel.load_from_checkpoint(checkpoint_path=CKPT_PATH)
  
  
  def translate(sentences, lightning_module):
      src_ids = [i.ids for i in lightning_module.src_tokenizer.encode_batch(sentences)]
      src_ids = torch.LongTensor(src_ids).to(lightning_module.device) # (B, seq_len)
      src_mask = lightning_module.model.make_src_mask(src_ids)
  
      encoder_output = lightning_module.model.encode(src_ids, src_mask)
      tgt_init_str = [""] * len(sentences)
      # tgt_tokenizer will return seq_len + 1 tokens, so we strip one here
      tgt_ids = [i.ids[:-1] for i in lightning_module.tgt_tokenizer.encode_batch(tgt_init_str)]
      tgt_ids = np.array(tgt_ids)
      tgt_eos_idx = lightning_module.tgt_tokenizer.get_vocab()["[EOS]"]
  
      tgt_mask = lightning_module.model.casual_mask.unsqueeze(0).to(lightning_module.device) # (1, seq_len, seq_len)
      completed = np.array([False for _ in sentences])
  
      for step_idx in range(lightning_module.hparams.seq_len - 1):
          tgt_ids_inp = torch.LongTensor(tgt_ids).to(lightning_module.device)
          logits, attentions = lightning_module.model.decode(encoder_output, src_mask, tgt_ids_inp, tgt_mask)
          # logits: B, seq_len, vocab_size
          current_decode_step_logits = logits[:, step_idx] # (B, V)
          _, best_ids = current_decode_step_logits.topk(1, dim=-1) # argmax on dim vocab
          best_ids = best_ids.detach().cpu().numpy().flatten() # (V,)
          tgt_ids[:, step_idx+1] = best_ids # Set best tokens (output of current step) to input of next step
          completed = completed | (best_ids == tgt_eos_idx) # update `completed` when best_ids[i] == EOS
          if all(completed):
              break
  
      # Strip [EOS] token for all tgt_ids
      predict_sentences = []
      strip_tgt_ids = []
      for tgt_sentence in tgt_ids:
          # Strip all token after EOS
          strip_ids = tgt_sentence
          if tgt_eos_idx in tgt_sentence:
              eos_idx = tgt_sentence.tolist().index(tgt_eos_idx)
              strip_ids = tgt_sentence[:eos_idx + 1] # Include <EOS>
          strip_tgt_ids.append(strip_ids.tolist())
          predict_sentences.append( lightning_module.tgt_tokenizer.decode(strip_ids) )
      return strip_tgt_ids, predict_sentences, attentions
  
  
  if __name__ == "__main__":
      src_sentence = [
          "Hi",
      ]
  
      tgt_ids, predict_sentences, attentions = translate(src_sentence, litmodel)
      for src, pred in zip(src_sentence, predict_sentences):
          print(src, "--->", pred)
  
  