


'''simple docstring'''






from manim import *






class  __magic_name__							(		lowerCAmelCase    ):






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




          '''simple docstring'''
          _UpperCAmelCase       :       Optional[int]							   =Rectangle(height=0.5  ,							width=0.5)
          _UpperCAmelCase       :       int							   =Rectangle(height=0.46  ,							width=0.46).set_stroke(width=0)

          _UpperCAmelCase       :       Union[str, Any]							   =Rectangle(height=0.25  ,							width=0.25)

          _UpperCAmelCase       :       List[str]							   =[mem.copy() for i in range(6)]
          _UpperCAmelCase       :       int							   =[mem.copy() for i in range(6)]
          _UpperCAmelCase       :       str							   =VGroup(*snake_case).arrange(snake_case  ,							buff=0)
          _UpperCAmelCase       :       int							   =VGroup(*snake_case).arrange(snake_case  ,							buff=0)
          _UpperCAmelCase       :       Optional[Any]							   =VGroup(snake_case  ,							snake_case).arrange(snake_case  ,							buff=0)
          _UpperCAmelCase       :       Dict							   =Text('CPU'  ,							font_size=2_4)
          _UpperCAmelCase       :       int							   =Group(snake_case  ,							snake_case).arrange(snake_case  ,							buff=0.5  ,							aligned_edge=snake_case)
          cpu.move_to([-2.5, -0.5, 0])
          self.add(snake_case)

          _UpperCAmelCase       :       List[str]							   =[mem.copy() for i in range(4)]
          _UpperCAmelCase       :       int							   =VGroup(*snake_case).arrange(snake_case  ,							buff=0)
          _UpperCAmelCase       :       Union[str, Any]							   =Text('GPU'  ,							font_size=2_4)
          _UpperCAmelCase       :       Dict							   =Group(snake_case  ,							snake_case).arrange(snake_case  ,							buff=0.5  ,							aligned_edge=snake_case)
          gpu.move_to([-1, -1, 0])
          self.add(snake_case)

          _UpperCAmelCase       :       Any							   =[mem.copy() for i in range(6)]
          _UpperCAmelCase       :       List[str]							   =VGroup(*snake_case).arrange(snake_case  ,							buff=0)

          _UpperCAmelCase       :       int							   =Text('Model'  ,							font_size=2_4)
          _UpperCAmelCase       :       Optional[Any]							   =Group(snake_case  ,							snake_case).arrange(snake_case  ,							buff=0.5  ,							aligned_edge=snake_case)
          model.move_to([3, -1.0, 0])
          self.add(snake_case)

          _UpperCAmelCase       :       Optional[Any]							   =[]
          _UpperCAmelCase       :       Optional[int]							   =[]

          for i, rect in enumerate(snake_case):
             _UpperCAmelCase       :       List[str]							   =fill.copy().set_fill(snake_case  ,							opacity=0.8)
             target.move_to(snake_case)
             model_arr.append(snake_case)

             _UpperCAmelCase       :       List[Any]							   =Rectangle(height=0.46  ,							width=0.46).set_stroke(width=0.0).set_fill(snake_case  ,							opacity=0.8)
             cpu_target.move_to(cpu_left_col_base[i])
             model_cpu_arr.append(snake_case)

          self.add(*snake_case  ,							*snake_case)

          _UpperCAmelCase       :       Any							   =[meta_mem.copy() for i in range(6)]
          _UpperCAmelCase       :       List[Any]							   =[meta_mem.copy() for i in range(6)]
          _UpperCAmelCase       :       Dict							   =VGroup(*snake_case).arrange(snake_case  ,							buff=0)
          _UpperCAmelCase       :       Union[str, Any]							   =VGroup(*snake_case).arrange(snake_case  ,							buff=0)
          _UpperCAmelCase       :       List[Any]							   =VGroup(snake_case  ,							snake_case).arrange(snake_case  ,							buff=0)
          _UpperCAmelCase       :       int							   =Text('Disk'  ,							font_size=2_4)
          _UpperCAmelCase       :       Any							   =Group(snake_case  ,							snake_case).arrange(snake_case  ,							buff=0.5  ,							aligned_edge=snake_case)
          disk.move_to([-4, -1.25, 0])
          self.add(snake_case  ,							snake_case)

          _UpperCAmelCase       :       List[str]							   =Square(side_length=2.2)
          key.move_to([-5, 2, 0])

          _UpperCAmelCase       :       List[str]							   =MarkupText(
              f"<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model"  ,							font_size=1_8  ,							)

          key_text.move_to([-5, 2.4, 0])

          self.add(snake_case  ,							snake_case)

          _UpperCAmelCase       :       Union[str, Any]							   =MarkupText(
              f"<span fgcolor='{BLUE}'>●</span> Checkpoint"  ,							font_size=1_8  ,							)

          blue_text.next_to(snake_case  ,							DOWN * 2.4  ,							aligned_edge=key_text.get_left())
          self.add(snake_case)

          _UpperCAmelCase       :       Optional[int]							   =MarkupText(
              f"Now watch as an input is passed through the model\nand how the memory is utilized and handled."  ,							font_size=2_4  ,							)
          step_a.move_to([2, 2, 0])

          self.play(Write(snake_case))

          _UpperCAmelCase       :       Optional[int]							   =Square(0.3)
          input.set_fill(snake_case  ,							opacity=1.0)
          input.set_stroke(width=0.0)
          input.next_to(model_base[0]  ,							snake_case  ,							buff=0.5)

          self.play(Write(snake_case))

          input.generate_target()
          input.target.next_to(model_arr[0]  ,							direction=snake_case  ,							buff=0.02)
          self.play(MoveToTarget(snake_case))

          self.play(FadeOut(snake_case))

          _UpperCAmelCase       :       Optional[Any]							   =Arrow(start=snake_case  ,							end=snake_case  ,							color=snake_case  ,							buff=0.5)
          a.next_to(model_arr[0].get_left()  ,							snake_case  ,							buff=0.2)

          model_cpu_arr[0].generate_target()
          model_cpu_arr[0].target.move_to(gpu_rect[0])

          _UpperCAmelCase       :       Dict							   =MarkupText(
              f"As the input reaches a layer, the hook triggers\nand weights are moved from the CPU\nto the GPU and back."  ,							font_size=2_4  ,							)
          step_a.move_to([2, 2, 0])

          self.play(Write(snake_case  ,							run_time=3))

          _UpperCAmelCase       :       List[str]							   ={'run_time': 1, 'fade_in': True, 'fade_out': True, 'buff': 0.02}

          self.play(
              Write(snake_case)  ,							Circumscribe(model_arr[0]  ,							color=snake_case  ,							**snake_case)  ,							Circumscribe(model_cpu_arr[0]  ,							color=snake_case  ,							**snake_case)  ,							Circumscribe(gpu_rect[0]  ,							color=snake_case  ,							**snake_case)  ,							)
          self.play(MoveToTarget(model_cpu_arr[0]))

          _UpperCAmelCase       :       Tuple							   =a.copy()
          for i in range(6):
             a_c.next_to(model_arr[i].get_right() + 0.02  ,							snake_case  ,							buff=0.2)

             input.generate_target()
             input.target.move_to(model_arr[i].get_right() + 0.02)

             _UpperCAmelCase       :       List[str]							   =AnimationGroup(
                 FadeOut(snake_case  ,							run_time=0.5)  ,							MoveToTarget(snake_case  ,							run_time=0.5)  ,							FadeIn(snake_case  ,							run_time=0.5)  ,							lag_ratio=0.2)

             self.play(snake_case)

             model_cpu_arr[i].generate_target()
             model_cpu_arr[i].target.move_to(cpu_left_col_base[i])

             if i < 5:
                model_cpu_arr[i + 1].generate_target()
                model_cpu_arr[i + 1].target.move_to(gpu_rect[0])
                if i >= 1:
                   _UpperCAmelCase       :       Dict							   =0.7

                self.play(
                    Circumscribe(model_arr[i]  ,							**snake_case)  ,							Circumscribe(cpu_left_col_base[i]  ,							**snake_case)  ,							Circumscribe(cpu_left_col_base[i + 1]  ,							color=snake_case  ,							**snake_case)  ,							Circumscribe(gpu_rect[0]  ,							color=snake_case  ,							**snake_case)  ,							Circumscribe(model_arr[i + 1]  ,							color=snake_case  ,							**snake_case)  ,							)
                if i < 1:
                   self.play(
                       MoveToTarget(model_cpu_arr[i])  ,							MoveToTarget(model_cpu_arr[i + 1])  ,							)
                else:
                   self.play(
                       MoveToTarget(model_cpu_arr[i]  ,							run_time=0.7)  ,							MoveToTarget(model_cpu_arr[i + 1]  ,							run_time=0.7)  ,							)
             else:
                model_cpu_arr[i].generate_target()
                model_cpu_arr[i].target.move_to(cpu_left_col_base[-1])
                input.generate_target()
                input.target.next_to(model_arr[-1].get_right()  ,							RIGHT + 0.02  ,							buff=0.2)

                self.play(
                    Circumscribe(model_arr[-1]  ,							color=snake_case  ,							**snake_case)  ,							Circumscribe(cpu_left_col_base[-1]  ,							color=snake_case  ,							**snake_case)  ,							Circumscribe(gpu_rect[0]  ,							color=snake_case  ,							**snake_case)  ,							)

                self.play(MoveToTarget(model_cpu_arr[i]))

             _UpperCAmelCase       :       str							   =a_c
             _UpperCAmelCase       :       List[Any]							   =a_c.copy()

          input.generate_target()
          input.target.next_to(model_base[-1]  ,							RIGHT + 0.02  ,							buff=0.5)
          self.play(
              FadeOut(snake_case)  ,							FadeOut(snake_case  ,							run_time=0.5)  ,							)

          _UpperCAmelCase       :       List[Any]							   =MarkupText(f"Inference on a model too large for GPU memory\nis successfully completed."  ,							font_size=2_4)
          step_a.move_to([2, 2, 0])

          self.play(Write(snake_case  ,							run_time=3)  ,							MoveToTarget(snake_case))

          self.wait()





'''simple docstring'''






from collections import UserDict
from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
 from PIL import Image

 from ..image_utils import load_image

if is_torch_available():
 from ..models.auto.modeling_auto import MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING

if is_tf_available():
 from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING
 from ..tf_utils import stable_softmax

lowercase							=logging.get_logger(__name__)






@add_end_docstrings(lowerCAmelCase    )
class  __magic_name__							(		lowerCAmelCase    ):






       def __init__(      self  ,							**snake_case)						->  Optional[int]:




          '''simple docstring'''
          super().__init__(**snake_case)

          requires_backends(self  ,							'vision')
          self.check_model_type(
              TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING
              if self.framework == 'tf'
              else MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING)






       def __call__(      self  ,							snake_case  ,							**snake_case)						->  str:




          '''simple docstring'''
          return super().__call__(snake_case  ,							**snake_case)






       def       lowerCAmelCase		(      self  ,							**snake_case)						->  int:




          '''simple docstring'''
          _UpperCAmelCase       :       str							   ={}
          if "candidate_labels" in kwargs:
             _UpperCAmelCase       :       Union[str, Any]							   =kwargs['candidate_labels']
          if "hypothesis_template" in kwargs:
             _UpperCAmelCase       :       List[Any]							   =kwargs['hypothesis_template']

          return preprocess_params, {}, {}






       def       lowerCAmelCase		(      self  ,							snake_case  ,							snake_case=None  ,							snake_case="This is a photo of {}.")						->  Any:




          '''simple docstring'''
          _UpperCAmelCase       :       Optional[Any]							   =load_image(snake_case)
          _UpperCAmelCase       :       Union[str, Any]							   =self.image_processor(images=[image]  ,							return_tensors=self.framework)
          _UpperCAmelCase       :       Union[str, Any]							   =candidate_labels
          _UpperCAmelCase       :       List[Any]							   =[hypothesis_template.format(snake_case) for x in candidate_labels]
          _UpperCAmelCase       :       str							   =self.tokenizer(snake_case  ,							return_tensors=self.framework  ,							padding=snake_case)
          _UpperCAmelCase       :       Any							   =[text_inputs]
          return inputs






       def       lowerCAmelCase		(      self  ,							snake_case)						->  str:




          '''simple docstring'''
          _UpperCAmelCase       :       List[str]							   =model_inputs.pop('candidate_labels')
          _UpperCAmelCase       :       Tuple							   =model_inputs.pop('text_inputs')
          if isinstance(text_inputs[0]  ,							snake_case):
             _UpperCAmelCase       :       Any							   =text_inputs[0]
          else:
             # Batching case.
             _UpperCAmelCase       :       str							   =text_inputs[0][0]

          _UpperCAmelCase       :       Any							   =self.model(**snake_case  ,							**snake_case)

          _UpperCAmelCase       :       List[str]							   ={
              'candidate_labels': candidate_labels,
              'logits': outputs.logits_per_image,
          }
          return model_outputs






       def       lowerCAmelCase		(      self  ,							snake_case)						->  Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase       :       str							   =model_outputs.pop('candidate_labels')
          _UpperCAmelCase       :       Union[str, Any]							   =model_outputs['logits'][0]
          if self.framework == "pt":
             _UpperCAmelCase       :       Dict							   =logits.softmax(dim=-1).squeeze(-1)
             _UpperCAmelCase       :       Union[str, Any]							   =probs.tolist()
             if not isinstance(snake_case  ,							snake_case):
                _UpperCAmelCase       :       Union[str, Any]							   =[scores]
          elif self.framework == "tf":
             _UpperCAmelCase       :       Dict							   =stable_softmax(snake_case  ,							axis=-1)
             _UpperCAmelCase       :       str							   =probs.numpy().tolist()
          else:
             raise ValueError(f"Unsupported framework: {self.framework}")

          _UpperCAmelCase       :       List[str]							   =[
              {'score': score, 'label': candidate_label}
              for score, candidate_label in sorted(zip(snake_case  ,							snake_case)  ,							key=lambda snake_case: -x[0])
          ]
          return result

