import unittest

from transformers import MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING, AutoTokenizer, is_vision_available
from transformers.pipelines import pipeline
from transformers.pipelines.document_question_answering import apply_tesseract
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_detectrona,
    require_pytesseract,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
							from PIL import Image

							from transformers.image_utils import load_image
else:

							class  __A       :


											@staticmethod
											def _snake_case    (						*UpperCAmelCase_				,	**UpperCAmelCase_							):
															pass





							def        _lowercase   (  _UpperCAmelCase			)  ->					Optional[Any]:
											return None


# This is a pinned image from a specific revision of a document question answering space, hosted by HuggingFace,
# so we can expect it to be available.
UpperCAmelCase__  :     Optional[Any]   			=(
    '''https://huggingface.co/spaces/impira/docquery/resolve/2f6c96314dc84dfda62d40de9da55f2f5165d403/invoice.png'''
)

@is_pipeline_test
@require_torch
@require_vision
class  __A       (  unittest.TestCase      ):
				__A								=			MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING


				@require_pytesseract
				@require_vision
				def _snake_case    (						self				,	UpperCAmelCase_				,	UpperCAmelCase_				,	UpperCAmelCase_							):
								lowerCamelCase			=pipeline(
								    """document-question-answering"""				,	model=UpperCAmelCase_				,	tokenizer=UpperCAmelCase_				,	image_processor=UpperCAmelCase_							)

								lowerCamelCase			=INVOICE_URL
								lowerCamelCase			=list(zip(*apply_tesseract(load_image(UpperCAmelCase_							)				,	UpperCAmelCase_				,	""""""							)							)							)
								lowerCamelCase			="""What is the placebo?"""
								lowerCamelCase			=[
								    {
								        """image""": load_image(UpperCAmelCase_							),
								        """question""": question,
								    },
								    {
								        """image""": image,
								        """question""": question,
								    },
								    {
								        """image""": image,
								        """question""": question,
								        """word_boxes""": word_boxes,
								    },
								]
								return dqa_pipeline, examples


				def _snake_case    (						self				,	UpperCAmelCase_				,	UpperCAmelCase_							):
								lowerCamelCase			=dqa_pipeline(UpperCAmelCase_				,	top_k=2							)
								self.assertEqual(
								    UpperCAmelCase_				,	[
								        [
								            {"""score""": ANY(UpperCAmelCase_							), """answer""": ANY(UpperCAmelCase_							), """start""": ANY(UpperCAmelCase_							), """end""": ANY(UpperCAmelCase_							)},
								            {"""score""": ANY(UpperCAmelCase_							), """answer""": ANY(UpperCAmelCase_							), """start""": ANY(UpperCAmelCase_							), """end""": ANY(UpperCAmelCase_							)},
								        ]
								    ]
								    * 3				,	)


				@require_torch
				@require_detectrona
				@require_pytesseract
				def _snake_case    (						self							):
								lowerCamelCase			=pipeline("""document-question-answering"""				,	model="""hf-internal-testing/tiny-random-layoutlmv2"""							)
								lowerCamelCase			=INVOICE_URL
								lowerCamelCase			="""How many cats are there?"""

								lowerCamelCase			=[
								    {"""score""": 0.0_0_0_1, """answer""": """oy 2312/2019""", """start""": 38, """end""": 39},
								    {"""score""": 0.0_0_0_1, """answer""": """oy 2312/2019 DUE""", """start""": 38, """end""": 40},
								]
								lowerCamelCase			=dqa_pipeline(image=UpperCAmelCase_				,	question=UpperCAmelCase_				,	top_k=2							)
								self.assertEqual(nested_simplify(UpperCAmelCase_				,	decimals=4							)				,	UpperCAmelCase_							)

								lowerCamelCase			=dqa_pipeline({"""image""": image, """question""": question}				,	top_k=2							)
								self.assertEqual(nested_simplify(UpperCAmelCase_				,	decimals=4							)				,	UpperCAmelCase_							)

								# This image does not detect ANY text in it, meaning layoutlmv2 should fail.
								# Empty answer probably
								lowerCamelCase			="""./tests/fixtures/tests_samples/COCO/000000039769.png"""
								lowerCamelCase			=dqa_pipeline(image=UpperCAmelCase_				,	question=UpperCAmelCase_				,	top_k=2							)
								self.assertEqual(UpperCAmelCase_				,	[]							)

								# We can optionnally pass directly the words and bounding boxes
								lowerCamelCase			="""./tests/fixtures/tests_samples/COCO/000000039769.png"""
								lowerCamelCase			=[]
								lowerCamelCase			=[]
								lowerCamelCase			=dqa_pipeline(image=UpperCAmelCase_				,	question=UpperCAmelCase_				,	words=UpperCAmelCase_				,	boxes=UpperCAmelCase_				,	top_k=2							)
								self.assertEqual(UpperCAmelCase_				,	[]							)


				@slow
				@require_torch
				@require_detectrona
				@require_pytesseract
				def _snake_case    (						self							):
								lowerCamelCase			=pipeline(
								    """document-question-answering"""				,	model="""tiennvcs/layoutlmv2-base-uncased-finetuned-docvqa"""				,	revision="""9977165"""				,	)
								lowerCamelCase			=INVOICE_URL
								lowerCamelCase			="""What is the invoice number?"""

								lowerCamelCase			=dqa_pipeline(image=UpperCAmelCase_				,	question=UpperCAmelCase_				,	top_k=2							)
								self.assertEqual(
								    nested_simplify(UpperCAmelCase_				,	decimals=4							)				,	[
								        {"""score""": 0.9_9_4_4, """answer""": """us-001""", """start""": 16, """end""": 16},
								        {"""score""": 0.0_0_0_9, """answer""": """us-001""", """start""": 16, """end""": 16},
								    ]				,	)

								lowerCamelCase			=dqa_pipeline({"""image""": image, """question""": question}				,	top_k=2							)
								self.assertEqual(
								    nested_simplify(UpperCAmelCase_				,	decimals=4							)				,	[
								        {"""score""": 0.9_9_4_4, """answer""": """us-001""", """start""": 16, """end""": 16},
								        {"""score""": 0.0_0_0_9, """answer""": """us-001""", """start""": 16, """end""": 16},
								    ]				,	)

								lowerCamelCase			=dqa_pipeline(
								    [{"""image""": image, """question""": question}, {"""image""": image, """question""": question}]				,	top_k=2							)
								self.assertEqual(
								    nested_simplify(UpperCAmelCase_				,	decimals=4							)				,	[
								        [
								            {"""score""": 0.9_9_4_4, """answer""": """us-001""", """start""": 16, """end""": 16},
								            {"""score""": 0.0_0_0_9, """answer""": """us-001""", """start""": 16, """end""": 16},
								        ],
								    ]
								    * 2				,	)


				@slow
				@require_torch
				@require_detectrona
				@require_pytesseract
				def _snake_case    (						self							):
								lowerCamelCase			=pipeline(
								    """document-question-answering"""				,	model="""tiennvcs/layoutlmv2-base-uncased-finetuned-docvqa"""				,	revision="""9977165"""				,	max_seq_len=50				,	)
								lowerCamelCase			=INVOICE_URL
								lowerCamelCase			="""What is the invoice number?"""

								lowerCamelCase			=dqa_pipeline(image=UpperCAmelCase_				,	question=UpperCAmelCase_				,	top_k=2							)
								self.assertEqual(
								    nested_simplify(UpperCAmelCase_				,	decimals=4							)				,	[
								        {"""score""": 0.9_9_7_4, """answer""": """1110212019""", """start""": 23, """end""": 23},
								        {"""score""": 0.9_9_4_8, """answer""": """us-001""", """start""": 16, """end""": 16},
								    ]				,	)

								lowerCamelCase			=dqa_pipeline({"""image""": image, """question""": question}				,	top_k=2							)
								self.assertEqual(
								    nested_simplify(UpperCAmelCase_				,	decimals=4							)				,	[
								        {"""score""": 0.9_9_7_4, """answer""": """1110212019""", """start""": 23, """end""": 23},
								        {"""score""": 0.9_9_4_8, """answer""": """us-001""", """start""": 16, """end""": 16},
								    ]				,	)

								lowerCamelCase			=dqa_pipeline(
								    [{"""image""": image, """question""": question}, {"""image""": image, """question""": question}]				,	top_k=2							)
								self.assertEqual(
								    nested_simplify(UpperCAmelCase_				,	decimals=4							)				,	[
								        [
								            {"""score""": 0.9_9_7_4, """answer""": """1110212019""", """start""": 23, """end""": 23},
								            {"""score""": 0.9_9_4_8, """answer""": """us-001""", """start""": 16, """end""": 16},
								        ]
								    ]
								    * 2				,	)


				@slow
				@require_torch
				@require_pytesseract
				@require_vision
				def _snake_case    (						self							):
								lowerCamelCase			=AutoTokenizer.from_pretrained(
								    """impira/layoutlm-document-qa"""				,	revision="""3dc6de3"""				,	add_prefix_space=UpperCAmelCase_							)
								lowerCamelCase			=pipeline(
								    """document-question-answering"""				,	model="""impira/layoutlm-document-qa"""				,	tokenizer=UpperCAmelCase_				,	revision="""3dc6de3"""				,	)
								lowerCamelCase			=INVOICE_URL
								lowerCamelCase			="""What is the invoice number?"""

								lowerCamelCase			=dqa_pipeline(image=UpperCAmelCase_				,	question=UpperCAmelCase_				,	top_k=2							)
								self.assertEqual(
								    nested_simplify(UpperCAmelCase_				,	decimals=4							)				,	[
								        {"""score""": 0.4_2_5_1, """answer""": """us-001""", """start""": 16, """end""": 16},
								        {"""score""": 0.0_8_1_9, """answer""": """1110212019""", """start""": 23, """end""": 23},
								    ]				,	)

								lowerCamelCase			=dqa_pipeline({"""image""": image, """question""": question}				,	top_k=2							)
								self.assertEqual(
								    nested_simplify(UpperCAmelCase_				,	decimals=4							)				,	[
								        {"""score""": 0.4_2_5_1, """answer""": """us-001""", """start""": 16, """end""": 16},
								        {"""score""": 0.0_8_1_9, """answer""": """1110212019""", """start""": 23, """end""": 23},
								    ]				,	)

								lowerCamelCase			=dqa_pipeline(
								    [{"""image""": image, """question""": question}, {"""image""": image, """question""": question}]				,	top_k=2							)
								self.assertEqual(
								    nested_simplify(UpperCAmelCase_				,	decimals=4							)				,	[
								        [
								            {"""score""": 0.4_2_5_1, """answer""": """us-001""", """start""": 16, """end""": 16},
								            {"""score""": 0.0_8_1_9, """answer""": """1110212019""", """start""": 23, """end""": 23},
								        ]
								    ]
								    * 2				,	)

								lowerCamelCase			=list(zip(*apply_tesseract(load_image(UpperCAmelCase_							)				,	UpperCAmelCase_				,	""""""							)							)							)

								# This model should also work if `image` is set to None
								lowerCamelCase			=dqa_pipeline({"""image""": None, """word_boxes""": word_boxes, """question""": question}				,	top_k=2							)
								self.assertEqual(
								    nested_simplify(UpperCAmelCase_				,	decimals=4							)				,	[
								        {"""score""": 0.4_2_5_1, """answer""": """us-001""", """start""": 16, """end""": 16},
								        {"""score""": 0.0_8_1_9, """answer""": """1110212019""", """start""": 23, """end""": 23},
								    ]				,	)


				@slow
				@require_torch
				@require_pytesseract
				@require_vision
				def _snake_case    (						self							):
								lowerCamelCase			=AutoTokenizer.from_pretrained(
								    """impira/layoutlm-document-qa"""				,	revision="""3dc6de3"""				,	add_prefix_space=UpperCAmelCase_							)
								lowerCamelCase			=pipeline(
								    """document-question-answering"""				,	model="""impira/layoutlm-document-qa"""				,	tokenizer=UpperCAmelCase_				,	revision="""3dc6de3"""				,	max_seq_len=50				,	)
								lowerCamelCase			=INVOICE_URL
								lowerCamelCase			="""What is the invoice number?"""

								lowerCamelCase			=dqa_pipeline(image=UpperCAmelCase_				,	question=UpperCAmelCase_				,	top_k=2							)
								self.assertEqual(
								    nested_simplify(UpperCAmelCase_				,	decimals=4							)				,	[
								        {"""score""": 0.9_9_9_9, """answer""": """us-001""", """start""": 16, """end""": 16},
								        {"""score""": 0.9_9_9_8, """answer""": """us-001""", """start""": 16, """end""": 16},
								    ]				,	)

								lowerCamelCase			=dqa_pipeline(
								    [{"""image""": image, """question""": question}, {"""image""": image, """question""": question}]				,	top_k=2							)
								self.assertEqual(
								    nested_simplify(UpperCAmelCase_				,	decimals=4							)				,	[
								        [
								            {"""score""": 0.9_9_9_9, """answer""": """us-001""", """start""": 16, """end""": 16},
								            {"""score""": 0.9_9_9_8, """answer""": """us-001""", """start""": 16, """end""": 16},
								        ]
								    ]
								    * 2				,	)

								lowerCamelCase			=list(zip(*apply_tesseract(load_image(UpperCAmelCase_							)				,	UpperCAmelCase_				,	""""""							)							)							)

								# This model should also work if `image` is set to None
								lowerCamelCase			=dqa_pipeline({"""image""": None, """word_boxes""": word_boxes, """question""": question}				,	top_k=2							)
								self.assertEqual(
								    nested_simplify(UpperCAmelCase_				,	decimals=4							)				,	[
								        {"""score""": 0.9_9_9_9, """answer""": """us-001""", """start""": 16, """end""": 16},
								        {"""score""": 0.9_9_9_8, """answer""": """us-001""", """start""": 16, """end""": 16},
								    ]				,	)


				@slow
				@require_torch
				def _snake_case    (						self							):
								lowerCamelCase			=pipeline(
								    """document-question-answering"""				,	model="""naver-clova-ix/donut-base-finetuned-docvqa"""				,	tokenizer=AutoTokenizer.from_pretrained("""naver-clova-ix/donut-base-finetuned-docvqa"""							)				,	feature_extractor="""naver-clova-ix/donut-base-finetuned-docvqa"""				,	)

								lowerCamelCase			=INVOICE_URL
								lowerCamelCase			="""What is the invoice number?"""
								lowerCamelCase			=dqa_pipeline(image=UpperCAmelCase_				,	question=UpperCAmelCase_				,	top_k=2							)
								self.assertEqual(nested_simplify(UpperCAmelCase_				,	decimals=4							)				,	[{"""answer""": """us-001"""}]							)


				@require_tf
				@unittest.skip("""Document question answering not implemented in TF"""							)
				def _snake_case    (						self							):
								pass












import argparse
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.local_sgd import LocalSGD


########################################################################
# This is a fully working simple example to use Accelerate
# with LocalSGD, which is a method to synchronize model
# parameters every K batches. It is different, but complementary
# to gradient accumulation.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


UpperCAmelCase__  :     Union[str, Any]   			=16
UpperCAmelCase__  :     Any   			=32





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase = 16			)  ->					int:
				lowerCamelCase			=AutoTokenizer.from_pretrained("""bert-base-cased"""			)
				lowerCamelCase			=load_dataset("""glue"""     , """mrpc"""			)

				def tokenize_function(_UpperCAmelCase			):
								# max_length=None => use the model max length (it's actually the default)
								lowerCamelCase			=tokenizer(examples["""sentence1"""]     , examples["""sentence2"""]     , truncation=_UpperCAmelCase     , max_length=_UpperCAmelCase			)
								return outputs

				# Apply the method we just defined to all the examples in all the splits of the dataset
				# starting with the main process first:
				with accelerator.main_process_first():
								lowerCamelCase			=datasets.map(
								    _UpperCAmelCase     , batched=_UpperCAmelCase     , remove_columns=["""idx""", """sentence1""", """sentence2"""]     , )

				# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
				# transformers library
				lowerCamelCase			=tokenized_datasets.rename_column("""label"""     , """labels"""			)

				def collate_fn(_UpperCAmelCase			):
								# On TPU it's best to pad everything to the same length or training will be very slow.
								lowerCamelCase			=1_28 if accelerator.distributed_type == DistributedType.TPU else None
								# When using mixed precision we want round multiples of 8/16
								if accelerator.mixed_precision == "fp8":
												lowerCamelCase			=16
								elif accelerator.mixed_precision != "no":
												lowerCamelCase			=8
								else:
												lowerCamelCase			=None

								return tokenizer.pad(
								    _UpperCAmelCase     , padding="""longest"""     , max_length=_UpperCAmelCase     , pad_to_multiple_of=_UpperCAmelCase     , return_tensors="""pt"""     , )

				# Instantiate dataloaders.
				lowerCamelCase			=DataLoader(
				    tokenized_datasets["""train"""]     , shuffle=_UpperCAmelCase     , collate_fn=_UpperCAmelCase     , batch_size=_UpperCAmelCase			)
				lowerCamelCase			=DataLoader(
				    tokenized_datasets["""validation"""]     , shuffle=_UpperCAmelCase     , collate_fn=_UpperCAmelCase     , batch_size=_UpperCAmelCase			)

				return train_dataloader, eval_dataloader


# For testing only
if os.environ.get('''TESTING_MOCKED_DATALOADERS''', None) == "1":
							from accelerate.test_utils.training import mocked_dataloaders

							UpperCAmelCase__  :     Dict   			=mocked_dataloaders  # noqa: F811





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase			)  ->					Optional[int]:
				# For testing only
				if os.environ.get("""TESTING_MOCKED_DATALOADERS"""     , _UpperCAmelCase			) == "1":
								lowerCamelCase			=2
				# New Code #
				lowerCamelCase			=int(args.gradient_accumulation_steps			)
				lowerCamelCase			=int(args.local_sgd_steps			)
				# Initialize accelerator
				lowerCamelCase			=Accelerator(
				    cpu=args.cpu     , mixed_precision=args.mixed_precision     , gradient_accumulation_steps=_UpperCAmelCase			)
				if accelerator.distributed_type not in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
								raise NotImplementedError("""LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)"""			)
				# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
				lowerCamelCase			=config["""lr"""]
				lowerCamelCase			=int(config["""num_epochs"""]			)
				lowerCamelCase			=int(config["""seed"""]			)
				lowerCamelCase			=int(config["""batch_size"""]			)

				lowerCamelCase			=evaluate.load("""glue"""     , """mrpc"""			)

				set_seed(_UpperCAmelCase			)
				lowerCamelCase	,					lowerCamelCase			=get_dataloaders(_UpperCAmelCase     , _UpperCAmelCase			)
				# Instantiate the model (we build the model here so that the seed also control new weights initialization)
				lowerCamelCase			=AutoModelForSequenceClassification.from_pretrained("""bert-base-cased"""     , return_dict=_UpperCAmelCase			)

				# We could avoid this line since the accelerator is set with `device_placement=True` (default value).
				# Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
				# creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
				lowerCamelCase			=model.to(accelerator.device			)

				# Instantiate optimizer
				lowerCamelCase			=AdamW(params=model.parameters()     , lr=_UpperCAmelCase			)

				# Instantiate scheduler
				lowerCamelCase			=get_linear_schedule_with_warmup(
				    optimizer=_UpperCAmelCase     , num_warmup_steps=1_00     , num_training_steps=(len(_UpperCAmelCase			) * num_epochs)     , )

				# Prepare everything
				# There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
				# prepare method.
				lowerCamelCase	,					lowerCamelCase	,					lowerCamelCase	,					lowerCamelCase	,					lowerCamelCase			=accelerator.prepare(
				    _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase			)

				# Now we train the model
				for epoch in range(_UpperCAmelCase			):
								model.train()
								with LocalSGD(
								    accelerator=_UpperCAmelCase     , model=_UpperCAmelCase     , local_sgd_steps=_UpperCAmelCase     , enabled=local_sgd_steps is not None			) as local_sgd:
												for step, batch in enumerate(_UpperCAmelCase			):
																# We could avoid this line since we set the accelerator with `device_placement=True`.
																batch.to(accelerator.device			)
																# New code #
																# We use the new `accumulate` context manager to perform gradient accumulation
																# We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.
																with accelerator.accumulate(_UpperCAmelCase			):
																				lowerCamelCase			=model(**_UpperCAmelCase			)
																				lowerCamelCase			=output.loss
																				accelerator.backward(_UpperCAmelCase			)
																				optimizer.step()
																				lr_scheduler.step()
																				optimizer.zero_grad()
																				# LocalSGD-specific line
																				local_sgd.step()

								model.eval()
								for step, batch in enumerate(_UpperCAmelCase			):
												# We could avoid this line since we set the accelerator with `device_placement=True`.
												batch.to(accelerator.device			)
												with torch.no_grad():
																lowerCamelCase			=model(**_UpperCAmelCase			)
												lowerCamelCase			=outputs.logits.argmax(dim=-1			)
												lowerCamelCase	,					lowerCamelCase			=accelerator.gather_for_metrics((predictions, batch["""labels"""])			)
												metric.add_batch(
												    predictions=_UpperCAmelCase     , references=_UpperCAmelCase     , )

								lowerCamelCase			=metric.compute()
								# Use accelerator.print to print only on the main process.
								accelerator.print(F"""epoch {epoch}:"""     , _UpperCAmelCase			)





def        _lowercase   (  )  ->					Any:
				lowerCamelCase			=argparse.ArgumentParser(description="""Simple example of training script."""			)
				parser.add_argument(
				    """--mixed_precision"""     , type=_UpperCAmelCase     , default=_UpperCAmelCase     , choices=["""no""", """fp16""", """bf16""", """fp8"""]     , help="""Whether to use mixed precision. Choose"""
				    """between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."""
				    """and an Nvidia Ampere GPU."""     , )
				# New Code #
				parser.add_argument(
				    """--gradient_accumulation_steps"""     , type=_UpperCAmelCase     , default=1     , help="""The number of minibatches to be ran before gradients are accumulated."""     , )
				parser.add_argument(
				    """--local_sgd_steps"""     , type=_UpperCAmelCase     , default=8     , help="""Number of local SGD steps or None to disable local SGD"""			)
				parser.add_argument("""--cpu"""     , action="""store_true"""     , help="""If passed, will train on the CPU."""			)
				lowerCamelCase			=parser.parse_args()
				lowerCamelCase			={"""lr""": 2e-5, """num_epochs""": 3, """seed""": 42, """batch_size""": 16}
				training_function(_UpperCAmelCase     , _UpperCAmelCase			)


if __name__ == "__main__":
							main()





