


"""simple docstring"""






import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
	import jax
	import jax.numpy as jnp

	from transformers.generation import (
	    FlaxForcedBOSTokenLogitsProcessor,
	    FlaxForcedEOSTokenLogitsProcessor,
	    FlaxLogitsProcessorList,
	    FlaxMinLengthLogitsProcessor,
	    FlaxTemperatureLogitsWarper,
	    FlaxTopKLogitsWarper,
	    FlaxTopPLogitsWarper,
	)


@require_flax
class 			lowerCamelCase_		(			unittest.TestCase   ):





				"""simple docstring"""







				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__     ):
								SCREAMING_SNAKE_CASE__     							=				jnp.ones((batch_size, length)     ) / length
								return scores



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				None
								SCREAMING_SNAKE_CASE__     							=				20

								SCREAMING_SNAKE_CASE__     							=				self._get_uniform_logits(batch_size=2       ,		length=UpperCAmelCase__     )

								# tweak scores to not be uniform anymore
								SCREAMING_SNAKE_CASE__     							=				scores.at[1, 5].set((1 / length) + 0.1     )  # peak, 1st batch
								SCREAMING_SNAKE_CASE__     							=				scores.at[1, 10].set((1 / length) - 0.4     )  # valley, 1st batch

								# compute softmax
								SCREAMING_SNAKE_CASE__     							=				jax.nn.softmax(UpperCAmelCase__       ,		axis=-1     )

								SCREAMING_SNAKE_CASE__     							=				FlaxTemperatureLogitsWarper(temperature=0.5     )
								SCREAMING_SNAKE_CASE__     							=				FlaxTemperatureLogitsWarper(temperature=1.3     )

								SCREAMING_SNAKE_CASE__     							=				jax.nn.softmax(temp_dist_warper_sharper(UpperCAmelCase__       ,		scores.copy()       ,		cur_len=UpperCAmelCase__     )       ,		axis=-1     )
								SCREAMING_SNAKE_CASE__     							=				jax.nn.softmax(temp_dist_warper_smoother(UpperCAmelCase__       ,		scores.copy()       ,		cur_len=UpperCAmelCase__     )       ,		axis=-1     )

								# uniform distribution stays uniform
								self.assertTrue(jnp.allclose(probs[0, :]       ,		warped_prob_sharp[0, :]       ,		atol=1e-3     )     )
								self.assertTrue(jnp.allclose(probs[0, :]       ,		warped_prob_smooth[0, :]       ,		atol=1e-3     )     )

								# sharp peaks get higher, valleys get lower
								self.assertLess(probs[1, :].max()       ,		warped_prob_sharp[1, :].max()     )
								self.assertGreater(probs[1, :].min()       ,		warped_prob_sharp[1, :].min()     )

								# smooth peaks get lower, valleys get higher
								self.assertGreater(probs[1, :].max()       ,		warped_prob_smooth[1, :].max()     )
								self.assertLess(probs[1, :].min()       ,		warped_prob_smooth[1, :].min()     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				None
								SCREAMING_SNAKE_CASE__     							=				10
								SCREAMING_SNAKE_CASE__     							=				2

								# create ramp distribution
								SCREAMING_SNAKE_CASE__     							=				np.broadcast_to(np.arange(UpperCAmelCase__     )[None, :]       ,		(batch_size, vocab_size)     ).copy()
								SCREAMING_SNAKE_CASE__     							=				ramp_logits[1:, : vocab_size // 2] + vocab_size

								SCREAMING_SNAKE_CASE__     							=				FlaxTopKLogitsWarper(3     )

								SCREAMING_SNAKE_CASE__     							=				top_k_warp(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )

								# check that correct tokens are filtered
								self.assertListEqual(jnp.isinf(scores[0]     ).tolist()       ,		7 * [True] + 3 * [False]     )
								self.assertListEqual(jnp.isinf(scores[1]     ).tolist()       ,		2 * [True] + 3 * [False] + 5 * [True]     )

								# check special case
								SCREAMING_SNAKE_CASE__     							=				5
								SCREAMING_SNAKE_CASE__     							=				FlaxTopKLogitsWarper(top_k=1       ,		filter_value=0.0       ,		min_tokens_to_keep=3     )

								SCREAMING_SNAKE_CASE__     							=				np.broadcast_to(np.arange(UpperCAmelCase__     )[None, :]       ,		(batch_size, length)     ).copy()
								SCREAMING_SNAKE_CASE__     							=				top_k_warp_safety_check(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )

								# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
								self.assertListEqual((scores == 0.0).sum(axis=-1     ).tolist()       ,		[2, 2]     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				None
								SCREAMING_SNAKE_CASE__     							=				10
								SCREAMING_SNAKE_CASE__     							=				2

								# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
								SCREAMING_SNAKE_CASE__     							=				np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]     )     )

								SCREAMING_SNAKE_CASE__     							=				FlaxTopPLogitsWarper(0.8     )
								SCREAMING_SNAKE_CASE__     							=				np.exp(top_p_warp(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )     )

								# dist should be filtered to keep min num values so that sum is >= top_p
								# exp (-inf) => 0
								SCREAMING_SNAKE_CASE__     							=				np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]     )
								self.assertTrue(np.allclose(UpperCAmelCase__       ,		UpperCAmelCase__       ,		atol=1e-3     )     )

								# check edge cases with negative and extreme logits
								SCREAMING_SNAKE_CASE__     							=				np.broadcast_to(np.arange(UpperCAmelCase__     )[None, :]       ,		(batch_size, vocab_size)     ).copy() - (
								    vocab_size // 2
								)

								# make ramp_logits more extreme
								SCREAMING_SNAKE_CASE__     							=				ramp_logits[1] * 100.0

								# make sure at least 2 tokens are kept
								SCREAMING_SNAKE_CASE__     							=				FlaxTopPLogitsWarper(0.9       ,		min_tokens_to_keep=2       ,		filter_value=0.0     )
								SCREAMING_SNAKE_CASE__     							=				top_p_warp(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )

								# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
								self.assertListEqual((filtered_dist != 0.0).sum(axis=-1     ).tolist()       ,		[3, 2]     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				20
								SCREAMING_SNAKE_CASE__     							=				4
								SCREAMING_SNAKE_CASE__     							=				0

								SCREAMING_SNAKE_CASE__     							=				FlaxMinLengthLogitsProcessor(min_length=10       ,		eos_token_id=UpperCAmelCase__     )

								# check that min length is applied at length 5
								SCREAMING_SNAKE_CASE__     							=				ids_tensor((batch_size, 20)       ,		vocab_size=20     )
								SCREAMING_SNAKE_CASE__     							=				5
								SCREAMING_SNAKE_CASE__     							=				self._get_uniform_logits(UpperCAmelCase__       ,		UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				min_dist_processor(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
								self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()       ,		4 * [-float("inf"     )]     )

								# check that min length is not applied anymore at length 15
								SCREAMING_SNAKE_CASE__     							=				self._get_uniform_logits(UpperCAmelCase__       ,		UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				15
								SCREAMING_SNAKE_CASE__     							=				min_dist_processor(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
								self.assertFalse(jnp.isinf(UpperCAmelCase__     ).any()     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				20
								SCREAMING_SNAKE_CASE__     							=				4
								SCREAMING_SNAKE_CASE__     							=				0

								SCREAMING_SNAKE_CASE__     							=				FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase__     )

								# check that all scores are -inf except the bos_token_id score
								SCREAMING_SNAKE_CASE__     							=				ids_tensor((batch_size, 1)       ,		vocab_size=20     )
								SCREAMING_SNAKE_CASE__     							=				1
								SCREAMING_SNAKE_CASE__     							=				self._get_uniform_logits(UpperCAmelCase__       ,		UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				logits_processor(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
								self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]     ).all()     )
								self.assertListEqual(scores[:, bos_token_id].tolist()       ,		4 * [0]     )  # score for bos_token_id shold be zero

								# check that bos_token_id is not forced if current length is greater than 1
								SCREAMING_SNAKE_CASE__     							=				3
								SCREAMING_SNAKE_CASE__     							=				self._get_uniform_logits(UpperCAmelCase__       ,		UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				logits_processor(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
								self.assertFalse(jnp.isinf(UpperCAmelCase__     ).any()     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				20
								SCREAMING_SNAKE_CASE__     							=				4
								SCREAMING_SNAKE_CASE__     							=				0
								SCREAMING_SNAKE_CASE__     							=				5

								SCREAMING_SNAKE_CASE__     							=				FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase__       ,		eos_token_id=UpperCAmelCase__     )

								# check that all scores are -inf except the eos_token_id when max_length is reached
								SCREAMING_SNAKE_CASE__     							=				ids_tensor((batch_size, 4)       ,		vocab_size=20     )
								SCREAMING_SNAKE_CASE__     							=				4
								SCREAMING_SNAKE_CASE__     							=				self._get_uniform_logits(UpperCAmelCase__       ,		UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				logits_processor(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
								self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]     ).all()     )
								self.assertListEqual(scores[:, eos_token_id].tolist()       ,		4 * [0]     )  # score for eos_token_id should be zero

								# check that eos_token_id is not forced if max_length is not reached
								SCREAMING_SNAKE_CASE__     							=				3
								SCREAMING_SNAKE_CASE__     							=				self._get_uniform_logits(UpperCAmelCase__       ,		UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				logits_processor(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
								self.assertFalse(jnp.isinf(UpperCAmelCase__     ).any()     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				4
								SCREAMING_SNAKE_CASE__     							=				10
								SCREAMING_SNAKE_CASE__     							=				15
								SCREAMING_SNAKE_CASE__     							=				2
								SCREAMING_SNAKE_CASE__     							=				1
								SCREAMING_SNAKE_CASE__     							=				15

								# dummy input_ids and scores
								SCREAMING_SNAKE_CASE__     							=				ids_tensor((batch_size, sequence_length)       ,		UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				input_ids.copy()

								SCREAMING_SNAKE_CASE__     							=				self._get_uniform_logits(UpperCAmelCase__       ,		UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				scores.copy()

								# instantiate all dist processors
								SCREAMING_SNAKE_CASE__     							=				FlaxTemperatureLogitsWarper(temperature=0.5     )
								SCREAMING_SNAKE_CASE__     							=				FlaxTopKLogitsWarper(3     )
								SCREAMING_SNAKE_CASE__     							=				FlaxTopPLogitsWarper(0.8     )

								# instantiate all logits processors
								SCREAMING_SNAKE_CASE__     							=				FlaxMinLengthLogitsProcessor(min_length=10       ,		eos_token_id=UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase__       ,		eos_token_id=UpperCAmelCase__     )

								SCREAMING_SNAKE_CASE__     							=				10

								# no processor list
								SCREAMING_SNAKE_CASE__     							=				temp_dist_warp(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				top_k_warp(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				top_p_warp(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				min_dist_proc(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				bos_dist_proc(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				eos_dist_proc(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )

								# with processor list
								SCREAMING_SNAKE_CASE__     							=				FlaxLogitsProcessorList(
								    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]     )
								SCREAMING_SNAKE_CASE__     							=				processor(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )

								# scores should be equal
								self.assertTrue(jnp.allclose(UpperCAmelCase__       ,		UpperCAmelCase__       ,		atol=1e-3     )     )

								# input_ids should never be changed
								self.assertListEqual(input_ids.tolist()       ,		input_ids_comp.tolist()     )






				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				4
								SCREAMING_SNAKE_CASE__     							=				10
								SCREAMING_SNAKE_CASE__     							=				15
								SCREAMING_SNAKE_CASE__     							=				2
								SCREAMING_SNAKE_CASE__     							=				1
								SCREAMING_SNAKE_CASE__     							=				15

								# dummy input_ids and scores
								SCREAMING_SNAKE_CASE__     							=				ids_tensor((batch_size, sequence_length)       ,		UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				input_ids.copy()

								SCREAMING_SNAKE_CASE__     							=				self._get_uniform_logits(UpperCAmelCase__       ,		UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				scores.copy()

								# instantiate all dist processors
								SCREAMING_SNAKE_CASE__     							=				FlaxTemperatureLogitsWarper(temperature=0.5     )
								SCREAMING_SNAKE_CASE__     							=				FlaxTopKLogitsWarper(3     )
								SCREAMING_SNAKE_CASE__     							=				FlaxTopPLogitsWarper(0.8     )

								# instantiate all logits processors
								SCREAMING_SNAKE_CASE__     							=				FlaxMinLengthLogitsProcessor(min_length=10       ,		eos_token_id=UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase__       ,		eos_token_id=UpperCAmelCase__     )

								SCREAMING_SNAKE_CASE__     							=				10

								# no processor list
								def run_no_processor_list(UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__     ):
												SCREAMING_SNAKE_CASE__     							=				temp_dist_warp(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
												SCREAMING_SNAKE_CASE__     							=				top_k_warp(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
												SCREAMING_SNAKE_CASE__     							=				top_p_warp(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
												SCREAMING_SNAKE_CASE__     							=				min_dist_proc(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
												SCREAMING_SNAKE_CASE__     							=				bos_dist_proc(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
												SCREAMING_SNAKE_CASE__     							=				eos_dist_proc(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
												return scores

								# with processor list
								def run_processor_list(UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__     ):
												SCREAMING_SNAKE_CASE__     							=				FlaxLogitsProcessorList(
												    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]     )
												SCREAMING_SNAKE_CASE__     							=				processor(UpperCAmelCase__       ,		UpperCAmelCase__       ,		cur_len=UpperCAmelCase__     )
												return scores

								SCREAMING_SNAKE_CASE__     							=				jax.jit(UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				jax.jit(UpperCAmelCase__     )

								SCREAMING_SNAKE_CASE__     							=				jitted_run_no_processor_list(UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				jitted_run_processor_list(UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__     )

								# scores should be equal
								self.assertTrue(jnp.allclose(UpperCAmelCase__       ,		UpperCAmelCase__       ,		atol=1e-3     )     )

								# input_ids should never be changed
								self.assertListEqual(input_ids.tolist()       ,		input_ids_comp.tolist()     )




"""simple docstring"""






import heapq


def __lowercase     (      lowerCamelCase_  :      dict  ):
				SCREAMING_SNAKE_CASE__     							=				[]

				# for each node and his adjacency list add them and the rank of the node to queue
				# using heapq module the queue will be filled like a Priority Queue
				# heapq works with a min priority queue, so I used -1*len(v) to build it
				for key, value in graph.items():
								# O(log(n))
								heapq.heappush(lowerCamelCase_						,					[-1 * len(lowerCamelCase_  ), (key, value)]  )

				# chosen_vertices = set of chosen vertices
				SCREAMING_SNAKE_CASE__     							=				set()

				# while queue isn't empty and there are still edges
				#   (queue[0][0] is the rank of the node with max rank)
				while queue and queue[0][0] != 0:
								# extract vertex with max rank from queue and add it to chosen_vertices
								SCREAMING_SNAKE_CASE__     							=				heapq.heappop(lowerCamelCase_  )[1][0]
								chosen_vertices.add(lowerCamelCase_  )

								# Remove all arcs adjacent to argmax
								for elem in queue:
												# if v haven't adjacent node, skip
												if elem[0] == 0:
																continue
												# if argmax is reachable from elem
												# remove argmax from elem's adjacent list and update his rank
												if argmax in elem[1][1]:
																SCREAMING_SNAKE_CASE__     							=				elem[1][1].index(lowerCamelCase_  )
																del elem[1][1][index]
																elem[0] += 1
        # re-order the queue
								heapq.heapify(lowerCamelCase_  )
				return chosen_vertices


if __name__ == "__main__":
	import doctest

	doctest.testmod()

	_lowerCamelCase     =  {0: [1, 3], 1: [0, 3], 2: [0, 3, 4], 3: [0, 1, 2], 4: [2, 3]}
	print(f"""Minimum vertex cover:\n{greedy_min_vertex_cover(graph)}""")
