


"""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 _lowerCAmelCase    :







     @staticmethod
     def 			lowerCamelCase		(     *UpperCamelCase__						,     **UpperCamelCase__						)     ->		Union[str, Any]:

            '''simple docstring'''




            pass







  def   __lowerCAmelCase						(						lowercase  :  str			)     ->						Dict:


         """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.
__snake_case             =   (
    """https://huggingface.co/spaces/impira/docquery/resolve/2f6c96314dc84dfda62d40de9da55f2f5165d403/invoice.png"""
)







@is_pipeline_test
@require_torch
@require_vision
class _lowerCAmelCase    (				unittest.TestCase     ):
   __UpperCAmelCase			:	Any							  =  MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING







   @require_pytesseract
   @require_vision
   def 			lowerCamelCase		(     self						,     UpperCamelCase__						,     UpperCamelCase__						,     UpperCamelCase__						)     ->		int:

          '''simple docstring'''




          snake_case   :				Dict				 =   pipeline(
              "document-question-answering"						,     model=UpperCamelCase__						,     tokenizer=UpperCamelCase__						,     image_processor=UpperCamelCase__						)

          snake_case   :				List[Any]				 =   INVOICE_URL
          snake_case   :				int				 =   list(zip(*apply_tesseract(load_image(UpperCamelCase__						)						,     UpperCamelCase__						,     ""						)						)						)
          snake_case   :				Tuple				 =   "What is the placebo?"
          snake_case   :				Optional[Any]				 =   [
              {
                  "image": load_image(UpperCamelCase__						),
                  "question": question,
              },
              {
                  "image": image,
                  "question": question,
              },
              {
                  "image": image,
                  "question": question,
                  "word_boxes": word_boxes,
              },
          ]
          return dqa_pipeline, examples







   def 			lowerCamelCase		(     self						,     UpperCamelCase__						,     UpperCamelCase__						)     ->		Optional[Any]:

          '''simple docstring'''




          snake_case   :				List[Any]				 =   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 			lowerCamelCase		(     self						)     ->		int:

          '''simple docstring'''




          snake_case   :				Any				 =   pipeline("document-question-answering"						,     model="hf-internal-testing/tiny-random-layoutlmv2"						)
          snake_case   :				Optional[Any]				 =   INVOICE_URL
          snake_case   :				Optional[int]				 =   "How many cats are there?"

          snake_case   :				Dict				 =   [
              {"score": 0.0001, "answer": "oy 2312/2019", "start": 38, "end": 39},
              {"score": 0.0001, "answer": "oy 2312/2019 DUE", "start": 38, "end": 40},
          ]
          snake_case   :				str				 =   dqa_pipeline(image=UpperCamelCase__						,     question=UpperCamelCase__						,     top_k=2						)
          self.assertEqual(nested_simplify(UpperCamelCase__						,     decimals=4						)						,     UpperCamelCase__						)

          snake_case   :				Optional[Any]				 =   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
          snake_case   :				List[str]				 =   "./tests/fixtures/tests_samples/COCO/000000039769.png"
          snake_case   :				Union[str, Any]				 =   dqa_pipeline(image=UpperCamelCase__						,     question=UpperCamelCase__						,     top_k=2						)
          self.assertEqual(UpperCamelCase__						,     []						)

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







   @slow
   @require_torch
   @require_detectrona
   @require_pytesseract
   def 			lowerCamelCase		(     self						)     ->		Tuple:

          '''simple docstring'''




          snake_case   :				List[Any]				 =   pipeline(
              "document-question-answering"						,     model="tiennvcs/layoutlmv2-base-uncased-finetuned-docvqa"						,     revision="9977165"						,     )
          snake_case   :				Optional[int]				 =   INVOICE_URL
          snake_case   :				List[str]				 =   "What is the invoice number?"

          snake_case   :				Dict				 =   dqa_pipeline(image=UpperCamelCase__						,     question=UpperCamelCase__						,     top_k=2						)
          self.assertEqual(
              nested_simplify(UpperCamelCase__						,     decimals=4						)						,     [
                  {"score": 0.9944, "answer": "us-001", "start": 16, "end": 16},
                  {"score": 0.0009, "answer": "us-001", "start": 16, "end": 16},
              ]						,     )

          snake_case   :				str				 =   dqa_pipeline({"image": image, "question": question}						,     top_k=2						)
          self.assertEqual(
              nested_simplify(UpperCamelCase__						,     decimals=4						)						,     [
                  {"score": 0.9944, "answer": "us-001", "start": 16, "end": 16},
                  {"score": 0.0009, "answer": "us-001", "start": 16, "end": 16},
              ]						,     )

          snake_case   :				int				 =   dqa_pipeline(
              [{"image": image, "question": question}, {"image": image, "question": question}]						,     top_k=2						)
          self.assertEqual(
              nested_simplify(UpperCamelCase__						,     decimals=4						)						,     [
                  [
                      {"score": 0.9944, "answer": "us-001", "start": 16, "end": 16},
                      {"score": 0.0009, "answer": "us-001", "start": 16, "end": 16},
                  ],
              ]
              * 2						,     )







   @slow
   @require_torch
   @require_detectrona
   @require_pytesseract
   def 			lowerCamelCase		(     self						)     ->		Any:

          '''simple docstring'''




          snake_case   :				List[Any]				 =   pipeline(
              "document-question-answering"						,     model="tiennvcs/layoutlmv2-base-uncased-finetuned-docvqa"						,     revision="9977165"						,     max_seq_len=50						,     )
          snake_case   :				Any				 =   INVOICE_URL
          snake_case   :				Tuple				 =   "What is the invoice number?"

          snake_case   :				List[str]				 =   dqa_pipeline(image=UpperCamelCase__						,     question=UpperCamelCase__						,     top_k=2						)
          self.assertEqual(
              nested_simplify(UpperCamelCase__						,     decimals=4						)						,     [
                  {"score": 0.9974, "answer": "1110212019", "start": 23, "end": 23},
                  {"score": 0.9948, "answer": "us-001", "start": 16, "end": 16},
              ]						,     )

          snake_case   :				Union[str, Any]				 =   dqa_pipeline({"image": image, "question": question}						,     top_k=2						)
          self.assertEqual(
              nested_simplify(UpperCamelCase__						,     decimals=4						)						,     [
                  {"score": 0.9974, "answer": "1110212019", "start": 23, "end": 23},
                  {"score": 0.9948, "answer": "us-001", "start": 16, "end": 16},
              ]						,     )

          snake_case   :				Optional[Any]				 =   dqa_pipeline(
              [{"image": image, "question": question}, {"image": image, "question": question}]						,     top_k=2						)
          self.assertEqual(
              nested_simplify(UpperCamelCase__						,     decimals=4						)						,     [
                  [
                      {"score": 0.9974, "answer": "1110212019", "start": 23, "end": 23},
                      {"score": 0.9948, "answer": "us-001", "start": 16, "end": 16},
                  ]
              ]
              * 2						,     )







   @slow
   @require_torch
   @require_pytesseract
   @require_vision
   def 			lowerCamelCase		(     self						)     ->		Dict:

          '''simple docstring'''




          snake_case   :				Optional[Any]				 =   AutoTokenizer.from_pretrained(
              "impira/layoutlm-document-qa"						,     revision="3dc6de3"						,     add_prefix_space=UpperCamelCase__						)
          snake_case   :				List[Any]				 =   pipeline(
              "document-question-answering"						,     model="impira/layoutlm-document-qa"						,     tokenizer=UpperCamelCase__						,     revision="3dc6de3"						,     )
          snake_case   :				str				 =   INVOICE_URL
          snake_case   :				Union[str, Any]				 =   "What is the invoice number?"

          snake_case   :				int				 =   dqa_pipeline(image=UpperCamelCase__						,     question=UpperCamelCase__						,     top_k=2						)
          self.assertEqual(
              nested_simplify(UpperCamelCase__						,     decimals=4						)						,     [
                  {"score": 0.4251, "answer": "us-001", "start": 16, "end": 16},
                  {"score": 0.0819, "answer": "1110212019", "start": 23, "end": 23},
              ]						,     )

          snake_case   :				Optional[Any]				 =   dqa_pipeline({"image": image, "question": question}						,     top_k=2						)
          self.assertEqual(
              nested_simplify(UpperCamelCase__						,     decimals=4						)						,     [
                  {"score": 0.4251, "answer": "us-001", "start": 16, "end": 16},
                  {"score": 0.0819, "answer": "1110212019", "start": 23, "end": 23},
              ]						,     )

          snake_case   :				Dict				 =   dqa_pipeline(
              [{"image": image, "question": question}, {"image": image, "question": question}]						,     top_k=2						)
          self.assertEqual(
              nested_simplify(UpperCamelCase__						,     decimals=4						)						,     [
                  [
                      {"score": 0.4251, "answer": "us-001", "start": 16, "end": 16},
                      {"score": 0.0819, "answer": "1110212019", "start": 23, "end": 23},
                  ]
              ]
              * 2						,     )

          snake_case   :				List[str]				 =   list(zip(*apply_tesseract(load_image(UpperCamelCase__						)						,     UpperCamelCase__						,     ""						)						)						)

          # This model should also work if `image` is set to None
          snake_case   :				Optional[int]				 =   dqa_pipeline({"image": None, "word_boxes": word_boxes, "question": question}						,     top_k=2						)
          self.assertEqual(
              nested_simplify(UpperCamelCase__						,     decimals=4						)						,     [
                  {"score": 0.4251, "answer": "us-001", "start": 16, "end": 16},
                  {"score": 0.0819, "answer": "1110212019", "start": 23, "end": 23},
              ]						,     )







   @slow
   @require_torch
   @require_pytesseract
   @require_vision
   def 			lowerCamelCase		(     self						)     ->		Union[str, Any]:

          '''simple docstring'''




          snake_case   :				List[str]				 =   AutoTokenizer.from_pretrained(
              "impira/layoutlm-document-qa"						,     revision="3dc6de3"						,     add_prefix_space=UpperCamelCase__						)
          snake_case   :				List[Any]				 =   pipeline(
              "document-question-answering"						,     model="impira/layoutlm-document-qa"						,     tokenizer=UpperCamelCase__						,     revision="3dc6de3"						,     max_seq_len=50						,     )
          snake_case   :				Any				 =   INVOICE_URL
          snake_case   :				Optional[int]				 =   "What is the invoice number?"

          snake_case   :				List[Any]				 =   dqa_pipeline(image=UpperCamelCase__						,     question=UpperCamelCase__						,     top_k=2						)
          self.assertEqual(
              nested_simplify(UpperCamelCase__						,     decimals=4						)						,     [
                  {"score": 0.9999, "answer": "us-001", "start": 16, "end": 16},
                  {"score": 0.9998, "answer": "us-001", "start": 16, "end": 16},
              ]						,     )

          snake_case   :				Any				 =   dqa_pipeline(
              [{"image": image, "question": question}, {"image": image, "question": question}]						,     top_k=2						)
          self.assertEqual(
              nested_simplify(UpperCamelCase__						,     decimals=4						)						,     [
                  [
                      {"score": 0.9999, "answer": "us-001", "start": 16, "end": 16},
                      {"score": 0.9998, "answer": "us-001", "start": 16, "end": 16},
                  ]
              ]
              * 2						,     )

          snake_case   :				List[Any]				 =   list(zip(*apply_tesseract(load_image(UpperCamelCase__						)						,     UpperCamelCase__						,     ""						)						)						)

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







   @slow
   @require_torch
   def 			lowerCamelCase		(     self						)     ->		Dict:

          '''simple docstring'''




          snake_case   :				List[Any]				 =   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"						,     )

          snake_case   :				Union[str, Any]				 =   INVOICE_URL
          snake_case   :				Dict				 =   "What is the invoice number?"
          snake_case   :				Union[str, Any]				 =   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 			lowerCamelCase		(     self						)     ->		Optional[int]:

          '''simple docstring'''




          pass









"""simple docstring"""



import argparse
import logging
import pickle
from collections import Counter


logging.basicConfig(
    format="""%(asctime)s - %(levelname)s - %(name)s - %(message)s""", datefmt="""%m/%d/%Y %H:%M:%S""", level=logging.INFO
)
__snake_case             =   logging.getLogger(__name__)

if __name__ == "__main__":
  __snake_case             =   argparse.ArgumentParser(
      description="""Token Counts for smoothing the masking probabilities in MLM (cf XLM/word2vec)"""
  )
  parser.add_argument(
      """--data_file""", type=str, default="""data/dump.bert-base-uncased.pickle""", help="""The binarized dataset."""
  )
  parser.add_argument(
      """--token_counts_dump""", type=str, default="""data/token_counts.bert-base-uncased.pickle""", help="""The dump file."""
  )
  parser.add_argument("""--vocab_size""", default=30522, type=int)
  __snake_case             =   parser.parse_args()

  logger.info(F'''Loading data from {args.data_file}''')
  with open(args.data_file, """rb""") as fp:
    __snake_case             =   pickle.load(fp)

  logger.info("""Counting occurrences for MLM.""")
  __snake_case             =   Counter()
  for tk_ids in data:
    counter.update(tk_ids)
  __snake_case             =   [0] * args.vocab_size
  for k, v in counter.items():
    __snake_case             =   v

  logger.info(F'''Dump to {args.token_counts_dump}''')
  with open(args.token_counts_dump, """wb""") as handle:
    pickle.dump(counts, handle, protocol=pickle.HIGHEST_PROTOCOL)





