




"""simple docstring"""




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


__A							:    List[Any]  							=							pytest.mark.integration






@require_faiss
class 					__magic_name__							( UpperCAmelCase__       ):






				'''simple docstring'''





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      Dataset.from_dict({"""filename""": ["""my_name-train""" + """_""" + str(_a				) for x in np.arange(30				).tolist()]}				)
							return dset





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							import faiss

							lowerCamelCase       =      self._create_dummy_dataset()
							lowerCamelCase       =      dset.map(
							    lambda _a	,				_a				: {"vecs": i * np.ones(5	,				dtype=np.floataa				)}	,				with_indices=_a	,				keep_in_memory=_a				)
							lowerCamelCase       =      dset.add_faiss_index("""vecs"""	,				batch_size=100	,				metric_type=faiss.METRIC_INNER_PRODUCT				)
							lowerCamelCase   ,     lowerCamelCase       =      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      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							import faiss

							lowerCamelCase       =      self._create_dummy_dataset()
							dset.add_faiss_index_from_external_arrays(
							    external_arrays=np.ones((30, 5)				) * np.arange(30				).reshape(-1	,				1				)	,				index_name="""vecs"""	,				batch_size=100	,				metric_type=faiss.METRIC_INNER_PRODUCT	,				)
							lowerCamelCase   ,     lowerCamelCase       =      dset.get_nearest_examples("""vecs"""	,				np.ones(5	,				dtype=np.floataa				)				)
							self.assertEqual(examples["""filename"""][0]	,				"""my_name-train_29"""				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							import faiss

							lowerCamelCase       =      self._create_dummy_dataset()
							dset.add_faiss_index_from_external_arrays(
							    external_arrays=np.ones((30, 5)				) * np.arange(30				).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=_a				) as tmp_file:
										dset.save_faiss_index("""vecs"""	,				tmp_file.name				)
										dset.load_faiss_index("""vecs2"""	,				tmp_file.name				)
							os.unlink(tmp_file.name				)

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self._create_dummy_dataset()
							dset.add_faiss_index_from_external_arrays(
							    external_arrays=np.ones((30, 5)				) * np.arange(30				).reshape(-1	,				1				)	,				index_name="""vecs"""				)
							dset.drop_index("""vecs"""				)
							self.assertRaises(_a	,				partial(dset.get_nearest_examples	,				"""vecs2"""	,				np.ones(5	,				dtype=np.floataa				)				)				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							from elasticsearch import Elasticsearch

							lowerCamelCase       =      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:
										lowerCamelCase       =      {"""acknowledged""": True}
										mocked_bulk.return_value([(True, None)] * 30				)
										lowerCamelCase       =      {"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 29}]}}
										lowerCamelCase       =      Elasticsearch()

										dset.add_elasticsearch_index("""filename"""	,				es_client=_a				)
										lowerCamelCase   ,     lowerCamelCase       =      dset.get_nearest_examples("""filename"""	,				"""my_name-train_29"""				)
										self.assertEqual(examples["""filename"""][0]	,				"""my_name-train_29"""				)






@require_faiss
class 					__magic_name__							( UpperCAmelCase__       ):






				'''simple docstring'''





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							import faiss

							lowerCamelCase       =      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	,				10				)

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

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							import faiss

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							import faiss

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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							import faiss

							lowerCamelCase       =      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=_a				) as tmp_file:
										index.save(tmp_file.name				)
										lowerCamelCase       =      FaissIndex.load(tmp_file.name				)
							os.unlink(tmp_file.name				)

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




@require_faiss
def 					a__				(       snake_case__       )		->      Tuple:
			import faiss

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

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

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






@require_elasticsearch
class 					__magic_name__							( UpperCAmelCase__       ):






				'''simple docstring'''





				def      _lowerCAmelCase       ( 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:
										lowerCamelCase       =      Elasticsearch()
										lowerCamelCase       =      {"""acknowledged""": True}
										lowerCamelCase       =      ElasticSearchIndex(es_client=_a				)
										mocked_bulk.return_value([(True, None)] * 3				)
										index.add_documents(["""foo""", """bar""", """foobar"""]				)

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

										# single query with timeout
										lowerCamelCase       =      """foo"""
										lowerCamelCase       =      {"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 0}]}}
										lowerCamelCase   ,     lowerCamelCase       =      index.search(_a	,				request_timeout=30				)
										self.assertEqual(scores[0]	,				1				)
										self.assertEqual(indices[0]	,				0				)

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

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









"""simple docstring"""




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

from ..utils import add_end_docstrings, is_tf_available, is_torch_available, logging
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_tf_available():
					import tensorflow as tf

if is_torch_available():
					import torch


lowerCAmelCase							:    List[str]  							=							logging.get_logger(__name__)






class 					__magic_name__							:






				'''simple docstring'''





				def __init__( self	,				_a = None	,				_a = None	,				_a=None	,				_a=None				):





							"""simple docstring"""
							if not conversation_id:
										lowerCamelCase       =      uuid.uuida()
							if past_user_inputs is None:
										lowerCamelCase       =      []
							if generated_responses is None:
										lowerCamelCase       =      []

							lowerCamelCase       =      conversation_id
							lowerCamelCase       =      past_user_inputs
							lowerCamelCase       =      generated_responses
							lowerCamelCase       =      text





				def __eq__( self	,				_a				):





							"""simple docstring"""
							if not isinstance(_a	,				_a				):
										return False
							if self.uuid == other.uuid:
										return True
							return (
							    self.new_user_input == other.new_user_input
							    and self.past_user_inputs == other.past_user_inputs
							    and self.generated_responses == other.generated_responses
							)





				def      _lowerCAmelCase       ( self	,				_a	,				_a = False				):





							"""simple docstring"""
							if self.new_user_input:
										if overwrite:
													logger.warning(
													    f'User input added while unprocessed input was existing: "{self.new_user_input}" was overwritten '
													    f'with: "{text}".'				)
													lowerCamelCase       =      text
										else:
													logger.warning(
													    f'User input added while unprocessed input was existing: "{self.new_user_input}" new input '
													    f'ignored: "{text}". Set `overwrite` to True to overwrite unprocessed user input'				)
							else:
										lowerCamelCase       =      text





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							if self.new_user_input:
										self.past_user_inputs.append(self.new_user_input				)
							lowerCamelCase       =      None





				def      _lowerCAmelCase       ( self	,				_a				):





							"""simple docstring"""
							self.generated_responses.append(_a				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							for user_input, generated_response in zip(self.past_user_inputs	,				self.generated_responses				):
										yield True, user_input
										yield False, generated_response
							if self.new_user_input:
										yield True, self.new_user_input





				def __repr__( self				):





							"""simple docstring"""
							lowerCamelCase       =      f'Conversation id: {self.uuid} \n'
							for is_user, text in self.iter_texts():
										lowerCamelCase       =      """user""" if is_user else """bot"""
										output += f'{name} >> {text} \n'
							return output






@add_end_docstrings(
    UpperCAmelCase__		,      R"\n        min_length_for_response (`int`, *optional*, defaults to 32):\n            The minimum length (in number of tokens) for a response.\n        minimum_tokens (`int`, *optional*, defaults to 10):\n            The minimum length of tokens to leave for a response.\n    "		,      )
class 					__magic_name__							( UpperCAmelCase__       ):






				'''simple docstring'''





				def __init__( self	,				*_a	,				**_a				):





							"""simple docstring"""
							super().__init__(*_a	,				**_a				)
							if self.tokenizer.pad_token_id is None:
										lowerCamelCase       =      self.tokenizer.eos_token





				def      _lowerCAmelCase       ( self	,				_a=None	,				_a=None	,				_a=None	,				**_a				):





							"""simple docstring"""
							lowerCamelCase       =      {}
							lowerCamelCase       =      {}
							lowerCamelCase       =      {}

							if min_length_for_response is not None:
										lowerCamelCase       =      min_length_for_response
							if minimum_tokens is not None:
										lowerCamelCase       =      minimum_tokens

							if "max_length" in generate_kwargs:
										lowerCamelCase       =      generate_kwargs["""max_length"""]
										# self.max_length = generate_kwargs.get("max_length", self.model.config.max_length)
							if clean_up_tokenization_spaces is not None:
										lowerCamelCase       =      clean_up_tokenization_spaces

							if generate_kwargs:
										forward_params.update(_a				)
							return preprocess_params, forward_params, postprocess_params





				def __call__( self	,				_a	,				_a=0	,				**_a				):





							"""simple docstring"""
							lowerCamelCase       =      super().__call__(_a	,				num_workers=_a	,				**_a				)
							if isinstance(_a	,				_a				) and len(_a				) == 1:
										return outputs[0]
							return outputs





				def      _lowerCAmelCase       ( self	,				_a	,				_a=32				):





							"""simple docstring"""
							if not isinstance(_a	,				_a				):
										raise ValueError("""ConversationalPipeline, expects Conversation as inputs"""				)
							if conversation.new_user_input is None:
										raise ValueError(
										    f'Conversation with UUID {type(conversation.uuid				)} does not contain new user input to process. '
										    """Add user inputs with the conversation's `add_user_input` method"""				)
							if hasattr(self.tokenizer	,				"""_build_conversation_input_ids"""				):
										lowerCamelCase       =      self.tokenizer._build_conversation_input_ids(_a				)
							else:
										# If the tokenizer cannot handle conversations, we default to only the old version
										lowerCamelCase       =      self._legacy_parse_and_tokenize(_a				)

							if self.framework == "pt":
										lowerCamelCase       =      torch.LongTensor([input_ids]				)
							elif self.framework == "tf":
										lowerCamelCase       =      tf.constant([input_ids]				)
							return {"input_ids": input_ids, "conversation": conversation}





				def      _lowerCAmelCase       ( self	,				_a	,				_a=10	,				**_a				):





							"""simple docstring"""
							lowerCamelCase       =      generate_kwargs.get("""max_length"""	,				self.model.config.max_length				)

							lowerCamelCase       =      model_inputs["""input_ids"""].shape[1]
							if max_length - minimum_tokens < n:
										logger.warning(f'Conversation input is to long ({n}), trimming it to ({max_length} - {minimum_tokens})'				)
										lowerCamelCase       =      max_length - minimum_tokens
										lowerCamelCase       =      model_inputs["""input_ids"""][:, -trim:]
										if "attention_mask" in model_inputs:
													lowerCamelCase       =      model_inputs["""attention_mask"""][:, -trim:]
							lowerCamelCase       =      model_inputs.pop("""conversation"""				)
							lowerCamelCase       =      max_length
							lowerCamelCase       =      self.model.generate(**_a	,				**_a				)
							if self.model.config.is_encoder_decoder:
										lowerCamelCase       =      1
							else:
										lowerCamelCase       =      n
							return {"output_ids": output_ids[:, start_position:], "conversation": conversation}





				def      _lowerCAmelCase       ( self	,				_a	,				_a=True				):





							"""simple docstring"""
							lowerCamelCase       =      model_outputs["""output_ids"""]
							lowerCamelCase       =      self.tokenizer.decode(
							    output_ids[0]	,				skip_special_tokens=_a	,				clean_up_tokenization_spaces=_a	,				)
							lowerCamelCase       =      model_outputs["""conversation"""]
							conversation.mark_processed()
							conversation.append_response(_a				)
							return conversation





				def      _lowerCAmelCase       ( self	,				_a				):





							"""simple docstring"""
							lowerCamelCase       =      self.tokenizer.eos_token_id
							lowerCamelCase       =      []
							for is_user, text in conversation.iter_texts():
										if eos_token_id is not None:
													input_ids.extend(self.tokenizer.encode(_a	,				add_special_tokens=_a				) + [eos_token_id]				)
										else:
													input_ids.extend(self.tokenizer.encode(_a	,				add_special_tokens=_a				)				)

							if len(_a				) > self.tokenizer.model_max_length:
										lowerCamelCase       =      input_ids[-self.tokenizer.model_max_length :]
							return input_ids



