
import os
import tempfile
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
 import torch
 from torch import nn

 from transformers import (
     Adafactor,
     AdamW,
     get_constant_schedule,
     get_constant_schedule_with_warmup,
     get_cosine_schedule_with_warmup,
     get_cosine_with_hard_restarts_schedule_with_warmup,
     get_inverse_sqrt_schedule,
     get_linear_schedule_with_warmup,
     get_polynomial_decay_schedule_with_warmup,
 )


def     UpperCAmelCase_  (						__lowerCAmelCase	,					__lowerCAmelCase=10	) ->						Tuple:
    __lowercase		:			Optional[int]	    =       []
    for _ in range(__lowerCAmelCase	):
        lrs.append(scheduler.get_lr()[0]	)
        scheduler.step()
    return lrs


def     UpperCAmelCase_  (						__lowerCAmelCase	,					__lowerCAmelCase=10	) ->						Any:
    __lowercase		:			Union[str, Any]	    =       []
    for step in range(__lowerCAmelCase	):
        lrs.append(scheduler.get_lr()[0]	)
        scheduler.step()
        if step == num_steps // 2:
            with tempfile.TemporaryDirectory() as tmpdirname:
                __lowercase		:			Union[str, Any]	    =       os.path.join(__lowerCAmelCase	,					'''schedule.bin'''	)
                torch.save(scheduler.state_dict()	,					__lowerCAmelCase	)

                __lowercase		:			List[Any]	    =       torch.load(__lowerCAmelCase	)
                scheduler.load_state_dict(__lowerCAmelCase	)
    return lrs


@require_torch
class 						__lowerCAmelCase	(	unittest.TestCase      ):



      """simple docstring"""







      def 		snake_case_			(      self		:       List[Any]				,       _snake_case		:       int				,       _snake_case		:       List[Any]				,       _snake_case		:       List[Any]				):
          self.assertEqual(len(_snake_case				)				,       len(_snake_case				)				)
          for a, b in zip(_snake_case				,       _snake_case				):
              self.assertAlmostEqual(_snake_case				,       _snake_case				,       delta=_snake_case				)




      def 		snake_case_			(      self		:       Union[str, Any]				):
          __lowercase		:			List[Any]	    =       torch.tensor([0.1, -0.2, -0.1]				,       requires_grad=_snake_case				)
          __lowercase		:			Union[str, Any]	    =       torch.tensor([0.4, 0.2, -0.5]				)
          __lowercase		:			Optional[Any]	    =       nn.MSELoss()
          # No warmup, constant schedule, no gradient clipping
          __lowercase		:			Optional[int]	    =       AdamW(params=[w]				,       lr=2E-1				,       weight_decay=0.0				)
          for _ in range(100				):
              __lowercase		:			List[Any]	    =       criterion(_snake_case				,       _snake_case				)
              loss.backward()
              optimizer.step()
              w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
              w.grad.zero_()
          self.assertListAlmostEqual(w.tolist()				,       [0.4, 0.2, -0.5]				,       tol=1E-2				)




      def 		snake_case_			(      self		:       str				):
          __lowercase		:			int	    =       torch.tensor([0.1, -0.2, -0.1]				,       requires_grad=_snake_case				)
          __lowercase		:			List[str]	    =       torch.tensor([0.4, 0.2, -0.5]				)
          __lowercase		:			List[str]	    =       nn.MSELoss()
          # No warmup, constant schedule, no gradient clipping
          __lowercase		:			int	    =       Adafactor(
              params=[w]				,       lr=1E-2				,       eps=(1E-3_0, 1E-3)				,       clip_threshold=1.0				,       decay_rate=-0.8				,       betaa=_snake_case				,       weight_decay=0.0				,       relative_step=_snake_case				,       scale_parameter=_snake_case				,       warmup_init=_snake_case				,       )
          for _ in range(1000				):
              __lowercase		:			Dict	    =       criterion(_snake_case				,       _snake_case				)
              loss.backward()
              optimizer.step()
              w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
              w.grad.zero_()
          self.assertListAlmostEqual(w.tolist()				,       [0.4, 0.2, -0.5]				,       tol=1E-2				)






@require_torch
class 						__lowerCAmelCase	(	unittest.TestCase      ):



      """simple docstring"""







      A__	:			int        =    nn.Linear(5_0 , 5_0      ) if is_torch_available() else None
      A__	:			int        =    AdamW(m.parameters() , lr=10.0      ) if is_torch_available() else None
      A__	:			List[str]        =    1_0




      def 		snake_case_			(      self		:       int				,       _snake_case		:       List[Any]				,       _snake_case		:       str				,       _snake_case		:       Optional[Any]				,       _snake_case		:       Optional[int]=None				):
          self.assertEqual(len(_snake_case				)				,       len(_snake_case				)				)
          for a, b in zip(_snake_case				,       _snake_case				):
              self.assertAlmostEqual(_snake_case				,       _snake_case				,       delta=_snake_case				,       msg=_snake_case				)




      def 		snake_case_			(      self		:       List[Any]				):
          __lowercase		:			Union[str, Any]	    =       {'''num_warmup_steps''': 2, '''num_training_steps''': 10}
          # schedulers doct format
          # function: (sched_args_dict, expected_learning_rates)
          __lowercase		:			int	    =       {
              get_constant_schedule: ({}, [10.0] * self.num_steps),
              get_constant_schedule_with_warmup: (
                  {'''num_warmup_steps''': 4},
                  [0.0, 2.5, 5.0, 7.5, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0],
              ),
              get_linear_schedule_with_warmup: (
                  {**common_kwargs},
                  [0.0, 5.0, 10.0, 8.75, 7.5, 6.25, 5.0, 3.75, 2.5, 1.25],
              ),
              get_cosine_schedule_with_warmup: (
                  {**common_kwargs},
                  [0.0, 5.0, 10.0, 9.61, 8.53, 6.91, 5.0, 3.08, 1.46, 0.38],
              ),
              get_cosine_with_hard_restarts_schedule_with_warmup: (
                  {**common_kwargs, '''num_cycles''': 2},
                  [0.0, 5.0, 10.0, 8.53, 5.0, 1.46, 10.0, 8.53, 5.0, 1.46],
              ),
              get_polynomial_decay_schedule_with_warmup: (
                  {**common_kwargs, '''power''': 2.0, '''lr_end''': 1E-7},
                  [0.0, 5.0, 10.0, 7.6_56, 5.6_25, 3.9_06, 2.5, 1.4_06, 0.6_25, 0.1_56],
              ),
              get_inverse_sqrt_schedule: (
                  {'''num_warmup_steps''': 2},
                  [0.0, 5.0, 10.0, 8.1_65, 7.0_71, 6.3_25, 5.7_74, 5.3_45, 5.0, 4.7_14],
              ),
          }

          for scheduler_func, data in scheds.items():
              __lowercase    ,      __lowercase		:			Any	    =       data

              __lowercase		:			Dict	    =       scheduler_func(self.optimizer				,       **_snake_case				)
              self.assertEqual(len([scheduler.get_lr()[0]]				)				,       1				)
              __lowercase		:			Optional[int]	    =       unwrap_schedule(_snake_case				,       self.num_steps				)
              self.assertListAlmostEqual(
                  _snake_case				,       _snake_case				,       tol=1E-2				,       msg=F'failed for {scheduler_func} in normal scheduler'				,       )

              __lowercase		:			Optional[int]	    =       scheduler_func(self.optimizer				,       **_snake_case				)
              if scheduler_func.__name__ != "get_constant_schedule":
                  LambdaScheduleWrapper.wrap_scheduler(_snake_case				)  # wrap to test picklability of the schedule
              __lowercase		:			Union[str, Any]	    =       unwrap_and_save_reload_schedule(_snake_case				,       self.num_steps				)
              self.assertListEqual(_snake_case				,       _snake_case				,       msg=F'failed for {scheduler_func} in save and reload'				)






class 						__lowerCAmelCase	:



      """simple docstring"""







      def __init__(      self		:       Optional[int]				,       _snake_case		:       Optional[int]				):
          __lowercase		:			str	    =       fn




      def __call__(      self		:       List[Any]				,       *_snake_case		:       Tuple				,       **_snake_case		:       Any				):
          return self.fn(*_snake_case				,       **_snake_case				)




      @classmethod
      def 		snake_case_			(      self		:       Union[str, Any]				,       _snake_case		:       List[Any]				):
          __lowercase		:			Optional[Any]	    =       list(map(self				,       scheduler.lr_lambdas				)				)




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 						__lowerCAmelCase	:



       """simple docstring"""







       @staticmethod
       def 		snake_case_			(      *_snake_case		:       Union[str, Any]				,       **_snake_case		:       str				):
           pass






 def     UpperCAmelCase_  (						__lowerCAmelCase	) ->						Union[str, 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.
__lowerCAmelCase					:			Dict			      =		(
    "https://huggingface.co/spaces/impira/docquery/resolve/2f6c96314dc84dfda62d40de9da55f2f5165d403/invoice.png"
)


@is_pipeline_test
@require_torch
@require_vision
class 						__lowerCAmelCase	(	unittest.TestCase      ):



      """simple docstring"""







      A__	:			int        =    MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING




      @require_pytesseract
      @require_vision
      def 		snake_case_			(      self		:       List[Any]				,       _snake_case		:       Any				,       _snake_case		:       Dict				,       _snake_case		:       Any				):
          __lowercase		:			List[str]	    =       pipeline(
              '''document-question-answering'''				,       model=_snake_case				,       tokenizer=_snake_case				,       image_processor=_snake_case				)

          __lowercase		:			int	    =       INVOICE_URL
          __lowercase		:			Optional[int]	    =       list(zip(*apply_tesseract(load_image(_snake_case				)				,       _snake_case				,       ''''''				)				)				)
          __lowercase		:			List[str]	    =       '''What is the placebo?'''
          __lowercase		:			Tuple	    =       [
              {
                  '''image''': load_image(_snake_case				),
                  '''question''': question,
              },
              {
                  '''image''': image,
                  '''question''': question,
              },
              {
                  '''image''': image,
                  '''question''': question,
                  '''word_boxes''': word_boxes,
              },
          ]
          return dqa_pipeline, examples




      def 		snake_case_			(      self		:       List[str]				,       _snake_case		:       int				,       _snake_case		:       List[Any]				):
          __lowercase		:			Optional[Any]	    =       dqa_pipeline(_snake_case				,       top_k=2				)
          self.assertEqual(
              _snake_case				,       [
                  [
                      {'''score''': ANY(_snake_case				), '''answer''': ANY(_snake_case				), '''start''': ANY(_snake_case				), '''end''': ANY(_snake_case				)},
                      {'''score''': ANY(_snake_case				), '''answer''': ANY(_snake_case				), '''start''': ANY(_snake_case				), '''end''': ANY(_snake_case				)},
                  ]
              ]
              * 3				,       )




      @require_torch
      @require_detectrona
      @require_pytesseract
      def 		snake_case_			(      self		:       Optional[int]				):
          __lowercase		:			List[str]	    =       pipeline('''document-question-answering'''				,       model='''hf-internal-testing/tiny-random-layoutlmv2'''				)
          __lowercase		:			Any	    =       INVOICE_URL
          __lowercase		:			str	    =       '''How many cats are there?'''

          __lowercase		:			Optional[int]	    =       [
              {'''score''': 0.00_01, '''answer''': '''oy 2312/2019''', '''start''': 38, '''end''': 39},
              {'''score''': 0.00_01, '''answer''': '''oy 2312/2019 DUE''', '''start''': 38, '''end''': 40},
          ]
          __lowercase		:			Optional[Any]	    =       dqa_pipeline(image=_snake_case				,       question=_snake_case				,       top_k=2				)
          self.assertEqual(nested_simplify(_snake_case				,       decimals=4				)				,       _snake_case				)

          __lowercase		:			Optional[Any]	    =       dqa_pipeline({'''image''': image, '''question''': question}				,       top_k=2				)
          self.assertEqual(nested_simplify(_snake_case				,       decimals=4				)				,       _snake_case				)

          # This image does not detect ANY text in it, meaning layoutlmv2 should fail.
          # Empty answer probably
          __lowercase		:			Optional[int]	    =       '''./tests/fixtures/tests_samples/COCO/000000039769.png'''
          __lowercase		:			str	    =       dqa_pipeline(image=_snake_case				,       question=_snake_case				,       top_k=2				)
          self.assertEqual(_snake_case				,       []				)

          # We can optionnally pass directly the words and bounding boxes
          __lowercase		:			List[Any]	    =       '''./tests/fixtures/tests_samples/COCO/000000039769.png'''
          __lowercase		:			Dict	    =       []
          __lowercase		:			List[str]	    =       []
          __lowercase		:			Tuple	    =       dqa_pipeline(image=_snake_case				,       question=_snake_case				,       words=_snake_case				,       boxes=_snake_case				,       top_k=2				)
          self.assertEqual(_snake_case				,       []				)




      @slow
      @require_torch
      @require_detectrona
      @require_pytesseract
      def 		snake_case_			(      self		:       int				):
          __lowercase		:			List[Any]	    =       pipeline(
              '''document-question-answering'''				,       model='''tiennvcs/layoutlmv2-base-uncased-finetuned-docvqa'''				,       revision='''9977165'''				,       )
          __lowercase		:			List[str]	    =       INVOICE_URL
          __lowercase		:			Optional[Any]	    =       '''What is the invoice number?'''

          __lowercase		:			int	    =       dqa_pipeline(image=_snake_case				,       question=_snake_case				,       top_k=2				)
          self.assertEqual(
              nested_simplify(_snake_case				,       decimals=4				)				,       [
                  {'''score''': 0.99_44, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
                  {'''score''': 0.00_09, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
              ]				,       )

          __lowercase		:			Union[str, Any]	    =       dqa_pipeline({'''image''': image, '''question''': question}				,       top_k=2				)
          self.assertEqual(
              nested_simplify(_snake_case				,       decimals=4				)				,       [
                  {'''score''': 0.99_44, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
                  {'''score''': 0.00_09, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
              ]				,       )

          __lowercase		:			List[str]	    =       dqa_pipeline(
              [{'''image''': image, '''question''': question}, {'''image''': image, '''question''': question}]				,       top_k=2				)
          self.assertEqual(
              nested_simplify(_snake_case				,       decimals=4				)				,       [
                  [
                      {'''score''': 0.99_44, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
                      {'''score''': 0.00_09, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
                  ],
              ]
              * 2				,       )




      @slow
      @require_torch
      @require_detectrona
      @require_pytesseract
      def 		snake_case_			(      self		:       int				):
          __lowercase		:			Optional[Any]	    =       pipeline(
              '''document-question-answering'''				,       model='''tiennvcs/layoutlmv2-base-uncased-finetuned-docvqa'''				,       revision='''9977165'''				,       max_seq_len=50				,       )
          __lowercase		:			Union[str, Any]	    =       INVOICE_URL
          __lowercase		:			Any	    =       '''What is the invoice number?'''

          __lowercase		:			Any	    =       dqa_pipeline(image=_snake_case				,       question=_snake_case				,       top_k=2				)
          self.assertEqual(
              nested_simplify(_snake_case				,       decimals=4				)				,       [
                  {'''score''': 0.99_74, '''answer''': '''1110212019''', '''start''': 23, '''end''': 23},
                  {'''score''': 0.99_48, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
              ]				,       )

          __lowercase		:			Dict	    =       dqa_pipeline({'''image''': image, '''question''': question}				,       top_k=2				)
          self.assertEqual(
              nested_simplify(_snake_case				,       decimals=4				)				,       [
                  {'''score''': 0.99_74, '''answer''': '''1110212019''', '''start''': 23, '''end''': 23},
                  {'''score''': 0.99_48, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
              ]				,       )

          __lowercase		:			List[Any]	    =       dqa_pipeline(
              [{'''image''': image, '''question''': question}, {'''image''': image, '''question''': question}]				,       top_k=2				)
          self.assertEqual(
              nested_simplify(_snake_case				,       decimals=4				)				,       [
                  [
                      {'''score''': 0.99_74, '''answer''': '''1110212019''', '''start''': 23, '''end''': 23},
                      {'''score''': 0.99_48, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
                  ]
              ]
              * 2				,       )




      @slow
      @require_torch
      @require_pytesseract
      @require_vision
      def 		snake_case_			(      self		:       str				):
          __lowercase		:			List[Any]	    =       AutoTokenizer.from_pretrained(
              '''impira/layoutlm-document-qa'''				,       revision='''3dc6de3'''				,       add_prefix_space=_snake_case				)
          __lowercase		:			Union[str, Any]	    =       pipeline(
              '''document-question-answering'''				,       model='''impira/layoutlm-document-qa'''				,       tokenizer=_snake_case				,       revision='''3dc6de3'''				,       )
          __lowercase		:			Tuple	    =       INVOICE_URL
          __lowercase		:			Any	    =       '''What is the invoice number?'''

          __lowercase		:			Optional[Any]	    =       dqa_pipeline(image=_snake_case				,       question=_snake_case				,       top_k=2				)
          self.assertEqual(
              nested_simplify(_snake_case				,       decimals=4				)				,       [
                  {'''score''': 0.42_51, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
                  {'''score''': 0.08_19, '''answer''': '''1110212019''', '''start''': 23, '''end''': 23},
              ]				,       )

          __lowercase		:			List[Any]	    =       dqa_pipeline({'''image''': image, '''question''': question}				,       top_k=2				)
          self.assertEqual(
              nested_simplify(_snake_case				,       decimals=4				)				,       [
                  {'''score''': 0.42_51, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
                  {'''score''': 0.08_19, '''answer''': '''1110212019''', '''start''': 23, '''end''': 23},
              ]				,       )

          __lowercase		:			int	    =       dqa_pipeline(
              [{'''image''': image, '''question''': question}, {'''image''': image, '''question''': question}]				,       top_k=2				)
          self.assertEqual(
              nested_simplify(_snake_case				,       decimals=4				)				,       [
                  [
                      {'''score''': 0.42_51, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
                      {'''score''': 0.08_19, '''answer''': '''1110212019''', '''start''': 23, '''end''': 23},
                  ]
              ]
              * 2				,       )

          __lowercase		:			Union[str, Any]	    =       list(zip(*apply_tesseract(load_image(_snake_case				)				,       _snake_case				,       ''''''				)				)				)

          # This model should also work if `image` is set to None
          __lowercase		:			Union[str, Any]	    =       dqa_pipeline({'''image''': None, '''word_boxes''': word_boxes, '''question''': question}				,       top_k=2				)
          self.assertEqual(
              nested_simplify(_snake_case				,       decimals=4				)				,       [
                  {'''score''': 0.42_51, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
                  {'''score''': 0.08_19, '''answer''': '''1110212019''', '''start''': 23, '''end''': 23},
              ]				,       )




      @slow
      @require_torch
      @require_pytesseract
      @require_vision
      def 		snake_case_			(      self		:       List[str]				):
          __lowercase		:			Optional[int]	    =       AutoTokenizer.from_pretrained(
              '''impira/layoutlm-document-qa'''				,       revision='''3dc6de3'''				,       add_prefix_space=_snake_case				)
          __lowercase		:			List[str]	    =       pipeline(
              '''document-question-answering'''				,       model='''impira/layoutlm-document-qa'''				,       tokenizer=_snake_case				,       revision='''3dc6de3'''				,       max_seq_len=50				,       )
          __lowercase		:			Tuple	    =       INVOICE_URL
          __lowercase		:			Optional[int]	    =       '''What is the invoice number?'''

          __lowercase		:			Any	    =       dqa_pipeline(image=_snake_case				,       question=_snake_case				,       top_k=2				)
          self.assertEqual(
              nested_simplify(_snake_case				,       decimals=4				)				,       [
                  {'''score''': 0.99_99, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
                  {'''score''': 0.99_98, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
              ]				,       )

          __lowercase		:			Any	    =       dqa_pipeline(
              [{'''image''': image, '''question''': question}, {'''image''': image, '''question''': question}]				,       top_k=2				)
          self.assertEqual(
              nested_simplify(_snake_case				,       decimals=4				)				,       [
                  [
                      {'''score''': 0.99_99, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
                      {'''score''': 0.99_98, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
                  ]
              ]
              * 2				,       )

          __lowercase		:			int	    =       list(zip(*apply_tesseract(load_image(_snake_case				)				,       _snake_case				,       ''''''				)				)				)

          # This model should also work if `image` is set to None
          __lowercase		:			Tuple	    =       dqa_pipeline({'''image''': None, '''word_boxes''': word_boxes, '''question''': question}				,       top_k=2				)
          self.assertEqual(
              nested_simplify(_snake_case				,       decimals=4				)				,       [
                  {'''score''': 0.99_99, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
                  {'''score''': 0.99_98, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
              ]				,       )




      @slow
      @require_torch
      def 		snake_case_			(      self		:       List[Any]				):
          __lowercase		:			int	    =       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'''				,       )

          __lowercase		:			List[Any]	    =       INVOICE_URL
          __lowercase		:			List[Any]	    =       '''What is the invoice number?'''
          __lowercase		:			int	    =       dqa_pipeline(image=_snake_case				,       question=_snake_case				,       top_k=2				)
          self.assertEqual(nested_simplify(_snake_case				,       decimals=4				)				,       [{'''answer''': '''us-001'''}]				)




      @require_tf
      @unittest.skip('''Document question answering not implemented in TF'''				)
      def 		snake_case_			(      self		:       int				):
          pass


