





from typing import Dict

import numpy as np

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


if is_tf_available():
      import tensorflow as tf

      from ..tf_utils import stable_softmax


if is_torch_available():
      import torch


UpperCAmelCase_					:   str 							=						logging.get_logger(__name__)





@add_end_docstrings(
    __A     ,					R"""
        top_k (`int`, defaults to 5):
            The number of predictions to return.
        targets (`str` or `List[str]`, *optional*):
            When passed, the model will limit the scores to the passed targets instead of looking up in the whole
            vocab. If the provided targets are not in the model vocab, they will be tokenized and the first resulting
            token will be used (with a warning, and that might be slower).

    """     ,					)
class      lowercase__    (							__A  ):







  def 	UpperCAmelCase__		(	self    ,						_lowercase						):
       if self.framework == "tf":
            lowerCAmelCase_					:       int    =		tf.where(input_ids == self.tokenizer.mask_token_id						).numpy()
       elif self.framework == "pt":
            lowerCAmelCase_					:       List[Any]    =		torch.nonzero(input_ids == self.tokenizer.mask_token_id    ,						as_tuple=_lowercase						)
       else:
            raise ValueError("""Unsupported framework"""						)
       return masked_index







  def 	UpperCAmelCase__		(	self    ,						_lowercase						):
       lowerCAmelCase_					:       Dict    =		self.get_masked_index(_lowercase						)
       lowerCAmelCase_					:       Optional[Any]    =		np.prod(masked_index.shape						)
       if numel < 1:
            raise PipelineException(
                """fill-mask"""    ,						self.model.base_model_prefix    ,						F'No mask_token ({self.tokenizer.mask_token}) found on the input'    ,						)







  def 	UpperCAmelCase__		(	self    ,						_lowercase						):
       if isinstance(_lowercase    ,						_lowercase						):
            for model_input in model_inputs:
                 self._ensure_exactly_one_mask_token(model_input["""input_ids"""][0]						)
       else:
            for input_ids in model_inputs["input_ids"]:
                 self._ensure_exactly_one_mask_token(_lowercase						)







  def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase=None    ,						**_lowercase						):
       if return_tensors is None:
            lowerCAmelCase_					:       List[Any]    =		self.framework
       lowerCAmelCase_					:       Any    =		self.tokenizer(_lowercase    ,						return_tensors=_lowercase						)
       self.ensure_exactly_one_mask_token(_lowercase						)
       return model_inputs







  def 	UpperCAmelCase__		(	self    ,						_lowercase						):
       lowerCAmelCase_					:       str    =		self.model(**_lowercase						)
       lowerCAmelCase_					:       Optional[Any]    =		model_inputs["""input_ids"""]
       return model_outputs







  def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase=5    ,						_lowercase=None						):
       # Cap top_k if there are targets
       if target_ids is not None and target_ids.shape[0] < top_k:
            lowerCAmelCase_					:       List[Any]    =		target_ids.shape[0]
       lowerCAmelCase_					:       Tuple    =		model_outputs["""input_ids"""][0]
       lowerCAmelCase_					:       Optional[int]    =		model_outputs["""logits"""]

       if self.framework == "tf":
            lowerCAmelCase_					:       str    =		tf.where(input_ids == self.tokenizer.mask_token_id						).numpy()[:, 0]

            lowerCAmelCase_					:       int    =		outputs.numpy()

            lowerCAmelCase_					:       int    =		outputs[0, masked_index, :]
            lowerCAmelCase_					:       Dict    =		stable_softmax(_lowercase    ,						axis=-1						)
            if target_ids is not None:
                 lowerCAmelCase_					:       Union[str, Any]    =		tf.gather_nd(tf.squeeze(_lowercase    ,						0						)    ,						target_ids.reshape(-1    ,						1						)						)
                 lowerCAmelCase_					:       Optional[Any]    =		tf.expand_dims(_lowercase    ,						0						)

            lowerCAmelCase_					:       Optional[Any]    =		tf.math.top_k(_lowercase    ,						k=_lowercase						)
            lowerCAmelCase_   ,  lowerCAmelCase_					:       Optional[Any]    =		topk.values.numpy(), topk.indices.numpy()
       else:
            lowerCAmelCase_					:       Union[str, Any]    =		torch.nonzero(input_ids == self.tokenizer.mask_token_id    ,						as_tuple=_lowercase						).squeeze(-1						)
            # Fill mask pipeline supports only one ${mask_token} per sample

            lowerCAmelCase_					:       List[Any]    =		outputs[0, masked_index, :]
            lowerCAmelCase_					:       List[str]    =		logits.softmax(dim=-1						)
            if target_ids is not None:
                 lowerCAmelCase_					:       Dict    =		probs[..., target_ids]

            lowerCAmelCase_   ,  lowerCAmelCase_					:       Dict    =		probs.topk(_lowercase						)

       lowerCAmelCase_					:       Optional[int]    =		[]
       lowerCAmelCase_					:       List[str]    =		values.shape[0] == 1
       for i, (_values, _predictions) in enumerate(zip(values.tolist()    ,						predictions.tolist()						)						):
            lowerCAmelCase_					:       List[str]    =		[]
            for v, p in zip(_values    ,						_predictions						):
                 # Copy is important since we're going to modify this array in place
                 lowerCAmelCase_					:       Optional[Any]    =		input_ids.numpy().copy()
                 if target_ids is not None:
                      lowerCAmelCase_					:       str    =		target_ids[p].tolist()

                 lowerCAmelCase_					:       Dict    =		p
                 # Filter padding out:
                 lowerCAmelCase_					:       Union[str, Any]    =		tokens[np.where(tokens != self.tokenizer.pad_token_id						)]
                 # Originally we skip special tokens to give readable output.
                 # For multi masks though, the other [MASK] would be removed otherwise
                 # making the output look odd, so we add them back
                 lowerCAmelCase_					:       Dict    =		self.tokenizer.decode(_lowercase    ,						skip_special_tokens=_lowercase						)
                 lowerCAmelCase_					:       Optional[Any]    =		{"""score""": v, """token""": p, """token_str""": self.tokenizer.decode([p]						), """sequence""": sequence}
                 row.append(_lowercase						)
            result.append(_lowercase						)
       if single_mask:
            return result[0]
       return result







  def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase=None						):
       if isinstance(_lowercase    ,						_lowercase						):
            lowerCAmelCase_					:       Tuple    =		[targets]
       try:
            lowerCAmelCase_					:       Tuple    =		self.tokenizer.get_vocab()
       except Exception:
            lowerCAmelCase_					:       Optional[Any]    =		{}
       lowerCAmelCase_					:       List[Any]    =		[]
       for target in targets:
            lowerCAmelCase_					:       Dict    =		vocab.get(_lowercase    ,						_lowercase						)
            if id_ is None:
                 lowerCAmelCase_					:       Union[str, Any]    =		self.tokenizer(
                     _lowercase    ,						add_special_tokens=_lowercase    ,						return_attention_mask=_lowercase    ,						return_token_type_ids=_lowercase    ,						max_length=1    ,						truncation=_lowercase    ,						)["""input_ids"""]
                 if len(_lowercase						) == 0:
                      logger.warning(
                          F'The specified target token `{target}` does not exist in the model vocabulary. '
                          """We cannot replace it with anything meaningful, ignoring it"""						)
                      continue
                 lowerCAmelCase_					:       Optional[int]    =		input_ids[0]
                 # XXX: If users encounter this pass
                 # it becomes pretty slow, so let's make sure
                 # The warning enables them to fix the input to
                 # get faster performance.
                 logger.warning(
                     F'The specified target token `{target}` does not exist in the model vocabulary. '
                     F'Replacing with `{self.tokenizer.convert_ids_to_tokens(id_						)}`.'						)
            target_ids.append(id_						)
       lowerCAmelCase_					:       List[str]    =		list(set(_lowercase						)						)
       if len(_lowercase						) == 0:
            raise ValueError("""At least one target must be provided when passed."""						)
       lowerCAmelCase_					:       str    =		np.array(_lowercase						)
       return target_ids







  def 	UpperCAmelCase__		(	self    ,						_lowercase=None    ,						_lowercase=None						):
       lowerCAmelCase_					:       Tuple    =		{}

       if targets is not None:
            lowerCAmelCase_					:       Union[str, Any]    =		self.get_target_ids(_lowercase    ,						_lowercase						)
            lowerCAmelCase_					:       int    =		target_ids

       if top_k is not None:
            lowerCAmelCase_					:       str    =		top_k

       if self.tokenizer.mask_token_id is None:
            raise PipelineException(
                """fill-mask"""    ,						self.model.base_model_prefix    ,						"""The tokenizer does not define a `mask_token`."""						)
       return {}, {}, postprocess_params







  def __call__(	self    ,						_lowercase    ,						*_lowercase    ,						**_lowercase						):
       lowerCAmelCase_					:       Any    =		super().__call__(_lowercase    ,						**_lowercase						)
       if isinstance(_lowercase    ,						_lowercase						) and len(_lowercase						) == 1:
            return outputs[0]
       return outputs







import argparse
import glob
import importlib.util
import os
import re

import black
from doc_builder.style_doc import style_docstrings_in_code


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_copies.py
UpperCAmelCase_					:   Tuple 							=						"""src/diffusers"""
UpperCAmelCase_					:   str 							=						"""."""


# This is to make sure the diffusers module imported is the one in the repo.
UpperCAmelCase_					:   Any 							=						importlib.util.spec_from_file_location(
    """diffusers""",
    os.path.join(DIFFUSERS_PATH, """__init__.py"""),
    submodule_search_locations=[DIFFUSERS_PATH],
)
UpperCAmelCase_					:   Union[str, Any] 							=						spec.loader.load_module()






def 							_lowerCAmelCase (  _a   :					Optional[int]					,   _a   :					Optional[int]			)	->			Tuple:
     return line.startswith(_a			) or len(_a			) <= 1 or re.search(R"""^\s*\)(\s*->.*:|:)\s*$"""					,   _a			) is not None






def 							_lowerCAmelCase (  _a   :					List[Any]			)	->			int:
     lowerCAmelCase_					:       List[str]    =		object_name.split("""."""			)
     lowerCAmelCase_					:       List[Any]    =		0

     # First let's find the module where our object lives.
     lowerCAmelCase_					:       Optional[int]    =		parts[i]
     while i < len(_a			) and not os.path.isfile(os.path.join(_a					,   F'{module}.py'			)			):
          i += 1
          if i < len(_a			):
               lowerCAmelCase_					:       Dict    =		os.path.join(_a					,   parts[i]			)
     if i >= len(_a			):
          raise ValueError(F'`object_name` should begin with the name of a module of diffusers but got {object_name}.'			)

     with open(os.path.join(_a					,   F'{module}.py'			)					,   """r"""					,   encoding="""utf-8"""					,   newline="""\n"""			) as f:
          lowerCAmelCase_					:       Union[str, Any]    =		f.readlines()

     # Now let's find the class / func in the code!
     lowerCAmelCase_					:       List[Any]    =		""""""
     lowerCAmelCase_					:       Optional[int]    =		0
     for name in parts[i + 1 :]:
          while (
              line_index < len(_a			) and re.search(RF'^{indent}(class|def)\s+{name}(\(|\:)'					,   lines[line_index]			) is None
          ):
               line_index += 1
          indent += "    "
          line_index += 1

     if line_index >= len(_a			):
          raise ValueError(F' {object_name} does not match any function or class in {module}.'			)

     # We found the beginning of the class / func, now let's find the end (when the indent diminishes).
     lowerCAmelCase_					:       str    =		line_index
     while line_index < len(_a			) and _should_continue(lines[line_index]					,   _a			):
          line_index += 1
     # Clean up empty lines at the end (if any).
     while len(lines[line_index - 1]			) <= 1:
          line_index -= 1

     lowerCAmelCase_					:       int    =		lines[start_index:line_index]
     return "".join(_a			)


UpperCAmelCase_					:   str 							=						re.compile(r"""^(\s*)#\s*Copied from\s+diffusers\.(\S+\.\S+)\s*($|\S.*$)""")
UpperCAmelCase_					:   int 							=						re.compile(r"""^\s*(\S+)->(\S+)(\s+.*|$)""")
UpperCAmelCase_					:   List[str] 							=						re.compile(r"""<FILL\s+[^>]*>""")






def 							_lowerCAmelCase (  _a   :					Dict			)	->			List[Any]:
     lowerCAmelCase_					:       int    =		code.split("""\n"""			)
     lowerCAmelCase_					:       Any    =		0
     while idx < len(_a			) and len(lines[idx]			) == 0:
          idx += 1
     if idx < len(_a			):
          return re.search(R"""^(\s*)\S"""					,   lines[idx]			).groups()[0]
     return ""






def 							_lowerCAmelCase (  _a   :					Optional[int]			)	->			Union[str, Any]:
     lowerCAmelCase_					:       Tuple    =		len(get_indent(_a			)			) > 0
     if has_indent:
          lowerCAmelCase_					:       Union[str, Any]    =		F'class Bla:\n{code}'
     lowerCAmelCase_					:       Tuple    =		black.Mode(target_versions={black.TargetVersion.PYaa}					,   line_length=1_19					,   preview=_a			)
     lowerCAmelCase_					:       str    =		black.format_str(_a					,   mode=_a			)
     lowerCAmelCase_   ,  lowerCAmelCase_					:       List[str]    =		style_docstrings_in_code(_a			)
     return result[len("""class Bla:\n"""			) :] if has_indent else result






def 							_lowerCAmelCase (  _a   :					Tuple					,   _a   :					str=False			)	->			Tuple:
     with open(_a					,   """r"""					,   encoding="""utf-8"""					,   newline="""\n"""			) as f:
          lowerCAmelCase_					:       List[str]    =		f.readlines()
     lowerCAmelCase_					:       Union[str, Any]    =		[]
     lowerCAmelCase_					:       Dict    =		0
     # Not a for loop cause `lines` is going to change (if `overwrite=True`).
     while line_index < len(_a			):
          lowerCAmelCase_					:       str    =		_re_copy_warning.search(lines[line_index]			)
          if search is None:
               line_index += 1
               continue

          # There is some copied code here, let's retrieve the original.
          lowerCAmelCase_   ,  lowerCAmelCase_   ,  lowerCAmelCase_					:       str    =		search.groups()
          lowerCAmelCase_					:       List[Any]    =		find_code_in_diffusers(_a			)
          lowerCAmelCase_					:       Tuple    =		get_indent(_a			)

          lowerCAmelCase_					:       Union[str, Any]    =		line_index + 1 if indent == theoretical_indent else line_index + 2
          lowerCAmelCase_					:       str    =		theoretical_indent
          lowerCAmelCase_					:       Tuple    =		start_index

          # Loop to check the observed code, stop when indentation diminishes or if we see a End copy comment.
          lowerCAmelCase_					:       List[Any]    =		True
          while line_index < len(_a			) and should_continue:
               line_index += 1
               if line_index >= len(_a			):
                    break
               lowerCAmelCase_					:       List[Any]    =		lines[line_index]
               lowerCAmelCase_					:       Tuple    =		_should_continue(_a					,   _a			) and re.search(F'^{indent}# End copy'					,   _a			) is None
          # Clean up empty lines at the end (if any).
          while len(lines[line_index - 1]			) <= 1:
               line_index -= 1

          lowerCAmelCase_					:       str    =		lines[start_index:line_index]
          lowerCAmelCase_					:       Optional[Any]    =		"""""".join(_a			)

          # Remove any nested `Copied from` comments to avoid circular copies
          lowerCAmelCase_					:       Optional[int]    =		[line for line in theoretical_code.split("""\n"""			) if _re_copy_warning.search(_a			) is None]
          lowerCAmelCase_					:       str    =		"""\n""".join(_a			)

          # Before comparing, use the `replace_pattern` on the original code.
          if len(_a			) > 0:
               lowerCAmelCase_					:       List[str]    =		replace_pattern.replace("""with"""					,   """"""			).split(""","""			)
               lowerCAmelCase_					:       Any    =		[_re_replace_pattern.search(_a			) for p in patterns]
               for pattern in patterns:
                    if pattern is None:
                         continue
                    lowerCAmelCase_   ,  lowerCAmelCase_   ,  lowerCAmelCase_					:       List[Any]    =		pattern.groups()
                    lowerCAmelCase_					:       Union[str, Any]    =		re.sub(_a					,   _a					,   _a			)
                    if option.strip() == "all-casing":
                         lowerCAmelCase_					:       Dict    =		re.sub(obja.lower()					,   obja.lower()					,   _a			)
                         lowerCAmelCase_					:       Dict    =		re.sub(obja.upper()					,   obja.upper()					,   _a			)

            # Blackify after replacement. To be able to do that, we need the header (class or function definition)
            # from the previous line
               lowerCAmelCase_					:       str    =		blackify(lines[start_index - 1] + theoretical_code			)
               lowerCAmelCase_					:       Optional[Any]    =		theoretical_code[len(lines[start_index - 1]			) :]

          # Test for a diff and act accordingly.
          if observed_code != theoretical_code:
               diffs.append([object_name, start_index]			)
               if overwrite:
                    lowerCAmelCase_					:       Any    =		lines[:start_index] + [theoretical_code] + lines[line_index:]
                    lowerCAmelCase_					:       Dict    =		start_index + 1

     if overwrite and len(_a			) > 0:
          # Warn the user a file has been modified.
          print(F'Detected changes, rewriting {filename}.'			)
          with open(_a					,   """w"""					,   encoding="""utf-8"""					,   newline="""\n"""			) as f:
               f.writelines(_a			)
     return diffs






def 							_lowerCAmelCase (  _a   :					bool = False			)	->			List[Any]:
     lowerCAmelCase_					:       Optional[Any]    =		glob.glob(os.path.join(_a					,   """**/*.py"""			)					,   recursive=_a			)
     lowerCAmelCase_					:       Union[str, Any]    =		[]
     for filename in all_files:
          lowerCAmelCase_					:       Union[str, Any]    =		is_copy_consistent(_a					,   _a			)
          diffs += [F'- {filename}: copy does not match {d[0]} at line {d[1]}' for d in new_diffs]
     if not overwrite and len(_a			) > 0:
          lowerCAmelCase_					:       List[Any]    =		"""\n""".join(_a			)
          raise Exception(
              """Found the following copy inconsistencies:\n"""
              + diff
              + """\nRun `make fix-copies` or `python utils/check_copies.py --fix_and_overwrite` to fix them."""			)


if __name__ == "__main__":
      UpperCAmelCase_					:   Optional[int] 							=						argparse.ArgumentParser()
      parser.add_argument("""--fix_and_overwrite""", action="""store_true""", help="""Whether to fix inconsistencies.""")
      UpperCAmelCase_					:   List[str] 							=						parser.parse_args()

      check_copies(args.fix_and_overwrite)
