def read_classifier_dataset(
    predict_data,
    data_format,
    compression_type,
    max_seq_length,
    max_predictions_per_seq,
    add_aggregation_function_id,
    add_classification_labels,
    add_answer,
):
def _get_in_memory_input_fn(examples):
  def _input_fn(params):
def compute_prediction_sequence(estimator, examples_by_position):
def _get_question_id(features):
def get_cell_token_probs(prediction):
def get_mean_cell_probs(prediction):
def _geometric_mean(values):
  def from_text(cls, json_str):
def _to_token_answer(
    prediction,
    begin_token_index,
    end_token_index,
    score,
):
def _get_token_answers(
    prediction,
    cell_classification_threshold,
):
def _write_prediction(
    prediction,
    cell_classification_threshold,
    output_token_probabilities,
    do_model_aggregation,
    do_model_classification,
    output_token_answers,
    writer,
):
def token_answers_to_text(answers):
def token_answers_from_text(text):
def write_json(predictions,
               output_predict_file):
def write_predictions(
    predictions,
    output_predict_file,
    do_model_aggregation,
    do_model_classification,
    cell_classification_threshold,
    output_token_probabilities,
    output_token_answers,
):
class TokenAnswer: _EPSILON_ZERO_DIVISION = utils.EPSILON_ZERO_DIVISION
_CLOSE_ENOUGH_TO_LOG_ZERO = utils.CLOSE_ENOUGH_TO_LOG_ZERO
  RATIO = "ratio"
  FIRST_ORDER = "first_order"
  SECOND_ORDER = "second_order"
class AverageApproximationFunction(str, enum.Enum):
class TapasClassifierConfig:
    class EnhancedJSONEncoder(json.JSONEncoder):
class Outputs:
  def to_json_string(self):
      def default(self, o):
  def to_json_file(self, json_file):
  def from_dict(cls, json_object, for_prediction=False):
  def from_json_file(cls, json_file, for_prediction=False):
def _get_probs(dist):
def _calculate_aggregation_logits(output_layer_aggregation, output_weights_agg,
                                  output_bias_agg):
def _calculate_aggregation_loss_known(logits_aggregation, aggregate_mask,
                                      aggregation_function_id,
                                      config):
def _calculate_aggregation_loss_unknown(logits_aggregation, aggregate_mask):
def _calculate_aggregation_loss(logits_aggregation, aggregate_mask,
                                aggregation_function_id,
                                config):
def _calculate_expected_result(dist_per_cell, numeric_values,
                               numeric_values_scale, input_mask_float,
                               logits_aggregation,
                               config):
def _calculate_regression_loss(answer, aggregate_mask, dist_per_cell,
                               numeric_values, numeric_values_scale,
                               input_mask_float, logits_aggregation,
                               config):
def _calculate_aggregate_mask(answer, output_layer_aggregation, output_bias_agg,
                              output_weights_agg, cell_select_pref, label_ids):
def compute_classification_logits(num_classification_labels, output_layer):
def _single_column_cell_selection_loss(token_logits, column_logits, label_ids,
                                       cell_index, col_index, cell_mask):
def _get_classification_outputs(
    config,
    is_training,
    output_layer,
    output_layer_aggregation,
    label_ids,
    input_mask,
    table_mask,
    aggregation_function_id,
    answer,
    numeric_values,
    numeric_values_scale,
    row_ids,
    column_ids,
    classification_class_index,
):
def _calculate_eval_metrics_fn(
    loss,
    label_ids,
    logits,
    input_mask,
    aggregation_function_id,
    logits_aggregation,
    classification_class_index,
    logits_cls,
    table_pruning_loss,
):
def model_fn_builder(
    config,
    custom_prediction_keys = None,
):
  def model_fn(features, labels, mode, params):
    def add_init_checkpoint(init_checkpoint, scope=None):
        def tpu_scaffold():
def input_fn(
    name,
    file_patterns,
    data_format,
    compression_type,
    is_training,
    max_seq_length,
    max_predictions_per_seq,
    add_aggregation_function_id,
    add_classification_labels,
    add_answer,
    include_id,
    params,
): class BertConfig(object):
class BertModel(object):
  def __init__(self,
               vocab_size,
               hidden_size=768,
               num_hidden_layers=12,
               num_attention_heads=12,
               intermediate_size=3072,
               hidden_act="gelu",
               hidden_dropout_prob=0.1,
               attention_probs_dropout_prob=0.1,
               max_position_embeddings=512,
               type_vocab_size=16,
               initializer_range=0.02,
               softmax_temperature=1.0):
  def from_dict(cls, json_object):
  def from_json_file(cls, json_file):
  def to_dict(self):
  def to_json_string(self):
  def to_json_file(self, json_file):
  def __init__(self,
               config,
               is_training,
               input_ids,
               input_mask=None,
               attention_mask=None,
               token_weights=None,
               custom_attention_layer=None,
               custom_transformer_layer=None,
               token_type_ids=None,
               disabled_ids=None,
               extra_embeddings=None,
               use_position_embeddings=True,
               reset_position_index_per_cell=False,
               proj_value_length=None,
               scope=None):
  def get_pooled_output(self):
  def get_sequence_output(self):
  def get_all_encoder_layers(self):
  def get_all_attention_probs(self):
  def get_embedding_output(self):
  def get_embedding_table(self):
def gelu(x):
def get_activation(activation_string):
def get_assignment_map_from_checkpoint(tvars, init_checkpoint, scope=None):
def dropout(input_tensor, dropout_prob):
def layer_norm(input_tensor, name=None):
def layer_norm_and_dropout(input_tensor, dropout_prob, name=None):
def create_initializer(initializer_range=0.02):
def embedding_lookup(input_ids,
                     vocab_size,
                     embedding_size=128,
                     initializer_range=0.02,
                     word_embedding_name="word_embeddings"):
def _get_absolute_position_embeddings(full_position_embeddings, seq_length,
                                      width, num_dims):
def _get_relative_position_embeddings(
    full_position_embeddings,
    token_type_ids,
    token_type_vocab_size,
    seq_length,
    batch_size,
    max_position_embeddings,
):
def embedding_postprocessor(input_tensor,
                            use_token_type=False,
                            token_type_ids=None,
                            disabled_ids=None,
                            token_type_vocab_size=None,
                            token_type_embedding_name="token_type_embeddings",
                            use_position_embeddings=True,
                            reset_position_index_per_cell=False,
                            position_embedding_name="position_embeddings",
                            initializer_range=0.02,
                            max_position_embeddings=512,
                            extra_embeddings=None,
                            dropout_prob=0.1):
def create_attention_mask_from_input_mask(
    from_tensor,
    to_mask,
):
def dense_layer_3d(input_tensor,
                   num_attention_heads,
                   size_per_head,
                   initializer,
                   activation,
                   name=None):
def dense_layer_3d_proj(input_tensor,
                        hidden_size,
                        num_attention_heads,
                        head_size,
                        initializer,
                        activation,
                        name=None):
def dense_layer_2d(input_tensor,
                   output_size,
                   initializer,
                   activation,
                   name=None):
def attention_layer(from_tensor,
                    to_tensor,
                    attention_mask=None,
                    input_mask=None,
                    num_attention_heads=1,
                    size_per_head=512,
                    query_act=None,
                    key_act=None,
                    value_act=None,
                    attention_probs_dropout_prob=0.0,
                    initializer_range=0.02,
                    softmax_temperature=1.0,
                    batch_size=None,
                    from_seq_length=None,
                    to_seq_length=None,
                    to_proj_length=None):
def transformer_model(input_tensor,
                      attention_mask=None,
                      input_mask=None,
                      custom_attention_layer=None,
                      hidden_size=768,
                      num_hidden_layers=12,
                      num_attention_heads=12,
                      intermediate_size=3072,
                      intermediate_act_fn=gelu,
                      hidden_dropout_prob=0.1,
                      attention_probs_dropout_prob=0.1,
                      initializer_range=0.02,
                      softmax_temperature=1.0,
                      do_return_all_layers=False,
                      do_return_attention_probs=False,
                      proj_value_length=None):
def get_shape_list(tensor, expected_rank=None, name=None):
def reshape_to_matrix(input_tensor):
def reshape_from_matrix(output_tensor, orig_shape_list):
def assert_rank(tensor, expected_rank, name=None): class EvalResult:
class ScoredCandidate:
class Wordpiece:
class DeTokenizer:
  def to_dict(self):
def _find(subpieces, pieces):
  def __init__(self, vocab_file, do_lower_case=True, split_on_punc=True):
  def _tokenize_text(self, text):
  def _convert_ids_to_text(
      self,
      token_ids,
  ):
  def _get_unique_texts(
      self,
      table,
  ):
  def _detokenize_text(
      self,
      text,
      token_ids,
  ):
  def detokenize(
      self,
      table,
      pieces,
  ):
  def detokenize_token_answer(
      self,
      table,
      token_answer,
  ):
def _get_example_id(question_id):
def _guess_vocab_file(interaction_file):
def _is_valid_answer(token_ids):
def _get_reconstructed_answer_text(
    detokenizer,
    table,
    prediction,
):
def _get_scored_candidates(
    detokenizer,
    interactions,
    predictions_by_qid,
):
def _is_correct_table(interaction):
def _list_to_answer(texts):
def _get_answer_texts(
    interactions):
def _remove_articles(text):
def _fix_whitespace(text):
def _remove_punctuation(text):
def _lower(text):
def _normalize_answer(s):
def _tokenize(s):
def _compute_exact(a_gold, a_pred):
def _compute_f1(a_gold, a_pred):
def _mean(values):
def _get_best_metrics(
    reference_answer_texts,
    best_answer_text,
):
def _get_question(
    candidates,):
def _get_oracle_metrics(
    candidates,
    reference_answer_texts,
    metric,
):
def is_table_answer(question_id,
                    references):
def is_passage_answer(question_id,
                      references):
def _evaluate_retrieval_e2e(
    vocab_file,
    interaction_list,
    predictions,
    references = None):
def evaluate_retrieval_e2e(
    interaction_file,
    prediction_file,
    references_file = None,
    vocab_file = None,
):
def generate_hybridqa_codalab_predictions(
    interaction_file,
    prediction_file): class Example:
class StructuredMetrics:
class DenotationResult:
class DenotationStats:
def write_to_tensorboard(
    metrics,
    global_step,
    logdir,
):
def read_data_examples_from_interactions(
    interactions_path):
def example_from_question(
    interaction,
    question,
):
def read_predictions(predictions_path, examples):
def _calc_acc(correct):
def calc_structure_metrics(
    examples,
    denotation_errors_path = None):
def _collect_cells_from_table(cell_coos,
                              table):
def _safe_convert_to_float(value):
def _parse_value(value):
def _to_float32s(elements):
def execute(aggregation_type, cell_coos,
            table):
def _highlight_cells(coordinates, table):
def _get_debug_row(result, table):
def _get_gold_denotation_result(example):
def _get_pred_denotation_result(example):
def get_denotation_stats(example):
def calc_weighted_denotation_accuracy(examples,
                                      denotation_errors_path,
                                      predictions_file_name,
                                      add_weights):
def calc_denotation_accuracy(examples,
                             denotation_errors_path,
                             predictions_file_name):
def calc_classification_accuracy(examples): def parse_coordinates(raw_coordinates):
def iterate_predictions(
    prediction_file):
def is_tfrecord(filename):
def iterate_interactions(
    interactions_file):
def parse_interaction_id(text):
def table_to_panda_frame(table): def make_directories(path):
def list_directory(path): def get_sqa_hparams():
def get_wtq_hparams():
def get_wikisql_super_hparams():
def get_wikisql_hparams():
def get_hybridqa_hparams():
def get_hybridqa_rc_hparams():
def get_tabfact_hparams():
def get_nq_e2e_hparams():
def get_hparams(task): _DATE_TUPLE_SIZE = 3
def _get_value_type(numeric_value):
def _get_value_as_primitive_value(
    numeric_value):
def _get_all_types(
    numeric_values):
def get_numeric_sort_key_fn(
    numeric_values):
  def _sort_key_fn(numeric_value):
def _consolidate_numeric_values(
    row_index_to_values,
    min_consolidation_fraction,
    debug_info):
def _get_numeric_values(text):
def _get_column_values(
    table,
    col_index):
def get_numeric_relation(
    value,
    other_value,
    sort_key_fn):
def add_numeric_table_values(table,
                             min_consolidation_fraction=0.7,
                             debug_info = None):
def add_numeric_values_to_questions(interaction
                                   ):
def add_numeric_values(interaction): _NS = "main"
  COLUMN = 1
  CELL = 2
def get_column_texts(
    table,
    column_index,
):
def _get_column_cost(
    tokenizer,
    table,
    column_index,
):
def _get_question_cost(
    tokenizer,
    question,
):
def get_cleaned_seq_tokens(str_tokens):
def from_selected_columns_to_selected_tokens(
    interaction,
    selected_columns,
    tokenizer):
  def _add_cell_tokens(column_index, row_index, cell_text, selected_tokens):
def _get_question_column_similarity(column_tokens,
                                    question_tokens,
                                    epsilon = 0.0):
  def __init__(self, vocab_file, use_previous_questions):
  def prepare(self):
  def select_tokens(
      self,
      interaction,
      question,
  ):
  def annotated_interaction(
      self,
      interaction,
  ):
  def __init__(self, vocab_file, max_nb_tokens,
               selection_level, use_previous_answer,
               use_previous_questions):
  def _get_column_tokens(self, interaction,
                         column_index):
  def _get_columns_tokens(
      self, interaction):
  def _get_columns_similarity_scores(
      self, interaction,
      question_tokens):
  def _get_headers_tokens(
      self, interaction):
  def _get_headers_similarity_scores(
      self, interaction,
      question_tokens):
  def _get_cells(
      self, interaction
  ):
  def _get_row_tokens(self, interaction,
                      row_index):
  def _get_header_tokens(
      self, interaction):
  def _get_row_similarity_scores(
      self, interaction,
      question_tokens):
  def _update_priority_queue_from_previous_answer(
      self,
      selection_type,
      priority_queue,
      interaction,
      question,
  ):
  def _get_columns_from_priority_queue(
      self,
      columns_queue,
      interaction,
      question,
  ):
  def _get_tokens_from_priority_queue(
      self,
      tokens_queue,
      interaction,
      question,
  ):
  def _select_columns(
      self,
      interaction,
      question,
  ):
  def _select_cells(
      self,
      interaction,
      question,
  ):
  def _select_columns(
      self,
      interaction,
      question,
  ):
  def _select_cells(
      self,
      interaction,
      question,
  ):
  def select_tokens(self, interaction,
                    question):
class SelectionType(enum.Enum):
class TokenScoreDebugInfo:
class TokenSelection:
class TokenSelector:
class MaxTokenSelector(TokenSelector):
class HeuristicExactMatchTokenSelector(MaxTokenSelector):
class SelectAllTokensFn(TokenSelector): def get_interaction_dir(output_dir):
def _to_tfrecord(
    interactions,
    output_dir,
    token_selector,
):
def get_train_filename(task):
def get_dev_filename(task):
def get_test_filename(task):
def get_supervision_modes(task):
def create_interactions(
    task,
    input_dir,
    output_dir,
    token_selector,
):  # pylint: disable=g-doc-args
  def to_tfrecord(
      interactions): class Task(enum.Enum):
  SQA = 0
  WTQ = 1
  WIKISQL = 2
  WIKISQL_SUPERVISED = 3
  TABFACT = 4
  NQ_RETRIEVAL = 7
  HYBRIDQA_RC = 8  # Reading comprehension (RC) of HybridQA.
  HYBRIDQA_E2E = 9  # HybridQA RC using preds of HybridQA Cell Selection.
  HYBRIDQA = 6
  SEM_TAB_FACT = 10  # https://sites.google.com/corp/view/sem-tab-facts _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
class Token:
class TrainingInstance:
class TokenCoordinates:
class TokenizedTable:
class MaskedLmInstance:
class ConversionConfig:
class PretrainConversionConfig(ConversionConfig):
class TrimmedConversionConfig(ConversionConfig):
class ClassifierConversionConfig(TrimmedConversionConfig):
class RetrievalConversionConfig(TrimmedConversionConfig):
class SerializedExample:
class TapasTokenizer:
class ToTensorflowExampleBase:
class ToPretrainingTensorflowExample(ToTensorflowExampleBase):
class ToTrimmedTensorflowExample(ToTensorflowExampleBase):
class ToClassifierTensorflowExample(ToTrimmedTensorflowExample):
def copy_vocab(input_vocab, output_vocab):
def _get_pieces(tokens):
def fingerprint(text):
def create_int_feature(values):
def create_float_feature(values):
def create_string_feature(values):
def _is_inner_wordpiece(token):
def _get_cell_token_indexes(column_ids, row_ids,
                            column_id, row_id):
def _get_buckets(value, buckets, name):
def _get_all_answer_ids_from_coordinates(
    column_ids,
    row_ids,
    answers_list,
):
def _get_all_answer_ids(
    column_ids,
    row_ids,
    questions,
):
  def _to_coordinates(
      question,):
def _find_tokens(text, segment):
def _find_answer_coordinates_from_answer_text(
    tokenized_table,
    answer_text,
):
def _find_answer_ids_from_answer_texts(
    column_ids,
    row_ids,
    tokenized_table,
    answer_texts,
):
def _get_answer_ids(column_ids, row_ids,
                    question):
def _get_annotation_name(identifier):
def _add_entity_descriptions_to_table(
    question,
    descriptions,
    table,
    use_entity_title,
    num_results,
):
  def __init__(self, vocab_file):
  def get_vocab(self):
  def tokenize(self, text):
  def convert_tokens_to_ids(self, word_pieces):
  def question_encoding_cost(self, question_tokens):
  def __init__(self, config):
  def _tokenize_table(
      self,
      table,
  ):
  def _get_table_values(self, table, num_columns,
                        num_rows,
                        num_tokens):
  def _serialize_text(
      self, question_tokens
  ):
  def _serialize(
      self,
      question_tokens,
      table,
      num_columns,
      num_rows,
      num_tokens,
  ):
  def _tokenize(self, text):
  def _get_token_budget(self, question_tokens):
  def _get_table_boundaries(self,
                            table):
  def _get_table_cost(self, table, num_columns,
                      num_rows, num_tokens):
  def _get_column_values(
      self, table,
      col_index):
  def _add_numeric_column_ranks(self, column_ids, row_ids,
                                table,
                                features):
  def _get_numeric_sort_key_fn(self, table_numeric_values, value):
  def _add_numeric_relations(self, question,
                             column_ids, row_ids,
                             table,
                             features):
  def _add_numeric_values(self, table,
                          token_ids_dict,
                          features):
  def _add_numeric_values_scale(self, table, token_ids_dict, features):
  def _pad_to_seq_length(self,
                         inputs,
                         override_max_seq_length = None):
  def _to_token_ids(self, tokens):
  def _to_features(
      self, tokens, token_ids_dict,
      table,
      question):
  def __init__(self, config):
  def _to_example(self, table,
                  instance):
  def convert(
      self,
      rng,
      interaction,
      random_table,
  ):
  def _create_masked_lm_predictions(
      self, interaction, tokens,
      column_ids, row_ids,
      rng):
  def _get_question_tokens(self, interaction,
                           rng):
  def _get_table_sizes(self, token_budget, table,
                       rng):
  def __init__(self, config):
  def _get_num_columns(self, table):
  def _get_num_rows(self, table,
                    drop_rows_to_fit):
  def _to_trimmed_features(
      self,
      question,
      table,
      question_tokens,
      tokenized_table,
      num_columns,
      num_rows,
      drop_rows_to_fit = False,
  ):
  def _get_max_num_tokens(
      self,
      question_tokens,
      tokenized_table,
      num_columns,
      num_rows,
  ):
  def __init__(self, config):
  def _tokenize_extended_question(
      self,
      question,
      table,
  ):
  def convert(self, interaction,
              index):
    def get_answer_ids(question):
  def get_empty_example(self):
import dataclasses
import enum
import functools
import os
import random
import time
import tensorflow.compat.v1 as tf
from typing import Mapping, Optional, Text
from absl import app
from absl import flags
from absl import logging
from tapas.experiments import prediction_utils as exp_prediction_utils
from tapas.models import tapas_classifier_model
from tapas.models.bert import modeling
from tapas.retrieval import e2e_eval_utils
from tapas.scripts import calc_metrics_utils
from tapas.scripts import prediction_utils
from tapas.utils import file_utils
from tapas.utils import hparam_utils
from tapas.utils import number_annotation_utils
from tapas.utils import pruning_utils
from tapas.utils import task_utils
from tapas.utils import tasks
from tapas.utils import tf_example_utils
  if batch_size is None:
    random.shuffle(examples)
  else:
    # Make sure the eval sets are divisible by the test batch size since
    # otherwise examples will be dropped on TPU.
    # These examples will later be ignored when writing the predictions.
    originial_num_examples = len(examples)
    while len(examples) % batch_size != 0:
      examples.append(converter.get_empty_example())
    if originial_num_examples != len(examples):
      _print(f'Padded with {len(examples) - originial_num_examples} examples.')

  with tf.io.TFRecordWriter(
      example_path,
      options=_to_tf_compression_type(FLAGS.compression_type),
  ) as writer:
    for example in examples:
      writer.write(example.SerializeToString())


def _get_train_examples_file(task, output_dir):
  return os.path.join(output_dir, 'tf_examples',
                      f'{task_utils.get_train_filename(task)}.tfrecord')


def _get_test_filename(task, test_set):
  if test_set == TestSet.TEST:
    return task_utils.get_test_filename(task)
  if test_set == TestSet.DEV:
    return task_utils.get_dev_filename(task)
  raise ValueError(f'Unknown test set: {test_set}')


def _get_test_examples_file(
    task,
    output_dir,
    test_set,
):
  filename = _get_test_filename(task, test_set)
  return os.path.join(output_dir, 'tf_examples', f'{filename}.tfrecord')


def _get_test_interactions_file(
    task,
    output_dir,
    test_set,
):
  filename = _get_test_filename(task, test_set)
  return os.path.join(output_dir, 'interactions', f'{filename}.tfrecord')


def _get_test_prediction_file(
    task,
    model_dir,
    test_set,
    is_sequence,
    global_step,
):
  """Get prediction filename for different tasks and setups."""
  suffix = '' if global_step is None else f'_{global_step}'
  if is_sequence:
    suffix = f'_sequence{suffix}'
  filename = _get_test_filename(task, test_set)
  return os.path.join(model_dir, f'{filename}{suffix}.tsv')


def _get_token_selector():
  if not FLAGS.prune_columns:
    return None
  return pruning_utils.HeuristicExactMatchTokenSelector(
      FLAGS.bert_vocab_file,
      FLAGS.max_seq_length,
      pruning_utils.SelectionType.COLUMN,
      # Only relevant for SQA where questions come in sequence
      use_previous_answer=True,
      use_previous_questions=True,
  )


def _train_and_predict(
    task,
    tpu_options,
    test_batch_size,
    train_batch_size,
    gradient_accumulation_steps,
    bert_config_file,
    init_checkpoint,
    test_mode,
    mode,
    output_dir,
    model_dir,
    loop_predict,
):
  """Trains, produces test predictions and eval metric."""
  file_utils.make_directories(model_dir)

  if task in (tasks.Task.SQA, tasks.Task.HYBRIDQA, tasks.Task.HYBRIDQA_RC):
    num_aggregation_labels = 0
    num_classification_labels = 0
    use_answer_as_supervision = False
  elif task in [
      tasks.Task.WTQ, tasks.Task.WIKISQL, tasks.Task.WIKISQL_SUPERVISED
  ]:
    num_aggregation_labels = 4
    num_classification_labels = 0
    use_answer_as_supervision = task != tasks.Task.WIKISQL_SUPERVISED
  elif task == tasks.Task.TABFACT:
    num_classification_labels = 2
    num_aggregation_labels = 0
    use_answer_as_supervision = True
  elif task == tasks.Task.NQ_RETRIEVAL:
    num_aggregation_labels = 0
    num_classification_labels = 2
    use_answer_as_supervision = False
  else:
    raise ValueError(f'Unknown task: {task.name}')

  do_model_aggregation = num_aggregation_labels > 0
  do_model_classification = num_classification_labels > 0


  hparams = hparam_utils.get_hparams(task)