code: # coding=utf-8
  # Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
  # Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
  # Modifications copyright (C) 2020 Zi-Yi Dou
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  
  import argparse
  import random
  import itertools
  import os
  
  import numpy as np
  import torch
  from tqdm import trange
  from torch.nn.utils.rnn import pad_sequence
  from torch.utils.data import DataLoader, Dataset, SequentialSampler
  
  import modeling
  from configuration_bert import BertConfig
  from modeling import BertForMaskedLM
  from tokenization_bert import BertTokenizer
  from tokenization_utils import PreTrainedTokenizer
  from modeling_utils import PreTrainedModel
  
  
  
  def set_seed(args):
      if args.seed >= 0:
          random.seed(args.seed)
          np.random.seed(args.seed)
          torch.manual_seed(args.seed)
          torch.cuda.manual_seed_all(args.seed)
  
  class LineByLineTextDataset(Dataset):
      def __init__(self, tokenizer: PreTrainedTokenizer, args, file_path):
          assert os.path.isfile(file_path)
          print('Loading the dataset...')
          self.examples = []
          with open(file_path, encoding="utf-8") as f:
              for idx, line in enumerate(f.readlines()):
                  if len(line) == 0 or line.isspace() or not len(line.split(' ||| ')) == 2:
                      raise ValueError(f'Line {idx+1} is not in the correct format!')
                  
                  src, tgt = line.split(' ||| ')
                  if src.rstrip() == '' or tgt.rstrip() == '':
                      raise ValueError(f'Line {idx+1} is not in the correct format!')
              
                  sent_src, sent_tgt = src.strip().split(), tgt.strip().split()
                  token_src, token_tgt = [tokenizer.tokenize(word) for word in sent_src], [tokenizer.tokenize(word) for word in sent_tgt]
                  wid_src, wid_tgt = [tokenizer.convert_tokens_to_ids(x) for x in token_src], [tokenizer.convert_tokens_to_ids(x) for x in token_tgt]
  
                  ids_src, ids_tgt = tokenizer.prepare_for_model(list(itertools.chain(*wid_src)), return_tensors='pt', max_length=tokenizer.max_len)['input_ids'], tokenizer.prepare_for_model(list(itertools.chain(*wid_tgt)), return_tensors='pt', max_length=tokenizer.max_len)['input_ids']
  
                  bpe2word_map_src = []
                  for i, word_list in enumerate(token_src):
                      bpe2word_map_src += [i for x in word_list]
                  bpe2word_map_tgt = []
                  for i, word_list in enumerate(token_tgt):
                      bpe2word_map_tgt += [i for x in word_list]
  
                  self.examples.append( (ids_src[0], ids_tgt[0], bpe2word_map_src, bpe2word_map_tgt) )
  
      def __len__(self):
          return len(self.examples)
  
      def __getitem__(self, i):
          return self.examples[i]
  
  def word_align(args, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, output_word_alignments = False):
      def collate(examples):
          ids_src, ids_tgt, bpe2word_map_src, bpe2word_map_tgt = zip(*examples)
          ids_src = pad_sequence(ids_src, batch_first=True, padding_value=tokenizer.pad_token_id)
          ids_tgt = pad_sequence(ids_tgt, batch_first=True, padding_value=tokenizer.pad_token_id)
          return ids_src, ids_tgt, bpe2word_map_src, bpe2word_map_tgt
  
      dataset = LineByLineTextDataset(tokenizer, args, file_path=args.data_file)
      sampler = SequentialSampler(dataset)
      dataloader = DataLoader(
          dataset, sampler=sampler, batch_size=args.batch_size, collate_fn=collate
      )
  
      model.to(args.device)
      model.eval()
      tqdm_iterator = trange(dataset.__len__(), desc="Extracting")
      with open(args.output_file, 'w') as writer:
          for batch in dataloader:
              with torch.no_grad():
                  ids_src, ids_tgt, bpe2word_map_src, bpe2word_map_tgt = batch
                  word_aligns_list = model.get_aligned_word(ids_src, ids_tgt, bpe2word_map_src, bpe2word_map_tgt, args.device, 0, 0, align_layer=args.align_layer, extraction=args.extraction, softmax_threshold=args.softmax_threshold, test=True)
                  for word_aligns in word_aligns_list:
                      output_str = []
                      for word_align in word_aligns:
                          output_str.append(f'{word_align[0]}-{word_align[1]}')
                      writer.write(' '.join(output_str)+'\n')
                  tqdm_iterator.update(len(ids_src))
  
      if output_word_alignments:
          with open(args.output_file, 'r') as fh:
              outputf = (fh.read()).split("\n")
          with open(args.data_file, 'r') as fh:
              datalines = (fh.read()).split("\n")
  
          with open(args.output_file+".outtxt", 'w') as fwriter:
              for indices, line in zip(outputf, datalines):
                  srcline, tgtline = line.split(' ||| ')
                  indices = indices.split()
                  srcwrds = srcline.split()
                  tgtwrds = tgtline.split()
                  output_wrds = []
                  for wrd in indices:
                      srcix,tgtix = wrd.split("-")
                      srcix, tgtix = int(srcix), int(tgtix)
                      output_wrds.append(f"{srcwrds[srcix]}-{tgtwrds[tgtix]}")
                  fwriter.write(' '.join(output_wrds)+'\n')
  
  def main():
      parser = argparse.ArgumentParser()
  
      # Required parameters
      parser.add_argument(
          "--data_file", default=None, type=str, required=True, help="The input data file (a text file)."
      )
      parser.add_argument(
          "--output_file",
          type=str,
          required=True,
          help="The output directory where the model predictions and checkpoints will be written.",
      )
      parser.add_argument("--align_layer", type=int, default=8, help="layer for alignment extraction")
      parser.add_argument(
          "--extraction", default='softmax', type=str, help='softmax or entmax15'
      )
      parser.add_argument(
          "--softmax_threshold", type=float, default=0.001
      )
      parser.add_argument(
          "--model_name_or_path",
          default=None,
          type=str,
          help="The model checkpoint for weights initialization. Leave None if you want to train a model from scratch.",
      )
      parser.add_argument(
          "--config_name",
          default=None,
          type=str,
          help="Optional pretrained config name or path if not the same as model_name_or_path. If both are None, initialize a new config.",
      )
      parser.add_argument(
          "--tokenizer_name",
          default=None,
          type=str,
          help="Optional pretrained tokenizer name or path if not the same as model_name_or_path. If both are None, initialize a new tokenizer.",
      )
      parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")
      parser.add_argument("--batch_size", default=32, type=int)
      parser.add_argument(
          "--cache_dir",
          default='cache_dir',
          type=str,
          help="Optional directory to store the pre-trained models downloaded from s3 (instead of the default one)",
      )
      parser.add_argument("--no_cuda", action="store_true", help="Avoid using CUDA when available")
      args = parser.parse_args()
      device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
      args.device = device
  
      # Set seed
      set_seed(args)
      config_class, model_class, tokenizer_class = BertConfig, BertForMaskedLM, BertTokenizer
      if args.config_name:
          config = config_class.from_pretrained(args.config_name, cache_dir=args.cache_dir)
      elif args.model_name_or_path:
          config = config_class.from_pretrained(args.model_name_or_path, cache_dir=args.cache_dir)
      else:
          config = config_class()
  
      if args.tokenizer_name:
          tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name, cache_dir=args.cache_dir)
      elif args.model_name_or_path:
          tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path, cache_dir=args.cache_dir)
      else:
          raise ValueError(
              "You are instantiating a new {} tokenizer. This is not supported, but you can do it from another script, save it,"
              "and load it from here, using --tokenizer_name".format(tokenizer_class.__name__)
          )
  
      modeling.PAD_ID = tokenizer.pad_token_id
      modeling.CLS_ID = tokenizer.cls_token_id
      modeling.SEP_ID = tokenizer.sep_token_id
  
      if args.model_name_or_path:
          model = model_class.from_pretrained(
              args.model_name_or_path,
              from_tf=bool(".ckpt" in args.model_name_or_path),
              config=config,
              cache_dir=args.cache_dir,
          )
      else:
          model = model_class(config=config)
  
      word_align(args, model, tokenizer)
  
  if __name__ == "__main__":
      main()
  
apis: torch.utils.data.DataLoader torch.nn.utils.rnn.pad_sequence torch.cuda.manual_seed_all torch.manual_seed torch.utils.data.SequentialSampler numpy.random.seed torch.no_grad torch.cuda.is_available