def read_dataset(
    parse_examples_fn,
    name,
    file_patterns,
    data_format,
    compression_type,
    is_training,
    params,
):
    def fetch_dataset(filename):
def build_parser_function(feature_types,
                          params):
  def parse_examples(serialized_examples): class TableTask(enum.Enum):
  CLASSIFICATION = 0
  PRETRAINING = 1
  RETRIEVAL = 2
  RETRIEVAL_NEGATIVES = 3
def parse_table_examples(
    max_seq_length,
    max_predictions_per_seq,
    task_type,
    add_aggregation_function_id,
    add_classification_labels,
    add_answer,
    include_id,
    add_candidate_answers,
    max_num_candidates,
    params,
):
  def _parse_fn(serialized_example):
def _preprocess_candidate_answers(features, max_num_candidates, max_seq_length): class IndexMap(object):
class ProductIndexMap(IndexMap):
  def __init__(self, indices, num_segments, batch_dims=0):
  def batch_shape(self):
  def __init__(self, outer_index, inner_index):
  def project_outer(self, index):
  def project_inner(self, index):
def gather(values, index, name='segmented_gather'):
def flatten(index, name='segmented_flatten'):
def range_index_map(batch_shape, num_segments, name='range_index_map'):
def _segment_reduce(values, index, segment_reduce_fn, name):
def reduce_mean(values, index, name='segmented_reduce_mean'):
def reduce_sum(values, index, name='segmented_reduce_sum'):
def reduce_max(values, index, name='segmented_reduce_max'):
def reduce_min(values, index, name='segmented_reduce_min'): EPSILON_ZERO_DIVISION = 1e-10
CLOSE_ENOUGH_TO_LOG_ZERO = -10000.0
def classification_initializer():
def extract_answer_from_features(
    features, use_answer_as_supervision
):
def compute_token_logits(output_layer, temperature,
                         init_cell_selection_weights_to_zero):
def compute_column_logits(output_layer,
                          cell_index,
                          cell_mask,
                          init_cell_selection_weights_to_zero,
                          allow_empty_column_selection): 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): def create_optimizer(loss,
                     init_lr,
                     num_train_steps,
                     num_warmup_steps,
                     use_tpu,
                     optimizer="adamw",
                     poly_power=1.0,
                     start_warmup_step=0,
                     gradient_accumulation_steps=1,
                     grad_clipping=None):
  def __init__(self,
               opt,
               steps,
               grad_clipping = None,
               name="GradientAccumulationOptimizer"):
  def _get_variable_name(self, param_name):
  def apply_gradients(self, grads_and_vars, global_step=None, name=None):
    def _apply_and_zero():
    def _accum():
  def __init__(self,
               learning_rate,
               weight_decay_rate=0.0,
               beta_1=0.9,
               beta_2=0.999,
               epsilon=1e-6,
               exclude_from_weight_decay=None,
               grad_clipping=None,
               name="AdamWeightDecayOptimizer"):
  def apply_gradients(self, grads_and_vars, global_step=None, name=None):
  def _do_use_weight_decay(self, param_name):
  def _get_variable_name(self, param_name):
class GradientAccumulationOptimizer(tf.train.Optimizer):
class AdamWeightDecayOptimizer(tf.train.Optimizer): class CustomAttention:
class CustomAttentionConfig:
def get_token_type_features():
def get_token_type_ids(features, token_type_features, disabled_features):
def run_custom_attention(
    is_training,
    num_attention_heads,
    config,
    features,
):
def create_model(
    features,
    mode,
    bert_config,
    restrict_attention_mode=_AttentionMode.FULL,
    restrict_attention_bucket_size=0,
    restrict_attention_header_size=None,
    restrict_attention_row_heads_ratio=0.5,
    restrict_attention_sort_after_projection=True,
    token_weights=None,
    disabled_features=None,
    disable_position_embeddings=False,
    reset_position_index_per_cell=False,
    proj_value_length=None,
): _INFINITY = 10000  # Same value used in BERT
  FULL = 'full'
  SAME_COLUMN_OR_ROW = 'same_colum_or_row'
  HEADWISE_SAME_COLUMN_OR_ROW = 'headwise_same_colum_or_row'
  HEADWISE_EFFICIENT = 'headwise_efficient'
class RestrictAttentionMode(str, enum.Enum):
class _BucketedTensor():
def _matches_token_type_id(tensor):
def compute_sparse_attention_mask(segment_ids, column_ids, row_ids, input_mask,
                                  **_):
def _compute_bucket_id(bucket_size, header_size, token_type_id,
                       input_mask):
def _compute_cross_bucket_attention_mask(bucket_size, header_size,
                                         token_type_id, input_mask):
def compute_headwise_sparse_attention_mask(num_row_heads,
                                           num_column_heads,
                                           bucket_size, header_size,
                                           segment_ids, column_ids, row_ids,
                                           input_mask, **_):
  def expand_from(self, view):
  def expand_to(self, view):
def _create_bucketed_tensor(tensor,
                            name,
                            order_indices,
                            bucket_size,
                            header_size,
                            num_heads,
                            sort_after_projection,
                            embedding_size = None):
    def _rotate(t, offset):
def _additive_mask(mask):
def _dense_layer_3d(input_tensor, activation,
                    name, num_heads, size_per_head,
                    initializer_range, sort_after_projection):
def _dropout(tensor, dropout_prob):
def _compute_bucketed_attention_mask(from_view, to_view,
                                     ids,
                                     attend_always):
def _attention_scores(from_view, to_view, additive_mask,
                      queries, keys):
def create_bucketed_attention_layer(
    input_mask, input_header, bucket_size, header_size,
    token_type_ids,
    sort_after_projection):
  def _gather_and_repeat(tensors, name):
def _bucketed_attention_layer(
    from_tensor,
    to_tensor,
    bucketed_tensor_args,
    header_size,
    gather_and_repeat,
    sort_after_projection,
    type_ranks,
    head_full_mask,
    tail_head_mask,
    tail_window_mask,
    num_attention_heads,
    size_per_head = 512,
    query_act = None,
    key_act = None,
    value_act = None,
    attention_mask=None,
    input_mask=None,
    attention_probs_dropout_prob = 0.0,
    initializer_range = 0.02,
    softmax_temperature = 1.0,
    to_proj_length = None): _NEG_INFINITY = -10000.0
  NONE = "none"
  BOUNDARY = "boundary"
  SPAN = "span"
class SpanPredictionMode(str, enum.Enum):
def _get_indexes(seq_length, batch_size):
def _get_valid_indexes(
    indexes,
    offset,
):
def _gather_indexes(
    indexes,
    flat_spans_2d,
    start_or_end,
):
def _get_mask(condition, dtype=tf.float32):
def _get_id_mask(
    ids,
    flat_spans_2d,
    dtype=tf.float32,
):
def _get_start_and_end_ids(label_ids):
def compute_token_span_logits(
    embeddings,
    name,
):
def _get_spans(
    seq_length,
    batch_size,
    max_span_length,
):
def _get_span_mask(column_ids, row_ids, flat_spans_2d):
def _mask_spans(span_mask, span_logits):
def _get_span_logits(
    embeddings,
    start_ids,
    end_ids,
    column_ids,
    row_ids,
    max_span_length,
):
  def _get_loss(masked_span_logits):
def get_span_logits(
    embeddings,
    label_ids,
    column_ids,
    row_ids,
    max_span_length,
):
def _get_span_loss(logits, label_ids):
def get_boundary_logits(
    embeddings,
    label_ids,
    column_ids,
    row_ids,
    max_span_length,
):
def get_span_logits_by_mode(
    mode, embeddings, label_ids, column_ids, row_ids,
    max_span_length):
def compute_masked_example_loss(
    label_ids,
    per_example_loss,
): PRUNING_SCOPE = "pruning"
_SEQUENCE_OUTPUT_KEEP_PROB = 0.9
class Scores:
class TablePruningSelector:
class NoTablePruning(TablePruningSelector):
class ModelPruningSelector(TablePruningSelector):
class OnesTablePruning(ModelPruningSelector):
class TapasPruningSelector(ModelPruningSelector):
class AverageCosineSimilaritySelector(ModelPruningSelector):
class LossSelector:
class Unsupervised(LossSelector):
class HardSelection:
class NoHardSelection(HardSelection):
class MaskTopK(HardSelection):
class TopK(MaskTopK):
  def __init__(self, config, max_num_tokens, max_num_columns):
  def select_columns(self, mode, features):
  def _select_columns(self, mode, features):
  def compute_loss(self, required_columns, column_scores,
                   column_score_mask,
                   token_scores):
  def apply_hard_selection(self, mode, scores):
  def compute_scores(self, mode, features):
  def get_sequence_output(self):
  def get_pooled_output(self):
  def should_add_classification_loss(self):
  def __init__(self):
  def _select_columns(self, mode, features):
  def _init_create_model_selector(self, config, max_num_tokens):
  def __init__(self, config, max_num_columns, max_num_tokens):
  def compute_loss(self, required_columns, column_scores,
                   column_score_mask,
                   token_scores):
  def apply_hard_selection(self, mode, scores):
  def _gather_op(
      self,
  ):
    def gather_top_k(features, scores):
  def _get_index_gather(self, scores,
                        input_mask):
  def _gather_scores(self, scores,
                     input_mask):
  def compute_scores(self, mode, features):
  def _compute_columns_stats(self, column_scores,
                             column_ids):
  def _computes_column_and_token_scores(
      self, mode, features):
  def should_add_classification_loss(self):
  def _init_create_model_selector(self, config, max_num_tokens):
  def __init__(self, config, max_num_columns,
               max_num_tokens):
  def _select_columns(self, mode, features):
  def apply_hard_selection(self, mode, scores):
  def compute_loss(self, required_columns, column_scores,
                   column_score_mask,
                   token_scores):
  def should_add_classification_loss(self):
  def __init__(self, config, max_num_columns, max_num_rows,
               max_num_tokens):
  def _compute_column_scores_from_token_scores(
      self, mode, output_layer,
      features):
  def _select_columns(self, mode, features):
  def _compute_token_scores(self, mode,
                            features):
  def _computes_column_and_token_scores(
      self, mode, features):
  def get_sequence_output(self):
  def get_pooled_output(self):
  def should_add_classification_loss(self):
  def __init__(self, vocab_size, hidden_size,
               initializer_range, max_num_columns,
               type_vocab_size, disabled_features,
               disable_position_embeddings, max_position_embeddings,
               config, max_num_tokens):
  def _select_columns(self, mode, features):
  def should_add_classification_loss(self):
def create_selector(
    table_pruning_config_file,
    vocab_size,
    hidden_size,
    initializer_range,
    max_num_columns,
    max_num_rows,
    type_vocab_size,
    disabled_features,
    disable_position_embeddings,
    max_position_embeddings,
):
  def __init__(self, config, max_num_tokens):
  def compute_loss(self, required_columns, column_scores,
                   column_score_mask,
                   token_scores):
  def apply_hard_selection(self, mode, scores):
  def should_add_classification_loss(self):
  def __init__(self, config, max_num_tokens):
  def compute_loss(self, required_columns, column_scores,
                   column_score_mask,
                   token_scores):
def create_loss_selector(config, max_num_tokens):
  def apply_hard_selection(self, scores):
  def apply_hard_selection(self, scores):
  def __init__(self, max_num_tokens):
  def apply_hard_selection(self, scores):
  def apply_hard_selection(self, scores):
def create_hard_selection_selector(config,
                                   max_num_tokens):
def get_mask_columns_scores(max_num_columns, column_ids,
                            scores):
def get_token_scores_from_column_scores(
    column_ids,
    column_probs,
    input_mask,
    max_num_columns,
):
def get_table_pruning_loss(
    table_selector,
    table_selector_output,
    do_model_aggregation,
    do_model_classification,
    initial_features,
    is_training,
    config,
    classification_fun):
import dataclasses
import enum
import json
import tensorflow.compat.v1 as tf
import tensorflow_probability as tfp
from typing import Iterable, List, Mapping, Optional, Set, Text
from tapas.datasets import dataset
from tapas.datasets import table_dataset
from tapas.models import segmented_tensor
from tapas.models import tapas_classifier_model_utils as utils
from tapas.models.bert import modeling
from tapas.models.bert import optimization
from tapas.models.bert import table_bert
from tapas.utils import attention_utils
from tapas.utils import span_prediction_utils
from tapas.utils import table_pruning
      per_example_additional_loss = _calculate_aggregation_loss(
          logits_aggregation, aggregate_mask, aggregation_function_id, config)

      if config.use_answer_as_supervision:
        # Add regression loss for numeric answers which require aggregation.
        answer_loss, large_answer_loss_mask = _calculate_regression_loss(
            answer, aggregate_mask, dist_per_token, numeric_values,
            numeric_values_scale, table_mask_float, logits_aggregation, config)
        per_example_additional_loss += answer_loss
        # Zero loss for examples with answer_loss > cutoff.
        per_example_additional_loss *= large_answer_loss_mask

      total_loss += tf.reduce_mean(per_example_additional_loss)

    return Outputs(
        total_loss=total_loss,
        logits=logits,
        probs=_get_probs(dist_per_token) * input_mask_float,
        logits_aggregation=logits_aggregation,
        logits_cls=logits_cls,
        span_indexes=span_indexes,
        span_logits=span_logits,
    )


def _calculate_eval_metrics_fn(
    loss,
    label_ids,
    logits,
    input_mask,
    aggregation_function_id,
    logits_aggregation,
    classification_class_index,
    logits_cls,
    table_pruning_loss,
):
  """Calculates metrics for both cells and aggregation functions."""
  logits.shape.assert_has_rank(2)
  label_ids.shape.assert_has_rank(2)

  # <int32>[batch size, seq_length]
  predictions = tf.where(logits >= 0, tf.ones_like(logits, dtype=tf.int32),
                         tf.zeros_like(logits, dtype=tf.int32))
  input_mask_float = tf.cast(input_mask, tf.float32)

  loss = tf.metrics.mean(values=loss)

  # <bool>[batch size, seq_length]
  token_correct = tf.logical_or(
      tf.equal(label_ids, predictions),
      tf.logical_not(tf.cast(input_mask, tf.bool)))
  # <bool>[batch size]
  per_sequence_accuracy = tf.reduce_all(token_correct, axis=1)
  sequence_accuracy = tf.metrics.mean(values=per_sequence_accuracy)
  mean_label = tf.metrics.mean(
      values=tf.cast(label_ids, tf.float32), weights=input_mask_float)

  metrics = {
      "eval_loss": loss,
      "eval_sequence_accuracy": sequence_accuracy,
      "eval_mean_label": mean_label,
  }

  if logits_cls is not None:
    # <int32>[batch size]
    predictions_cls = tf.argmax(logits_cls, axis=-1, output_type=tf.int32)
    accuracy_cls = tf.metrics.accuracy(
        labels=classification_class_index, predictions=predictions_cls)
    mean_per_class_accuracy_cls = tf.metrics.mean_per_class_accuracy(
        labels=classification_class_index,
        predictions=predictions_cls,
        num_classes=logits_cls.shape[-1].value)
    metrics.update({
        "eval_classification_accuracy":
            accuracy_cls,
        "eval_mean_per_class_classification_accuracy":
            mean_per_class_accuracy_cls,
    })

  if logits_aggregation is not None:
    # <int32>[batch size]
    predictions_agg = tf.argmax(
        logits_aggregation, axis=-1, output_type=tf.int32)
    accuracy_agg = tf.metrics.accuracy(
        labels=aggregation_function_id, predictions=predictions_agg)
    # <bool>[batch size]
    per_sequence_agg_accuracy = tf.equal(aggregation_function_id,
                                         predictions_agg)
    # Whether cells and aggregation function predictions are both correct.
    per_sequence_joint_accuracy = tf.logical_and(per_sequence_agg_accuracy,
                                                 per_sequence_accuracy)
    joint_accuracy = tf.metrics.mean(values=per_sequence_joint_accuracy)
    metrics.update({
        "eval_aggregation_accuracy": accuracy_agg,
        "eval_joint_accuracy": joint_accuracy,
    })
  if table_pruning_loss is not None:
    metrics["pruning_loss"] = tf.metrics.mean(table_pruning_loss)
  return metrics


def model_fn_builder(
    config,
    custom_prediction_keys = None,
):
  """Returns `model_fn` closure for TPUEstimator.

  Args:
   config: Tapas config.
   custom_prediction_keys: Names of the features to export in predict mode. This
     allows fine-grained control over which features to export.

  Returns:
    A model_fn
  """

  def model_fn(features, labels, mode, params):
    """The `model_fn` for TPUEstimator."""

    del labels  # Unused.

    table_selector = table_pruning.create_selector(