



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 							_a							(      _lowerCAmelCase				):

			def __call__(self,      SCREAMING_SNAKE_CASE_=None,      SCREAMING_SNAKE_CASE_=None,      **SCREAMING_SNAKE_CASE_				)					->      List[str]:
						# Returns tqa_pipeline_inputs of shape:
						# [
						#   {"table": pd.DataFrame, "query": List[str]},
						#   ...,
						#   {"table": pd.DataFrame, "query" : List[str]}
						# ]
						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(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_				) and table.get("""query"""				) is not None and table.get("""table"""				) is not None:
												UpperCAmelCase_:		Any       			=    [table]
									elif isinstance(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_				) and len(SCREAMING_SNAKE_CASE_				) > 0:
												if not all(isinstance(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_				) for d in table				):
															raise ValueError(
															    f'Keyword argument `table` should be a list of dict, but is {(type(SCREAMING_SNAKE_CASE_				) for d in table)}'				)

												if table[0].get("""query"""				) is not None and table[0].get("""table"""				) is not None:
															UpperCAmelCase_:		Union[str, 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(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_				) or isinstance(SCREAMING_SNAKE_CASE_,      types.GeneratorType				):
												return table
									else:
												raise ValueError(
												    """Invalid input. Keyword argument `table` should be either of type `dict` or `list`, but """
												    f'is {type(SCREAMING_SNAKE_CASE_				)})'				)
						else:
									UpperCAmelCase_:		int       			=    [{"""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."""				)

												UpperCAmelCase_:		Optional[int]       			=    pd.DataFrame(tqa_pipeline_input["""table"""]				)

						return tqa_pipeline_inputs







@add_end_docstrings(_lowerCAmelCase				)
class 							_a							(      _lowerCAmelCase				):
			A					=  '''table,query'''

			def __init__(self,      SCREAMING_SNAKE_CASE_=TableQuestionAnsweringArgumentHandler(),      *SCREAMING_SNAKE_CASE_,      **SCREAMING_SNAKE_CASE_				)					->      Union[str, Any]:
						super().__init__(*SCREAMING_SNAKE_CASE_,      **SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Optional[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()				)				)

						UpperCAmelCase_:		List[Any]       			=    bool(getattr(self.model.config,      """aggregation_labels""",      SCREAMING_SNAKE_CASE_				)				) and bool(
						    getattr(self.model.config,      """num_aggregation_labels""",      SCREAMING_SNAKE_CASE_				)				)
						UpperCAmelCase_:		int       			=    """tapas""" if hasattr(self.model.config,      """aggregation_labels"""				) else None

			def     __snake_case					(self,      **SCREAMING_SNAKE_CASE_				)					->      str:
						return self.model(**SCREAMING_SNAKE_CASE_				)

			def     __snake_case					(self,      **SCREAMING_SNAKE_CASE_				)					->      Dict:
						if self.framework == "pt":
									UpperCAmelCase_:		Any       			=    []
									UpperCAmelCase_:		Optional[int]       			=    []
									UpperCAmelCase_:		List[str]       			=    None
									UpperCAmelCase_:		int       			=    inputs["""input_ids"""].shape[0]

									UpperCAmelCase_:		str       			=    inputs["""input_ids"""].to(self.device				)
									UpperCAmelCase_:		List[Any]       			=    inputs["""attention_mask"""].to(self.device				)
									UpperCAmelCase_:		Optional[int]       			=    inputs["""token_type_ids"""].to(self.device				)
									UpperCAmelCase_:		List[str]       			=    None

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

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

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

															UpperCAmelCase_:		Optional[int]       			=    torch.from_numpy(SCREAMING_SNAKE_CASE_				).type(torch.long				).to(self.device				)

												UpperCAmelCase_:		str       			=    input_ids[index]
												UpperCAmelCase_:		Any       			=    attention_mask[index]  # shape (seq_len,)
												UpperCAmelCase_:		List[Any]       			=    token_type_ids[index]  # shape (seq_len, 7)
												UpperCAmelCase_:		str       			=    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				),      )
												UpperCAmelCase_:		Tuple       			=    outputs.logits

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

												all_logits.append(SCREAMING_SNAKE_CASE_				)

												UpperCAmelCase_:		Dict       			=    torch.distributions.Bernoulli(logits=SCREAMING_SNAKE_CASE_				)
												UpperCAmelCase_:		Optional[int]       			=    dist_per_token.probs * attention_mask_example.type(torch.floataa				).to(
												    dist_per_token.probs.device				)

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

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

									UpperCAmelCase_:		Union[str, Any]       			=    torch.cat(tuple(SCREAMING_SNAKE_CASE_				),      0				)

									return (logits_batch,) if not self.aggregate else (logits_batch, torch.cat(tuple(SCREAMING_SNAKE_CASE_				),      0				))
						else:
									UpperCAmelCase_:		Tuple       			=    []
									UpperCAmelCase_:		int       			=    []
									UpperCAmelCase_:		List[Any]       			=    None
									UpperCAmelCase_:		str       			=    inputs["""input_ids"""].shape[0]

									UpperCAmelCase_:		Tuple       			=    inputs["""input_ids"""]
									UpperCAmelCase_:		Tuple       			=    inputs["""attention_mask"""]
									UpperCAmelCase_:		Optional[Any]       			=    inputs["""token_type_ids"""].numpy()
									UpperCAmelCase_:		List[str]       			=    None

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

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

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

															UpperCAmelCase_:		Tuple       			=    model_labels

												UpperCAmelCase_:		Tuple       			=    input_ids[index]
												UpperCAmelCase_:		Dict       			=    attention_mask[index]  # shape (seq_len,)
												UpperCAmelCase_:		int       			=    token_type_ids[index]  # shape (seq_len, 7)
												UpperCAmelCase_:		List[str]       			=    self.model(
												    input_ids=np.expand_dims(SCREAMING_SNAKE_CASE_,      axis=0				),      attention_mask=np.expand_dims(SCREAMING_SNAKE_CASE_,      axis=0				),      token_type_ids=np.expand_dims(SCREAMING_SNAKE_CASE_,      axis=0				),      )
												UpperCAmelCase_:		Union[str, Any]       			=    outputs.logits

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

												all_logits.append(SCREAMING_SNAKE_CASE_				)

												UpperCAmelCase_:		List[Any]       			=    tfp.distributions.Bernoulli(logits=SCREAMING_SNAKE_CASE_				)
												UpperCAmelCase_:		str       			=    dist_per_token.probs_parameter() * tf.cast(SCREAMING_SNAKE_CASE_,      tf.floataa				)

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

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

									UpperCAmelCase_:		int       			=    tf.concat(tuple(SCREAMING_SNAKE_CASE_				),      0				)

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

			def __call__(self,      *SCREAMING_SNAKE_CASE_,      **SCREAMING_SNAKE_CASE_				)					->      List[str]:
						UpperCAmelCase_:		List[str]       			=    self._args_parser(*SCREAMING_SNAKE_CASE_,      **SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		Optional[Any]       			=    super().__call__(SCREAMING_SNAKE_CASE_,      **SCREAMING_SNAKE_CASE_				)
						if len(SCREAMING_SNAKE_CASE_				) == 1:
									return results[0]
						return results

			def     __snake_case					(self,      SCREAMING_SNAKE_CASE_=None,      SCREAMING_SNAKE_CASE_=None,      SCREAMING_SNAKE_CASE_=None,      **SCREAMING_SNAKE_CASE_				)					->      Union[str, Any]:
						UpperCAmelCase_:		Optional[int]       			=    {}
						if padding is not None:
									UpperCAmelCase_:		str       			=    padding
						if truncation is not None:
									UpperCAmelCase_:		List[str]       			=    truncation

						UpperCAmelCase_:		List[str]       			=    {}
						if sequential is not None:
									UpperCAmelCase_:		List[Any]       			=    sequential
						return preprocess_params, forward_params, {}

			def     __snake_case					(self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_=None,      SCREAMING_SNAKE_CASE_=True,      SCREAMING_SNAKE_CASE_=None				)					->      str:
						if truncation is None:
									if self.type == "tapas":
												UpperCAmelCase_:		Optional[Any]       			=    """drop_rows_to_fit"""
									else:
												UpperCAmelCase_:		Dict       			=    """do_not_truncate"""

						UpperCAmelCase_			,       UpperCAmelCase_:		Optional[Any]       			=    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"""				)
						UpperCAmelCase_:		Dict       			=    self.tokenizer(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      return_tensors=self.framework,      truncation=SCREAMING_SNAKE_CASE_,      padding=SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Dict       			=    table
						return inputs

			def     __snake_case					(self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_=False				)					->      Union[str, Any]:
						UpperCAmelCase_:		Optional[Any]       			=    model_inputs.pop("""table"""				)

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

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

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

												answers.append(SCREAMING_SNAKE_CASE_				)
									if len(SCREAMING_SNAKE_CASE_				) == 0:
												raise PipelineException("""Empty answer"""				)
						else:
									UpperCAmelCase_:		int       			=    [{"""answer""": answer} for answer in self.tokenizer.batch_decode(SCREAMING_SNAKE_CASE_,      skip_special_tokens=SCREAMING_SNAKE_CASE_				)]

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











def    lowerCAmelCase_  (lowerCAmelCase__:		str     ):





			"""simple docstring"""


			return "".join(chr(ord(lowerCAmelCase__     ) + 3_2     ) if """A""" <= char <= """Z""" else char for char in word     )


if __name__ == "__main__":
	from doctest import testmod

	testmod()






