


'''simple docstring'''




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 _snake_case   ( a__   ):
							"""simple docstring"""





							def __call__(   self			,				__A=None			,				__A=None			,				**__A	)     ->Optional[int]:



											"""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:
																			__lowerCamelCase		:      List[str]											=					[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:
																							__lowerCamelCase		:      Dict											=					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:
															__lowerCamelCase		:      Tuple											=					[{"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."	)

																			__lowerCamelCase		:      List[str]											=					pd.DataFrame(tqa_pipeline_input["table"]	)

											return tqa_pipeline_inputs





@add_end_docstrings(a__   )
class _snake_case   ( a__   ):
							"""simple docstring"""





							__a												=     '''table,query'''







							def __init__(   self			,				__A=TableQuestionAnsweringArgumentHandler()			,				*__A			,				**__A	)     ->Any:



											"""simple docstring"""







											super().__init__(*__A			,				**__A	)
											__lowerCamelCase		:      Tuple											=					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()	)	)

											__lowerCamelCase		:      int											=					bool(getattr(self.model.config			,				"aggregation_labels"			,				__A	)	) and bool(
											    getattr(self.model.config			,				"num_aggregation_labels"			,				__A	)	)
											__lowerCamelCase		:      Tuple											=					"tapas" if hasattr(self.model.config			,				"aggregation_labels"	) else None







							def 		__SCREAMING_SNAKE_CASE(   self			,				**__A	)     ->List[Any]:



											"""simple docstring"""







											return self.model(**__A	)







							def 		__SCREAMING_SNAKE_CASE(   self			,				**__A	)     ->Dict:



											"""simple docstring"""







											if self.framework == "pt":
															__lowerCamelCase		:      Union[str, Any]											=					[]
															__lowerCamelCase		:      List[str]											=					[]
															__lowerCamelCase		:      Optional[int]											=					None
															__lowerCamelCase		:      Tuple											=					inputs["input_ids"].shape[0]

															__lowerCamelCase		:      Any											=					inputs["input_ids"].to(self.device	)
															__lowerCamelCase		:      Union[str, Any]											=					inputs["attention_mask"].to(self.device	)
															__lowerCamelCase		:      Any											=					inputs["token_type_ids"].to(self.device	)
															__lowerCamelCase		:      Dict											=					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:
																							__lowerCamelCase		:      List[str]											=					token_type_ids_example[:, 3]  # shape (seq_len,)
																							__lowerCamelCase		:      List[str]											=					np.zeros_like(prev_labels_example.cpu().numpy()	)  # shape (seq_len,)

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

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

																							__lowerCamelCase		:      List[Any]											=					torch.from_numpy(__A	).type(torch.long	).to(self.device	)

																			__lowerCamelCase		:      Dict											=					input_ids[index]
																			__lowerCamelCase		:      Union[str, Any]											=					attention_mask[index]  # shape (seq_len,)
																			__lowerCamelCase		:      Any											=					token_type_ids[index]  # shape (seq_len, 7)
																			__lowerCamelCase		:      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	)			,				)
																			__lowerCamelCase		:      Optional[int]											=					outputs.logits

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

																			all_logits.append(__A	)

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

																			__lowerCamelCase		:      Dict											=					collections.defaultdict(__A	)
																			for i, p in enumerate(probabilities.squeeze().tolist()	):
																							__lowerCamelCase		:      int											=					token_type_ids_example[:, 0].tolist()[i]
																							__lowerCamelCase		:      List[str]											=					token_type_ids_example[:, 1].tolist()[i] - 1
																							__lowerCamelCase		:      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	)

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

															__lowerCamelCase		:      List[str]											=					torch.cat(tuple(__A	)			,				0	)

															return (logits_batch,) if not self.aggregate else (logits_batch, torch.cat(tuple(__A	)			,				0	))
											else:
															__lowerCamelCase		:      List[Any]											=					[]
															__lowerCamelCase		:      str											=					[]
															__lowerCamelCase		:      Optional[Any]											=					None
															__lowerCamelCase		:      Optional[Any]											=					inputs["input_ids"].shape[0]

															__lowerCamelCase		:      List[Any]											=					inputs["input_ids"]
															__lowerCamelCase		:      Dict											=					inputs["attention_mask"]
															__lowerCamelCase		:      Union[str, Any]											=					inputs["token_type_ids"].numpy()
															__lowerCamelCase		:      Optional[int]											=					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:
																							__lowerCamelCase		:      Optional[Any]											=					token_type_ids_example[:, 3]  # shape (seq_len,)
																							__lowerCamelCase		:      List[Any]											=					np.zeros_like(__A			,				dtype=np.intaa	)  # shape (seq_len,)

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

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

																							__lowerCamelCase		:      Optional[Any]											=					model_labels

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

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

																			all_logits.append(__A	)

																			__lowerCamelCase		:      Union[str, Any]											=					tfp.distributions.Bernoulli(logits=__A	)
																			__lowerCamelCase		:      Optional[int]											=					dist_per_token.probs_parameter() * tf.cast(__A			,				tf.floataa	)

																			__lowerCamelCase		:      str											=					collections.defaultdict(__A	)
																			__lowerCamelCase		:      int											=					token_type_ids_example
																			for i, p in enumerate(tf.squeeze(__A	).numpy().tolist()	):
																							__lowerCamelCase		:      Union[str, Any]											=					token_type_ids_example[:, 0].tolist()[i]
																							__lowerCamelCase		:      int											=					token_type_ids_example[:, 1].tolist()[i] - 1
																							__lowerCamelCase		:      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(__A	)

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

															__lowerCamelCase		:      List[Any]											=					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	)     ->Any:



											"""simple docstring"""







											__lowerCamelCase		:      Dict											=					self._args_parser(*__A			,				**__A	)

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







							def 		__SCREAMING_SNAKE_CASE(   self			,				__A=None			,				__A=None			,				__A=None			,				**__A	)     ->Optional[Any]:



											"""simple docstring"""







											__lowerCamelCase		:      Any											=					{}
											if padding is not None:
															__lowerCamelCase		:      int											=					padding
											if truncation is not None:
															__lowerCamelCase		:      List[str]											=					truncation

											__lowerCamelCase		:      Optional[Any]											=					{}
											if sequential is not None:
															__lowerCamelCase		:      Any											=					sequential
											return preprocess_params, forward_params, {}







							def 		__SCREAMING_SNAKE_CASE(   self			,				__A			,				__A=None			,				__A=True			,				__A=None	)     ->Union[str, Any]:



											"""simple docstring"""







											if truncation is None:
															if self.type == "tapas":
																			__lowerCamelCase		:      int											=					"drop_rows_to_fit"
															else:
																			__lowerCamelCase		:      str											=					"do_not_truncate"

											__lowerCamelCase						,					__lowerCamelCase		:      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"	)
											__lowerCamelCase		:      Any											=					self.tokenizer(__A			,				__A			,				return_tensors=self.framework			,				truncation=__A			,				padding=__A	)
											__lowerCamelCase		:      Tuple											=					table
											return inputs







							def 		__SCREAMING_SNAKE_CASE(   self			,				__A			,				__A=False	)     ->int:



											"""simple docstring"""







											__lowerCamelCase		:      List[str]											=					model_inputs.pop("table"	)

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







							def 		__SCREAMING_SNAKE_CASE(   self			,				__A	)     ->List[Any]:



											"""simple docstring"""







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

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

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

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









'''simple docstring'''




import unittest
from functools import partial

import torch

from accelerate import Accelerator, debug_launcher
from accelerate.state import AcceleratorState, GradientState
from accelerate.test_utils import require_cpu, require_huggingface_suite
from accelerate.utils import GradientAccumulationPlugin





def 							a_       (						lowercase=2							,						lowercase=True							,						lowercase=False    )			->Optional[Any]:






				'''simple docstring'''

				__lowerCamelCase		:      List[str]											=					Accelerator(step_scheduler_with_optimizer=lowercase							,						split_batches=lowercase    )
				__lowerCamelCase		:      Union[str, Any]											=					torch.nn.Linear(2							,						4    )
				__lowerCamelCase		:      str											=					torch.optim.AdamW(model.parameters()							,						lr=1.0    )
				__lowerCamelCase		:      List[str]											=					torch.optim.lr_scheduler.OneCycleLR(lowercase							,						max_lr=0.0_1							,						steps_per_epoch=2							,						epochs=1    )
				__lowerCamelCase						,					__lowerCamelCase						,					__lowerCamelCase		:      List[Any]											=					accelerator.prepare(lowercase							,						lowercase							,						lowercase    )

				# Optimizer has stepped
				scheduler.step()
				if step_scheduler_with_optimizer or (num_processes == 1):
								assert (
								    scheduler.scheduler.last_epoch == num_processes
								), F'Last Epoch ({scheduler.scheduler.last_epoch}) != Num Processes ({num_processes})'
				else:
								assert (
								    scheduler.scheduler.last_epoch != num_processes
								), F'Last Epoch ({scheduler.scheduler.last_epoch}) == Num Processes ({num_processes})'





def 							a_       (						lowercase=2							,						lowercase=True							,						lowercase=False    )			->Dict:






				'''simple docstring'''

				__lowerCamelCase		:      Dict											=					Accelerator(step_scheduler_with_optimizer=lowercase							,						split_batches=lowercase    )
				__lowerCamelCase		:      Tuple											=					torch.nn.Linear(2							,						4    )
				__lowerCamelCase		:      Optional[Any]											=					torch.optim.AdamW(model.parameters()							,						lr=1.0    )
				__lowerCamelCase		:      Optional[Any]											=					torch.optim.lr_scheduler.LambdaLR(lowercase							,						lr_lambda=lambda lowercase    : 1 - n / 10    )
				__lowerCamelCase						,					__lowerCamelCase						,					__lowerCamelCase		:      Optional[int]											=					accelerator.prepare(lowercase							,						lowercase							,						lowercase    )

				# Optimizer has stepped
				__lowerCamelCase		:      List[Any]											=					False
				scheduler.step()
				__lowerCamelCase		:      List[Any]											=					1 - (num_processes if (step_scheduler_with_optimizer and not split_batches) else 1) / 10
				assert (
				    scheduler.get_last_lr()[0] == expected_lr
				), F'Wrong lr found at first step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}'

				# Optimizer has not stepped
				__lowerCamelCase		:      Any											=					True
				scheduler.step()
				if not step_scheduler_with_optimizer:
								__lowerCamelCase		:      int											=					1 - 2 / 10
				assert (
				    scheduler.get_last_lr()[0] == expected_lr
				), F'Wrong lr found at second step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}'





def 							a_       (						lowercase = 2    )			->Dict:






				'''simple docstring'''

				from transformers import get_linear_schedule_with_warmup

				__lowerCamelCase		:      str											=					[1, 2, 4]
				for num_steps in steps:
								__lowerCamelCase		:      List[Any]											=					GradientAccumulationPlugin(num_steps=lowercase							,						adjust_scheduler=num_steps > 1    )
								__lowerCamelCase		:      Dict											=					Accelerator(gradient_accumulation_plugin=lowercase    )
								__lowerCamelCase		:      Any											=					torch.nn.Linear(2							,						4    )
								__lowerCamelCase		:      List[Any]											=					torch.optim.AdamW(model.parameters()							,						lr=1_0.0    )
								__lowerCamelCase		:      Optional[Any]											=					get_linear_schedule_with_warmup(optimizer=lowercase							,						num_warmup_steps=0							,						num_training_steps=20    )

								__lowerCamelCase						,					__lowerCamelCase						,					__lowerCamelCase		:      Optional[int]											=					accelerator.prepare(lowercase							,						lowercase							,						lowercase    )

								for i in range(10 * num_steps    ):
												with accelerator.accumulate(lowercase    ):
																optimizer.step()
																scheduler.step()

												if i == (10 * num_steps - 2):
																assert (
																    scheduler.get_last_lr()[0] != 0
																), F'Wrong lr found at second-to-last step, expected non-zero, got {scheduler.get_last_lr()[0]}. num_steps: {num_steps}'
								assert (
								    scheduler.get_last_lr()[0] == 0
								), F'Wrong lr found at last step, expected 0, got {scheduler.get_last_lr()[0]}'
								GradientState._reset_state()





@require_cpu
class _snake_case   ( unittest.TestCase   ):
							"""simple docstring"""





							def 		__SCREAMING_SNAKE_CASE(   self	)     ->Tuple:



											"""simple docstring"""







											debug_launcher(partial(__A			,				num_processes=1	)			,				num_processes=1	)
											debug_launcher(partial(__A			,				num_processes=1			,				split_batches=__A	)			,				num_processes=1	)







							def 		__SCREAMING_SNAKE_CASE(   self	)     ->Tuple:



											"""simple docstring"""







											debug_launcher(partial(__A			,				num_processes=1	)			,				num_processes=1	)
											debug_launcher(partial(__A			,				num_processes=1			,				split_batches=__A	)			,				num_processes=1	)







							def 		__SCREAMING_SNAKE_CASE(   self	)     ->Optional[Any]:



											"""simple docstring"""







											debug_launcher(partial(__A			,				num_processes=1			,				step_scheduler_with_optimizer=__A	)			,				num_processes=1	)







							def 		__SCREAMING_SNAKE_CASE(   self	)     ->Dict:



											"""simple docstring"""







											debug_launcher(partial(__A			,				num_processes=1			,				step_scheduler_with_optimizer=__A	)			,				num_processes=1	)







							def 		__SCREAMING_SNAKE_CASE(   self	)     ->Optional[int]:



											"""simple docstring"""







											AcceleratorState._reset_state(__A	)
											debug_launcher(__A	)
											debug_launcher(partial(__A			,				num_processes=1			,				split_batches=__A	)			,				num_processes=1	)







							def 		__SCREAMING_SNAKE_CASE(   self	)     ->Union[str, Any]:



											"""simple docstring"""







											AcceleratorState._reset_state(__A	)
											debug_launcher(__A	)
											debug_launcher(partial(__A			,				num_processes=1			,				split_batches=__A	)			,				num_processes=1	)







							def 		__SCREAMING_SNAKE_CASE(   self	)     ->Any:



											"""simple docstring"""







											AcceleratorState._reset_state(__A	)
											debug_launcher(partial(__A			,				step_scheduler_with_optimizer=__A	)	)







							def 		__SCREAMING_SNAKE_CASE(   self	)     ->List[Any]:



											"""simple docstring"""







											AcceleratorState._reset_state(__A	)
											debug_launcher(partial(__A			,				step_scheduler_with_optimizer=__A	)	)







							@require_huggingface_suite
							def 		__SCREAMING_SNAKE_CASE(   self	)     ->Dict:



											"""simple docstring"""







											AcceleratorState._reset_state(__A	)
											debug_launcher(partial(__A			,				num_processes=1	)	)
											debug_launcher(__A	)





