




import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase__    			=			logging.get_logger(__name__)

UpperCamelCase__    			=			{
    """BAAI/AltCLIP""": """https://huggingface.co/BAAI/AltCLIP/resolve/main/config.json""",
    # See all AltCLIP models at https://huggingface.co/models?filter=altclip
}
class 			a__    (					snake_case__      ):
						_a							:    List[Any]					=    """altclip_text_model"""






						def __init__(   self		,			_A=2_5_0_0_0_2		,			_A=1_0_2_4		,			_A=2_4		,			_A=1_6		,			_A=4_0_9_6		,			_A="gelu"		,			_A=0.1		,			_A=0.1		,			_A=5_1_4		,			_A=1		,			_A=0.02		,			_A=0.02		,			_A=1E-0_5		,			_A=1		,			_A=0		,			_A=2		,			_A="absolute"		,			_A=True		,			_A=7_6_8		,			**_A		,			):




											"""simple docstring"""



											super().__init__(pad_token_id=_A		,			bos_token_id=_A		,			eos_token_id=_A		,			**_A			)

											__lowerCAmelCase				   =       vocab_size
											__lowerCAmelCase				   =       hidden_size
											__lowerCAmelCase				   =       num_hidden_layers
											__lowerCAmelCase				   =       num_attention_heads
											__lowerCAmelCase				   =       hidden_act
											__lowerCAmelCase				   =       intermediate_size
											__lowerCAmelCase				   =       hidden_dropout_prob
											__lowerCAmelCase				   =       attention_probs_dropout_prob
											__lowerCAmelCase				   =       max_position_embeddings
											__lowerCAmelCase				   =       type_vocab_size
											__lowerCAmelCase				   =       initializer_range
											__lowerCAmelCase				   =       initializer_factor
											__lowerCAmelCase				   =       layer_norm_eps
											__lowerCAmelCase				   =       position_embedding_type
											__lowerCAmelCase				   =       use_cache
											__lowerCAmelCase				   =       project_dim






class 			a__    (					snake_case__      ):
						_a							:    List[Any]					=    """altclip_vision_model"""






						def __init__(   self		,			_A=7_6_8		,			_A=3_0_7_2		,			_A=5_1_2		,			_A=1_2		,			_A=1_2		,			_A=3		,			_A=2_2_4		,			_A=3_2		,			_A="quick_gelu"		,			_A=1E-5		,			_A=0.0		,			_A=0.02		,			_A=1.0		,			**_A		,			):




											"""simple docstring"""



											super().__init__(**_A			)

											__lowerCAmelCase				   =       hidden_size
											__lowerCAmelCase				   =       intermediate_size
											__lowerCAmelCase				   =       projection_dim
											__lowerCAmelCase				   =       num_hidden_layers
											__lowerCAmelCase				   =       num_attention_heads
											__lowerCAmelCase				   =       num_channels
											__lowerCAmelCase				   =       patch_size
											__lowerCAmelCase				   =       image_size
											__lowerCAmelCase				   =       initializer_range
											__lowerCAmelCase				   =       initializer_factor
											__lowerCAmelCase				   =       attention_dropout
											__lowerCAmelCase				   =       layer_norm_eps
											__lowerCAmelCase				   =       hidden_act






						@classmethod
						def 				__SCREAMING_SNAKE_CASE(   cls		,			_A		,			**_A			):




											"""simple docstring"""



											cls._set_token_in_kwargs(_A			)

											__lowerCAmelCase   ,							__lowerCAmelCase				   =       cls.get_config_dict(_A		,			**_A			)

											# get the vision config dict if we are loading from AltCLIPConfig
											if config_dict.get("model_type"			) == "altclip":
																__lowerCAmelCase				   =       config_dict["vision_config"]

											if "model_type" in config_dict and hasattr(cls		,			"model_type"			) and config_dict["model_type"] != cls.model_type:
																logger.warning(
																    f"""You are using a model of type {config_dict['model_type']} to instantiate a model of type """
																    f"""{cls.model_type}. This is not supported for all configurations of models and can yield errors."""			)

											return cls.from_dict(_A		,			**_A			)






class 			a__    (					snake_case__      ):
						_a							:    Any					=    """altclip"""
						_a							:    Optional[int]					=    True






						def __init__(   self		,			_A=None		,			_A=None		,			_A=7_6_8		,			_A=2.65_92		,			**_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       kwargs.pop("text_config_dict"		,			_A			)
											__lowerCAmelCase				   =       kwargs.pop("vision_config_dict"		,			_A			)

											super().__init__(**_A			)

											# Instead of simply assigning `[text|vision]_config_dict` to `[text|vision]_config`, we use the values in
											# `[text|vision]_config_dict` to update the values in `[text|vision]_config`. The values should be same in most
											# cases, but we don't want to break anything regarding `_config_dict` that existed before commit `8827e1b2`.
											if text_config_dict is not None:
																if text_config is None:
																					__lowerCAmelCase				   =       {}

																# This is the complete result when using `text_config_dict`.
																__lowerCAmelCase				   =       AltCLIPTextConfig(**_A			).to_dict()

																# Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
																for key, value in _text_config_dict.items():
																					if key in text_config and value != text_config[key] and key not in ["transformers_version"]:
																										# If specified in `text_config_dict`
																										if key in text_config_dict:
																															__lowerCAmelCase				   =       (
																															    f"""`{key}` is found in both `text_config_dict` and `text_config` but with different values. """
																															    f"""The value `text_config_dict[\"{key}\"]` will be used instead."""
																															)
																										# If inferred from default argument values (just to be super careful)
																										else:
																															__lowerCAmelCase				   =       (
																															    f"""`text_config_dict` is provided which will be used to initialize `AltCLIPTextConfig`. The """
																															    f"""value `text_config[\"{key}\"]` will be overriden."""
																															)
																										logger.warning(_A			)

            # Update all values in `text_config` with the ones in `_text_config_dict`.
																text_config.update(_text_config_dict			)

											if vision_config_dict is not None:
																if vision_config is None:
																					__lowerCAmelCase				   =       {}

																# This is the complete result when using `vision_config_dict`.
																__lowerCAmelCase				   =       AltCLIPVisionConfig(**_A			).to_dict()
																# convert keys to string instead of integer
																if "id2label" in _vision_config_dict:
																					__lowerCAmelCase				   =       {
																					    str(_A			): value for key, value in _vision_config_dict["id2label"].items()
																					}

																# Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different.
																for key, value in _vision_config_dict.items():
																					if key in vision_config and value != vision_config[key] and key not in ["transformers_version"]:
																										# If specified in `vision_config_dict`
																										if key in vision_config_dict:
																															__lowerCAmelCase				   =       (
																															    f"""`{key}` is found in both `vision_config_dict` and `vision_config` but with different """
																															    f"""values. The value `vision_config_dict[\"{key}\"]` will be used instead."""
																															)
																										# If inferred from default argument values (just to be super careful)
																										else:
																															__lowerCAmelCase				   =       (
																															    f"""`vision_config_dict` is provided which will be used to initialize `AltCLIPVisionConfig`. """
																															    f"""The value `vision_config[\"{key}\"]` will be overriden."""
																															)
																										logger.warning(_A			)

            # Update all values in `vision_config` with the ones in `_vision_config_dict`.
																vision_config.update(_vision_config_dict			)

											if text_config is None:
																__lowerCAmelCase				   =       {}
																logger.info("`text_config` is `None`. Initializing the `AltCLIPTextConfig` with default values."			)

											if vision_config is None:
																__lowerCAmelCase				   =       {}
																logger.info("`vision_config` is `None`. initializing the `AltCLIPVisionConfig` with default values."			)

											__lowerCAmelCase				   =       AltCLIPTextConfig(**_A			)
											__lowerCAmelCase				   =       AltCLIPVisionConfig(**_A			)

											__lowerCAmelCase				   =       projection_dim
											__lowerCAmelCase				   =       logit_scale_init_value
											__lowerCAmelCase				   =       1.0






						@classmethod
						def 				__SCREAMING_SNAKE_CASE(   cls		,			_A		,			_A		,			**_A			):




											"""simple docstring"""



											return cls(text_config=text_config.to_dict()		,			vision_config=vision_config.to_dict()		,			**_A			)






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase				   =       copy.deepcopy(self.__dict__			)
											__lowerCAmelCase				   =       self.text_config.to_dict()
											__lowerCAmelCase				   =       self.vision_config.to_dict()
											__lowerCAmelCase				   =       self.__class__.model_type
											return output










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 	_a      (						SCREAMING_SNAKE_CASE_							:				Optional[int]     ):
					__lowerCAmelCase				   =       filter(lambda SCREAMING_SNAKE_CASE_     : p.requires_grad      ,      model.parameters()     )
					__lowerCAmelCase				   =       sum([np.prod(p.size()     ) for p in model_parameters]     )
					return params


UpperCamelCase__    			=			logging.getLogger(__name__)







def 	_a      (						SCREAMING_SNAKE_CASE_							:				Union[str, Any]      ,      SCREAMING_SNAKE_CASE_							:				Any     ):
					if metric == "rouge2":
										__lowerCAmelCase				   =       "{val_avg_rouge2:.4f}-{step_count}"
					elif metric == "bleu":
										__lowerCAmelCase				   =       "{val_avg_bleu:.4f}-{step_count}"
					elif metric == "em":
										__lowerCAmelCase				   =       "{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."     )

					__lowerCAmelCase				   =       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 	_a      (						SCREAMING_SNAKE_CASE_							:				int      ,      SCREAMING_SNAKE_CASE_							:				Union[str, Any]     ):
					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      ):






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       {f"""lr_group_{i}""": param["lr"] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups			)}
											pl_module.logger.log_metrics(_A			)






						@rank_zero_only
						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A		,			_A		,			_A=True			):




											"""simple docstring"""



											logger.info(f"""***** {type_path} results at step {trainer.global_step:05d} *****"""			)
											__lowerCAmelCase				   =       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
											__lowerCAmelCase				   =       Path(pl_module.hparams.output_dir			)
											if type_path == "test":
																__lowerCAmelCase				   =       od / "test_results.txt"
																__lowerCAmelCase				   =       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.
																__lowerCAmelCase				   =       od / f"""{type_path}_results/{trainer.global_step:05d}.txt"""
																__lowerCAmelCase				   =       od / f"""{type_path}_generations/{trainer.global_step:05d}.txt"""
																results_file.parent.mkdir(exist_ok=_A			)
																generations_file.parent.mkdir(exist_ok=_A			)
											with open(_A		,			"a+"			) as writer:
																for key in sorted(_A			):
																					if key in ["log", "progress_bar", "preds"]:
																										continue
																					__lowerCAmelCase				   =       metrics[key]
																					if isinstance(_A		,			torch.Tensor			):
																										__lowerCAmelCase				   =       val.item()
																					__lowerCAmelCase				   =       f"""{key}: {val:.6f}\n"""
																					writer.write(_A			)

											if not save_generations:
																return

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






						@rank_zero_only
						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A			):




											"""simple docstring"""



											try:
																__lowerCAmelCase				   =       pl_module.model.model.num_parameters()
											except AttributeError:
																__lowerCAmelCase				   =       pl_module.model.num_parameters()

											__lowerCAmelCase				   =       count_trainable_parameters(_A			)
											# 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 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A			):




											"""simple docstring"""



											save_json(pl_module.metrics		,			pl_module.metrics_save_path			)
											return self._write_logs(_A		,			_A		,			"test"			)






						@rank_zero_only
						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A			):




											"""simple docstring"""



											save_json(pl_module.metrics		,			pl_module.metrics_save_path			)
											# Uncommenting this will save val generations
											# return self._write_logs(trainer, pl_module, "valid")




