



"""simple docstring"""





import inspect

import jax
import jax.lax as lax
import jax.numpy as jnp

from ..utils import add_start_docstrings
from ..utils.logging import get_logger


lowerCamelCase  			=		get_logger(__name__)


lowerCamelCase  			=		r"""
    Args:
        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
            Indices of input sequence tokens in the vocabulary.

            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            [What are input IDs?](../glossary#input-ids)
        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):
            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam
            search or log softmax for each vocabulary token when using beam search
        kwargs (`Dict[str, Any]`, *optional*):
            Additional logits processor specific kwargs.

    Return:
        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.

"""
class  lowercase__  :







       '''simple docstring'''


       @add_start_docstrings(_UpperCAmelCase       )
       def __call__(				self     :    str						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    jnp.ndarray       )					->		jnp.ndarray:


             '''simple docstring'''




             raise NotImplementedError(
                 F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""       )



class  lowercase__  :







       '''simple docstring'''


       @add_start_docstrings(_UpperCAmelCase       )
       def __call__(				self     :    Union[str, Any]						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    jnp.ndarray       )					->		jnp.ndarray:


             '''simple docstring'''




             raise NotImplementedError(
                 F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""       )



class  lowercase__  (       SCREAMING_SNAKE_CASE ):







       '''simple docstring'''


       @add_start_docstrings(_UpperCAmelCase       )
       def __call__(				self     :    Any						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    int						,							**_UpperCAmelCase     :    List[str]       )					->		jnp.ndarray:


             '''simple docstring'''




             for processor in self:
                   UpperCAmelCase_     					=				inspect.signature(processor.__call__       ).parameters
                   if len(_UpperCAmelCase       ) > 3:
                         if not all(arg in kwargs for arg in list(function_args.keys()       )[2:]       ):
                               raise ValueError(
                                   F"""Make sure that all the required parameters: {list(function_args.keys()       )} for """
                                   F"""{processor.__class__} are passed to the logits processor."""       )
                         UpperCAmelCase_     					=				processor(_UpperCAmelCase						,							_UpperCAmelCase						,							_UpperCAmelCase						,							**_UpperCAmelCase       )
                   else:
                         UpperCAmelCase_     					=				processor(_UpperCAmelCase						,							_UpperCAmelCase						,							_UpperCAmelCase       )
             return scores



class  lowercase__  (       SCREAMING_SNAKE_CASE ):







       '''simple docstring'''


       def __init__(				self     :    Optional[int]						,							_UpperCAmelCase     :    float       )					->		List[Any]:


             '''simple docstring'''




             if not isinstance(_UpperCAmelCase						,							_UpperCAmelCase       ) or not (temperature > 0):
                   raise ValueError(F"""`temperature` has to be a strictly positive float, but is {temperature}"""       )

             UpperCAmelCase_     					=				temperature







       def __call__(				self     :    Dict						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    int       )					->		jnp.ndarray:


             '''simple docstring'''




             UpperCAmelCase_     					=				scores / self.temperature
             return scores



class  lowercase__  (       SCREAMING_SNAKE_CASE ):







       '''simple docstring'''


       def __init__(				self     :    Union[str, Any]						,							_UpperCAmelCase     :    float						,							_UpperCAmelCase     :    float = -float("Inf"       )						,							_UpperCAmelCase     :    int = 1       )					->		Optional[int]:


             '''simple docstring'''




             if not isinstance(_UpperCAmelCase						,							_UpperCAmelCase       ) or (top_p < 0 or top_p > 1.0):
                   raise ValueError(F"""`top_p` has to be a float > 0 and < 1, but is {top_p}"""       )
             if not isinstance(_UpperCAmelCase						,							_UpperCAmelCase       ) or (min_tokens_to_keep < 1):
                   raise ValueError(F"""`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}"""       )

             UpperCAmelCase_     					=				top_p
             UpperCAmelCase_     					=				filter_value
             UpperCAmelCase_     					=				min_tokens_to_keep







       def __call__(				self     :    Optional[int]						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    int       )					->		jnp.ndarray:


             '''simple docstring'''




             UpperCAmelCase_					,       UpperCAmelCase_     					=				lax.top_k(_UpperCAmelCase						,							scores.shape[-1]       )

             UpperCAmelCase_     					=				jnp.full_like(_UpperCAmelCase						,							self.filter_value       )
             UpperCAmelCase_     					=				jax.nn.softmax(_UpperCAmelCase						,							axis=-1       ).cumsum(axis=-1       )
             UpperCAmelCase_     					=				cumulative_probs < self.top_p

             # include the token that is higher than top_p as well
             UpperCAmelCase_     					=				jnp.roll(_UpperCAmelCase						,							1       )
             score_mask |= score_mask.at[:, 0].set(_UpperCAmelCase       )

             # min tokens to keep
             UpperCAmelCase_     					=				score_mask.at[:, : self.min_tokens_to_keep].set(_UpperCAmelCase       )

             UpperCAmelCase_     					=				jnp.where(_UpperCAmelCase						,							_UpperCAmelCase						,							_UpperCAmelCase       )
             UpperCAmelCase_     					=				jax.lax.sort_key_val(_UpperCAmelCase						,							_UpperCAmelCase       )[-1]

             return next_scores



class  lowercase__  (       SCREAMING_SNAKE_CASE ):







       '''simple docstring'''


       def __init__(				self     :    Dict						,							_UpperCAmelCase     :    int						,							_UpperCAmelCase     :    float = -float("Inf"       )						,							_UpperCAmelCase     :    int = 1       )					->		Optional[int]:


             '''simple docstring'''




             if not isinstance(_UpperCAmelCase						,							_UpperCAmelCase       ) or top_k <= 0:
                   raise ValueError(F"""`top_k` has to be a strictly positive integer, but is {top_k}"""       )

             UpperCAmelCase_     					=				max(_UpperCAmelCase						,							_UpperCAmelCase       )
             UpperCAmelCase_     					=				filter_value







       def __call__(				self     :    Union[str, Any]						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    int       )					->		jnp.ndarray:


             '''simple docstring'''




             UpperCAmelCase_					,       UpperCAmelCase_     					=				scores.shape
             UpperCAmelCase_     					=				jnp.full(batch_size * vocab_size						,							self.filter_value       )

             UpperCAmelCase_     					=				min(self.top_k						,							scores.shape[-1]       )  # Safety check
             UpperCAmelCase_					,       UpperCAmelCase_     					=				lax.top_k(_UpperCAmelCase						,							_UpperCAmelCase       )
             UpperCAmelCase_     					=				jnp.broadcast_to((jnp.arange(_UpperCAmelCase       ) * vocab_size)[:, None]						,							(batch_size, topk)       ).flatten()
             UpperCAmelCase_     					=				topk_scores.flatten()
             UpperCAmelCase_     					=				topk_indices.flatten() + shift

             UpperCAmelCase_     					=				next_scores_flat.at[topk_indices_flat].set(_UpperCAmelCase       )
             UpperCAmelCase_     					=				next_scores_flat.reshape(_UpperCAmelCase						,							_UpperCAmelCase       )
             return next_scores



class  lowercase__  (       SCREAMING_SNAKE_CASE ):







       '''simple docstring'''


       def __init__(				self     :    str						,							_UpperCAmelCase     :    int       )					->		Union[str, Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				bos_token_id







       def __call__(				self     :    Optional[Any]						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    int       )					->		jnp.ndarray:


             '''simple docstring'''




             UpperCAmelCase_     					=				jnp.full(scores.shape						,							-float("inf"       )       )

             UpperCAmelCase_     					=				1 - jnp.bool_(cur_len - 1       )

             UpperCAmelCase_     					=				jnp.where(_UpperCAmelCase						,							new_scores.at[:, self.bos_token_id].set(0       )						,							_UpperCAmelCase       )

             return scores



class  lowercase__  (       SCREAMING_SNAKE_CASE ):







       '''simple docstring'''


       def __init__(				self     :    Any						,							_UpperCAmelCase     :    int						,							_UpperCAmelCase     :    int       )					->		Tuple:


             '''simple docstring'''




             UpperCAmelCase_     					=				max_length
             UpperCAmelCase_     					=				eos_token_id







       def __call__(				self     :    Optional[int]						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    int       )					->		jnp.ndarray:


             '''simple docstring'''




             UpperCAmelCase_     					=				jnp.full(scores.shape						,							-float("inf"       )       )

             UpperCAmelCase_     					=				1 - jnp.bool_(cur_len - self.max_length + 1       )

             UpperCAmelCase_     					=				jnp.where(_UpperCAmelCase						,							new_scores.at[:, self.eos_token_id].set(0       )						,							_UpperCAmelCase       )

             return scores



class  lowercase__  (       SCREAMING_SNAKE_CASE ):







       '''simple docstring'''


       def __init__(				self     :    Tuple						,							_UpperCAmelCase     :    int						,							_UpperCAmelCase     :    int       )					->		List[Any]:


             '''simple docstring'''




             if not isinstance(_UpperCAmelCase						,							_UpperCAmelCase       ) or min_length < 0:
                   raise ValueError(F"""`min_length` has to be a positive integer, but is {min_length}"""       )

             if not isinstance(_UpperCAmelCase						,							_UpperCAmelCase       ) or eos_token_id < 0:
                   raise ValueError(F"""`eos_token_id` has to be a positive integer, but is {eos_token_id}"""       )

             UpperCAmelCase_     					=				min_length
             UpperCAmelCase_     					=				eos_token_id







       def __call__(				self     :    List[str]						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    int       )					->		jnp.ndarray:


             '''simple docstring'''




             UpperCAmelCase_     					=				1 - jnp.clip(cur_len - self.min_length						,							0						,							1       )

             UpperCAmelCase_     					=				jnp.where(_UpperCAmelCase						,							scores.at[:, self.eos_token_id].set(-float("inf"       )       )						,							_UpperCAmelCase       )

             return scores



class  lowercase__  (       SCREAMING_SNAKE_CASE ):







       '''simple docstring'''


       def __init__(				self     :    List[str]						,							_UpperCAmelCase     :    List[Any]						,							_UpperCAmelCase     :    Optional[Any]       )					->		List[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				list(_UpperCAmelCase       )
             UpperCAmelCase_     					=				begin_index







       def __call__(				self     :    str						,							_UpperCAmelCase     :    Union[str, Any]						,							_UpperCAmelCase     :    int						,							_UpperCAmelCase     :    int       )					->		Any:


             '''simple docstring'''




             UpperCAmelCase_     					=				1 - jnp.bool_(cur_len - self.begin_index       )

             UpperCAmelCase_     					=				jnp.where(_UpperCAmelCase						,							scores.at[:, self.begin_suppress_tokens].set(-float("inf"       )       )						,							_UpperCAmelCase       )

             return scores



class  lowercase__  (       SCREAMING_SNAKE_CASE ):







       '''simple docstring'''


       def __init__(				self     :    List[str]						,							_UpperCAmelCase     :    list       )					->		Tuple:


             '''simple docstring'''




             UpperCAmelCase_     					=				list(_UpperCAmelCase       )







       def __call__(				self     :    Dict						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    int       )					->		jnp.ndarray:


             '''simple docstring'''




             UpperCAmelCase_     					=				scores.at[..., self.suppress_tokens].set(-float("inf"       )       )

             return scores



class  lowercase__  (       SCREAMING_SNAKE_CASE ):







       '''simple docstring'''


       def __init__(				self     :    Optional[int]						,							_UpperCAmelCase     :    str       )					->		Any:


             '''simple docstring'''




             UpperCAmelCase_     					=				dict(_UpperCAmelCase       )
             # Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the
             # index of the array corresponds to the index of the token to be forced, for XLA compatibility.
             # Indexes without forced tokens will have a negative value.
             UpperCAmelCase_     					=				jnp.ones((max(force_token_map.keys()       ) + 1)						,							dtype=jnp.intaa       ) * -1
             for index, token in force_token_map.items():
                   if token is not None:
                         UpperCAmelCase_     					=				force_token_array.at[index].set(_UpperCAmelCase       )
             UpperCAmelCase_     					=				jnp.intaa(_UpperCAmelCase       )







       def __call__(				self     :    str						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    jnp.ndarray						,							_UpperCAmelCase     :    int       )					->		jnp.ndarray:


             '''simple docstring'''




             def _force_token(_UpperCAmelCase     :    Optional[Any]       ):
                   UpperCAmelCase_     					=				scores.shape[0]
                   UpperCAmelCase_     					=				self.force_token_array[generation_idx]

                   UpperCAmelCase_     					=				jnp.ones_like(_UpperCAmelCase						,							dtype=scores.dtype       ) * -float("inf"       )
                   UpperCAmelCase_     					=				jnp.zeros((batch_size, 1)						,							dtype=scores.dtype       )
                   UpperCAmelCase_     					=				lax.dynamic_update_slice(_UpperCAmelCase						,							_UpperCAmelCase						,							(0, current_token)       )
                   return new_scores

             UpperCAmelCase_     					=				lax.cond(
                 cur_len >= self.force_token_array.shape[0]						,							lambda: scores						,							lambda: lax.cond(
                     self.force_token_array[cur_len] >= 0						,							lambda: _force_token(_UpperCAmelCase       )						,							lambda: scores						,							)						,							)
             return scores



class  lowercase__  (       SCREAMING_SNAKE_CASE ):







       '''simple docstring'''


       def __init__(				self     :    Tuple						,							_UpperCAmelCase     :    Optional[Any]						,							_UpperCAmelCase     :    Dict						,							_UpperCAmelCase     :    Optional[int]       )					->		Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				generate_config.eos_token_id
             UpperCAmelCase_     					=				generate_config.no_timestamps_token_id
             UpperCAmelCase_     					=				generate_config.no_timestamps_token_id + 1

             UpperCAmelCase_     					=				decoder_input_length + 1

             if generate_config.is_multilingual:
                   # room for language token and task token
                   self.begin_index += 2
             if hasattr(_UpperCAmelCase						,							"max_initial_timestamp_index"       ):
                   UpperCAmelCase_     					=				generate_config.max_initial_timestamp_index
             else:
                   UpperCAmelCase_     					=				model_config.vocab_size
             if self.max_initial_timestamp_index is None:
                   UpperCAmelCase_     					=				model_config.vocab_size







       def __call__(				self     :    Optional[int]						,							_UpperCAmelCase     :    Tuple						,							_UpperCAmelCase     :    Optional[int]						,							_UpperCAmelCase     :    Optional[int]       )					->		int:


             '''simple docstring'''




             UpperCAmelCase_     					=				scores.at[:, self.no_timestamps_token_id].set(-float("inf"       )       )

             def handle_pairs(_UpperCAmelCase     :    int						,							_UpperCAmelCase     :    int       ):
                   UpperCAmelCase_     					=				jnp.where((cur_len - self.begin_index) >= 1						,							_UpperCAmelCase						,							_UpperCAmelCase       )
                   UpperCAmelCase_     					=				jnp.where(
                       input_ids_k[cur_len - 1] >= self.timestamp_begin						,							True and last_was_timestamp						,							_UpperCAmelCase						,							)

                   UpperCAmelCase_     					=				jnp.where((cur_len - self.begin_index) < 2						,							_UpperCAmelCase						,							_UpperCAmelCase       )
                   UpperCAmelCase_     					=				jnp.where(
                       input_ids_k[cur_len - 2] >= self.timestamp_begin						,							_UpperCAmelCase						,							_UpperCAmelCase						,							)

                   return jnp.where(
                       _UpperCAmelCase						,							jnp.where(
                           penultimate_was_timestamp > 0						,							scores_k.at[self.timestamp_begin :].set(-float("inf"       )       )						,							scores_k.at[: self.eos_token_id].set(-float("inf"       )       )						,							)						,							_UpperCAmelCase						,							)

             UpperCAmelCase_     					=				jax.vmap(_UpperCAmelCase       )(_UpperCAmelCase						,							_UpperCAmelCase       )

             UpperCAmelCase_     					=				jnp.where(cur_len == self.begin_index						,							_UpperCAmelCase						,							_UpperCAmelCase       )
             UpperCAmelCase_     					=				jnp.where(
                 self.max_initial_timestamp_index is not None						,							True and apply_max_initial_timestamp						,							_UpperCAmelCase						,							)

             UpperCAmelCase_     					=				self.timestamp_begin + self.max_initial_timestamp_index

             UpperCAmelCase_     					=				jnp.where(
                 _UpperCAmelCase						,							scores.at[:, last_allowed + 1 :].set(-float("inf"       )       )						,							_UpperCAmelCase						,							)

             # if sum of probability over timestamps is above any other token, sample timestamp
             UpperCAmelCase_     					=				jax.nn.log_softmax(_UpperCAmelCase						,							axis=-1       )

             def handle_cumulative_probs(_UpperCAmelCase     :    Optional[Any]						,							_UpperCAmelCase     :    Union[str, Any]       ):
                   UpperCAmelCase_     					=				jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]						,							axis=-1       )
                   UpperCAmelCase_     					=				jnp.max(logprobs_k[: self.timestamp_begin]       )
                   return jnp.where(
                       timestamp_logprob > max_text_token_logprob						,							scores_k.at[: self.timestamp_begin].set(-float("inf"       )       )						,							_UpperCAmelCase						,							)

             UpperCAmelCase_     					=				jax.vmap(_UpperCAmelCase       )(_UpperCAmelCase						,							_UpperCAmelCase       )

             return scores





import json
import logging
import os
import sys
from pathlib import Path

import finetune_rag

from transformers.file_utils import is_apex_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    require_ray,
    require_torch_gpu,
    require_torch_multi_gpu,
)


logging.basicConfig(level=logging.DEBUG)
SCREAMING_SNAKE_CASE			:						List[Any]				  =		logging.getLogger()

SCREAMING_SNAKE_CASE			:						List[str]				  =		logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)
class  A_			(							a_			):

  def 		_UpperCAmelCase (		self			:   str  ,				__SCREAMING_SNAKE_CASE			:   int							):
   os.makedirs(__SCREAMING_SNAKE_CASE  ,				exist_ok=__SCREAMING_SNAKE_CASE							)
   __a							=	{"source": "What is love ?", "target": "life"}
   __a							=	{"train": 12, "val": 2, "test": 2}
   for split in ["train", "test", "val"]:
    for field in ["source", "target"]:
     __a							=	"\n".join([contents[field]] * n_lines[split]							)
     with open(os.path.join(__SCREAMING_SNAKE_CASE  ,				f"""{split}.{field}"""							)  ,				"w"							) as f:
      f.write(__SCREAMING_SNAKE_CASE							)

  def 		_UpperCAmelCase (		self			:   str  ,				__SCREAMING_SNAKE_CASE			:   int  ,				__SCREAMING_SNAKE_CASE			:   str = "pytorch"							):
   __a							=	self.get_auto_remove_tmp_dir()
   __a							=	os.path.join(__SCREAMING_SNAKE_CASE  ,				"output"							)
   __a							=	os.path.join(__SCREAMING_SNAKE_CASE  ,				"data"							)
   self._create_dummy_data(data_dir=__SCREAMING_SNAKE_CASE							)

   __a							=	f"""
                --data_dir {data_dir} \
                --output_dir {output_dir} \
                --model_name_or_path facebook/rag-sequence-base \
                --model_type rag_sequence \
                --do_train \
                --do_predict \
                --n_val -1 \
                --val_check_interval 1.0 \
                --train_batch_size 2 \
                --eval_batch_size 1 \
                --max_source_length 25 \
                --max_target_length 25 \
                --val_max_target_length 25 \
                --test_max_target_length 25 \
                --label_smoothing 0.1 \
                --dropout 0.1 \
                --attention_dropout 0.1 \
                --weight_decay 0.001 \
                --adam_epsilon 1e-08 \
                --max_grad_norm 0.1 \
                --lr_scheduler polynomial \
                --learning_rate 3e-04 \
                --num_train_epochs 1 \
                --warmup_steps 4 \
                --gradient_accumulation_steps 1 \
                --distributed-port 8787 \
                --use_dummy_dataset 1 \
                --distributed_retriever {distributed_retriever} \
            """.split()

   if gpus > 0:
    testargs.append(f"""--gpus={gpus}"""							)
    if is_apex_available():
     testargs.append("--fp16"							)
   else:
    testargs.append("--gpus=0"							)
    testargs.append("--distributed_backend=ddp_cpu"							)
    testargs.append("--num_processes=2"							)

   __a							=	[sys.executable, str(Path(finetune_rag.__file__							).resolve()							)] + testargs
   execute_subprocess_async(__SCREAMING_SNAKE_CASE  ,				env=self.get_env()							)

   __a							=	os.path.join(__SCREAMING_SNAKE_CASE  ,				"metrics.json"							)
   with open(__SCREAMING_SNAKE_CASE							) as f:
    __a							=	json.load(__SCREAMING_SNAKE_CASE							)
   return result

  @require_torch_gpu
  def 		_UpperCAmelCase (		self			:   Dict							):
   __a							=	self._run_finetune(gpus=1							)
   self.assertGreaterEqual(result["test"][0]["test_avg_em"]  ,				0.2							)

  @require_torch_multi_gpu
  def 		_UpperCAmelCase (		self			:   Optional[int]							):
   __a							=	self._run_finetune(gpus=2							)
   self.assertGreaterEqual(result["test"][0]["test_avg_em"]  ,				0.2							)

  @require_torch_gpu
  @require_ray
  def 		_UpperCAmelCase (		self			:   Optional[Any]							):
   __a							=	self._run_finetune(gpus=1  ,				distributed_retriever="ray"							)
   self.assertGreaterEqual(result["test"][0]["test_avg_em"]  ,				0.2							)

  @require_torch_multi_gpu
  @require_ray
  def 		_UpperCAmelCase (		self			:   Any							):
   __a							=	self._run_finetune(gpus=1  ,				distributed_retriever="ray"							)
   self.assertGreaterEqual(result["test"][0]["test_avg_em"]  ,				0.2							)


