

"""simple docstring"""





import unittest

import requests

from transformers import MODEL_FOR_VISION_2_SEQ_MAPPING, TF_MODEL_FOR_VISION_2_SEQ_MAPPING, is_vision_available
from transformers.pipelines import pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    is_torch_available,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_torch_available():
   from transformers.pytorch_utils import is_torch_greater_or_equal_than_1_11
else:
   _a				   =							False


if is_vision_available():
   from PIL import Image
else:




   class       _lowerCAmelCase						:







    """simple docstring"""







    @staticmethod
    def       _lowercase				(    *UpperCAmelCase__							:    Tuple,      **UpperCAmelCase__							:    Optional[Any]			):
       pass




@is_pipeline_test
@require_vision
class       _lowerCAmelCase						(			unittest.TestCase     ):







 """simple docstring"""







 __UpperCAmelCase			:      Union[str, Any]     	=			MODEL_FOR_VISION_2_SEQ_MAPPING
 __UpperCAmelCase			:      Dict     	=			TF_MODEL_FOR_VISION_2_SEQ_MAPPING


 def       _lowercase				(    self							:    Any,      UpperCAmelCase__							:    int,      UpperCAmelCase__							:    Optional[int],      UpperCAmelCase__							:    List[str]			):
    __lowercase			=		pipeline("image-to-text",      model=UpperCAmelCase__,      tokenizer=UpperCAmelCase__,      image_processor=UpperCAmelCase__			)
    __lowercase			=		[
        Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"			),
        "./tests/fixtures/tests_samples/COCO/000000039769.png",
    ]
    return pipe, examples


 def       _lowercase				(    self							:    List[Any],      UpperCAmelCase__							:    Optional[int],      UpperCAmelCase__							:    Tuple			):
    __lowercase			=		pipe(UpperCAmelCase__			)
    self.assertEqual(
        UpperCAmelCase__,      [
            [{"generated_text": ANY(UpperCAmelCase__			)}],
            [{"generated_text": ANY(UpperCAmelCase__			)}],
        ],      )


 @require_tf
 def       _lowercase				(    self							:    List[str]			):
    __lowercase			=		pipeline("image-to-text",      model="hf-internal-testing/tiny-random-vit-gpt2",      framework="tf"			)
    __lowercase			=		"./tests/fixtures/tests_samples/COCO/000000039769.png"

    __lowercase			=		pipe(UpperCAmelCase__			)
    self.assertEqual(
        UpperCAmelCase__,      [
            {
                "generated_text": "growthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthGOGO"
            },
        ],      )

    __lowercase			=		pipe([image, image]			)
    self.assertEqual(
        UpperCAmelCase__,      [
            [
                {
                    "generated_text": "growthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthGOGO"
                }
            ],
            [
                {
                    "generated_text": "growthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthGOGO"
                }
            ],
        ],      )

    __lowercase			=		pipe(UpperCAmelCase__,      max_new_tokens=1			)
    self.assertEqual(
        UpperCAmelCase__,      [{"generated_text": "growth"}],      )


 @require_torch
 def       _lowercase				(    self							:    Any			):
    __lowercase			=		pipeline("image-to-text",      model="hf-internal-testing/tiny-random-vit-gpt2"			)
    __lowercase			=		"./tests/fixtures/tests_samples/COCO/000000039769.png"

    __lowercase			=		pipe(UpperCAmelCase__			)
    self.assertEqual(
        UpperCAmelCase__,      [
            {
                "generated_text": "growthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthGOGO"
            },
        ],      )

    __lowercase			=		pipe([image, image]			)
    self.assertEqual(
        UpperCAmelCase__,      [
            [
                {
                    "generated_text": "growthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthGOGO"
                }
            ],
            [
                {
                    "generated_text": "growthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthgrowthGOGO"
                }
            ],
        ],      )


 @require_torch
 def       _lowercase				(    self							:    int			):
    __lowercase			=		pipeline("image-to-text",      model="hf-internal-testing/tiny-random-BlipForConditionalGeneration"			)
    __lowercase			=		"./tests/fixtures/tests_samples/COCO/000000039769.png"
    __lowercase			=		"a photo of"

    __lowercase			=		pipe(UpperCAmelCase__,      prompt=UpperCAmelCase__			)
    self.assertTrue(outputs[0]["generated_text"].startswith(UpperCAmelCase__			)			)


 @slow
 @require_torch
 def       _lowercase				(    self							:    Optional[Any]			):
    __lowercase			=		pipeline("image-to-text",      model="ydshieh/vit-gpt2-coco-en"			)
    __lowercase			=		"./tests/fixtures/tests_samples/COCO/000000039769.png"

    __lowercase			=		pipe(UpperCAmelCase__			)
    self.assertEqual(UpperCAmelCase__,      [{"generated_text": "a cat laying on a blanket next to a cat laying on a bed "}]			)

    __lowercase			=		pipe([image, image]			)
    self.assertEqual(
        UpperCAmelCase__,      [
            [{"generated_text": "a cat laying on a blanket next to a cat laying on a bed "}],
            [{"generated_text": "a cat laying on a blanket next to a cat laying on a bed "}],
        ],      )


 @slow
 @require_torch
 def       _lowercase				(    self							:    Tuple			):
    __lowercase			=		pipeline("image-to-text",      model="Salesforce/blip-image-captioning-base"			)
    __lowercase			=		"https://huggingface.co/datasets/sayakpaul/sample-datasets/resolve/main/pokemon.png"
    __lowercase			=		Image.open(requests.get(UpperCAmelCase__,      stream=UpperCAmelCase__			).raw			)

    __lowercase			=		pipe(UpperCAmelCase__			)
    self.assertEqual(UpperCAmelCase__,      [{"generated_text": "a pink pokemon pokemon with a blue shirt and a blue shirt"}]			)


 @slow
 @require_torch
 def       _lowercase				(    self							:    List[str]			):
    __lowercase			=		pipeline("image-to-text",      model="microsoft/git-base-coco"			)
    __lowercase			=		"https://huggingface.co/datasets/sayakpaul/sample-datasets/resolve/main/pokemon.png"
    __lowercase			=		Image.open(requests.get(UpperCAmelCase__,      stream=UpperCAmelCase__			).raw			)

    __lowercase			=		pipe(UpperCAmelCase__			)
    self.assertEqual(UpperCAmelCase__,      [{"generated_text": "a cartoon of a purple character."}]			)


 @slow
 @require_torch
 def       _lowercase				(    self							:    Dict			):
    __lowercase			=		pipeline("image-to-text",      model="Salesforce/blip-image-captioning-base"			)
    __lowercase			=		"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/ai2d-demo.jpg"
    __lowercase			=		Image.open(requests.get(UpperCAmelCase__,      stream=UpperCAmelCase__			).raw			)

    __lowercase			=		"a photography of"

    __lowercase			=		pipe(UpperCAmelCase__,      prompt=UpperCAmelCase__			)
    self.assertEqual(UpperCAmelCase__,      [{"generated_text": "a photography of a volcano"}]			)

    with self.assertRaises(UpperCAmelCase__			):
       __lowercase			=		pipe([image, image],      prompt=[prompt, prompt]			)


 @slow
 @require_torch
 def       _lowercase				(    self							:    List[str]			):
    __lowercase			=		pipeline("image-to-text",      model="microsoft/git-base-coco"			)
    __lowercase			=		"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/ai2d-demo.jpg"
    __lowercase			=		Image.open(requests.get(UpperCAmelCase__,      stream=UpperCAmelCase__			).raw			)

    __lowercase			=		"a photo of a"

    __lowercase			=		pipe(UpperCAmelCase__,      prompt=UpperCAmelCase__			)
    self.assertEqual(UpperCAmelCase__,      [{"generated_text": "a photo of a tent with a tent and a tent in the background."}]			)

    with self.assertRaises(UpperCAmelCase__			):
       __lowercase			=		pipe([image, image],      prompt=[prompt, prompt]			)


 @unittest.skipIf(
     not is_torch_greater_or_equal_than_1_11,      reason="`Pix2StructImageProcessor` requires `torch>=1.11.0`."			)
 @slow
 @require_torch
 def       _lowercase				(    self							:    Any			):
    __lowercase			=		pipeline("image-to-text",      model="google/pix2struct-ai2d-base"			)
    __lowercase			=		"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/ai2d-demo.jpg"
    __lowercase			=		Image.open(requests.get(UpperCAmelCase__,      stream=UpperCAmelCase__			).raw			)

    __lowercase			=		"What does the label 15 represent? (1) lava (2) core (3) tunnel (4) ash cloud"

    __lowercase			=		pipe(UpperCAmelCase__,      prompt=UpperCAmelCase__			)
    self.assertEqual(UpperCAmelCase__,      [{"generated_text": "ash cloud"}]			)

    with self.assertRaises(UpperCAmelCase__			):
       __lowercase			=		pipe([image, image],      prompt=[prompt, prompt]			)



 @slow
 @require_tf
 def       _lowercase				(    self							:    Union[str, Any]			):
    __lowercase			=		pipeline("image-to-text",      model="ydshieh/vit-gpt2-coco-en",      framework="tf"			)
    __lowercase			=		"./tests/fixtures/tests_samples/COCO/000000039769.png"

    __lowercase			=		pipe(UpperCAmelCase__			)
    self.assertEqual(UpperCAmelCase__,      [{"generated_text": "a cat laying on a blanket next to a cat laying on a bed "}]			)

    __lowercase			=		pipe([image, image]			)
    self.assertEqual(
        UpperCAmelCase__,      [
            [{"generated_text": "a cat laying on a blanket next to a cat laying on a bed "}],
            [{"generated_text": "a cat laying on a blanket next to a cat laying on a bed "}],
        ],      )







"""simple docstring"""





import os
from dataclasses import replace

import jax
import wandb
from bigbird_flax import Args, DataCollator, FlaxBigBirdForNaturalQuestions, Trainer, build_tx, train_step, val_step
from datasets import load_dataset
from flax import jax_utils

from transformers import BigBirdTokenizerFast


if __name__ == "__main__":
   print('#################### AVAILABLE DEVICES ####################')
   print(jax.devices())
   print('###########################################################')

   # setup for wandb sweep
   _a				   =							Args()
   _a				   =							wandb.init(project='bigbird-natural-questions', config=args.__dict__)
   _a				   =							dict(logger.config)
   del wandb_args["batch_size"]
   _a				   =							replace(args, **wandb_args)
   _a				   =							args.base_dir + '-' + wandb.run.id
   _a				   =							replace(args, base_dir=base_dir)
   print(args)

   _a				   =							load_dataset('json', data_files=args.tr_data_path)['train']
   _a				   =							load_dataset('json', data_files=args.val_data_path)['train']

   # drop extra batch for now
   _a				   =							range(len(tr_dataset) - len(tr_dataset) % args.batch_size)
   _a				   =							tr_dataset.shuffle().select(indices)
   _a				   =							range(len(val_dataset) - len(val_dataset) % args.batch_size)
   _a				   =							val_dataset.shuffle().select(indices)

   if os.environ.get('TRAIN_ON_SMALL', 'false') == "true":
      _a				   =							tr_dataset.shuffle().select(range(8_00_00))
      _a				   =							val_dataset.shuffle().select(range(80_00))

   print(tr_dataset)
   print(val_dataset)

   _a				   =							FlaxBigBirdForNaturalQuestions.from_pretrained(
       args.model_id, block_size=args.block_size, num_random_blocks=args.num_random_blocks
   )
   _a				   =							BigBirdTokenizerFast.from_pretrained(args.model_id)
   _a				   =							DataCollator(pad_id=tokenizer.pad_token_id, max_length=40_96)

   _a				   =							{
       'lr': args.lr,
       'init_lr': args.init_lr,
       'warmup_steps': args.warmup_steps,
       'num_train_steps': args.max_epochs * (len(tr_dataset) // args.batch_size),
       'weight_decay': args.weight_decay,
   }
   _a   ,		_a				   =							build_tx(**tx_args)

   _a				   =							Trainer(
       args=args,
       data_collator=data_collator,
       model_save_fn=model.save_pretrained,
       train_step_fn=train_step,
       val_step_fn=val_step,
       logger=logger,
       scheduler_fn=lr,
   )

   _a				   =							None
   _a				   =							trainer.create_state(model, tx, num_train_steps=tx_args['num_train_steps'], ckpt_dir=ckpt_dir)
   try:
      trainer.train(state, tr_dataset, val_dataset)
   except KeyboardInterrupt:
      print('Oooops; TRAINING STOPPED UNFORTUNATELY')

   print('SAVING WEIGHTS IN `final-weights`')
   _a				   =							jax_utils.unreplicate(state.params)
   model.save_pretrained(os.path.join(args.base_dir, 'final-weights'), params=params)




