

'''simple docstring'''
import os
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers.models.realm.configuration_realm import RealmConfig
from transformers.models.realm.retrieval_realm import _REALM_BLOCK_RECORDS_FILENAME, RealmRetriever
from transformers.models.realm.tokenization_realm import VOCAB_FILES_NAMES, RealmTokenizer







class 							A__    (       A__  ):





  def 						A  (		self      :	Any )	->    Optional[int]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =tempfile.mkdtemp()
       _SCREAMING_SNAKE_CASE							    =5

       # Realm tok
       _SCREAMING_SNAKE_CASE							    =[
           '[UNK]',
           '[CLS]',
           '[SEP]',
           '[PAD]',
           '[MASK]',
           'test',
           'question',
           'this',
           'is',
           'the',
           'first',
           'second',
           'third',
           'fourth',
           'fifth',
           'record',
           'want',
           '##want',
           '##ed',
           'wa',
           'un',
           'runn',
           '##ing',
           ',',
           'low',
           'lowest',
       ]
       _SCREAMING_SNAKE_CASE							    =os.path.join(self.tmpdirname							,					'realm_tokenizer' )
       os.makedirs(_a							,					exist_ok=_a )
       _SCREAMING_SNAKE_CASE							    =os.path.join(_a							,					VOCAB_FILES_NAMES['vocab_file'] )
       with open(self.vocab_file							,					'w'							,					encoding='utf-8' ) as vocab_writer:
            vocab_writer.write(''.join([x + '\n' for x in vocab_tokens] ) )

       _SCREAMING_SNAKE_CASE							    =os.path.join(self.tmpdirname							,					'realm_block_records' )
       os.makedirs(_a							,					exist_ok=_a )





  def 						A  (		self      :	Union[str, Any] )	->    RealmTokenizer:





       '''simple docstring'''






       return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname							,					'realm_tokenizer' ) )





  def 						A  (		self      :	int )	->    Any:





       '''simple docstring'''






       shutil.rmtree(self.tmpdirname )





  def 						A  (		self      :	Optional[int] )	->    List[str]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =RealmConfig(num_block_records=self.num_block_records )
       return config





  def 						A  (		self      :	Optional[Any] )	->    List[Any]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =Dataset.from_dict(
           {
               'id': ['0', '1'],
               'question': ['foo', 'bar'],
               'answers': [['Foo', 'Bar'], ['Bar']],
           } )
       return dataset





  def 						A  (		self      :	int )	->    Optional[int]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =np.array(
           [
               b'This is the first record',
               b'This is the second record',
               b'This is the third record',
               b'This is the fourth record',
               b'This is the fifth record',
               b'This is a longer longer longer record',
           ]							,					dtype=_a							,					)
       return block_records





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





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =RealmRetriever(
           block_records=self.get_dummy_block_records()							,					tokenizer=self.get_tokenizer()							,					)
       return retriever





  def 						A  (		self      :	str )	->    int:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =self.get_config()
       _SCREAMING_SNAKE_CASE							    =self.get_dummy_retriever()
       _SCREAMING_SNAKE_CASE							    =retriever.tokenizer

       _SCREAMING_SNAKE_CASE							    =np.array([0, 3]							,					dtype='long' )
       _SCREAMING_SNAKE_CASE							    =tokenizer(['Test question'] ).input_ids
       _SCREAMING_SNAKE_CASE							    =tokenizer(
           ['the fourth']							,					add_special_tokens=_a							,					return_token_type_ids=_a							,					return_attention_mask=_a							,					).input_ids
       _SCREAMING_SNAKE_CASE							    =config.reader_seq_len

       _SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE							    =retriever(
           _a							,					_a							,					answer_ids=_a							,					max_length=_a							,					return_tensors='np' )

       self.assertEqual(len(_a )							,					2 )
       self.assertEqual(len(_a )							,					2 )
       self.assertEqual(len(_a )							,					2 )
       self.assertEqual(concat_inputs.input_ids.shape							,					(2, 10) )
       self.assertEqual(concat_inputs.attention_mask.shape							,					(2, 10) )
       self.assertEqual(concat_inputs.token_type_ids.shape							,					(2, 10) )
       self.assertEqual(concat_inputs.special_tokens_mask.shape							,					(2, 10) )
       self.assertEqual(
           tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[0] )							,					['[CLS]', 'test', 'question', '[SEP]', 'this', 'is', 'the', 'first', 'record', '[SEP]']							,					)
       self.assertEqual(
           tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[1] )							,					['[CLS]', 'test', 'question', '[SEP]', 'this', 'is', 'the', 'fourth', 'record', '[SEP]']							,					)





  def 						A  (		self      :	List[Any] )	->    Tuple:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =self.get_config()
       _SCREAMING_SNAKE_CASE							    =self.get_dummy_retriever()
       _SCREAMING_SNAKE_CASE							    =retriever.tokenizer

       _SCREAMING_SNAKE_CASE							    =np.array([0, 3, 5]							,					dtype='long' )
       _SCREAMING_SNAKE_CASE							    =tokenizer(['Test question'] ).input_ids
       _SCREAMING_SNAKE_CASE							    =tokenizer(
           ['the fourth', 'longer longer']							,					add_special_tokens=_a							,					return_token_type_ids=_a							,					return_attention_mask=_a							,					).input_ids
       _SCREAMING_SNAKE_CASE							    =config.reader_seq_len

       _SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE							    =retriever(
           _a							,					_a							,					answer_ids=_a							,					max_length=_a							,					return_tensors='np' )

       self.assertEqual([False, True, True]							,					_a )
       self.assertEqual([[-1, -1, -1], [6, -1, -1], [6, 7, 8]]							,					_a )
       self.assertEqual([[-1, -1, -1], [7, -1, -1], [7, 8, 9]]							,					_a )





  def 						A  (		self      :	Tuple )	->    Union[str, Any]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =self.get_dummy_retriever()
       retriever.save_pretrained(os.path.join(self.tmpdirname							,					'realm_block_records' ) )

       # Test local path
       _SCREAMING_SNAKE_CASE							    =retriever.from_pretrained(os.path.join(self.tmpdirname							,					'realm_block_records' ) )
       self.assertEqual(retriever.block_records[0]							,					b'This is the first record' )

       # Test mocked remote path
       with patch('transformers.models.realm.retrieval_realm.hf_hub_download' ) as mock_hf_hub_download:
            _SCREAMING_SNAKE_CASE							    =os.path.join(
                os.path.join(self.tmpdirname							,					'realm_block_records' )							,					_REALM_BLOCK_RECORDS_FILENAME )
            _SCREAMING_SNAKE_CASE							    =RealmRetriever.from_pretrained('google/realm-cc-news-pretrained-openqa' )

       self.assertEqual(retriever.block_records[0]							,					b'This is the first record' )







"""simple docstring"""



# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch

from ..models.auto import AutoModelForSequenceClassification, AutoTokenizer
from .base import PipelineTool






class 		__snake_case					(		SCREAMING_SNAKE_CASE__     ):





 """simple docstring"""


 _lowerCamelCase        =     """facebook/bart-large-mnli"""
 _lowerCamelCase        =     (
     """This is a tool that classifies an English text using provided labels. It takes two inputs: `text`, which """
     """should be the text to classify, and `labels`, which should be the list of labels to use for classification. """
     """It returns the most likely label in the list of provided `labels` for the input text."""
 )
 _lowerCamelCase        =     """text_classifier"""
 _lowerCamelCase        =     AutoTokenizer
 _lowerCamelCase        =     AutoModelForSequenceClassification

 _lowerCamelCase        =     ["""text""", ["""text"""]]
 _lowerCamelCase        =     ["""text"""]

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     super().setup()
     __A    :						List[str]						 =     self.model.config
     __A    :						int						 =     -1
     for idx, label in config.idalabel.items():
         if label.lower().startswith('''entail'''  ):
             __A    :						List[str]						 =     int(__lowerCamelCase  )
     if self.entailment_id == -1:
         raise ValueError('''Could not determine the entailment ID from the model config, please pass it at init.'''  )

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						Union[str, Any]						 =     labels
     return self.pre_processor(
         [text] * len(__lowerCamelCase  )					,						[F"""This example is {label}""" for label in labels]					,						return_tensors='''pt'''					,						padding='''max_length'''					,						)







 def 					UpperCamelCase__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						List[Any]						 =     outputs.logits
     __A    :						List[str]						 =     torch.argmax(logits[:, 2]  ).item()
     return self._labels[label_id]




