text: #coding=utf8
  import torch
  import numpy as np
  from utils.example import Example, get_position_ids
  from utils.constants import PAD, UNK
  from model.model_utils import lens2mask, cached_property
  import torch.nn.functional as F
  
  def from_example_list_base(ex_list, device='cpu', train=True):
      """
          question_lens: torch.long, bsize
          questions: torch.long, bsize x max_question_len, include [CLS] if add_cls
          table_lens: torch.long, bsize, number of tables for each example
          table_word_lens: torch.long, number of words for each table name
          tables: torch.long, sum_of_tables x max_table_word_len
          column_lens: torch.long, bsize, number of columns for each example
          column_word_lens: torch.long, number of words for each column name
          columns: torch.long, sum_of_columns x max_column_word_len
      """
      batch = Batch(ex_list, device)
      plm = Example.plm
      pad_idx = Example.word_vocab[PAD] if plm is None else Example.tokenizer.pad_token_id
  
      question_lens = [len(ex.question) for ex in ex_list]
      batch.question_lens = torch.tensor(question_lens, dtype=torch.long, device=device)
      batch.table_lens = torch.tensor([len(ex.table) for ex in ex_list], dtype=torch.long, device=device)
      table_word_lens = [len(t) for ex in ex_list for t in ex.table]
      batch.table_word_lens = torch.tensor(table_word_lens, dtype=torch.long, device=device)
      batch.column_lens = torch.tensor([len(ex.column) for ex in ex_list], dtype=torch.long, device=device)
      column_word_lens = [len(c) for ex in ex_list for c in ex.column]
      batch.column_word_lens = torch.tensor(column_word_lens, dtype=torch.long, device=device)
  
      if plm is None: # glove.42B.300d
          questions = [ex.question_id + [pad_idx] * (batch.max_question_len - len(ex.question_id)) for ex in ex_list]
          batch.questions = torch.tensor(questions, dtype=torch.long, device=device)
          tables = [t + [pad_idx] * (batch.max_table_word_len - len(t)) for ex in ex_list for t in ex.table_id]
          batch.tables = torch.tensor(tables, dtype=torch.long, device=device)
          columns = [c + [pad_idx] * (batch.max_column_word_len - len(c)) for ex in ex_list for c in ex.column_id]
          batch.columns = torch.tensor(columns, dtype=torch.long, device=device)
      else:
          # prepare inputs for pretrained models
          batch.inputs = {"input_ids": None, "attention_mask": None, "token_type_ids": None, "position_ids": None}
          input_lens = [len(ex.input_id) for ex in ex_list]
          max_len = max(input_lens)
          input_ids = [ex.input_id + [pad_idx] * (max_len - len(ex.input_id)) for ex in ex_list]
          batch.inputs["input_ids"] = torch.tensor(input_ids, dtype=torch.long, device=device)
          attention_mask = [[1] * l + [0] * (max_len - l) for l in input_lens]
          batch.inputs["attention_mask"] = torch.tensor(attention_mask, dtype=torch.float, device=device)
          token_type_ids = [ex.segment_id + [0] * (max_len - len(ex.segment_id)) for ex in ex_list]
          batch.inputs["token_type_ids"] = torch.tensor(token_type_ids, dtype=torch.long, device=device)
          position_ids = [get_position_ids(ex, shuffle=train) + [0] * (max_len - len(ex.input_id)) for ex in ex_list]
          batch.inputs["position_ids"] = torch.tensor(position_ids, dtype=torch.long, device=device)
          # extract representations after plm, remove [SEP]
          question_mask_plm = [ex.question_mask_plm + [0] * (max_len - len(ex.question_mask_plm)) for ex in ex_list]
          batch.question_mask_plm = torch.tensor(question_mask_plm, dtype=torch.bool, device=device)
          table_mask_plm = [ex.table_mask_plm + [0] * (max_len - len(ex.table_mask_plm)) for ex in ex_list]
          batch.table_mask_plm = torch.tensor(table_mask_plm, dtype=torch.bool, device=device)
          column_mask_plm = [ex.column_mask_plm + [0] * (max_len - len(ex.column_mask_plm)) for ex in ex_list]
          batch.column_mask_plm = torch.tensor(column_mask_plm, dtype=torch.bool, device=device)
          # subword aggregation
          question_subword_lens = [l for ex in ex_list for l in ex.question_subword_len]
          batch.question_subword_lens = torch.tensor(question_subword_lens, dtype=torch.long, device=device)
          table_subword_lens = [l for ex in ex_list for l in ex.table_subword_len]
          batch.table_subword_lens = torch.tensor(table_subword_lens, dtype=torch.long, device=device)
          column_subword_lens = [l for ex in ex_list for l in ex.column_subword_len]
          batch.column_subword_lens = torch.tensor(column_subword_lens, dtype=torch.long, device=device)
  
      batch.question_unk_mask, batch.table_unk_mask, batch.column_unk_mask = None, None, None
      if not train and plm is None:
          # during evaluation, for words not in vocab but in glove vocab, extract its correpsonding embedding
          word2vec, unk_idx = Example.word2vec, Example.word_vocab[UNK]
  
          question_unk_mask = (batch.questions == unk_idx).cpu()
          if question_unk_mask.any().item():
              raw_questions = np.array([ex.question + [PAD] * (batch.max_question_len - len(ex.question)) for ex in ex_list], dtype='<U100')
              unk_words = raw_questions[question_unk_mask.numpy()].tolist()
              unk_word_embeddings = [word2vec.emb(w) for w in unk_words]
              oov_flag = torch.tensor([True if e is not None else False for e in unk_word_embeddings], dtype=torch.bool)
              if oov_flag.any().item():
                  batch.question_unk_mask = question_unk_mask.masked_scatter_(torch.clone(question_unk_mask), oov_flag).to(device)
                  batch.question_unk_embeddings = torch.tensor([e for e in unk_word_embeddings if e is not None], dtype=torch.float, device=device)
  
          table_unk_mask = (batch.tables == unk_idx).cpu()
          if table_unk_mask.any().item():
              raw_tables = np.array([t + [PAD] * (batch.max_table_word_len - len(t)) for ex in ex_list for t in ex.table], dtype='<U100')
              unk_words = raw_tables[table_unk_mask.numpy()].tolist()
              unk_word_embeddings = [word2vec.emb(w) for w in unk_words]
              oov_flag = torch.tensor([True if e is not None else False for e in unk_word_embeddings], dtype=torch.bool)
              if oov_flag.any().item():
                  batch.table_unk_mask = table_unk_mask.masked_scatter_(torch.clone(table_unk_mask), oov_flag).to(device)
                  batch.table_unk_embeddings = torch.tensor([e for e in unk_word_embeddings if e is not None], dtype=torch.float, device=device)
  
          column_unk_mask = (batch.columns == unk_idx).cpu()
          if column_unk_mask.any().item():
              raw_columns = np.array([c + [PAD] * (batch.max_column_word_len - len(c)) for ex in ex_list for c in ex.column], dtype='<U100')
              unk_words = raw_columns[column_unk_mask.numpy()].tolist()
              unk_word_embeddings = [word2vec.emb(w) for w in unk_words]
              oov_flag = torch.tensor([True if e is not None else False for e in unk_word_embeddings], dtype=torch.bool)
              if oov_flag.any().item():
                  batch.column_unk_mask = column_unk_mask.masked_scatter_(torch.clone(column_unk_mask), oov_flag).to(device)
                  batch.column_unk_embeddings = torch.tensor([e for e in unk_word_embeddings if e is not None], dtype=torch.float, device=device)
      return batch
  
  def from_example_list_text2sql(ex_list, device='cpu', train=True, **kwargs):
      """ New fields: batch.lens, batch.max_len, batch.relations, batch.relations_mask
      """
      batch = from_example_list_base(ex_list, device, train)
      batch.graph = Example.graph_factory.batch_graphs(ex_list, device, train=train, **kwargs)
      if train:
          batch.max_action_num = max([len(ex.tgt_action) for ex in ex_list])
      return batch
  
  class Batch():
  
      def __init__(self, examples, device='cpu'):
          super(Batch, self).__init__()
          self.examples = examples
          self.device = device
  
      @classmethod
      def from_example_list(cls, ex_list, device='cpu', train=True, method='text2sql', **kwargs):
          method_dict = {
              "text2sql": from_example_list_text2sql,
          }
          return method_dict[method](ex_list, device, train=train, **kwargs)
  
      def __len__(self):
          return len(self.examples)
  
      def __getitem__(self, idx):
          return self.examples[idx]
  
      @cached_property
      def max_question_len(self):
          return torch.max(self.question_lens).item()
  
      @cached_property
      def max_table_len(self):
          return torch.max(self.table_lens).item()
  
      @cached_property
      def max_column_len(self):
          return torch.max(self.column_lens).item()
  
      @cached_property
      def max_table_word_len(self):
          return torch.max(self.table_word_lens).item()
  
      @cached_property
      def max_column_word_len(self):
          return torch.max(self.column_word_lens).item()
  
      @cached_property
      def max_question_subword_len(self):
          return torch.max(self.question_subword_lens).item()
  
      @cached_property
      def max_table_subword_len(self):
          return torch.max(self.table_subword_lens).item()
  
      @cached_property
      def max_column_subword_len(self):
          return torch.max(self.column_subword_lens).item()
  
      """ Different types of nodes are seperated instead of concatenated together """
      @cached_property
      def mask(self):
          return torch.cat([self.question_mask, self.table_mask, self.column_mask], dim=1)
  
      @cached_property
      def question_mask(self):
          return lens2mask(self.question_lens)
  
      @cached_property
      def table_mask(self):
          return lens2mask(self.table_lens)
  
      @cached_property
      def column_mask(self):
          return lens2mask(self.column_lens)
  
      @cached_property
      def table_word_mask(self):
          return lens2mask(self.table_word_lens)
  
      @cached_property
      def column_word_mask(self):
          return lens2mask(self.column_word_lens)
  
      @cached_property
      def question_subword_mask(self):
          return lens2mask(self.question_subword_lens)
  
      @cached_property
      def table_subword_mask(self):
          return lens2mask(self.table_subword_lens)
  
      @cached_property
      def column_subword_mask(self):
          return lens2mask(self.column_subword_lens)
  
      def get_frontier_field_idx(self, t):
          ids = []
          for e in self.examples:
              if t < len(e.tgt_action):
                  ids.append(Example.grammar.field2id[e.tgt_action[t].frontier_field])
                  # assert self.grammar.id2field[ids[-1]] == e.tgt_action[t].frontier_field
              else:
                  ids.append(0)
          return torch.tensor(ids, dtype=torch.long, device=self.device)
  
      def get_frontier_prod_idx(self, t):
          ids = []
          for e in self.examples:
              if t < len(e.tgt_action):
                  ids.append(Example.grammar.prod2id[e.tgt_action[t].frontier_prod])
                  # assert self.grammar.id2prod[ids[-1]] == e.tgt_action[t].frontier_prod
              else:
                  ids.append(0)
          return torch.tensor(ids, dtype=torch.long, device=self.device)
  
      def get_frontier_field_type_idx(self, t):
          ids = []
          for e in self.examples:
              if t < len(e.tgt_action):
                  ids.append(Example.grammar.type2id[e.tgt_action[t].frontier_field.type])
                  # assert self.grammar.id2type[ids[-1]] == e.tgt_action[t].frontier_field.type
              else:
                  ids.append(0)
          return torch.tensor(ids, dtype=torch.long, device=self.device)
  