code: # coding=utf-8
  # Copyright 2019 The Google AI Language Team Authors.
  #
  # 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.
  # Lint as: python3
  """Utilities for converting interactions to TF examples."""
  
  import collections
  import dataclasses
  import hashlib
  import random
  from typing import Iterable, List, Mapping, Optional, Text, Tuple
  
  from absl import logging
  from apache_beam import metrics as beam_metrics
  from tapas.protos import annotated_text_pb2
  from tapas.protos import interaction_pb2
  from tapas.protos import table_selection_pb2
  from tapas.utils import constants
  from tapas.utils import interpretation_utils
  from tapas.utils import number_annotation_utils
  from tapas.utils import sentence_tokenizer
  from tapas.utils import text_index
  from tapas.utils import text_utils
  import tensorflow.compat.v1 as tf
  
  from official.nlp.bert import tokenization
  
  _NS = 'main'
  _CLS = '[CLS]'
  _EMPTY = '[EMPTY]'
  _MASK = '[MASK]'
  _SEP = '[SEP]'
  _NAN = float('nan')
  _MAX_NUM_CANDIDATES = 1000
  _MAX_NUM_ROWS = 32
  _WP_PER_CELL = 1.5
  _MAX_INDEX_LENGTH = int(_MAX_NUM_CANDIDATES * _MAX_NUM_ROWS * _WP_PER_CELL)
  _MAX_INT = 2**32 - 1
  
  
  @dataclasses.dataclass(frozen=True)
  class Token:
    original_text: Text
    piece: Text
  
  
  @dataclasses.dataclass(frozen=True)
  class TrainingInstance:
    tokens: List[Token]
    segment_ids: List[int]
    column_ids: List[int]
    row_ids: List[int]
    masked_lm_positions: List[int]
    masked_lm_labels: List[Text]
    is_random_table: bool
  
  
  @dataclasses.dataclass(frozen=True)
  class TokenCoordinates:
    column_index: int
    row_index: int
    token_index: int
  
  
  @dataclasses.dataclass
  class TokenizedTable:
    rows: List[List[List[Token]]]
    selected_tokens: List[TokenCoordinates]
  
  
  @dataclasses.dataclass(frozen=True)
  class MaskedLmInstance:
    index: int
    label: Text
  
  
  @dataclasses.dataclass(frozen=True)
  class ConversionConfig:
    """Configues conversion to TF example.
  
    vocab_file: Bert vocab file
    max_seq_length: Max length of a sequence in word pieces.
    max_column_id: Max column id to extract.
    max_row_id: Max row id to extract.
    """
    vocab_file: Text
    max_seq_length: int
    max_column_id: int
    max_row_id: int
    strip_column_names: bool
  
  
  @dataclasses.dataclass(frozen=True)
  class PretrainConversionConfig(ConversionConfig):
    """Configures options speciic to pretraining data creation.
  
    max_predictions_per_seq: Max predictions per sequence for mask task.
    min_question_length: Min question length.
    max_question_length: Max question length.
    always_continue_cells: If true always mask entire cells.
    strip_column_names: If true, add empty strings instead of column names.
    random_seed: Random seed.
    masked_lm_prob: Percentage of tokens to mask.
    concatenate_snippets: If true concatenate snippets in a random fashion.
    """
    max_predictions_per_seq: int
    masked_lm_prob: float
    random_seed: int
    min_question_length: int
    max_question_length: int
    always_continue_cells: bool
    concatenate_snippets: bool = True
  
  
  @dataclasses.dataclass(frozen=True)
  class TrimmedConversionConfig(ConversionConfig):
    # if > 0: Trim cells so that the length is <= this value.
    # Also disables further cell trimming should thus be used with
    # 'drop_rows_to_fit' below.
    # TODO(thomasmueller) Make this a parameter of the base config.
    # TODO(thomasmueller) Consider giving this a better name.
    cell_trim_length: int = -1
  
  
  @dataclasses.dataclass(frozen=True)
  class ClassifierConversionConfig(TrimmedConversionConfig):
    """The config used to extract the tf examples for the classifier model."""
    add_aggregation_candidates: bool = False
    expand_entity_descriptions: bool = False
    use_entity_title: bool = False
    entity_descriptions_sentence_limit: int = 5
    use_document_title: bool = False
    # Re-computes answer coordinates from the answer text.
    update_answer_coordinates: bool = False
    # Drop last rows if table doesn't fit within max sequence length.
    drop_rows_to_fit: bool = False
    # If true adds the context heading of the table to the question.
    use_context_title: bool = False
    # For TPU prediction we serialize strings into a fix length.
    trim_question_ids: bool = False
    # For each data split how to up/down sample the dataset
    label_sampling_rate: Mapping[Tuple[Text, int],
                                 float] = dataclasses.field(default_factory=dict)
    is_multi_hop: bool = False
    # self._is_multi_hop = config.is_multi_hop
    use_bridge_entity: bool = False
    # self._use_bridge_entity = config.use_bridge_entity
    use_question_type: bool = False
    # self._use_question_type = config.use_question_type
  
  
  @dataclasses.dataclass(frozen=True)
  class RetrievalConversionConfig(TrimmedConversionConfig):
    use_document_title: bool = True
    use_section_title: bool = False
    use_caption: bool = False
    use_abbv: bool = False
    use_header: bool = True
    use_content: bool = True
    oracle_abbv_expansion: bool = False
  
  
  @dataclasses.dataclass(frozen=True)
  class SerializedExample:
    tokens: List[Token]
    column_ids: List[int]
    row_ids: List[int]
    segment_ids: List[int]
  
  
  def copy_vocab(input_vocab, output_vocab):
    """Copies vocabulary file and add [EMPTY] token."""
    with tf.io.gfile.GFile(input_vocab) as input_vocab_file:
      with tf.io.gfile.GFile(output_vocab, 'w') as output_vocab_file:
        for token in input_vocab_file:
          output_vocab_file.write('[EMPTY]\n' if token ==
                                  '[unused0]\n' else token)
  
  
  def _get_pieces(tokens):
    return (token.piece for token in tokens)
  
  
  def fingerprint(text):
    return int(hashlib.sha256(text.encode('utf-8')).hexdigest(), 16)
  
  
  def create_int_feature(values):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))
  
  
  def create_float_feature(values):
    return tf.train.Feature(float_list=tf.train.FloatList(value=list(values)))
  
  
  def create_string_feature(values):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=list(values)))
  
  
  def _is_inner_wordpiece(token):
    return token.piece.startswith('##')
  
  
  def _get_cell_token_indexes(column_ids, row_ids,
                              column_id, row_id):
    for index in range(len(column_ids)):
      if (column_ids[index] - 1 == column_id and row_ids[index] - 1 == row_id):
        yield index
  
  
  def _get_buckets(value, buckets, name):
    for bucket_value in buckets:
      if value <= bucket_value:
        return '%s: <= %d' % (name, bucket_value)
    return '%s: < inf' % (name)
  
  
  def _get_all_answer_ids_from_coordinates(
      column_ids,
      row_ids,
      answers_list,
  ):
    """Maps lists of answer coordinates to token indexes."""
    answer_ids = [0] * len(column_ids)
    found_answers = set()
    all_answers = set()
    for answers in answers_list:
      for column_index, row_index in answers:
        all_answers.add((column_index, row_index))
        for index in _get_cell_token_indexes(column_ids, row_ids, column_index,
                                             row_index):
          found_answers.add((column_index, row_index))
          answer_ids[index] = 1
  
    missing_count = len(all_answers) - len(found_answers)
    buckets = [1, 2, 3, 4, 5, 10, 25, 50, 100]
    if missing_count:
      beam_metrics.Metrics.counter(
          _NS, _get_buckets(missing_count, buckets, 'Missing answers')).inc()
    if found_answers:
      beam_metrics.Metrics.counter(
          _NS, _get_buckets(len(found_answers), buckets, 'Found answers')).inc()
    return answer_ids, missing_count
  
  
  def _get_all_answer_ids(
      column_ids,
      row_ids,
      questions,
  ):
    """Maps lists of questions with answer coordinates to token indexes."""
  
    def _to_coordinates(
        question,):
      return [(coords.column_index, coords.row_index)
              for coords in question.answer.answer_coordinates]
  
    return _get_all_answer_ids_from_coordinates(
        column_ids,
        row_ids,
        answers_list=(_to_coordinates(question) for question in questions),
    )
  
  
  def _find_tokens(text, segment):
    """Return start index of segment in text or None."""
    logging.info('text: %s %s', text, segment)
    for index in range(1 + len(text) - len(segment)):
      for seg_index, seg_token in enumerate(segment):
        if text[index + seg_index].piece != seg_token.piece:
          break
      else:
        return index
    return None
  
  
  def _find_answer_coordinates_from_answer_text(
      tokenized_table,
      answer_text,
  ):
    """Returns all occurrences of answer_text in the table."""
    logging.info('answer text: %s', answer_text)
    for row_index, row in enumerate(tokenized_table.rows):
      if row_index == 0:
        # We don't search for answers in the header.
        continue
      for col_index, cell in enumerate(row):
        token_index = _find_tokens(cell, answer_text)
        if token_index is not None:
          yield TokenCoordinates(
              row_index=row_index,
              column_index=col_index,
              token_index=token_index,
          )
  
  
  def _find_answer_ids_from_answer_texts(
      column_ids,
      row_ids,
      tokenized_table,
      answer_texts,
  ):
    """Maps question with answer texts to the first matching token indexes."""
    answer_ids = [0] * len(column_ids)
    for answer_text in answer_texts:
      found_answer_text = False
      found_answer_text_ids = False
      for coordinates in _find_answer_coordinates_from_answer_text(
          tokenized_table,
          answer_text,
      ):
        found_answer_text = True
        # Maps answer coordinates to indexes this can fail if tokens / rows have
        # been pruned.
        indexes = list(
            _get_cell_token_indexes(
                column_ids,
                row_ids,
                column_id=coordinates.column_index,
                row_id=coordinates.row_index - 1,
            ))
        indexes.sort()
        coordinate_answer_ids = []
        if indexes:
          begin_index = coordinates.token_index + indexes[0]
          end_index = begin_index + len(answer_text)
          for index in indexes:
            if index >= begin_index and index < end_index:
              coordinate_answer_ids.append(index)
        if len(coordinate_answer_ids) == len(answer_text):
          found_answer_text_ids = True
          for index in coordinate_answer_ids:
            answer_ids[index] = 1
          break
      beam_metrics.Metrics.counter(_NS, 'Answer texts: total').inc()
      if found_answer_text:
        beam_metrics.Metrics.counter(_NS, 'Answer texts: found').inc()
      if found_answer_text_ids:
        beam_metrics.Metrics.counter(_NS, 'Answer texts: found ids').inc()
  
    return answer_ids
  
  
  def _get_answer_ids(column_ids, row_ids,
                      question):
    """Maps answer coordinates to token indexes."""
    answer_ids, missing_count = _get_all_answer_ids(column_ids, row_ids,
                                                    [question])
  
    if missing_count:
      raise ValueError("Couldn't find all answers")
    return answer_ids
  
  
  def _get_annotation_name(identifier):
    """Extracts the clean title from a Wikipedia identifier."""
    # Example input: /wiki/New_York_City -> New York City
    return identifier.split('/')[-1].replace('_', ' ')
  
  
  def _add_entity_descriptions_to_table(
      question,
      descriptions,
      table,
      use_entity_title,
      num_results,
  ):
    """Expand table cells with the descriptions of the entities mentioned.
  
    This function will add entity descriptions inside the Table proto by expanding
    the content of each cell according to entities mentioned in that cell. The
    sentences in the descriptions will be ranked by similarity to the question and
    only the top results will be included.
  
    Args:
      question: Question proto containing the question text. The text will be used
        to filter only a subset of the descriptions using a similarity criteria.
      descriptions: A map that contains for entity id, its textual description.
      table: Table to be modified in-place. Some cells may contain annotation
        extensions with entity ids that will be expanded with their descriptions.
      use_entity_title: Prepend the entity title to entity descriptions.
      num_results: Limit on the number of entities to expand with a description.
    """
    descriptions = {
        key: sentence_tokenizer.tokenize(description)
        for key, description in descriptions.items()
    }
    documents = []
    for sentences in descriptions.values():
      for sentence in sentences:
        documents.append(sentence)
  
    search_results = text_index.TextIndex(documents).search(
        question.text, num_results=num_results)
    logging.log_first_n(logging.INFO,
                        '%s selected entity annotations for %s:  %s', 100,
                        question.id, question.text, search_results)
    search_results_set = {r.text for r in search_results}
  
    buckets = [1, 2, 3, 4, 5, 10, 25, 50, 100]
    sentences_kept = len(search_results_set)
    sentences_discarded = len(documents) - sentences_kept
    beam_metrics.Metrics.counter(
        _NS, _get_buckets(sentences_kept, buckets, 'Descriptions kept')).inc()
    beam_metrics.Metrics.counter(
        _NS, _get_buckets(sentences_discarded, buckets,
                          'Descriptions discarded')).inc()
  
    annotated_cell_ext = annotated_text_pb2.AnnotatedText.annotated_cell_ext
    for row in table.rows:
      for cell in row.cells:
        if annotated_cell_ext in cell.Extensions:
          for annotation in cell.Extensions[annotated_cell_ext].annotations:
            sentences = descriptions[annotation.identifier]
            filtered_sentences = ' '.join(
                sent for sent in sentences if sent in search_results_set)
            if filtered_sentences:
              if use_entity_title:
                annotation_name = _get_annotation_name(annotation.identifier)
                cell.text += f' ( {annotation_name} : {filtered_sentences} )'
              else:
                cell.text += f' ( {filtered_sentences} )'
  
  
  class TapasTokenizer:
    """Wraps a Bert tokenizer."""
  
    def __init__(self, vocab_file):
      self._basic_tokenizer = tokenization.BasicTokenizer(do_lower_case=True)
      self._wp_tokenizer = tokenization.FullTokenizer(
          vocab_file=vocab_file, do_lower_case=True)
  
    def get_vocab(self):
      return self._wp_tokenizer.vocab.keys()
  
    def tokenize(self, text):
      if text_utils.format_text(text) == constants.EMPTY_TEXT:
        return [Token(_EMPTY, _EMPTY)]
      tokens = []
      for token in self._basic_tokenizer.tokenize(text):
        for piece in self._wp_tokenizer.tokenize(token):
          tokens.append(Token(token, piece))
      return tokens
  
    def convert_tokens_to_ids(self, word_pieces):
      return self._wp_tokenizer.convert_tokens_to_ids(word_pieces)
  
    def question_encoding_cost(self, question_tokens):
      # Two extra spots of SEP and CLS.
      return len(question_tokens) + 2
  
  
  class ToTensorflowExampleBase:
    """Base class for converting interactions to TF examples."""
  
    def __init__(self, config):
      self._max_seq_length = config.max_seq_length
      self._max_column_id = config.max_column_id
      self._max_row_id = config.max_row_id
      self._strip_column_names = config.strip_column_names
      self._tokenizer = TapasTokenizer(config.vocab_file)
  
    def _tokenize_table(
        self,
        table,
    ):
      """Runs tokenizer over columns and table cell texts."""
      tokenized_rows = []
      tokenized_row = []
      for column in table.columns:
        if self._strip_column_names:
          tokenized_row.append(self._tokenizer.tokenize(''))
        else:
          tokenized_row.append(self._tokenizer.tokenize(column.text))
      tokenized_rows.append(tokenized_row)
  
      for row in table.rows:
        tokenized_row = []
        for cell in row.cells:
          tokenized_row.append(self._tokenizer.tokenize(cell.text))
        tokenized_rows.append(tokenized_row)
  
      token_coordinates = []
      for row_index, row in enumerate(tokenized_rows):
        for column_index, cell in enumerate(row):
          for token_index, _ in enumerate(cell):
            token_coordinates.append(
                TokenCoordinates(
                    row_index=row_index,
                    column_index=column_index,
                    token_index=token_index,
                ))
  
      return TokenizedTable(
          rows=tokenized_rows,
          selected_tokens=token_coordinates,
      )
  
    def _get_table_values(self, table, num_columns,
                          num_rows,
                          num_tokens):
      """Iterates over partial table and returns token, col. and row indexes."""
      for tc in table.selected_tokens:
        # First row is header row.
        if tc.row_index >= num_rows + 1:
          continue
        if tc.column_index >= num_columns:
          continue
        cell = table.rows[tc.row_index][tc.column_index]
        token = cell[tc.token_index]
        word_begin_index = tc.token_index
        # Don't add partial words. Find the starting word piece and check if it
        # fits in the token budget.
        while (word_begin_index >= 0 and
               _is_inner_wordpiece(cell[word_begin_index])):
          word_begin_index -= 1
        if word_begin_index >= num_tokens:
          continue
        yield token, tc.column_index + 1, tc.row_index
  
    def _serialize_text(
        self, question_tokens
    ):
      """Serialzes texts in index arrays."""
      tokens = []
      segment_ids = []
      column_ids = []
      row_ids = []
  
      tokens.append(Token(_CLS, _CLS))
      segment_ids.append(0)
      column_ids.append(0)
      row_ids.append(0)
  
      for token in question_tokens:
        tokens.append(token)
        segment_ids.append(0)
        column_ids.append(0)
        row_ids.append(0)
  
      return tokens, segment_ids, column_ids, row_ids
  
    def _serialize(
        self,
        question_tokens,
        table,
        num_columns,
        num_rows,
        num_tokens,
    ):
      """Serializes table and text."""
      tokens, segment_ids, column_ids, row_ids = self._serialize_text(
          question_tokens)
  
      tokens.append(Token(_SEP, _SEP))
      segment_ids.append(0)
      column_ids.append(0)
      row_ids.append(0)
  
      for token, column_id, row_id in self._get_table_values(
          table, num_columns, num_rows, num_tokens):
        tokens.append(token)
        segment_ids.append(1)
        column_ids.append(column_id)
        row_ids.append(row_id)
  
      return SerializedExample(
          tokens=tokens,
          segment_ids=segment_ids,
          column_ids=column_ids,
          row_ids=row_ids,
      )
  
    def _tokenize(self, text):
      return self._tokenizer.tokenize(text)
  
    def _get_token_budget(self, question_tokens):
      return self._max_seq_length - self._tokenizer.question_encoding_cost(
          question_tokens)
  
    def _get_table_boundaries(self,
                              table):
      """Return maximal number of rows, columns and tokens."""
      max_num_tokens = 0
      max_num_columns = 0
      max_num_rows = 0
      for tc in table.selected_tokens:
        max_num_columns = max(max_num_columns, tc.column_index + 1)
        max_num_rows = max(max_num_rows, tc.row_index + 1)
        max_num_tokens = max(max_num_tokens, tc.token_index + 1)
      max_num_columns = min(self._max_column_id, max_num_columns)
      max_num_rows = min(self._max_row_id, max_num_rows)
      return max_num_rows, max_num_columns, max_num_tokens
  
    def _get_table_cost(self, table, num_columns,
                        num_rows, num_tokens):
      return sum(1 for _ in self._get_table_values(table, num_columns, num_rows,
                                                   num_tokens))
  
    def _get_column_values(
        self, table,
        col_index):
      table_numeric_values = {}
      for row_index, row in enumerate(table.rows):
        cell = row.cells[col_index]
        if cell.HasField('numeric_value'):
          table_numeric_values[row_index] = cell.numeric_value
      return table_numeric_values
  
    def _add_numeric_column_ranks(self, column_ids, row_ids,
                                  table,
                                  features):
      """Adds column ranks for all numeric columns."""
  
      ranks = [0] * len(column_ids)
      inv_ranks = [0] * len(column_ids)
  
      if table:
        for col_index in range(len(table.columns)):
          table_numeric_values = self._get_column_values(table, col_index)
          if not table_numeric_values:
            continue
  
          try:
            key_fn = number_annotation_utils.get_numeric_sort_key_fn(
                table_numeric_values.values())
          except ValueError:
            continue
  
          table_numeric_values = {
              row_index: key_fn(value)
              for row_index, value in table_numeric_values.items()
          }
  
          table_numeric_values_inv = collections.defaultdict(list)
          for row_index, value in table_numeric_values.items():
            table_numeric_values_inv[value].append(row_index)
  
          unique_values = sorted(table_numeric_values_inv.keys())
  
          for rank, value in enumerate(unique_values):
            for row_index in table_numeric_values_inv[value]:
              for index in _get_cell_token_indexes(column_ids, row_ids, col_index,
                                                   row_index):
                ranks[index] = rank + 1
                inv_ranks[index] = len(unique_values) - rank
  
      features['column_ranks'] = create_int_feature(ranks)
      features['inv_column_ranks'] = create_int_feature(inv_ranks)
  
    def _get_numeric_sort_key_fn(self, table_numeric_values, value):
      """Returns the sort key function for comparing value to table values.
  
      The function returned will be a suitable input for the key param of the
      sort(). See number_annotation_utils._get_numeric_sort_key_fn for details.
  
      Args:
        table_numeric_values: Numeric values of a column
        value: Numeric value in the question.
  
      Returns:
        A function key function to compare column and question values.
  
      """
      if not table_numeric_values:
        return None
      all_values = list(table_numeric_values.values())
      all_values.append(value)
      try:
        return number_annotation_utils.get_numeric_sort_key_fn(all_values)
      except ValueError:
        return None
  
    def _add_numeric_relations(self, question,
                               column_ids, row_ids,
                               table,
                               features):
      """Adds numeric relation emebeddings to 'features'.
  
      Args:
        question: The question, numeric values are used.
        column_ids: Maps word piece position to column id.
        row_ids: Maps word piece position to row id.
        table: The table containing the numeric cell values.
        features: Output.
      """
  
      numeric_relations = [0 for _ in column_ids]
  
      # Create a dictionary that maps a table cell to the set of all relations
      # this cell has with any value in the question.
      cell_indices_to_relations = collections.defaultdict(set)
      if question is not None and table is not None:
        for numeric_value_span in question.annotations.spans:
          for value in numeric_value_span.values:
            for column_index in range(len(table.columns)):
              table_numeric_values = self._get_column_values(table, column_index)
              sort_key_fn = self._get_numeric_sort_key_fn(table_numeric_values,
                                                          value)
              if sort_key_fn is None:
                continue
              for row_index, cell_value in table_numeric_values.items():
                relation = number_annotation_utils.get_numeric_relation(
                    value, cell_value, sort_key_fn)
                if relation is not None:
                  cell_indices_to_relations[column_index, row_index].add(relation)
  
      # For each cell add a special feature for all its word pieces.
      for (column_index,
           row_index), relations in cell_indices_to_relations.items():
        relation_set_index = 0
        for relation in relations:
          assert relation.value >= constants.Relation.EQ.value
          relation_set_index += 2**(relation.value - constants.Relation.EQ.value)
        beam_metrics.Metrics.counter(
            _NS, 'Relation Set Index: %d' % relation_set_index).inc()
        for cell_token_index in _get_cell_token_indexes(column_ids, row_ids,
                                                        column_index, row_index):
          numeric_relations[cell_token_index] = relation_set_index
  
      features['numeric_relations'] = create_int_feature(numeric_relations)
  
    def _add_numeric_values(self, table,
                            token_ids_dict,
                            features):
      """Adds numeric values for computation of answer loss."""
      numeric_values = [_NAN] * self._max_seq_length
      if table:
        for col_index in range(len(table.columns)):
          for row_index in range(len(table.rows)):
  
            numeric_value = table.rows[row_index].cells[col_index].numeric_value
            if not numeric_value.HasField('float_value'):
              continue
  
            float_value = numeric_value.float_value
            if float_value == float('inf'):
              beam_metrics.Metrics.counter(
                  _NS, 'cell with numeric value of infinite').inc()
              continue
  
            for index in _get_cell_token_indexes(token_ids_dict['column_ids'],
                                                 token_ids_dict['row_ids'],
                                                 col_index, row_index):
              numeric_values[index] = float_value
      features['numeric_values'] = create_float_feature(numeric_values)
  
    def _add_numeric_values_scale(self, table, token_ids_dict, features):
      """Adds a scale to each token to down weigh the value of long words."""
      numeric_values_scale = [1.0] * self._max_seq_length
      if not table:
        return numeric_values_scale
      for col_index in range(len(table.columns)):
        for row_index in range(len(table.rows)):
          indices = [
              index for index in _get_cell_token_indexes(
                  token_ids_dict['column_ids'], token_ids_dict['row_ids'],
                  col_index, row_index)
          ]
          num_indices = len(indices)
          if num_indices > 1:
            for index in indices:
              numeric_values_scale[index] = float(num_indices)
      features['numeric_values_scale'] = create_float_feature(
          numeric_values_scale)
  
    def _pad_to_seq_length(self, inputs):
      while len(inputs) > self._max_seq_length:
        inputs.pop()
      while len(inputs) < self._max_seq_length:
        inputs.append(0)
  
    def _to_token_ids(self, tokens):
      return self._tokenizer.convert_tokens_to_ids(_get_pieces(tokens))
  
    def _to_features(
        self, tokens, token_ids_dict,
        table,
        question):
      """Produces a dict of TF features."""
      tokens = list(tokens)
      token_ids_dict = {
          key: list(values) for key, values in token_ids_dict.items()
      }
  
      length = len(tokens)
      for values in token_ids_dict.values():
        if len(values) != length:
          raise ValueError('Inconsistent length')
  
      input_ids = self._to_token_ids(tokens)
      input_mask = [1] * len(input_ids)
  
      self._pad_to_seq_length(input_ids)
      self._pad_to_seq_length(input_mask)
      for values in token_ids_dict.values():
        self._pad_to_seq_length(values)
  
      assert len(input_ids) == self._max_seq_length
      assert len(input_mask) == self._max_seq_length
      for values in token_ids_dict.values():
        assert len(values) == self._max_seq_length
  
      features = collections.OrderedDict()
      features['input_ids'] = create_int_feature(input_ids)
      features['input_mask'] = create_int_feature(input_mask)
      for key, values in sorted(token_ids_dict.items()):
        features[key] = create_int_feature(values)
  
      self._add_numeric_column_ranks(token_ids_dict['column_ids'],
                                     token_ids_dict['row_ids'], table, features)
  
      self._add_numeric_relations(question, token_ids_dict['column_ids'],
                                  token_ids_dict['row_ids'], table, features)
  
      self._add_numeric_values(table, token_ids_dict, features)
  
      self._add_numeric_values_scale(table, token_ids_dict, features)
  
      if table:
        features['table_id'] = create_string_feature(
            [table.table_id.encode('utf8')])
        features['table_id_hash'] = create_int_feature(
            [fingerprint(table.table_id) % _MAX_INT])
      return features
  
  
  class ToPretrainingTensorflowExample(ToTensorflowExampleBase):
    """Class for converting pretraining examples."""
  
    def __init__(self, config):
      super(ToPretrainingTensorflowExample, self).__init__(config)
      self._max_predictions_per_seq = config.max_predictions_per_seq
      self._masked_lm_prob = config.masked_lm_prob
      self._min_question_length = config.min_question_length
      self._max_question_length = config.max_question_length
      self._concatenate_snippets = config.concatenate_snippets
      self._always_continue_cells = config.always_continue_cells
      self._question_buckets = [
          self._min_question_length,
          (self._min_question_length + self._max_question_length) / 2,
          self._max_question_length
      ]
      self._vocab_words = list(self._tokenizer.get_vocab())
  
    def _to_example(self, table,
                    instance):
      """Creates TF example from TrainingInstance."""
  
      features = self._to_features(
          instance.tokens, {
              'column_ids': instance.column_ids,
              'prev_label_ids': [0] * len(instance.tokens),
              'row_ids': instance.row_ids,
              'segment_ids': instance.segment_ids,
          },
          table=table,
          question=None)
  
      masked_lm_positions = list(instance.masked_lm_positions)
      masked_lm_ids = self._tokenizer.convert_tokens_to_ids(
          instance.masked_lm_labels)
      masked_lm_weights = [1.0] * len(masked_lm_ids)
  
      while len(masked_lm_positions) < self._max_predictions_per_seq:
        masked_lm_positions.append(0)
        masked_lm_ids.append(0)
        masked_lm_weights.append(0.0)
  
      is_random_table = 1 if instance.is_random_table else 0
  
      features['masked_lm_positions'] = create_int_feature(masked_lm_positions)
      features['masked_lm_ids'] = create_int_feature(masked_lm_ids)
      features['masked_lm_weights'] = create_float_feature(masked_lm_weights)
      features['next_sentence_labels'] = create_int_feature([is_random_table])
      features['is_random_table'] = create_int_feature([is_random_table])
  
      return tf.train.Example(features=tf.train.Features(feature=features))
  
    def convert(
        self,
        rng,
        interaction,
        random_table,
    ):
      """Creates TF example from interaction."""
      question_tokens = self._get_question_tokens(interaction, rng)
  
      if random_table is not None and rng.random() < 0.5:
        is_random_table = True
        table = random_table
      else:
        is_random_table = False
        if interaction.HasField('table'):
          table = interaction.table
        else:
          table = None
  
      if table is None:
        beam_metrics.Metrics.counter(_NS, 'Examples without tables').inc()
        question_tokens = self._tokenizer.tokenize(
            interaction.questions[0].original_text)
        question_tokens = question_tokens[:self._max_seq_length - 1]
        tokens, segment_ids, column_ids, row_ids = self._serialize_text(
            question_tokens)
      else:
        if (not question_tokens or
            len(question_tokens) < self._min_question_length):
          beam_metrics.Metrics.counter(
              _NS,
              f'Remove question below the min length {self._min_question_length}'
          ).inc()
          return None
        beam_metrics.Metrics.counter(_NS, 'Examples with tables').inc()
        beam_metrics.Metrics.counter(
            _NS,
            _get_buckets(
                len(question_tokens), self._question_buckets,
                'Question Length')).inc()
        if random_table is not None:
          logging.log_every_n(logging.INFO,
                              'Table: %s Random Table: %s is_random_table: %s',
                              500000, interaction.table.table_id,
                              random_table.table_id, is_random_table)
  
        token_budget = self._get_token_budget(question_tokens)
        tokenized_table = self._tokenize_table(table)
        try:
          num_columns, num_rows, num_tokens = self._get_table_sizes(
              token_budget, tokenized_table, rng)
        except ValueError:
          return None
  
        serialized_example = self._serialize(question_tokens, tokenized_table,
                                             num_columns, num_rows, num_tokens)
        tokens = serialized_example.tokens
        segment_ids = serialized_example.segment_ids
        row_ids = serialized_example.row_ids
        column_ids = serialized_example.column_ids
  
      assert len(tokens) <= self._max_seq_length
  
      (tokens, masked_lm_positions,
       masked_lm_labels) = self._create_masked_lm_predictions(
           interaction, tokens, column_ids, row_ids, rng)
      instance = TrainingInstance(
          tokens=tokens,
          segment_ids=segment_ids,
          column_ids=column_ids,
          row_ids=row_ids,
          masked_lm_positions=masked_lm_positions,
          masked_lm_labels=masked_lm_labels,
          is_random_table=is_random_table)
      return self._to_example(table, instance)
  
    def _create_masked_lm_predictions(
        self, interaction, tokens,
        column_ids, row_ids,
        rng):
      """Creates the predictions for the masked LM objective."""
      cand_indexes = []
      for (i, token) in enumerate(tokens):
        if token.piece in [_CLS, _SEP]:
          continue
        column_id = column_ids[i]
        is_cell_continutation = column_id > 0 and column_id == column_ids[i - 1]
        if not self._always_continue_cells:
          is_cell_continutation = False
        if cand_indexes and (_is_inner_wordpiece(token) or is_cell_continutation):
          cand_indexes[-1].append(i)
        else:
          cand_indexes.append([i])
  
      rng.shuffle(cand_indexes)
  
      output_tokens = list(tokens)
  
      num_to_predict = min(self._max_predictions_per_seq,
                           max(1, int(round(len(tokens) * self._masked_lm_prob))))
  
      masked_lms = []
      covered_indexes = set()
      for index_set in cand_indexes:
        if len(masked_lms) >= num_to_predict:
          break
        # If adding a whole-word mask would exceed the maximum number of
        # predictions, then just skip this candidate.
        if len(masked_lms) + len(index_set) > num_to_predict:
          continue
  
        for index in index_set:
          assert index not in covered_indexes
          covered_indexes.add(index)
  
          masked_token = None
          # 80% of the time, replace with [MASK]
          if rng.random() < 0.8:
            masked_token = _MASK
          else:
            # 10% of the time, keep original
            if rng.random() < 0.5:
              masked_token = tokens[index].piece
            # 10% of the time, replace with random word
            else:
              masked_token = rng.choice(self._vocab_words)
  
          output_tokens[index] = Token(tokens[index].original_text, masked_token)
  
          masked_lms.append(
              MaskedLmInstance(index=index, label=tokens[index].piece))
      assert len(masked_lms) <= num_to_predict
      masked_lms = sorted(masked_lms, key=lambda x: x.index)
  
      masked_lm_positions = []
      masked_lm_labels = []
      for p in masked_lms:
        masked_lm_positions.append(p.index)
        masked_lm_labels.append(p.label)
  
      return (output_tokens, masked_lm_positions, masked_lm_labels)
  
    def _get_question_tokens(self, interaction,
                             rng):
      """Randomly gets a snippet of relevant text."""
      questions = [q.text for q in interaction.questions]
      rng.shuffle(questions)
      if not self._concatenate_snippets:
        # Find the first snippet that satisfies the requirements.
        for question in questions:
          tokens = self._tokenizer.tokenize(question)
          if len(tokens) > self._max_question_length:
            continue
          if len(tokens) < self._min_question_length:
            continue
          return tokens
        return None
      tokens = []
      for question in questions:
        tokens += self._tokenizer.tokenize(question)
  
      if len(tokens) < self._min_question_length:
        return None
  
      max_start_index = len(tokens) - self._min_question_length
      start_index = rng.randint(0, max_start_index)
      while start_index >= 0 and _is_inner_wordpiece(tokens[start_index]):
        start_index -= 1
  
      min_end_index = start_index + self._min_question_length
      max_end_index = min(len(tokens), self._max_question_length + start_index)
      assert min_end_index <= max_end_index
      end_index = rng.randint(min_end_index, max_end_index)
      assert (self._min_question_length <= end_index - start_index <=
              self._max_question_length)
      while end_index < len(tokens) and _is_inner_wordpiece(tokens[end_index]):
        end_index += 1
  
      return tokens[start_index:end_index]
  
    def _get_table_sizes(self, token_budget, table,
                         rng):
      """Computes column, row and token count for table."""
      num_columns = 1
      num_rows = 1
      num_tokens = 1
      table_cost = self._get_table_cost(table, num_columns, num_rows, num_tokens)
      if table_cost > token_budget:
        raise ValueError('Cannot create table that fits budget')
  
      max_num_rows, max_num_columns, max_num_tokens = self._get_table_boundaries(
          table)
  
      while (num_columns < max_num_columns or num_rows < max_num_rows or
             num_tokens < max_num_tokens):
        if num_columns < max_num_columns and rng.random() < 0.5:
          cost = self._get_table_cost(table, num_columns + 1, num_rows,
                                      num_tokens)
          if cost > token_budget:
            break
          num_columns += 1
        if num_rows < max_num_rows and rng.random() < 0.5:
          cost = self._get_table_cost(table, num_columns, num_rows + 1,
                                      num_tokens)
          if cost > token_budget:
            break
          num_rows += 1
        if num_tokens < max_num_tokens and rng.random() < 0.5:
          cost = self._get_table_cost(table, num_columns, num_rows,
                                      num_tokens + 1)
          if cost > token_budget:
            break
          num_tokens += 1
  
      buckets = [8, 16, 32, 64, 128, 256]
      beam_metrics.Metrics.counter(
          _NS, _get_buckets(num_columns * num_rows, buckets,
                            'Trimmed Table Size')).inc()
  
      # First row is the header.
      real_num_columns = len(table.rows[0]) if table.rows else 0
      # We don't count the header row.
      real_num_rows = len(table.rows) - 1
  
      beam_metrics.Metrics.counter(
          _NS,
          _get_buckets(real_num_columns * real_num_rows, buckets,
                       'Real Table Size')).inc()
      beam_metrics.Metrics.counter(
          _NS, _get_buckets(num_columns, buckets, 'Column Sizes')).inc()
      beam_metrics.Metrics.counter(_NS,
                                   _get_buckets(num_rows, buckets,
                                                'Row Sizes')).inc()
      beam_metrics.Metrics.counter(
          _NS, _get_buckets(num_tokens, buckets, 'Table Token Sizes')).inc()
  
      return num_columns, num_rows, num_tokens
  
  
  class ToTrimmedTensorflowExample(ToTensorflowExampleBase):
    """Helper that allows squeezing a table into the max seq length."""
  
    def __init__(self, config):
      super(ToTrimmedTensorflowExample, self).__init__(config)
      self._cell_trim_length = config.cell_trim_length
  
    def _get_num_columns(self, table):
      num_columns = len(table.columns)
      if num_columns >= self._max_column_id:
        raise ValueError('Too many columns')
      return num_columns
  
    def _get_num_rows(self, table,
                      drop_rows_to_fit):
      num_rows = len(table.rows)
      if num_rows >= self._max_row_id:
        if drop_rows_to_fit:
          num_rows = self._max_row_id - 1
        else:
          raise ValueError('Too many rows')
      return num_rows
  
    def _to_trimmed_features(
        self,
        question,
        table,
        question_tokens,
        tokenized_table,
        num_columns,
        num_rows,
        drop_rows_to_fit = False,
    ):
      """Finds optiomal number of table tokens to include and serializes."""
      init_num_rows = num_rows
      while True:
        num_tokens = self._get_max_num_tokens(
            question_tokens,
            tokenized_table,
            num_rows=num_rows,
            num_columns=num_columns,
        )
        if num_tokens is not None:
          # We could fit the table.
          break
        if not drop_rows_to_fit or num_rows == 0:
          raise ValueError('Sequence too long')
        # Try to drop a row to fit the table.
        num_rows -= 1
      if init_num_rows != num_rows:
        beam_metrics.Metrics.counter(_NS, 'Tables with trimmed rows').inc()
      serialized_example = self._serialize(question_tokens, tokenized_table,
                                           num_columns, num_rows, num_tokens)
  
      assert len(serialized_example.tokens) <= self._max_seq_length
  
      feature_dict = {
          'column_ids': serialized_example.column_ids,
          'row_ids': serialized_example.row_ids,
          'segment_ids': serialized_example.segment_ids,
      }
      features = self._to_features(
          serialized_example.tokens, feature_dict, table=table, question=question)
      return serialized_example, features
  
    def _get_max_num_tokens(
        self,
        question_tokens,
        tokenized_table,
        num_columns,
        num_rows,
    ):
      """Computes max number of tokens that can be squeezed into the budget."""
      token_budget = self._get_token_budget(question_tokens)
      _, _, max_num_tokens = self._get_table_boundaries(tokenized_table)
      if self._cell_trim_length >= 0 and max_num_tokens > self._cell_trim_length:
        max_num_tokens = self._cell_trim_length
      num_tokens = 0
      for num_tokens in range(max_num_tokens + 1):
        cost = self._get_table_cost(tokenized_table, num_columns, num_rows,
                                    num_tokens + 1)
        if cost > token_budget:
          break
      if num_tokens < max_num_tokens:
        if self._cell_trim_length >= 0:
          # We don't allow dynamic trimming if a cell_trim_length is set.
          return None
        if num_tokens == 0:
          return None
        beam_metrics.Metrics.counter(_NS, 'Tables with trimmed cells').inc()
      return num_tokens
  
  
  class ToClassifierTensorflowExample(ToTrimmedTensorflowExample):
    """Class for converting finetuning examples."""
  
    def __init__(self, config):
      super(ToClassifierTensorflowExample, self).__init__(config)
      self._add_aggregation_candidates = config.add_aggregation_candidates
      self._use_document_title = config.use_document_title
      self._use_context_title = config.use_context_title
      self._update_answer_coordinates = config.update_answer_coordinates
      self._drop_rows_to_fit = config.drop_rows_to_fit
      self._trim_question_ids = config.trim_question_ids
      self._expand_entity_descriptions = config.expand_entity_descriptions
      self._use_entity_title = config.use_entity_title
      self._entity_descriptions_sentence_limit = config.entity_descriptions_sentence_limit
      
      # TODO (Chia-Chun)
      self._is_multi_hop = config.is_multi_hop
      self._use_bridge_entity = config.use_bridge_entity
      self._use_question_type = config.use_question_type
  
    def _tokenize_extended_question(
        self,
        question,
        table,
    ):
      """Runs tokenizer over the question text and document title if it's used."""
      # (Chia-Chun): 
      text_tokens = []
      if self._use_question_type:
        question_type_tokens = self._tokenizer.tokenize(question.question_type)
        text_tokens.extend(question_type_tokens)
        text_tokens.append(Token(_SEP, _SEP))
      if self._is_multi_hop:
        hop_tokens = self._tokenizer.tokenize("Hop is " + str(question.hop))
        text_tokens.extend(hop_tokens)
        text_tokens.append(Token(_SEP, _SEP))
  
      if self._use_bridge_entity:
        for bridge_entity in question.bridge_entities:
          bridge_entity_tokens = self._tokenizer.tokenize(bridge_entity)
          text_tokens.extend(bridge_entity_tokens)
          text_tokens.append(Token(_SEP, _SEP))
  
      
  
  
      question_tokens = self._tokenizer.tokenize(question.text)
      text_tokens.extend(list(question_tokens))
      #import pdb;pdb.set_trace()
      if self._use_document_title and table.document_title:
        # TODO(thomasmueller) Consider adding a different segment id.
        document_title_tokens = self._tokenizer.tokenize(table.document_title)
        text_tokens.append(Token(_SEP, _SEP))
        text_tokens.extend(document_title_tokens)
      context_heading = table.context_heading
      if self._use_context_title and context_heading:
        context_title_tokens = self._tokenizer.tokenize(context_heading)
        text_tokens.append(Token(_SEP, _SEP))
        text_tokens.extend(context_title_tokens)
      
  
  
  
      return text_tokens
  
    def convert(self, interaction,
                index):
      """Converts question at 'index' to example."""
      table = interaction.table
  
      num_rows = self._get_num_rows(table, self._drop_rows_to_fit)
      num_columns = self._get_num_columns(table)
  
      question = interaction.questions[index]
      #debug, disable here
      if not interaction.questions[index].answer.is_valid:
        beam_metrics.Metrics.counter(
            _NS, 'Conversion skipped (answer not valid)').inc()
        raise ValueError('Invalid answer')
  
      annotation_descriptions_ext = (
          annotated_text_pb2.AnnotationDescription.annotation_descriptions_ext)
      if (self._expand_entity_descriptions and
          annotation_descriptions_ext in interaction.Extensions):
        descriptions = interaction.Extensions[
            annotation_descriptions_ext].descriptions
        _add_entity_descriptions_to_table(
            question,
            descriptions,
            table,
            use_entity_title=self._use_entity_title,
            num_results=self._entity_descriptions_sentence_limit)
  
      text_tokens = self._tokenize_extended_question(question, table)
      tokenized_table = self._tokenize_table(table)
      table_selection_ext = table_selection_pb2.TableSelection.table_selection_ext
      if table_selection_ext in question.Extensions:
        table_selection = question.Extensions[table_selection_ext]
        if not tokenized_table.selected_tokens:
          raise ValueError('No tokens selected')
        if table_selection.selected_tokens:
          selected_tokens = {(t.row_index, t.column_index, t.token_index)
                             for t in table_selection.selected_tokens}
          tokenized_table.selected_tokens = [
              t for t in tokenized_table.selected_tokens
              if (t.row_index, t.column_index, t.token_index) in selected_tokens
          ]
  
      serialized_example, features = self._to_trimmed_features(
          question=question,
          table=table,
          question_tokens=text_tokens,
          tokenized_table=tokenized_table,
          num_columns=num_columns,
          num_rows=num_rows,
          drop_rows_to_fit=self._drop_rows_to_fit)
  
      column_ids = serialized_example.column_ids
      row_ids = serialized_example.row_ids
  
      def get_answer_ids(question):
        if self._update_answer_coordinates:
          return _find_answer_ids_from_answer_texts(
              column_ids,
              row_ids,
              tokenized_table,
              answer_texts=[
                  self._tokenizer.tokenize(at)
                  for at in question.answer.answer_texts
              ],
          )
        return _get_answer_ids(column_ids, row_ids, question)
  
      answer_ids = get_answer_ids(question)
      self._pad_to_seq_length(answer_ids)
      features['label_ids'] = create_int_feature(answer_ids)
  
      if index > 0:
        prev_answer_ids = get_answer_ids(interaction.questions[index - 1],)
      else:
        prev_answer_ids = [0] * len(column_ids)
      self._pad_to_seq_length(prev_answer_ids)
      features['prev_label_ids'] = create_int_feature(prev_answer_ids)
      features['question_id'] = create_string_feature(
          [question.id.encode('utf8')])
      if self._trim_question_ids:
        question_id = question.id[-text_utils.DEFAULT_INTS_LENGTH:]
      else:
        question_id = question.id
      features['question_id_ints'] = create_int_feature(
          text_utils.str_to_ints(
              question_id, length=text_utils.DEFAULT_INTS_LENGTH))
      features['aggregation_function_id'] = create_int_feature(
          [question.answer.aggregation_function])
      features['classification_class_index'] = create_int_feature(
          [question.answer.class_index])
  
      answer = question.answer.float_value if question.answer.HasField(
          'float_value') else _NAN
      features['answer'] = create_float_feature([answer])
  
      if self._add_aggregation_candidates:
        rng = random.Random(fingerprint(question.id))
  
        candidates = interpretation_utils.find_candidates(rng, table, question)
        num_initial_candidates = len(candidates)
  
        candidates = [c for c in candidates if len(c.rows) < _MAX_NUM_ROWS]
        candidates = candidates[:_MAX_NUM_CANDIDATES]
  
        funs = [0] * _MAX_NUM_CANDIDATES
        sizes = [0] * _MAX_NUM_CANDIDATES
        indexes = []
  
        num_final_candidates = 0
        for index, candidate in enumerate(candidates):
          token_indexes = []
          for row in candidate.rows:
            token_indexes += _get_cell_token_indexes(column_ids, row_ids,
                                                     candidate.column, row)
          if len(indexes) + len(serialized_example.tokens) > _MAX_INDEX_LENGTH:
            break
          num_final_candidates += 1
          sizes[index] = len(token_indexes)
          funs[index] = candidate.agg_function
          indexes += token_indexes
  
        # <int>[1]
        features['cand_num'] = create_int_feature([num_final_candidates])
        # <int>[_MAX_NUM_CANDIDATES]
        features['can_aggregation_function_ids'] = create_int_feature(funs)
        # <int>[_MAX_NUM_CANDIDATES]
        features['can_sizes'] = create_int_feature(sizes)
        # <int>[_MAX_INDEX_LENGTH]
        # Actual length is sum(sizes).
        features['can_indexes'] = create_int_feature(indexes)
  
        if num_initial_candidates > 0:
          beam_metrics.Metrics.counter(
              _NS,
              _get_buckets(num_initial_candidates,
                           [10, 20, 50, 100, 200, 500, 1000, 1200, 1500],
                           'Candidates Size:')).inc()
  
          beam_metrics.Metrics.counter(_NS, 'Candidates: Input').inc()
          if num_final_candidates != num_initial_candidates:
            beam_metrics.Metrics.counter(_NS,
                                         'Candidates: Dropped candidates').inc()
  
      return tf.train.Example(features=tf.train.Features(feature=features))
  
    def get_empty_example(self):
      interaction = interaction_pb2.Interaction(questions=[
          interaction_pb2.Question(id=text_utils.get_padded_question_id())
      ])
      return self.convert(interaction, index=0)
  
apis: tensorflow.compat.v1.train.Features tensorflow.compat.v1.io.gfile.GFile