text: import argparse
  import logging
  import numpy as np
  import os
  import sys
  import time
  import torch
  from multiprocessing import Pool
  from typing import Dict, List, Tuple
  from utils import parsing
  from utils.data_utils import get_graph_features_from_smi, load_vocab, make_vocab, \
      tokenize_selfies_from_smiles, tokenize_smiles
  from utils.train_utils import log_tensor, set_seed, setup_logger
  
  
  def get_preprocess_parser():
      parser = argparse.ArgumentParser("preprocess")
      parsing.add_common_args(parser)
      parsing.add_preprocess_args(parser)
  
      return parser
  
  
  def tokenize(fns: Dict[str, List[Tuple[str, str]]], output_path: str, repr_start: str, repr_end: str):
      assert repr_start == "smiles", f"{repr_start} input provided. Only smiles inputs are supported!"
  
      if repr_end == "smiles":
          tokenize_line = tokenize_smiles
      elif repr_end == "selfies":
          tokenize_line = tokenize_selfies_from_smiles
      else:
          raise ValueError(f"{repr_end} output required. Only smiles and selfies outputs are supported!")
  
      ofns = {}
  
      for phase, file_list in fns.items():
          ofns[phase] = []
  
          for src_file, tgt_file in file_list:
              src_output = os.path.join(output_path, f"{repr_end}_tokenized_{os.path.basename(src_file)}")
              tgt_output = os.path.join(output_path, f"{repr_end}_tokenized_{os.path.basename(tgt_file)}")
  
              for fn, ofn in [(src_file, src_output),
                              (tgt_file, tgt_output)]:
                  if os.path.exists(ofn):
                      logging.info(f"Found {ofn}, skipping tokenization.")
                      continue
  
                  with open(fn, "r") as f, open(ofn, "w") as of:
                      logging.info(f"Tokenizing input {fn} into {ofn}")
                      for i, line in enumerate(f):
                          line = "".join(line.strip().split())
                          newline = tokenize_line(line)
                          of.write(f"{newline}\n")
                      logging.info(f"Done, total lines: {i + 1}")
  
              ofns[phase].append((src_output, tgt_output))
  
      return ofns
  
  
  def get_token_ids(tokens: list, vocab: Dict[str, int], max_len: int) -> Tuple[List, int]:
      # token_ids = [vocab["_SOS"]]               # shouldn't really need this
      token_ids = []
      token_ids.extend([vocab[token] for token in tokens])
      token_ids = token_ids[:max_len-1]
      token_ids.append(vocab["_EOS"])
  
      lens = len(token_ids)
      while len(token_ids) < max_len:
          token_ids.append(vocab["_PAD"])
  
      return token_ids, lens
  
  
  def get_seq_features_from_line(_args) -> Tuple[np.ndarray, int, np.ndarray, int]:
      i, src_line, tgt_line, max_src_len, max_tgt_len = _args
      assert isinstance(src_line, str) and isinstance(tgt_line, str)
      if i > 0 and i % 10000 == 0:
          logging.info(f"Processing {i}th SMILES")
  
      global G_vocab
  
      src_tokens = src_line.strip().split()
      if not src_tokens:
          src_tokens = ["C", "C"]             # hardcode to ignore
      tgt_tokens = tgt_line.strip().split()
      src_token_ids, src_lens = get_token_ids(src_tokens, G_vocab, max_len=max_src_len)
      tgt_token_ids, tgt_lens = get_token_ids(tgt_tokens, G_vocab, max_len=max_tgt_len)
  
      src_token_ids = np.array(src_token_ids, dtype=np.int32)
      tgt_token_ids = np.array(tgt_token_ids, dtype=np.int32)
  
      return src_token_ids, src_lens, tgt_token_ids, tgt_lens
  
  
  def binarize_s2s(src_file: str, tgt_file: str, prefix: str, output_path: str,
                   max_src_len: int, max_tgt_len: int, num_workers: int = 1):
      output_file = os.path.join(output_path, f"{prefix}.npz")
      logging.info(f"Binarizing (s2s) src {src_file} and tgt {tgt_file}, saving to {output_file}")
  
      with open(src_file, "r") as f:
          src_lines = f.readlines()
  
      with open(tgt_file, "r") as f:
          tgt_lines = f.readlines()
  
      logging.info("Getting seq features")
      start = time.time()
  
      p = Pool(num_workers)
      seq_features_and_lengths = p.imap(
          get_seq_features_from_line,
          ((i, src_line, tgt_line, max_src_len, max_tgt_len)
           for i, (src_line, tgt_line) in enumerate(zip(src_lines, tgt_lines)))
      )
  
      p.close()
      p.join()
  
      seq_features_and_lengths = list(seq_features_and_lengths)
  
      logging.info(f"Done seq featurization, time: {time.time() - start}. Collating")
      src_token_ids, src_lens, tgt_token_ids, tgt_lens = zip(*seq_features_and_lengths)
  
      src_token_ids = np.stack(src_token_ids, axis=0)
      src_lens = np.array(src_lens, dtype=np.int32)
      tgt_token_ids = np.stack(tgt_token_ids, axis=0)
      tgt_lens = np.array(tgt_lens, dtype=np.int32)
  
      np.savez(
          output_file,
          src_token_ids=src_token_ids,
          src_lens=src_lens,
          tgt_token_ids=tgt_token_ids,
          tgt_lens=tgt_lens
      )
  
  
  def binarize_g2s(src_file: str, tgt_file: str, prefix: str, output_path: str,
                   max_src_len: int, max_tgt_len: int, num_workers: int = 1):
      output_file = os.path.join(output_path, f"{prefix}.npz")
      logging.info(f"Binarizing (g2s) src {src_file} and tgt {tgt_file}, saving to {output_file}")
  
      with open(src_file, "r") as f:
          # lines = f.readlines()[164104:164106]
          src_lines = f.readlines()
  
      with open(tgt_file, "r") as f:
          tgt_lines = f.readlines()
  
      logging.info("Getting seq features")
      start = time.time()
  
      p = Pool(num_workers)
      seq_features_and_lengths = p.imap(
          get_seq_features_from_line,
          ((i, src_line, tgt_line, max_src_len, max_tgt_len)
           for i, (src_line, tgt_line) in enumerate(zip(src_lines, tgt_lines)))
      )
  
      p.close()
      p.join()
  
      seq_features_and_lengths = list(seq_features_and_lengths)
  
      logging.info(f"Done seq featurization, time: {time.time() - start}. Collating")
      src_token_ids, src_lens, tgt_token_ids, tgt_lens = zip(*seq_features_and_lengths)
  
      src_token_ids = np.stack(src_token_ids, axis=0)
      src_lens = np.array(src_lens, dtype=np.int32)
      tgt_token_ids = np.stack(tgt_token_ids, axis=0)
      tgt_lens = np.array(tgt_lens, dtype=np.int32)
  
      logging.info("Getting graph features")
      start = time.time()
  
      p = Pool(num_workers)
      graph_features_and_lengths = p.imap(
          get_graph_features_from_smi,
          ((i, "".join(line.split()), False) for i, line in enumerate(src_lines))
      )
  
      p.close()
      p.join()
  
      graph_features_and_lengths = list(graph_features_and_lengths)
      logging.info(f"Done graph featurization, time: {time.time() - start}. Collating and saving...")
      a_scopes, a_scopes_lens, b_scopes, b_scopes_lens, a_features, a_features_lens, \
          b_features, b_features_lens, a_graphs, b_graphs = zip(*graph_features_and_lengths)
  
      a_scopes = np.concatenate(a_scopes, axis=0)
      b_scopes = np.concatenate(b_scopes, axis=0)
      a_features = np.concatenate(a_features, axis=0)
      b_features = np.concatenate(b_features, axis=0)
      a_graphs = np.concatenate(a_graphs, axis=0)
      b_graphs = np.concatenate(b_graphs, axis=0)
  
      a_scopes_lens = np.array(a_scopes_lens, dtype=np.int32)
      b_scopes_lens = np.array(b_scopes_lens, dtype=np.int32)
      a_features_lens = np.array(a_features_lens, dtype=np.int32)
      b_features_lens = np.array(b_features_lens, dtype=np.int32)
  
      np.savez(
          output_file,
          src_token_ids=src_token_ids,
          src_lens=src_lens,
          tgt_token_ids=tgt_token_ids,
          tgt_lens=tgt_lens,
          a_scopes=a_scopes,
          b_scopes=b_scopes,
          a_features=a_features,
          b_features=b_features,
          a_graphs=a_graphs,
          b_graphs=b_graphs,
          a_scopes_lens=a_scopes_lens,
          b_scopes_lens=b_scopes_lens,
          a_features_lens=a_features_lens,
          b_features_lens=b_features_lens
      )
  
  
  def preprocess_main(args):
      parsing.log_args(args)
  
      os.makedirs(args.preprocess_output_path, exist_ok=True)
  
      fns = {
          "train": [(args.train_src, args.train_tgt)],
          "val": [(args.val_src, args.val_tgt)],
          "test": [(args.test_src, args.test_tgt)]
      }
  
      if not args.representation_start == args.representation_end:
          assert args.do_tokenize, f"Different representations, start: {args.representation_start}, " \
                                   f"end: {args.representation_end}. Please set '--do_tokenize'"
  
      if args.do_tokenize:
          ofns = tokenize(fns=fns,
                          output_path=args.preprocess_output_path,
                          repr_start=args.representation_start,
                          repr_end=args.representation_end)
          fns = ofns                          # just pass the handle of tokenized files
  
      vocab_file = os.path.join(args.preprocess_output_path,
                                f"vocab_{args.representation_end}.txt")
      if not os.path.exists(vocab_file):
          make_vocab(
              fns=fns,
              vocab_file=vocab_file,
              tokenized=True
          )
  
      if args.make_vocab_only:
          logging.info(f"--make_vocab_only flag detected. Skipping featurization")
          exit(0)
  
      global G_vocab
      G_vocab = load_vocab(vocab_file)
  
      if args.model == "s2s":
          binarize = binarize_s2s
      elif args.model.startswith("g2s"):
          binarize = binarize_g2s
      else:
          raise ValueError(f"Model {args.model} not supported!")
  
      for phase, file_list in fns.items():
          for i, (src_file, tgt_file) in enumerate(file_list):
              binarize(
                  src_file=src_file,
                  tgt_file=tgt_file,
                  prefix=f"{phase}_{i}",
                  output_path=args.preprocess_output_path,
                  max_src_len=args.max_src_len,
                  max_tgt_len=args.max_tgt_len,
                  num_workers=args.num_workers
              )
  
  
  if __name__ == "__main__":
      preprocess_parser = get_preprocess_parser()
      args = preprocess_parser.parse_args()
  
      # set random seed
      set_seed(args.seed)
  
      # logger setup
      logger = setup_logger(args)
  
      np.set_printoptions(threshold=sys.maxsize)
      torch.set_printoptions(profile="full")
  
      G_vocab = {}            # global vocab
  
      preprocess_main(args)
  