



"""simple docstring"""


from typing import Optional, Tuple

import flax.linen as nn
import jax
import jax.numpy as jnp
from configuration_hybrid_clip import HybridCLIPConfig
from flax.core.frozen_dict import FrozenDict

from transformers import FLAX_MODEL_MAPPING, FlaxCLIPVisionModel
from transformers.modeling_flax_utils import FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPOutput
from transformers.utils import logging


snake_case_								=			logging.get_logger(__name__)







class  A_     (	nn.Module							):



						"""simple docstring"""
						__UpperCamelCase													=				42
						__UpperCamelCase													=				jnp.floataa






						def       UpperCAmelCase__	(   self					:Optional[Any]   )							->     List[Any]:
												UpperCAmelCase			 =					self.config.text_config
												UpperCAmelCase			 =					self.config.vision_config

												UpperCAmelCase			 =					self.config.projection_dim
												UpperCAmelCase			 =					text_config.hidden_size
												UpperCAmelCase			 =					vision_config.hidden_size

												UpperCAmelCase			 =					FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
												UpperCAmelCase			 =					FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__  ,    lowercase_   ).module_class

												UpperCAmelCase			 =					text_module(lowercase_  ,    dtype=self.dtype   )
												UpperCAmelCase			 =					vision_module(lowercase_  ,    dtype=self.dtype   )

												UpperCAmelCase			 =					nn.Dense(
												    self.projection_dim  ,    dtype=self.dtype  ,    kernel_init=jax.nn.initializers.normal(0.02   )  ,    use_bias=lowercase_  ,    )
												UpperCAmelCase			 =					nn.Dense(
												    self.projection_dim  ,    dtype=self.dtype  ,    kernel_init=jax.nn.initializers.normal(0.02   )  ,    use_bias=lowercase_  ,    )
												UpperCAmelCase			 =					self.param('logit_scale'  ,    jax.nn.initializers.ones  ,    []   )






						def __call__(   self					:Dict  ,    lowercase_					:int=None  ,    lowercase_					:Dict=None  ,    lowercase_					:Dict=None  ,    lowercase_					:Optional[Any]=None  ,    lowercase_					:List[Any]=None  ,    lowercase_					:bool = True  ,    lowercase_					:List[Any]=None  ,    lowercase_					:str=None  ,    lowercase_					:Tuple=None  ,    )							->     Union[str, Any]:
												UpperCAmelCase			 =					return_dict if return_dict is not None else self.config.return_dict

												UpperCAmelCase			 =					self.vision_model(
												    pixel_values=lowercase_  ,    deterministic=lowercase_  ,    output_attentions=lowercase_  ,    output_hidden_states=lowercase_  ,    return_dict=lowercase_  ,    )

												UpperCAmelCase			 =					self.text_model(
												    input_ids=lowercase_  ,    attention_mask=lowercase_  ,    token_type_ids=lowercase_  ,    position_ids=lowercase_  ,    deterministic=lowercase_  ,    output_attentions=lowercase_  ,    output_hidden_states=lowercase_  ,    return_dict=lowercase_  ,    )

												UpperCAmelCase			 =					vision_outputs[1]
												UpperCAmelCase			 =					self.visual_projection(lowercase_   )

												UpperCAmelCase			 =					text_outputs[1]
												UpperCAmelCase			 =					self.text_projection(lowercase_   )

												# normalized features
												UpperCAmelCase			 =					image_embeds / jnp.linalg.norm(lowercase_  ,    axis=-1  ,    keepdims=lowercase_   )
												UpperCAmelCase			 =					text_embeds / jnp.linalg.norm(lowercase_  ,    axis=-1  ,    keepdims=lowercase_   )

												# cosine similarity as logits
												UpperCAmelCase			 =					jnp.exp(self.logit_scale   )
												UpperCAmelCase			 =					jnp.matmul(lowercase_  ,    image_embeds.T   ) * logit_scale
												UpperCAmelCase			 =					logits_per_text.T

												if not return_dict:
																		return (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)

												return FlaxCLIPOutput(
												    logits_per_image=lowercase_  ,    logits_per_text=lowercase_  ,    text_embeds=lowercase_  ,    image_embeds=lowercase_  ,    text_model_output=lowercase_  ,    vision_model_output=lowercase_  ,    )







class  A_     (	SCREAMING_SNAKE_CASE_							):



						"""simple docstring"""
						__UpperCamelCase													=				HybridCLIPConfig
						__UpperCamelCase													=				FlaxHybridCLIPModule






						def __init__(   self					:Optional[int]  ,    lowercase_					:HybridCLIPConfig  ,    lowercase_					:Optional[Tuple] = None  ,    lowercase_					:int = 0  ,    lowercase_					:jnp.dtype = jnp.floataa  ,    **lowercase_					:List[str]  ,    )							->     Union[str, Any]:
												if input_shape is None:
																		UpperCAmelCase			 =					((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

												UpperCAmelCase			 =					self.module_class(config=lowercase_  ,    dtype=lowercase_  ,    **lowercase_   )
												super().__init__(lowercase_  ,    lowercase_  ,    input_shape=lowercase_  ,    seed=lowercase_  ,    dtype=lowercase_   )






						def       UpperCAmelCase__	(   self					:Tuple  ,    lowercase_					:jax.random.PRNGKey  ,    lowercase_					:Tuple  ,    lowercase_					:FrozenDict = None   )							->     FrozenDict:
												# init input tensor
												UpperCAmelCase			 =					jnp.zeros(input_shape[0]  ,    dtype='i4'   )
												UpperCAmelCase			 =					jnp.broadcast_to(jnp.arange(jnp.atleast_ad(lowercase_   ).shape[-1]   )  ,    input_shape[0]   )
												UpperCAmelCase			 =					jnp.ones_like(lowercase_   )
												UpperCAmelCase			 =					jnp.ones_like(lowercase_   )

												UpperCAmelCase			 =					jax.random.normal(lowercase_  ,    input_shape[1]   )

												UpperCAmelCase ,       UpperCAmelCase			 =					jax.random.split(lowercase_   )
												UpperCAmelCase			 =					{'params': params_rng, 'dropout': dropout_rng}

												return self.module.init(lowercase_  ,    lowercase_  ,    lowercase_  ,    lowercase_  ,    lowercase_  ,    lowercase_   )["params"]






						def __call__(   self					:str  ,    lowercase_					:Dict  ,    lowercase_					:Optional[int]  ,    lowercase_					:List[Any]=None  ,    lowercase_					:Dict=None  ,    lowercase_					:Tuple=None  ,    lowercase_					:dict = None  ,    lowercase_					:jax.random.PRNGKey = None  ,    lowercase_					:bool = False  ,    lowercase_					:Optional[bool] = None  ,    lowercase_					:Optional[bool] = None  ,    lowercase_					:Optional[bool] = None  ,    )							->     Union[str, Any]:
												UpperCAmelCase			 =					output_attentions if output_attentions is not None else self.config.output_attentions
												UpperCAmelCase			 =					(
												    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
												)
												UpperCAmelCase			 =					return_dict if return_dict is not None else self.config.return_dict

												if position_ids is None:
																		UpperCAmelCase			 =					jnp.broadcast_to(jnp.arange(jnp.atleast_ad(lowercase_   ).shape[-1]   )  ,    input_ids.shape   )

												if token_type_ids is None:
																		UpperCAmelCase			 =					jnp.zeros_like(lowercase_   )

												if attention_mask is None:
																		UpperCAmelCase			 =					jnp.ones_like(lowercase_   )

												# Handle any PRNG if needed
												UpperCAmelCase			 =					{}
												if dropout_rng is not None:
																		UpperCAmelCase			 =					dropout_rng

												return self.module.apply(
												    {'params': params or self.params}  ,    jnp.array(lowercase_  ,    dtype='i4'   )  ,    jnp.array(lowercase_  ,    dtype=jnp.floataa   )  ,    jnp.array(lowercase_  ,    dtype='i4'   )  ,    jnp.array(lowercase_  ,    dtype='i4'   )  ,    jnp.array(lowercase_  ,    dtype='i4'   )  ,    not train  ,    lowercase_  ,    lowercase_  ,    lowercase_  ,    rngs=lowercase_  ,    )






						def       UpperCAmelCase__	(   self					:Tuple  ,    lowercase_					:Tuple  ,    lowercase_					:Optional[int]=None  ,    lowercase_					:Union[str, Any]=None  ,    lowercase_					:Tuple=None  ,    lowercase_					:dict = None  ,    lowercase_					:jax.random.PRNGKey = None  ,    lowercase_					:Optional[Any]=False  ,    )							->     str:
												if position_ids is None:
																		UpperCAmelCase			 =					jnp.broadcast_to(jnp.arange(jnp.atleast_ad(lowercase_   ).shape[-1]   )  ,    input_ids.shape   )

												if token_type_ids is None:
																		UpperCAmelCase			 =					jnp.zeros_like(lowercase_   )

												if attention_mask is None:
																		UpperCAmelCase			 =					jnp.ones_like(lowercase_   )

												# Handle any PRNG if needed
												UpperCAmelCase			 =					{}
												if dropout_rng is not None:
																		UpperCAmelCase			 =					dropout_rng

												def _get_features(lowercase_					:Dict  ,    lowercase_					:List[Any]  ,    lowercase_					:Union[str, Any]  ,    lowercase_					:Union[str, Any]  ,    lowercase_					:Any  ,    lowercase_					:List[Any]   ):
																		UpperCAmelCase			 =					module.text_model(
																		    input_ids=lowercase_  ,    attention_mask=lowercase_  ,    position_ids=lowercase_  ,    token_type_ids=lowercase_  ,    deterministic=lowercase_  ,    )
																		UpperCAmelCase			 =					text_outputs[1]
																		UpperCAmelCase			 =					module.text_projection(lowercase_   )
																		return text_features

												return self.module.apply(
												    {'params': params or self.params}  ,    jnp.array(lowercase_  ,    dtype='i4'   )  ,    jnp.array(lowercase_  ,    dtype='i4'   )  ,    jnp.array(lowercase_  ,    dtype='i4'   )  ,    jnp.array(lowercase_  ,    dtype='i4'   )  ,    not train  ,    method=_get_features  ,    rngs=lowercase_  ,    )






						def       UpperCAmelCase__	(   self					:Optional[Any]  ,    lowercase_					:Optional[int]  ,    lowercase_					:dict = None  ,    lowercase_					:jax.random.PRNGKey = None  ,    lowercase_					:int=False   )							->     List[Any]:
												UpperCAmelCase			 =					{}
												if dropout_rng is not None:
																		UpperCAmelCase			 =					dropout_rng

												def _get_features(lowercase_					:Any  ,    lowercase_					:Union[str, Any]  ,    lowercase_					:int   ):
																		UpperCAmelCase			 =					module.vision_model(pixel_values=lowercase_  ,    deterministic=lowercase_   )
																		UpperCAmelCase			 =					vision_outputs[1]  # pooled_output
																		UpperCAmelCase			 =					module.visual_projection(lowercase_   )
																		return image_features

												return self.module.apply(
												    {'params': params or self.params}  ,    jnp.array(lowercase_  ,    dtype=jnp.floataa   )  ,    not train  ,    method=_get_features  ,    rngs=lowercase_  ,    )






						@classmethod
						def       UpperCAmelCase__	(   cls					:str  ,    lowercase_					:str = None  ,    lowercase_					:str = None  ,    *lowercase_					:Tuple  ,    **lowercase_					:str  ,    )							->     FlaxPreTrainedModel:
												UpperCAmelCase			 =					{
												    argument[len('text_'   ) :]: value for argument, value in kwargs.items() if argument.startswith('text_'   )
												}

												UpperCAmelCase			 =					{
												    argument[len('vision_'   ) :]: value for argument, value in kwargs.items() if argument.startswith('vision_'   )
												}

												# remove text, vision kwargs from kwargs
												for key in kwargs_text.keys():
																		del kwargs["text_" + key]
												for key in kwargs_vision.keys():
																		del kwargs["vision_" + key]

												# Load and initialize the text and vision model
												UpperCAmelCase			 =					kwargs_text.pop('model'  ,    lowercase_   )
												if text_model is None:
																		assert (
																		    text_model_name_or_path is not None
																		), "If `model` is not defined as an argument, a `text_model_name_or_path` has to be defined"
																		from transformers import FlaxAutoModel

																		if "config" not in kwargs_text:
																								from transformers import AutoConfig

																								UpperCAmelCase			 =					AutoConfig.from_pretrained(lowercase_   )
																								UpperCAmelCase			 =					text_config

																		UpperCAmelCase			 =					FlaxAutoModel.from_pretrained(lowercase_  ,    *lowercase_  ,    **lowercase_   )

												UpperCAmelCase			 =					kwargs_vision.pop('model'  ,    lowercase_   )
												if vision_model is None:
																		assert (
																		    vision_model_name_or_path is not None
																		), "If `model` is not defined as an argument, a `vision_model_name_or_path` has to be defined"
																		from transformers import FlaxAutoModel

																		if "config" not in kwargs_vision:
																								from transformers import AutoConfig

																								UpperCAmelCase			 =					AutoConfig.from_pretrained(lowercase_   )
																								UpperCAmelCase			 =					vision_config

																		UpperCAmelCase			 =					FlaxAutoModel.from_pretrained(lowercase_  ,    *lowercase_  ,    **lowercase_   )

												# instantiate config with corresponding kwargs
												UpperCAmelCase			 =					kwargs.pop('dtype'  ,    jnp.floataa   )
												UpperCAmelCase			 =					HybridCLIPConfig.from_text_vision_configs(text_model.config  ,    vision_model.config  ,    **lowercase_   )

												# init model
												UpperCAmelCase			 =					cls(lowercase_  ,    *lowercase_  ,    dtype=lowercase_  ,    **lowercase_   )

												if vision_config.model_type == "clip":
																		UpperCAmelCase			 =					vision_model.params['vision_model']
																		UpperCAmelCase			 =					vision_model.params['visual_projection']['kernel']
												else:
																		UpperCAmelCase			 =					vision_model.params

												UpperCAmelCase			 =					text_model.params

												return model










"""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
				snake_case_								=			Args()
				snake_case_								=			wandb.init(project="""bigbird-natural-questions""", config=args.__dict__)
				snake_case_								=			dict(logger.config)
				del wandb_args["batch_size"]
				snake_case_								=			replace(args, **wandb_args)
				snake_case_								=			args.base_dir + """-""" + wandb.run.id
				snake_case_								=			replace(args, base_dir=base_dir)
				print(args)

				snake_case_								=			load_dataset("""json""", data_files=args.tr_data_path)["""train"""]
				snake_case_								=			load_dataset("""json""", data_files=args.val_data_path)["""train"""]

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

				if os.environ.get("""TRAIN_ON_SMALL""", """false""") == "true":
								snake_case_								=			tr_dataset.shuffle().select(range(8_0000))
								snake_case_								=			val_dataset.shuffle().select(range(8000))

				print(tr_dataset)
				print(val_dataset)

				snake_case_								=			FlaxBigBirdForNaturalQuestions.from_pretrained(
				    args.model_id, block_size=args.block_size, num_random_blocks=args.num_random_blocks
				)
				snake_case_								=			BigBirdTokenizerFast.from_pretrained(args.model_id)
				snake_case_								=			DataCollator(pad_id=tokenizer.pad_token_id, max_length=4096)

				snake_case_								=			{
				    """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,
				}
				snake_case_	,  snake_case_								=			build_tx(**tx_args)

				snake_case_								=			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,
				)

				snake_case_								=			None
				snake_case_								=			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`""")
				snake_case_								=			jax_utils.unreplicate(state.params)
				model.save_pretrained(os.path.join(args.base_dir, """final-weights"""), params=params)





