


import collections
import types

import numpy as np

from ..utils import (
    add_end_docstrings,
    is_tensorflow_probability_available,
    is_tf_available,
    is_torch_available,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, ArgumentHandler, Dataset, Pipeline, PipelineException


if is_torch_available():
  import torch

  from ..models.auto.modeling_auto import (
      MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
      MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING,
  )

if is_tf_available() and is_tensorflow_probability_available():
  import tensorflow as tf
  import tensorflow_probability as tfp

  from ..models.auto.modeling_tf_auto import (
      TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
      TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING,
  )

class   lowerCamelCase						(	SCREAMING_SNAKE_CASE  ):

       def __call__(      self						,    a__=None						,    a__=None						,    **a__       ):






        '''simple docstring'''






        requires_backends(self						,    '''pandas'''       )
        import pandas as pd

        if table is None:
         raise ValueError('''Keyword argument `table` cannot be None.'''       )
        elif query is None:
         if isinstance(a__						,    a__       ) and table.get('''query'''       ) is not None and table.get('''table'''       ) is not None:
          snake_case :			int							=   [table]
         elif isinstance(a__						,    a__       ) and len(a__       ) > 0:
          if not all(isinstance(a__						,    a__       ) for d in table       ):
           raise ValueError(
               f'''Keyword argument `table` should be a list of dict, but is {(type(a__       ) for d in table)}'''       )

          if table[0].get('''query'''       ) is not None and table[0].get('''table'''       ) is not None:
           snake_case :			Any							=   table
          else:
           raise ValueError(
               '''If keyword argument `table` is a list of dictionaries, each dictionary should have a `table`'''
               f''' and `query` key, but only dictionary has keys {table[0].keys()} `table` and `query` keys.'''       )
         elif Dataset is not None and isinstance(a__						,    a__       ) or isinstance(a__						,    types.GeneratorType       ):
          return table
         else:
          raise ValueError(
              '''Invalid input. Keyword argument `table` should be either of type `dict` or `list`, but '''
              f'''is {type(a__       )})'''       )
        else:
         snake_case :			Any							=   [{'''table''': table, '''query''': query}]

        for tqa_pipeline_input in tqa_pipeline_inputs:
         if not isinstance(tqa_pipeline_input['''table''']						,    pd.DataFrame       ):
          if tqa_pipeline_input["table"] is None:
           raise ValueError('''Table cannot be None.'''       )

          snake_case :			List[Any]							=   pd.DataFrame(tqa_pipeline_input['''table''']       )

        return tqa_pipeline_inputs







@add_end_docstrings(SCREAMING_SNAKE_CASE  )
class   lowerCamelCase						(	SCREAMING_SNAKE_CASE  ):
       _UpperCamelCase     		=   """table,query"""

       def __init__(      self						,    a__=TableQuestionAnsweringArgumentHandler()						,    *a__						,    **a__       ):






        '''simple docstring'''






        super().__init__(*a__						,    **a__       )
        snake_case :			Any							=   args_parser

        self.check_model_type(
            dict(
                TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING.items()
                + TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.items()       )
            if self.framework == '''tf'''
            else dict(
                MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING.items() + MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.items()       )       )

        snake_case :			int							=   bool(getattr(self.model.config						,    '''aggregation_labels'''						,    a__       )       ) and bool(
            getattr(self.model.config						,    '''num_aggregation_labels'''						,    a__       )       )
        snake_case :			Optional[Any]							=   '''tapas''' if hasattr(self.model.config						,    '''aggregation_labels'''       ) else None

       def 					UpperCAmelCase					(      self						,    **a__       ):






        '''simple docstring'''






        return self.model(**a__       )

       def 					UpperCAmelCase					(      self						,    **a__       ):






        '''simple docstring'''






        if self.framework == "pt":
         snake_case :			Union[str, Any]							=   []
         snake_case :			int							=   []
         snake_case :			Dict							=   None
         snake_case :			List[str]							=   inputs['''input_ids'''].shape[0]

         snake_case :			Dict							=   inputs['''input_ids'''].to(self.device       )
         snake_case :			int							=   inputs['''attention_mask'''].to(self.device       )
         snake_case :			Union[str, Any]							=   inputs['''token_type_ids'''].to(self.device       )
         snake_case :			List[Any]							=   None

         for index in range(a__       ):
          # If sequences have already been processed, the token type IDs will be created according to the previous
          # answer.
          if prev_answers is not None:
           snake_case :			Optional[Any]							=   token_type_ids_example[:, 3]  # shape (seq_len,)
           snake_case :			Tuple							=   np.zeros_like(prev_labels_example.cpu().numpy()       )  # shape (seq_len,)

           snake_case :			List[str]							=   token_type_ids[index]  # shape (seq_len, 7)
           for i in range(model_labels.shape[0]       ):
            snake_case :			int							=   token_type_ids_example[:, 0].tolist()[i]
            snake_case :			List[Any]							=   token_type_ids_example[:, 1].tolist()[i] - 1
            snake_case :			Tuple							=   token_type_ids_example[:, 2].tolist()[i] - 1

            if row_id >= 0 and col_id >= 0 and segment_id == 1:
             snake_case :			Union[str, Any]							=   int(prev_answers[(col_id, row_id)]       )

           snake_case :			int							=   torch.from_numpy(a__       ).type(torch.long       ).to(self.device       )

          snake_case :			Tuple							=   input_ids[index]
          snake_case :			int							=   attention_mask[index]  # shape (seq_len,)
          snake_case :			int							=   token_type_ids[index]  # shape (seq_len, 7)
          snake_case :			Any							=   self.model(
              input_ids=input_ids_example.unsqueeze(0       )						,    attention_mask=attention_mask_example.unsqueeze(0       )						,    token_type_ids=token_type_ids_example.unsqueeze(0       )						,    )
          snake_case :			Dict							=   outputs.logits

          if self.aggregate:
           all_aggregations.append(outputs.logits_aggregation       )

          all_logits.append(a__       )

          snake_case :			List[Any]							=   torch.distributions.Bernoulli(logits=a__       )
          snake_case :			List[Any]							=   dist_per_token.probs * attention_mask_example.type(torch.floataa       ).to(
              dist_per_token.probs.device       )

          snake_case :			Optional[int]							=   collections.defaultdict(a__       )
          for i, p in enumerate(probabilities.squeeze().tolist()       ):
           snake_case :			Any							=   token_type_ids_example[:, 0].tolist()[i]
           snake_case :			str							=   token_type_ids_example[:, 1].tolist()[i] - 1
           snake_case :			Optional[Any]							=   token_type_ids_example[:, 2].tolist()[i] - 1
           if col >= 0 and row >= 0 and segment_id == 1:
            coords_to_probs[(col, row)].append(a__       )

          snake_case :			Tuple							=   {key: np.array(coords_to_probs[key]       ).mean() > 0.5 for key in coords_to_probs}

         snake_case :			Union[str, Any]							=   torch.cat(tuple(a__       )						,    0       )

         return (logits_batch,) if not self.aggregate else (logits_batch, torch.cat(tuple(a__       )						,    0       ))
        else:
         snake_case :			Optional[int]							=   []
         snake_case :			Optional[Any]							=   []
         snake_case :			Optional[Any]							=   None
         snake_case :			List[Any]							=   inputs['''input_ids'''].shape[0]

         snake_case :			int							=   inputs['''input_ids''']
         snake_case :			int							=   inputs['''attention_mask''']
         snake_case :			Tuple							=   inputs['''token_type_ids'''].numpy()
         snake_case :			Any							=   None

         for index in range(a__       ):
          # If sequences have already been processed, the token type IDs will be created according to the previous
          # answer.
          if prev_answers is not None:
           snake_case :			Any							=   token_type_ids_example[:, 3]  # shape (seq_len,)
           snake_case :			Optional[Any]							=   np.zeros_like(a__						,    dtype=np.intaa       )  # shape (seq_len,)

           snake_case :			Union[str, Any]							=   token_type_ids[index]  # shape (seq_len, 7)
           for i in range(model_labels.shape[0]       ):
            snake_case :			int							=   token_type_ids_example[:, 0].tolist()[i]
            snake_case :			Union[str, Any]							=   token_type_ids_example[:, 1].tolist()[i] - 1
            snake_case :			Tuple							=   token_type_ids_example[:, 2].tolist()[i] - 1

            if row_id >= 0 and col_id >= 0 and segment_id == 1:
             snake_case :			Dict							=   int(prev_answers[(col_id, row_id)]       )

           snake_case :			List[str]							=   model_labels

          snake_case :			int							=   input_ids[index]
          snake_case :			int							=   attention_mask[index]  # shape (seq_len,)
          snake_case :			List[Any]							=   token_type_ids[index]  # shape (seq_len, 7)
          snake_case :			Optional[int]							=   self.model(
              input_ids=np.expand_dims(a__						,    axis=0       )						,    attention_mask=np.expand_dims(a__						,    axis=0       )						,    token_type_ids=np.expand_dims(a__						,    axis=0       )						,    )
          snake_case :			Optional[Any]							=   outputs.logits

          if self.aggregate:
           all_aggregations.append(outputs.logits_aggregation       )

          all_logits.append(a__       )

          snake_case :			Any							=   tfp.distributions.Bernoulli(logits=a__       )
          snake_case :			Any							=   dist_per_token.probs_parameter() * tf.cast(a__						,    tf.floataa       )

          snake_case :			Tuple							=   collections.defaultdict(a__       )
          snake_case :			List[str]							=   token_type_ids_example
          for i, p in enumerate(tf.squeeze(a__       ).numpy().tolist()       ):
           snake_case :			int							=   token_type_ids_example[:, 0].tolist()[i]
           snake_case :			Optional[Any]							=   token_type_ids_example[:, 1].tolist()[i] - 1
           snake_case :			List[Any]							=   token_type_ids_example[:, 2].tolist()[i] - 1
           if col >= 0 and row >= 0 and segment_id == 1:
            coords_to_probs[(col, row)].append(a__       )

          snake_case :			List[str]							=   {key: np.array(coords_to_probs[key]       ).mean() > 0.5 for key in coords_to_probs}

         snake_case :			Optional[int]							=   tf.concat(tuple(a__       )						,    0       )

         return (logits_batch,) if not self.aggregate else (logits_batch, tf.concat(tuple(a__       )						,    0       ))

       def __call__(      self						,    *a__						,    **a__       ):






        '''simple docstring'''






        snake_case :			Any							=   self._args_parser(*a__						,    **a__       )

        snake_case :			Dict							=   super().__call__(a__						,    **a__       )
        if len(a__       ) == 1:
         return results[0]
        return results

       def 					UpperCAmelCase					(      self						,    a__=None						,    a__=None						,    a__=None						,    **a__       ):






        '''simple docstring'''






        snake_case :			List[Any]							=   {}
        if padding is not None:
         snake_case :			Optional[Any]							=   padding
        if truncation is not None:
         snake_case :			Optional[Any]							=   truncation

        snake_case :			Optional[int]							=   {}
        if sequential is not None:
         snake_case :			Tuple							=   sequential
        return preprocess_params, forward_params, {}

       def 					UpperCAmelCase					(      self						,    a__						,    a__=None						,    a__=True						,    a__=None       ):






        '''simple docstring'''






        if truncation is None:
         if self.type == "tapas":
          snake_case :			List[str]							=   '''drop_rows_to_fit'''
         else:
          snake_case :			Dict							=   '''do_not_truncate'''

        snake_case	,       snake_case :			List[str]							=   pipeline_input['''table'''], pipeline_input['''query''']
        if table.empty:
         raise ValueError('''table is empty'''       )
        if query is None or query == "":
         raise ValueError('''query is empty'''       )
        snake_case :			List[str]							=   self.tokenizer(a__						,    a__						,    return_tensors=self.framework						,    truncation=a__						,    padding=a__       )
        snake_case :			str							=   table
        return inputs

       def 					UpperCAmelCase					(      self						,    a__						,    a__=False       ):






        '''simple docstring'''






        snake_case :			Dict							=   model_inputs.pop('''table'''       )

        if self.type == "tapas":
         if sequential:
          snake_case :			Union[str, Any]							=   self.sequential_inference(**a__       )
         else:
          snake_case :			Optional[Any]							=   self.batch_inference(**a__       )
        else:
         snake_case :			str							=   self.model.generate(**a__       )
        snake_case :			List[Any]							=   {'''model_inputs''': model_inputs, '''table''': table, '''outputs''': outputs}
        return model_outputs



       def 					UpperCAmelCase					(      self						,    a__       ):






        '''simple docstring'''






        snake_case :			List[str]							=   model_outputs['''model_inputs''']
        snake_case :			str							=   model_outputs['''table''']
        snake_case :			Tuple							=   model_outputs['''outputs''']
        if self.type == "tapas":
         if self.aggregate:
          snake_case	,       snake_case :			List[str]							=   outputs[:2]
          snake_case :			List[str]							=   self.tokenizer.convert_logits_to_predictions(a__						,    a__						,    a__       )
          snake_case	,       snake_case :			Any							=   predictions
          snake_case :			int							=   {i: self.model.config.aggregation_labels[pred] for i, pred in enumerate(a__       )}

          snake_case :			Tuple							=   self.model.config.no_aggregation_label_index
          snake_case :			Union[str, Any]							=   {
              i: aggregators[i] + ''' > ''' for i, pred in enumerate(a__       ) if pred != no_agg_label_index
          }
         else:
          snake_case :			int							=   outputs[0]
          snake_case :			int							=   self.tokenizer.convert_logits_to_predictions(a__						,    a__       )
          snake_case :			Optional[int]							=   predictions[0]
          snake_case :			str							=   {}
          snake_case :			List[str]							=   {}
         snake_case :			Optional[Any]							=   []
         for index, coordinates in enumerate(a__       ):
          snake_case :			List[str]							=   [table.iat[coordinate] for coordinate in coordinates]
          snake_case :			int							=   aggregators.get(a__						,    ''''''       )
          snake_case :			Any							=   aggregators_prefix.get(a__						,    ''''''       )
          snake_case :			str							=   {
              '''answer''': aggregator_prefix + ''', '''.join(a__       ),
              '''coordinates''': coordinates,
              '''cells''': [table.iat[coordinate] for coordinate in coordinates],
          }
          if aggregator:
           snake_case :			List[Any]							=   aggregator

          answers.append(a__       )
         if len(a__       ) == 0:
          raise PipelineException('''Empty answer'''       )
        else:
         snake_case :			Any							=   [{'''answer''': answer} for answer in self.tokenizer.batch_decode(a__						,    skip_special_tokens=a__       )]

        return answers if len(a__       ) > 1 else answers[0]










from typing import Any

import numpy as np






def  _UpperCAmelCase						( __lowerCamelCase							)		->bool:


 '''simple docstring'''
 assert np.shape(__lowerCamelCase							)[0] == np.shape(__lowerCamelCase							)[1]

 # If matrix not symmetric, exit right away.
 if np.allclose(__lowerCamelCase			,  matrix.T							) is False:
  return False

 # Get eigenvalues and eignevectors for a symmetric matrix.
 snake_case	,       snake_case :			int							=   np.linalg.eigh(__lowerCamelCase							)

 # Check sign of all eigenvalues.
 # np.all returns a value of type np.bool_
 return bool(np.all(eigen_values > 0							)							)






def  _UpperCAmelCase						( __lowerCamelCase							)		->Any:


 '''simple docstring'''
 snake_case :			Dict							=   np.random.randn(__lowerCamelCase			,  __lowerCamelCase							)
 snake_case :			List[Any]							=   np.dot(__lowerCamelCase			,  random_matrix.T							)
 assert _is_matrix_spd(__lowerCamelCase							)
 return spd_matrix






def  _UpperCAmelCase						( __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase = 1000			,  __lowerCamelCase = 1E-8			,  )		->Any:


 '''simple docstring'''
 assert np.shape(__lowerCamelCase							)[0] == np.shape(__lowerCamelCase							)[1]
 assert np.shape(__lowerCamelCase							)[0] == np.shape(__lowerCamelCase							)[0]
 assert _is_matrix_spd(__lowerCamelCase							)

 # Initialize solution guess, residual, search direction.
 snake_case :			Any							=   np.zeros((np.shape(__lowerCamelCase							)[0], 1)							)
 snake_case :			Optional[int]							=   np.copy(__lowerCamelCase							)
 snake_case :			Optional[Any]							=   np.copy(__lowerCamelCase							)

 # Set initial errors in solution guess and residual.
 snake_case :			Optional[Any]							=   1E9
 snake_case :			List[Any]							=   1E9
 snake_case :			List[str]							=   1E9

 # Set iteration counter to threshold number of iterations.
 snake_case :			List[str]							=   0

 while error > tol:
  # Save this value so we only calculate the matrix-vector product once.
  snake_case :			str							=   np.dot(__lowerCamelCase			,  __lowerCamelCase							)

  # The main algorithm.

  # Update search direction magnitude.
  snake_case :			List[str]							=   np.dot(ra.T			,  __lowerCamelCase							) / np.dot(pa.T			,  __lowerCamelCase							)
  # Update solution guess.
  snake_case :			Optional[int]							=   xa + alpha * pa
  # Calculate new residual.
  snake_case :			str							=   ra - alpha * w
  # Calculate new Krylov subspace scale.
  snake_case :			List[str]							=   np.dot(r.T			,  __lowerCamelCase							) / np.dot(ra.T			,  __lowerCamelCase							)
  # Calculate new A conjuage search direction.
  snake_case :			List[Any]							=   r + beta * pa

  # Calculate errors.
  snake_case :			List[Any]							=   np.linalg.norm(r - ra							)
  snake_case :			Any							=   np.linalg.norm(x - xa							)
  snake_case :			str							=   np.maximum(__lowerCamelCase			,  __lowerCamelCase							)

  # Update variables.
  snake_case :			Tuple							=   np.copy(__lowerCamelCase							)
  snake_case :			Tuple							=   np.copy(__lowerCamelCase							)
  snake_case :			List[str]							=   np.copy(__lowerCamelCase							)

  # Update number of iterations.
  iterations += 1
  if iterations > max_iterations:
   break

 return x






def  _UpperCAmelCase						( )		->None:


 '''simple docstring'''
 snake_case :			str							=   3
 snake_case :			Dict							=   _create_spd_matrix(__lowerCamelCase							)
 snake_case :			str							=   np.random.randn(__lowerCamelCase			,  1							)
 snake_case :			str							=   np.dot(__lowerCamelCase			,  __lowerCamelCase							)

 # Numpy solution.
 snake_case :			Any							=   np.linalg.solve(__lowerCamelCase			,  __lowerCamelCase							)

 # Our implementation.
 snake_case :			List[str]							=   conjugate_gradient(__lowerCamelCase			,  __lowerCamelCase							)

 # Ensure both solutions are close to x_true (and therefore one another).
 assert np.linalg.norm(x_numpy - x_true							) <= 1E-6
 assert np.linalg.norm(x_conjugate_gradient - x_true							) <= 1E-6


if __name__ == "__main__":
  import doctest

  doctest.testmod()
  test_conjugate_gradient()






