

from typing import Any, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
from torch.utils.data import DistributedSampler, RandomSampler

from transformers import PreTrainedModel, Trainer, logging
from transformers.integrations import is_fairscale_available
from transformers.models.fsmt.configuration_fsmt import FSMTConfig
from transformers.optimization import (
    Adafactor,
    AdamW,
    get_constant_schedule,
    get_constant_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.trainer_pt_utils import get_tpu_sampler
from transformers.training_args import ParallelMode
from transformers.utils import is_torch_tpu_available


if is_fairscale_available():
   from fairscale.optim import OSS


_lowerCamelCase      :					int					    =      logging.get_logger(__name__)

_lowerCamelCase      :					List[Any]					    =      {
    '''linear''': get_linear_schedule_with_warmup,
    '''cosine''': get_cosine_schedule_with_warmup,
    '''cosine_w_restarts''': get_cosine_with_hard_restarts_schedule_with_warmup,
    '''polynomial''': get_polynomial_decay_schedule_with_warmup,
    '''constant''': get_constant_schedule,
    '''constant_w_warmup''': get_constant_schedule_with_warmup,
}





class        lowerCAmelCase__  (					__magic_name__			):



       '''simple docstring'''







       def __init__(							self ,   lowercase__=None ,   lowercase__=None ,   *lowercase__ ,   **lowercase__     ):







             '''simple docstring'''



             super().__init__(*lowercase__ ,   **lowercase__     )

             if config is None:
                   assert isinstance(self.model ,   lowercase__     ), (
                       "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
                       f''' {self.model.__class__}'''
                   )
                   __A    =self.model.config
             else:
                   __A    =config

             __A    =data_args
             __A    =self.config.tgt_vocab_size if isinstance(self.config ,   lowercase__     ) else self.config.vocab_size

             if self.args.label_smoothing != 0 or (self.data_args is not None and self.data_args.ignore_pad_token_for_loss):
                   assert self.config.pad_token_id is not None, (
                       "Make sure that `config.pad_token_id` is correcly defined when ignoring `pad_token` for loss"
                       " calculation or doing label smoothing."
                   )

             if self.config.pad_token_id is None and self.config.eos_token_id is not None:
                   logger.warning(
                       f'''The `config.pad_token_id` is `None`. Using `config.eos_token_id` = {self.config.eos_token_id} for'''
                       ''' padding..'''     )

             if self.args.label_smoothing == 0:
                   __A    =torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id     )
             else:
                   # dynamically import label_smoothed_nll_loss
                   from utils import label_smoothed_nll_loss

                   __A    =label_smoothed_nll_loss

       def     __UpperCamelCase     (							self ,   lowercase__     ):







             '''simple docstring'''



             if self.optimizer is None:
                   __A    =['''bias''', '''LayerNorm.weight''']
                   __A    =[
                       {
                           '''params''': [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay     )],
                           '''weight_decay''': self.args.weight_decay,
                       },
                       {
                           '''params''': [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay     )],
                           '''weight_decay''': 0.0,
                       },
                   ]
                   __A    =Adafactor if self.args.adafactor else AdamW
                   if self.args.adafactor:
                         __A    =Adafactor
                         __A    ={'''scale_parameter''': False, '''relative_step''': False}
                   else:
                         __A    =AdamW
                         __A    ={
                             '''betas''': (self.args.adam_betaa, self.args.adam_betaa),
                             '''eps''': self.args.adam_epsilon,
                         }
                   __A    =self.args.learning_rate
                   if self.sharded_ddp:
                         __A    =OSS(
                             params=lowercase__ ,   optim=lowercase__ ,   **lowercase__ ,   )
                   else:
                         __A    =optimizer_cls(lowercase__ ,   **lowercase__     )

             if self.lr_scheduler is None:
                   __A    =self._get_lr_scheduler(lowercase__     )
             else:  # ignoring --lr_scheduler
                   logger.warning('''scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored.'''     )

       def     __UpperCamelCase     (							self ,   lowercase__     ):







             '''simple docstring'''



             __A    =arg_to_scheduler[self.args.lr_scheduler]
             if self.args.lr_scheduler == "constant":
                   __A    =schedule_func(self.optimizer     )
             elif self.args.lr_scheduler == "constant_w_warmup":
                   __A    =schedule_func(self.optimizer ,   num_warmup_steps=self.args.warmup_steps     )
             else:
                   __A    =schedule_func(
                       self.optimizer ,   num_warmup_steps=self.args.warmup_steps ,   num_training_steps=lowercase__     )
             return scheduler

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             if isinstance(self.train_dataset ,   torch.utils.data.IterableDataset     ):
                   return None
             elif is_torch_tpu_available():
                   return get_tpu_sampler(self.train_dataset     )
             else:
                   if self.args.sortish_sampler:
                         self.train_dataset.make_sortish_sampler(
                             self.args.per_device_train_batch_size ,   distributed=(self.args.parallel_mode == ParallelMode.DISTRIBUTED) ,   )

                   return (
                       RandomSampler(self.train_dataset     )
                       if self.args.local_rank == -1
                       else DistributedSampler(self.train_dataset     )
                   )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             if self.args.label_smoothing == 0:
                   if self.data_args is not None and self.data_args.ignore_pad_token_for_loss:
                         # force training to ignore pad token
                         __A    =model(**lowercase__ ,   use_cache=lowercase__     )[0]
                         __A    =self.loss_fn(logits.view(-1 ,   logits.shape[-1]     ) ,   labels.view(-1     )     )
                   else:
                         # compute usual loss via models
                         __A    ,		__A    =model(**lowercase__ ,   labels=lowercase__ ,   use_cache=lowercase__     )[:2]
             else:
                   # compute label smoothed loss
                   __A    =model(**lowercase__ ,   use_cache=lowercase__     )[0]
                   __A    =torch.nn.functional.log_softmax(lowercase__ ,   dim=-1     )
                   __A    ,		__A    =self.loss_fn(lowercase__ ,   lowercase__ ,   self.args.label_smoothing ,   ignore_index=self.config.pad_token_id     )
             return loss, logits

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =inputs.pop('''labels'''     )
             __A    ,		__A    =self._compute_loss(lowercase__ ,   lowercase__ ,   lowercase__     )
             return loss

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ = None ,   ):







             '''simple docstring'''



             __A    =self._prepare_inputs(lowercase__     )

             __A    ={
                 '''max_length''': self.data_args.val_max_target_length
                 if self.data_args is not None
                 else self.config.max_length,
                 '''num_beams''': self.data_args.eval_beams if self.data_args is not None else self.config.num_beams,
             }

             if self.args.predict_with_generate and not self.args.prediction_loss_only:
                   __A    =self.model.generate(
                       inputs['''input_ids'''] ,   attention_mask=inputs['''attention_mask'''] ,   **lowercase__ ,   )
                   # in case the batch is shorter than max length, the output should be padded
                   if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
                         __A    =self._pad_tensors_to_max_len(lowercase__ ,   gen_kwargs['''max_length''']     )

             __A    =inputs.pop('''labels'''     )
             with torch.no_grad():
                   # compute loss on predict data
                   __A    ,		__A    =self._compute_loss(lowercase__ ,   lowercase__ ,   lowercase__     )

             __A    =loss.mean().detach()
             if self.args.prediction_loss_only:
                   return (loss, None, None)

             __A    =generated_tokens if self.args.predict_with_generate else logits

             if labels.shape[-1] < gen_kwargs["max_length"]:
                   __A    =self._pad_tensors_to_max_len(lowercase__ ,   gen_kwargs['''max_length''']     )

             return (loss, logits, labels)


       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =self.config.pad_token_id if self.config.pad_token_id is not None else self.config.eos_token_id

             if pad_token_id is None:
                   raise ValueError(
                       '''Make sure that either `config.pad_token_id` or `config.eos_token_id` is defined if tensor has to be'''
                       f''' padded to `max_length`={max_length}'''     )

             __A    =pad_token_id * torch.ones(
                 (tensor.shape[0], max_length) ,   dtype=tensor.dtype ,   device=tensor.device     )
             __A    =tensor
             return padded_tensor





import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import numpy as np
import pytest

from datasets.arrow_dataset import Dataset
from datasets.search import ElasticSearchIndex, FaissIndex, MissingIndex

from .utils import require_elasticsearch, require_faiss


_lowerCamelCase      :					List[str]					    =      pytest.mark.integration





@require_faiss
class        lowerCAmelCase__  (					__magic_name__			):



       '''simple docstring'''







       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =Dataset.from_dict({'''filename''': ['''my_name-train''' + '''_''' + str(lowercase__     ) for x in np.arange(3_0     ).tolist()]}     )
             return dset

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

             __A    =self._create_dummy_dataset()
             __A    =dset.map(
                 lambda lowercase__ ,   lowercase__     : {"vecs": i * np.ones(5 ,   dtype=np.floataa     )} ,   with_indices=lowercase__ ,   keep_in_memory=lowercase__     )
             __A    =dset.add_faiss_index('''vecs''' ,   batch_size=1_0_0 ,   metric_type=faiss.METRIC_INNER_PRODUCT     )
             __A    ,		__A    =dset.get_nearest_examples('''vecs''' ,   np.ones(5 ,   dtype=np.floataa     )     )
             self.assertEqual(examples['''filename'''][0] ,   '''my_name-train_29'''     )
             dset.drop_index('''vecs'''     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

             __A    =self._create_dummy_dataset()
             dset.add_faiss_index_from_external_arrays(
                 external_arrays=np.ones((3_0, 5)     ) * np.arange(3_0     ).reshape(-1 ,   1     ) ,   index_name='''vecs''' ,   batch_size=1_0_0 ,   metric_type=faiss.METRIC_INNER_PRODUCT ,   )
             __A    ,		__A    =dset.get_nearest_examples('''vecs''' ,   np.ones(5 ,   dtype=np.floataa     )     )
             self.assertEqual(examples['''filename'''][0] ,   '''my_name-train_29'''     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

             __A    =self._create_dummy_dataset()
             dset.add_faiss_index_from_external_arrays(
                 external_arrays=np.ones((3_0, 5)     ) * np.arange(3_0     ).reshape(-1 ,   1     ) ,   index_name='''vecs''' ,   metric_type=faiss.METRIC_INNER_PRODUCT ,   )

             # Setting delete=False and unlinking manually is not pretty... but it is required on Windows to
             # ensure somewhat stable behaviour. If we don't, we get PermissionErrors. This is an age-old issue.
             # see https://bugs.python.org/issue14243 and
             # https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file/23212515
             with tempfile.NamedTemporaryFile(delete=lowercase__     ) as tmp_file:
                   dset.save_faiss_index('''vecs''' ,   tmp_file.name     )
                   dset.load_faiss_index('''vecs2''' ,   tmp_file.name     )
             os.unlink(tmp_file.name     )

             __A    ,		__A    =dset.get_nearest_examples('''vecs2''' ,   np.ones(5 ,   dtype=np.floataa     )     )
             self.assertEqual(examples['''filename'''][0] ,   '''my_name-train_29'''     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self._create_dummy_dataset()
             dset.add_faiss_index_from_external_arrays(
                 external_arrays=np.ones((3_0, 5)     ) * np.arange(3_0     ).reshape(-1 ,   1     ) ,   index_name='''vecs'''     )
             dset.drop_index('''vecs'''     )
             self.assertRaises(lowercase__ ,   partial(dset.get_nearest_examples ,   '''vecs2''' ,   np.ones(5 ,   dtype=np.floataa     )     )     )


       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             from elasticsearch import Elasticsearch

             __A    =self._create_dummy_dataset()
             with patch('''elasticsearch.Elasticsearch.search'''     ) as mocked_search, patch(
                 '''elasticsearch.client.IndicesClient.create'''     ) as mocked_index_create, patch('''elasticsearch.helpers.streaming_bulk'''     ) as mocked_bulk:
                   __A    ={'''acknowledged''': True}
                   mocked_bulk.return_value([(True, None)] * 3_0     )
                   __A    ={'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 2_9}]}}
                   __A    =Elasticsearch()

                   dset.add_elasticsearch_index('''filename''' ,   es_client=lowercase__     )
                   __A    ,		__A    =dset.get_nearest_examples('''filename''' ,   '''my_name-train_29'''     )
                   self.assertEqual(examples['''filename'''][0] ,   '''my_name-train_29'''     )





@require_faiss
class        lowerCAmelCase__  (					__magic_name__			):



       '''simple docstring'''







       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

             __A    =FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT     )

             # add vectors
             index.add_vectors(np.eye(5 ,   dtype=np.floataa     )     )
             self.assertIsNotNone(index.faiss_index     )
             self.assertEqual(index.faiss_index.ntotal ,   5     )
             index.add_vectors(np.zeros((5, 5) ,   dtype=np.floataa     )     )
             self.assertEqual(index.faiss_index.ntotal ,   1_0     )

             # single query
             __A    =np.zeros(5 ,   dtype=np.floataa     )
             __A    =1
             __A    ,		__A    =index.search(lowercase__     )
             self.assertRaises(lowercase__ ,   index.search ,   query.reshape(-1 ,   1     )     )
             self.assertGreater(scores[0] ,   0     )
             self.assertEqual(indices[0] ,   1     )

             # batched queries
             __A    =np.eye(5 ,   dtype=np.floataa     )[::-1]
             __A    ,		__A    =index.search_batch(lowercase__     )
             self.assertRaises(lowercase__ ,   index.search_batch ,   queries[0]     )
             __A    =[scores[0] for scores in total_scores]
             __A    =[indices[0] for indices in total_indices]
             self.assertGreater(np.min(lowercase__     ) ,   0     )
             self.assertListEqual([4, 3, 2, 1, 0] ,   lowercase__     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

             __A    =FaissIndex(string_factory='''Flat'''     )
             index.add_vectors(np.eye(5 ,   dtype=np.floataa     )     )
             self.assertIsInstance(index.faiss_index ,   faiss.IndexFlat     )
             __A    =FaissIndex(string_factory='''LSH'''     )
             index.add_vectors(np.eye(5 ,   dtype=np.floataa     )     )
             self.assertIsInstance(index.faiss_index ,   faiss.IndexLSH     )
             with self.assertRaises(lowercase__     ):
                   __A    =FaissIndex(string_factory='''Flat''' ,   custom_index=faiss.IndexFlat(5     )     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

             __A    =faiss.IndexFlat(5     )
             __A    =FaissIndex(custom_index=lowercase__     )
             index.add_vectors(np.eye(5 ,   dtype=np.floataa     )     )
             self.assertIsInstance(index.faiss_index ,   faiss.IndexFlat     )


       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             import faiss

             __A    =FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT     )
             index.add_vectors(np.eye(5 ,   dtype=np.floataa     )     )

             # Setting delete=False and unlinking manually is not pretty... but it is required on Windows to
             # ensure somewhat stable behaviour. If we don't, we get PermissionErrors. This is an age-old issue.
             # see https://bugs.python.org/issue14243 and
             # https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file/23212515
             with tempfile.NamedTemporaryFile(delete=lowercase__     ) as tmp_file:
                   index.save(tmp_file.name     )
                   __A    =FaissIndex.load(tmp_file.name     )
             os.unlink(tmp_file.name     )

             __A    =np.zeros(5 ,   dtype=np.floataa     )
             __A    =1
             __A    ,		__A    =index.search(lowercase__     )
             self.assertGreater(scores[0] ,   0     )
             self.assertEqual(indices[0] ,   1     )





@require_faiss
def      A__					(			__A  :						Dict			)					->List[Any]:
      import faiss

      __A    =FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT			)
      index.add_vectors(np.eye(5     ,				dtype=np.floataa			)			)

      __A    ='''index.faiss'''
      __A    =F'''mock://{index_name}'''
      index.save(__A     ,				storage_options=mockfs.storage_options			)
      __A    =FaissIndex.load(__A     ,				storage_options=mockfs.storage_options			)

      __A    =np.zeros(5     ,				dtype=np.floataa			)
      __A    =1
      __A    ,		__A    =index.search(__A			)
      assert scores[0] > 0
      assert indices[0] == 1





@require_elasticsearch
class        lowerCAmelCase__  (					__magic_name__			):



       '''simple docstring'''







       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             from elasticsearch import Elasticsearch

             with patch('''elasticsearch.Elasticsearch.search'''     ) as mocked_search, patch(
                 '''elasticsearch.client.IndicesClient.create'''     ) as mocked_index_create, patch('''elasticsearch.helpers.streaming_bulk'''     ) as mocked_bulk:
                   __A    =Elasticsearch()
                   __A    ={'''acknowledged''': True}
                   __A    =ElasticSearchIndex(es_client=lowercase__     )
                   mocked_bulk.return_value([(True, None)] * 3     )
                   index.add_documents(['''foo''', '''bar''', '''foobar''']     )

                   # single query
                   __A    ='''foo'''
                   __A    ={'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 0}]}}
                   __A    ,		__A    =index.search(lowercase__     )
                   self.assertEqual(scores[0] ,   1     )
                   self.assertEqual(indices[0] ,   0     )

                   # single query with timeout
                   __A    ='''foo'''
                   __A    ={'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 0}]}}
                   __A    ,		__A    =index.search(lowercase__ ,   request_timeout=3_0     )
                   self.assertEqual(scores[0] ,   1     )
                   self.assertEqual(indices[0] ,   0     )

                   # batched queries
                   __A    =['''foo''', '''bar''', '''foobar''']
                   __A    ={'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 1}]}}
                   __A    ,		__A    =index.search_batch(lowercase__     )
                   __A    =[scores[0] for scores in total_scores]
                   __A    =[indices[0] for indices in total_indices]
                   self.assertGreater(np.min(lowercase__     ) ,   0     )
                   self.assertListEqual([1, 1, 1] ,   lowercase__     )

                   # batched queries with timeout
                   __A    =['''foo''', '''bar''', '''foobar''']
                   __A    ={'''hits''': {'''hits''': [{'''_score''': 1, '''_id''': 1}]}}
                   __A    ,		__A    =index.search_batch(lowercase__ ,   request_timeout=3_0     )
                   __A    =[scores[0] for scores in total_scores]
                   __A    =[indices[0] for indices in total_indices]
                   self.assertGreater(np.min(lowercase__     ) ,   0     )
                   self.assertListEqual([1, 1, 1] ,   lowercase__     )


