

'''simple docstring'''



import logging
import math
from functools import partial
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import torch

from .tensor_utils import tensor_tree_map, tree_map



def   __lowerCAmelCase					(							lowerCamelCase   :						Union[dict, list, tuple, torch.Tensor]					):






	'''simple docstring'''




	__lowerCAmelCase												=  []
	if isinstance(lowerCamelCase				,		lowerCamelCase					):
		for v in tree.values():
			shapes.extend(_fetch_dims(lowerCamelCase					)					)
	elif isinstance(lowerCamelCase				,		(list, tuple)					):
		for t in tree:
			shapes.extend(_fetch_dims(lowerCamelCase					)					)
	elif isinstance(lowerCamelCase				,		torch.Tensor					):
		shapes.append(tree.shape					)
	else:
		raise ValueError("Not supported"					)

	return shapes



@torch.jit.ignore
def   __lowerCAmelCase					(							lowerCamelCase   :						int				,		lowerCamelCase   :						Tuple[int, ...]					):






	'''simple docstring'''




	__lowerCAmelCase												=  []
	for d in reversed(lowerCamelCase					):
		idx.append(flat_idx % d					)
		__lowerCAmelCase												=  flat_idx // d

	return tuple(reversed(lowerCamelCase					)					)



@torch.jit.ignore
def   __lowerCAmelCase					(							lowerCamelCase   :						Sequence[int]				,		lowerCamelCase   :						Sequence[int]				,		lowerCamelCase   :						Sequence[int]				,		lowerCamelCase   :						Optional[Sequence[bool]] = None				,		lowerCamelCase   :						Optional[Sequence[bool]] = None				,		):






	'''simple docstring'''




	def reduce_edge_list(lowerCamelCase   :						List[bool]					) -> None:
		__lowerCAmelCase												=  True
		for i in range(len(lowerCamelCase					)					):
			__lowerCAmelCase												=  -1 * (i + 1)
			l[reversed_idx] &= tally
			__lowerCAmelCase												=  l[reversed_idx]

	if start_edges is None:
		__lowerCAmelCase												=  [s == 0 for s in start]
		reduce_edge_list(lowerCamelCase					)
	if end_edges is None:
		__lowerCAmelCase												=  [e == (d - 1) for e, d in zip(lowerCamelCase				,		lowerCamelCase					)]
		reduce_edge_list(lowerCamelCase					)

	# Base cases. Either start/end are empty and we're done, or the final,
	# one-dimensional tensor can be simply sliced
	if len(lowerCamelCase					) == 0:
		return [()]
	elif len(lowerCamelCase					) == 1:
		return [(slice(start[0]				,		end[0] + 1					),)]

	__lowerCAmelCase												=  []
	__lowerCAmelCase												=  []

	# Dimensions common to start and end can be selected directly
	for s, e in zip(lowerCamelCase				,		lowerCamelCase					):
		if s == e:
			path_list.append(slice(lowerCamelCase				,		s + 1					)					)
		else:
			break

	__lowerCAmelCase												=  tuple(lowerCamelCase					)
	__lowerCAmelCase												=  len(lowerCamelCase					)

	# start == end, and we're done
	if divergence_idx == len(lowerCamelCase					):
		return [path]

	def upper() -> Tuple[Tuple[slice, ...], ...]:
		assert start_edges is not None
		assert end_edges is not None

		__lowerCAmelCase												=  start[divergence_idx]
		return tuple(
		    path + (slice(lowerCamelCase				,		sdi + 1					),) + s
		    for s in _get_minimal_slice_set(
		        start[divergence_idx + 1 :]				,		[d - 1 for d in dims[divergence_idx + 1 :]]				,		dims[divergence_idx + 1 :]				,		start_edges=start_edges[divergence_idx + 1 :]				,		end_edges=[True for _ in end_edges[divergence_idx + 1 :]]				,		)					)

	def lower() -> Tuple[Tuple[slice, ...], ...]:
		assert start_edges is not None
		assert end_edges is not None

		__lowerCAmelCase												=  end[divergence_idx]
		return tuple(
		    path + (slice(lowerCamelCase				,		edi + 1					),) + s
		    for s in _get_minimal_slice_set(
		        [0 for _ in start[divergence_idx + 1 :]]				,		end[divergence_idx + 1 :]				,		dims[divergence_idx + 1 :]				,		start_edges=[True for _ in start_edges[divergence_idx + 1 :]]				,		end_edges=end_edges[divergence_idx + 1 :]				,		)					)

	# If both start and end are at the edges of the subtree rooted at
	# divergence_idx, we can just select the whole subtree at once
	if start_edges[divergence_idx] and end_edges[divergence_idx]:
		slices.append(path + (slice(start[divergence_idx]				,		end[divergence_idx] + 1					),)					)
	# If just start is at the edge, we can grab almost all of the subtree,
	# treating only the ragged bottom edge as an edge case
	elif start_edges[divergence_idx]:
		slices.append(path + (slice(start[divergence_idx]				,		end[divergence_idx]					),)					)
		slices.extend(lower()					)
	# Analogous to the previous case, but the top is ragged this time
	elif end_edges[divergence_idx]:
		slices.extend(upper()					)
		slices.append(path + (slice(start[divergence_idx] + 1				,		end[divergence_idx] + 1					),)					)
	# If both sides of the range are ragged, we need to handle both sides
	# separately. If there's contiguous meat in between them, we can index it
	# in one big chunk
	else:
		slices.extend(upper()					)
		__lowerCAmelCase												=  end[divergence_idx] - start[divergence_idx]
		if middle_ground > 1:
			slices.append(path + (slice(start[divergence_idx] + 1				,		end[divergence_idx]					),)					)
		slices.extend(lower()					)

	return slices



@torch.jit.ignore
def   __lowerCAmelCase					(							lowerCamelCase   :						torch.Tensor				,		lowerCamelCase   :						int				,		lowerCamelCase   :						int				,		lowerCamelCase   :						int					):






	'''simple docstring'''




	__lowerCAmelCase												=  t.shape[:no_batch_dims]
	__lowerCAmelCase												=  list(_flat_idx_to_idx(lowerCamelCase				,		lowerCamelCase					)					)
	# _get_minimal_slice_set is inclusive
	__lowerCAmelCase												=  list(_flat_idx_to_idx(flat_end - 1				,		lowerCamelCase					)					)

	# Get an ordered list of slices to perform
	__lowerCAmelCase												=  _get_minimal_slice_set(
	    lowerCamelCase				,		lowerCamelCase				,		lowerCamelCase				,		)

	__lowerCAmelCase												=  [t[s] for s in slices]

	return torch.cat([s.view((-1,) + t.shape[no_batch_dims:]					) for s in sliced_tensors]					)



def   __lowerCAmelCase					(							lowerCamelCase   :						Callable				,		lowerCamelCase   :						Dict[str, Any]				,		lowerCamelCase   :						int				,		lowerCamelCase   :						int				,		lowerCamelCase   :						bool = False				,		lowerCamelCase   :						Any = None				,		lowerCamelCase   :						bool = False				,		):






	'''simple docstring'''




	if not (len(lowerCamelCase					) > 0):
		raise ValueError("Must provide at least one input"					)

	__lowerCAmelCase												=  [shape[:no_batch_dims] for shape in _fetch_dims(lowerCamelCase					)]
	__lowerCAmelCase												=  tuple([max(lowerCamelCase					) for s in zip(*lowerCamelCase					)]					)

	def _prep_inputs(lowerCamelCase   :						torch.Tensor					) -> torch.Tensor:
		if not low_mem:
			if not sum(t.shape[:no_batch_dims]					) == no_batch_dims:
				__lowerCAmelCase												=  t.expand(orig_batch_dims + t.shape[no_batch_dims:]					)
			__lowerCAmelCase												=  t.reshape(-1				,		*t.shape[no_batch_dims:]					)
		else:
			__lowerCAmelCase												=  t.expand(orig_batch_dims + t.shape[no_batch_dims:]					)
		return t

	__lowerCAmelCase												=  tensor_tree_map(_prep_inputs				,		lowerCamelCase					)
	__lowerCAmelCase												=  None
	if _out is not None:
		__lowerCAmelCase												=  tensor_tree_map(lambda lowerCamelCase					: t.view([-1] + list(t.shape[no_batch_dims:]					)					)				,		_out					)

	__lowerCAmelCase												=  1
	for d in orig_batch_dims:
		flat_batch_dim *= d

	__lowerCAmelCase												=  flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

	def _select_chunk(lowerCamelCase   :						torch.Tensor					) -> torch.Tensor:
		return t[i : i + chunk_size] if t.shape[0] != 1 else t

	__lowerCAmelCase												=  0
	__lowerCAmelCase												=  prepped_outputs
	for _ in range(lowerCamelCase					):
		# Chunk the input
		if not low_mem:
			__lowerCAmelCase												=  _select_chunk
		else:
			__lowerCAmelCase												=  partial(
			    _chunk_slice				,		flat_start=lowerCamelCase				,		flat_end=min(lowerCamelCase				,		i + chunk_size					)				,		no_batch_dims=len(lowerCamelCase					)				,		)

		__lowerCAmelCase												=  tensor_tree_map(lowerCamelCase				,		lowerCamelCase					)

		# Run the layer on the chunk
		__lowerCAmelCase												=  layer(**lowerCamelCase					)

		# Allocate space for the output
		if out is None:
			__lowerCAmelCase												=  tensor_tree_map(lambda lowerCamelCase					: t.new_zeros((flat_batch_dim,) + t.shape[1:]					)				,		lowerCamelCase					)

		# Put the chunk in its pre-allocated space
		if isinstance(lowerCamelCase				,		lowerCamelCase					):

			def assign(lowerCamelCase   :						dict				,		lowerCamelCase   :						dict					) -> None:
				for k, v in da.items():
					if isinstance(lowerCamelCase				,		lowerCamelCase					):
						assign(lowerCamelCase				,		da[k]					)
					else:
						if _add_into_out:
							v[i : i + chunk_size] += da[k]
						else:
							__lowerCAmelCase												=  da[k]

			assign(lowerCamelCase				,		lowerCamelCase					)
		elif isinstance(lowerCamelCase				,		lowerCamelCase					):
			for xa, xa in zip(lowerCamelCase				,		lowerCamelCase					):
				if _add_into_out:
					xa[i : i + chunk_size] += xa
				else:
					__lowerCAmelCase												=  xa
		elif isinstance(lowerCamelCase				,		torch.Tensor					):
			if _add_into_out:
				out[i : i + chunk_size] += output_chunk
			else:
				__lowerCAmelCase												=  output_chunk
		else:
			raise ValueError("Not supported"					)

		i += chunk_size

	__lowerCAmelCase												=  tensor_tree_map(lambda lowerCamelCase					: t.view(orig_batch_dims + t.shape[1:]					)				,		lowerCamelCase					)

	return out






class 					UpperCAmelCase__	:



					def __init__(    self			,				UpperCamelCase = 512			,				)     ->			Any:
						__lowerCAmelCase												=  max_chunk_size
						__lowerCAmelCase												=  None
						__lowerCAmelCase												=  None



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			int:
						logging.info("Tuning chunk size..."     )

						if min_chunk_size >= self.max_chunk_size:
							return min_chunk_size

						__lowerCAmelCase												=  [2**l for l in range(int(math.log(self.max_chunk_size			,				2     )     ) + 1     )]
						__lowerCAmelCase												=  [c for c in candidates if c > min_chunk_size]
						__lowerCAmelCase												=  [min_chunk_size] + candidates
						candidates[-1] += 4

						def test_chunk_size(UpperCamelCase     ) -> bool:
							try:
								with torch.no_grad():
									fn(*UpperCamelCase			,				chunk_size=UpperCamelCase     )
								return True
							except RuntimeError:
								return False

						__lowerCAmelCase												=  0
						__lowerCAmelCase												=  len(UpperCamelCase     ) - 1
						while i > min_viable_chunk_size_index:
							__lowerCAmelCase												=  test_chunk_size(candidates[i]     )
							if not viable:
								__lowerCAmelCase												=  (min_viable_chunk_size_index + i) // 2
							else:
								__lowerCAmelCase												=  i
								__lowerCAmelCase												=  (i + len(UpperCamelCase     ) - 1) // 2

						return candidates[min_viable_chunk_size_index]



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase     )     ->			bool:
						__lowerCAmelCase												=  True
						for aa, aa in zip(UpperCamelCase			,				UpperCamelCase     ):
							assert type(UpperCamelCase     ) == type(UpperCamelCase     )
							if isinstance(UpperCamelCase			,				(list, tuple)     ):
								consistent &= self._compare_arg_caches(UpperCamelCase			,				UpperCamelCase     )
							elif isinstance(UpperCamelCase			,				UpperCamelCase     ):
								__lowerCAmelCase												=  [v for _, v in sorted(aa.items()			,				key=lambda UpperCamelCase     : x[0]     )]
								__lowerCAmelCase												=  [v for _, v in sorted(aa.items()			,				key=lambda UpperCamelCase     : x[0]     )]
								consistent &= self._compare_arg_caches(UpperCamelCase			,				UpperCamelCase     )
							else:
								consistent &= aa == aa

						return consistent



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				)     ->			int:
						__lowerCAmelCase												=  True
						__lowerCAmelCase												=  tree_map(lambda UpperCamelCase     : a.shape if isinstance(UpperCamelCase			,				torch.Tensor     ) else a			,				UpperCamelCase			,				UpperCamelCase     )
						if self.cached_arg_data is not None:
							# If args have changed shape/value, we need to re-tune
							assert len(self.cached_arg_data     ) == len(UpperCamelCase     )
							__lowerCAmelCase												=  self._compare_arg_caches(self.cached_arg_data			,				UpperCamelCase     )
						else:
							# Otherwise, we can reuse the precomputed value
							__lowerCAmelCase												=  False

						if not consistent:
							__lowerCAmelCase												=  self._determine_favorable_chunk_size(
							    UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				)
							__lowerCAmelCase												=  arg_data

						assert self.cached_chunk_size is not None

						return self.cached_chunk_size




'''simple docstring'''



from argparse import ArgumentParser

from . import BaseTransformersCLICommand



def   __lowerCAmelCase					(							lowerCamelCase   :						List[str]					):






	'''simple docstring'''




	return DownloadCommand(args.model				,		args.cache_dir				,		args.force				,		args.trust_remote_code					)






class 					UpperCAmelCase__	(		UpperCamelCase__ ):



					@staticmethod
					def        UpperCAmelCase_ (    UpperCamelCase     )     ->			Tuple:
						__lowerCAmelCase												=  parser.add_parser("download"     )
						download_parser.add_argument(
						    "--cache-dir"			,				type=UpperCamelCase			,				default=UpperCamelCase			,				help="Path to location to store the models"     )
						download_parser.add_argument(
						    "--force"			,				action="store_true"			,				help="Force the model to be download even if already in cache-dir"     )
						download_parser.add_argument(
						    "--trust-remote-code"			,				action="store_true"			,				help="Whether or not to allow for custom models defined on the Hub in their own modeling files. Use only if you've reviewed the code as it will execute on your local machine"			,				)
						download_parser.add_argument("model"			,				type=UpperCamelCase			,				help="Name of the model to download"     )
						download_parser.set_defaults(func=UpperCamelCase     )



					def __init__(    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			List[str]:
						__lowerCAmelCase												=  model
						__lowerCAmelCase												=  cache
						__lowerCAmelCase												=  force
						__lowerCAmelCase												=  trust_remote_code



					def        UpperCAmelCase_ (    self     )     ->			Any:
						from ..models.auto import AutoModel, AutoTokenizer

						AutoModel.from_pretrained(
						    self._model			,				cache_dir=self._cache			,				force_download=self._force			,				trust_remote_code=self._trust_remote_code     )
						AutoTokenizer.from_pretrained(
						    self._model			,				cache_dir=self._cache			,				force_download=self._force			,				trust_remote_code=self._trust_remote_code     )