



"""simple docstring"""


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 :
								"""simple docstring"""



								@staticmethod
								def  lowercase__						(		*__UpperCamelCase     :    Dict				,   **__UpperCamelCase     :    Optional[Any]   )->List[str]:
												pass







					def  lowercase	(      _SCREAMING_SNAKE_CASE   :						Union[str, Any]				):






									'''simple docstring'''



									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.
__A   :							Dict										=   (
    "https://huggingface.co/spaces/impira/docquery/resolve/2f6c96314dc84dfda62d40de9da55f2f5165d403/invoice.png"
)
@is_pipeline_test
@require_torch
@require_vision
class 			_a (					unittest.TestCase):
			"""simple docstring"""
			UpperCamelCase__			      =	MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING



			@require_pytesseract
			@require_vision
			def  lowercase__						(		self     :    List[str]				,   __UpperCamelCase     :    Tuple				,   __UpperCamelCase     :    Any				,   __UpperCamelCase     :    str   )->Tuple:
							_UpperCAmelCase					=    pipeline(
							    '''document-question-answering'''				,   model=__UpperCamelCase				,   tokenizer=__UpperCamelCase				,   image_processor=__UpperCamelCase   )

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



			def  lowercase__						(		self     :    str				,   __UpperCamelCase     :    List[Any]				,   __UpperCamelCase     :    Optional[int]   )->Any:
							_UpperCAmelCase					=    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  lowercase__						(		self     :    str   )->Union[str, Any]:
							_UpperCAmelCase					=    pipeline('''document-question-answering'''				,   model='''hf-internal-testing/tiny-random-layoutlmv2'''   )
							_UpperCAmelCase					=    INVOICE_URL
							_UpperCAmelCase					=    '''How many cats are there?'''

							_UpperCAmelCase					=    [
							    {'''score''': 0.0_0_0_1, '''answer''': '''oy 2312/2019''', '''start''': 3_8, '''end''': 3_9},
							    {'''score''': 0.0_0_0_1, '''answer''': '''oy 2312/2019 DUE''', '''start''': 3_8, '''end''': 4_0},
							]
							_UpperCAmelCase					=    dqa_pipeline(image=__UpperCamelCase				,   question=__UpperCamelCase				,   top_k=2   )
							self.assertEqual(nested_simplify(__UpperCamelCase				,   decimals=4   )				,   __UpperCamelCase   )

							_UpperCAmelCase					=    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
							_UpperCAmelCase					=    '''./tests/fixtures/tests_samples/COCO/000000039769.png'''
							_UpperCAmelCase					=    dqa_pipeline(image=__UpperCamelCase				,   question=__UpperCamelCase				,   top_k=2   )
							self.assertEqual(__UpperCamelCase				,   []   )

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



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

							_UpperCAmelCase					=    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''': 1_6, '''end''': 1_6},
							        {'''score''': 0.0_0_0_9, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
							    ]				,   )

							_UpperCAmelCase					=    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''': 1_6, '''end''': 1_6},
							        {'''score''': 0.0_0_0_9, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
							    ]				,   )

							_UpperCAmelCase					=    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''': 1_6, '''end''': 1_6},
							            {'''score''': 0.0_0_0_9, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
							        ],
							    ]
							    * 2				,   )



			@slow
			@require_torch
			@require_detectrona
			@require_pytesseract
			def  lowercase__						(		self     :    int   )->Optional[int]:
							_UpperCAmelCase					=    pipeline(
							    '''document-question-answering'''				,   model='''tiennvcs/layoutlmv2-base-uncased-finetuned-docvqa'''				,   revision='''9977165'''				,   max_seq_len=5_0				,   )
							_UpperCAmelCase					=    INVOICE_URL
							_UpperCAmelCase					=    '''What is the invoice number?'''

							_UpperCAmelCase					=    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''': 2_3, '''end''': 2_3},
							        {'''score''': 0.9_9_4_8, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
							    ]				,   )

							_UpperCAmelCase					=    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''': 2_3, '''end''': 2_3},
							        {'''score''': 0.9_9_4_8, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
							    ]				,   )

							_UpperCAmelCase					=    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''': 2_3, '''end''': 2_3},
							            {'''score''': 0.9_9_4_8, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
							        ]
							    ]
							    * 2				,   )



			@slow
			@require_torch
			@require_pytesseract
			@require_vision
			def  lowercase__						(		self     :    Dict   )->Tuple:
							_UpperCAmelCase					=    AutoTokenizer.from_pretrained(
							    '''impira/layoutlm-document-qa'''				,   revision='''3dc6de3'''				,   add_prefix_space=__UpperCamelCase   )
							_UpperCAmelCase					=    pipeline(
							    '''document-question-answering'''				,   model='''impira/layoutlm-document-qa'''				,   tokenizer=__UpperCamelCase				,   revision='''3dc6de3'''				,   )
							_UpperCAmelCase					=    INVOICE_URL
							_UpperCAmelCase					=    '''What is the invoice number?'''

							_UpperCAmelCase					=    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''': 1_6, '''end''': 1_6},
							        {'''score''': 0.0_8_1_9, '''answer''': '''1110212019''', '''start''': 2_3, '''end''': 2_3},
							    ]				,   )

							_UpperCAmelCase					=    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''': 1_6, '''end''': 1_6},
							        {'''score''': 0.0_8_1_9, '''answer''': '''1110212019''', '''start''': 2_3, '''end''': 2_3},
							    ]				,   )

							_UpperCAmelCase					=    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''': 1_6, '''end''': 1_6},
							            {'''score''': 0.0_8_1_9, '''answer''': '''1110212019''', '''start''': 2_3, '''end''': 2_3},
							        ]
							    ]
							    * 2				,   )

							_UpperCAmelCase					=    list(zip(*apply_tesseract(load_image(__UpperCamelCase   )				,   __UpperCamelCase				,   ''''''   )   )   )

							# This model should also work if `image` is set to None
							_UpperCAmelCase					=    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''': 1_6, '''end''': 1_6},
							        {'''score''': 0.0_8_1_9, '''answer''': '''1110212019''', '''start''': 2_3, '''end''': 2_3},
							    ]				,   )



			@slow
			@require_torch
			@require_pytesseract
			@require_vision
			def  lowercase__						(		self     :    List[str]   )->Any:
							_UpperCAmelCase					=    AutoTokenizer.from_pretrained(
							    '''impira/layoutlm-document-qa'''				,   revision='''3dc6de3'''				,   add_prefix_space=__UpperCamelCase   )
							_UpperCAmelCase					=    pipeline(
							    '''document-question-answering'''				,   model='''impira/layoutlm-document-qa'''				,   tokenizer=__UpperCamelCase				,   revision='''3dc6de3'''				,   max_seq_len=5_0				,   )
							_UpperCAmelCase					=    INVOICE_URL
							_UpperCAmelCase					=    '''What is the invoice number?'''

							_UpperCAmelCase					=    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''': 1_6, '''end''': 1_6},
							        {'''score''': 0.9_9_9_8, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
							    ]				,   )

							_UpperCAmelCase					=    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''': 1_6, '''end''': 1_6},
							            {'''score''': 0.9_9_9_8, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
							        ]
							    ]
							    * 2				,   )

							_UpperCAmelCase					=    list(zip(*apply_tesseract(load_image(__UpperCamelCase   )				,   __UpperCamelCase				,   ''''''   )   )   )

							# This model should also work if `image` is set to None
							_UpperCAmelCase					=    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''': 1_6, '''end''': 1_6},
							        {'''score''': 0.9_9_9_8, '''answer''': '''us-001''', '''start''': 1_6, '''end''': 1_6},
							    ]				,   )



			@slow
			@require_torch
			def  lowercase__						(		self     :    int   )->str:
							_UpperCAmelCase					=    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'''				,   )

							_UpperCAmelCase					=    INVOICE_URL
							_UpperCAmelCase					=    '''What is the invoice number?'''
							_UpperCAmelCase					=    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  lowercase__						(		self     :    Tuple   )->Union[str, Any]:
							pass








"""simple docstring"""


import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only
from utils_rag import save_json







def  lowercase	(      _SCREAMING_SNAKE_CASE   :						Optional[Any]				):






				'''simple docstring'''



				_UpperCAmelCase					=    filter(lambda _SCREAMING_SNAKE_CASE				: p.requires_grad					,     model.parameters()				)
				_UpperCAmelCase					=    sum([np.prod(p.size()				) for p in model_parameters]				)
				return params


__A   :							Dict										=   logging.getLogger(__name__)







def  lowercase	(      _SCREAMING_SNAKE_CASE   :						Optional[Any]					,     _SCREAMING_SNAKE_CASE   :						Any				):






				'''simple docstring'''



				if metric == "rouge2":
								_UpperCAmelCase					=    '''{val_avg_rouge2:.4f}-{step_count}'''
				elif metric == "bleu":
								_UpperCAmelCase					=    '''{val_avg_bleu:.4f}-{step_count}'''
				elif metric == "em":
								_UpperCAmelCase					=    '''{val_avg_em:.4f}-{step_count}'''
				else:
								raise NotImplementedError(
								    f'seq2seq callbacks only support rouge2 and bleu, got {metric}, You can make your own by adding to this'
								    ''' function.'''				)

				_UpperCAmelCase					=    ModelCheckpoint(
				    dirpath=_SCREAMING_SNAKE_CASE					,     filename=_SCREAMING_SNAKE_CASE					,     monitor=f'val_{metric}'					,     mode='''max'''					,     save_top_k=3					,     every_n_epochs=1					,     )
				return checkpoint_callback







def  lowercase	(      _SCREAMING_SNAKE_CASE   :						str					,     _SCREAMING_SNAKE_CASE   :						List[str]				):






				'''simple docstring'''



				return EarlyStopping(
				    monitor=f'val_{metric}'					,     mode='''min''' if '''loss''' in metric else '''max'''					,     patience=_SCREAMING_SNAKE_CASE					,     verbose=_SCREAMING_SNAKE_CASE					,     )
class 			_a (					pl.Callback):
			"""simple docstring"""



			def  lowercase__						(		self     :    List[str]				,   __UpperCamelCase     :    str				,   __UpperCamelCase     :    str   )->Tuple:
							_UpperCAmelCase					=    {F'lr_group_{i}': param['''lr'''] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups   )}
							pl_module.logger.log_metrics(__UpperCamelCase   )



			@rank_zero_only
			def  lowercase__						(		self     :    Union[str, Any]				,   __UpperCamelCase     :    pl.Trainer				,   __UpperCamelCase     :    pl.LightningModule				,   __UpperCamelCase     :    str				,   __UpperCamelCase     :    Tuple=True   )->None:
							logger.info(F'***** {type_path} results at step {trainer.global_step:05d} *****'   )
							_UpperCAmelCase					=    trainer.callback_metrics
							trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ['''log''', '''progress_bar''', '''preds''']}   )
							# Log results
							_UpperCAmelCase					=    Path(pl_module.hparams.output_dir   )
							if type_path == "test":
											_UpperCAmelCase					=    od / '''test_results.txt'''
											_UpperCAmelCase					=    od / '''test_generations.txt'''
							else:
											# this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
											# If people want this it will be easy enough to add back.
											_UpperCAmelCase					=    od / F'{type_path}_results/{trainer.global_step:05d}.txt'
											_UpperCAmelCase					=    od / F'{type_path}_generations/{trainer.global_step:05d}.txt'
											results_file.parent.mkdir(exist_ok=__UpperCamelCase   )
											generations_file.parent.mkdir(exist_ok=__UpperCamelCase   )
							with open(__UpperCamelCase				,   '''a+'''   ) as writer:
											for key in sorted(__UpperCamelCase   ):
															if key in ["log", "progress_bar", "preds"]:
																			continue
															_UpperCAmelCase					=    metrics[key]
															if isinstance(__UpperCamelCase				,   torch.Tensor   ):
																			_UpperCAmelCase					=    val.item()
															_UpperCAmelCase					=    F'{key}: {val:.6f}\n'
															writer.write(__UpperCamelCase   )

							if not save_generations:
											return

							if "preds" in metrics:
											_UpperCAmelCase					=    '''\n'''.join(metrics['''preds''']   )
											generations_file.open('''w+'''   ).write(__UpperCamelCase   )



			@rank_zero_only
			def  lowercase__						(		self     :    Optional[int]				,   __UpperCamelCase     :    int				,   __UpperCamelCase     :    Dict   )->Union[str, Any]:
							try:
											_UpperCAmelCase					=    pl_module.model.model.num_parameters()
							except AttributeError:
											_UpperCAmelCase					=    pl_module.model.num_parameters()

							_UpperCAmelCase					=    count_trainable_parameters(__UpperCamelCase   )
							# mp stands for million parameters
							trainer.logger.log_metrics({'''n_params''': npars, '''mp''': npars / 1e6, '''grad_mp''': n_trainable_pars / 1e6}   )



			@rank_zero_only
			def  lowercase__						(		self     :    str				,   __UpperCamelCase     :    pl.Trainer				,   __UpperCamelCase     :    pl.LightningModule   )->List[Any]:
							save_json(pl_module.metrics				,   pl_module.metrics_save_path   )
							return self._write_logs(__UpperCamelCase				,   __UpperCamelCase				,   '''test'''   )



			@rank_zero_only
			def  lowercase__						(		self     :    Optional[Any]				,   __UpperCamelCase     :    pl.Trainer				,   __UpperCamelCase     :    List[str]   )->Optional[Any]:
							save_json(pl_module.metrics				,   pl_module.metrics_save_path   )
							# Uncommenting this will save val generations
							# return self._write_logs(trainer, pl_module, "valid")



