




"""simple docstring"""

from manim import *



class  _snake_case				(							A__       ):

				'''simple docstring'''




				def 				snake_case_ (							self      :  Optional[int]     ):
								UpperCAmelCase_					:Tuple      			=      Rectangle(height=0.5					,							width=0.5     )
								UpperCAmelCase_					:Dict      			=      Rectangle(height=0.46					,							width=0.46     ).set_stroke(width=0     )
								UpperCAmelCase_					:str      			=      Rectangle(height=0.25					,							width=0.25     )

								UpperCAmelCase_					:Union[str, Any]      			=      [mem.copy() for i in range(6     )]
								UpperCAmelCase_					:int      			=      [mem.copy() for i in range(6     )]
								UpperCAmelCase_					:List[str]      			=      VGroup(*__lowerCamelCase     ).arrange(__lowerCamelCase					,							buff=0     )
								UpperCAmelCase_					:Union[str, Any]      			=      VGroup(*__lowerCamelCase     ).arrange(__lowerCamelCase					,							buff=0     )
								UpperCAmelCase_					:Union[str, Any]      			=      VGroup(__lowerCamelCase					,							__lowerCamelCase     ).arrange(__lowerCamelCase					,							buff=0     )
								UpperCAmelCase_					:List[Any]      			=      Text('''CPU'''					,							font_size=24     )
								UpperCAmelCase_					:List[str]      			=      Group(__lowerCamelCase					,							__lowerCamelCase     ).arrange(__lowerCamelCase					,							buff=0.5					,							aligned_edge=__lowerCamelCase     )
								cpu.move_to([-2.5, -0.5, 0]     )
								self.add(__lowerCamelCase     )

								UpperCAmelCase_					:Optional[int]      			=      [mem.copy() for i in range(4     )]
								UpperCAmelCase_					:Dict      			=      VGroup(*__lowerCamelCase     ).arrange(__lowerCamelCase					,							buff=0     )
								UpperCAmelCase_					:Optional[int]      			=      Text('''GPU'''					,							font_size=24     )
								UpperCAmelCase_					:Optional[Any]      			=      Group(__lowerCamelCase					,							__lowerCamelCase     ).arrange(__lowerCamelCase					,							buff=0.5					,							aligned_edge=__lowerCamelCase     )
								gpu.move_to([-1, -1, 0]     )
								self.add(__lowerCamelCase     )

								UpperCAmelCase_					:int      			=      [mem.copy() for i in range(6     )]
								UpperCAmelCase_					:Any      			=      VGroup(*__lowerCamelCase     ).arrange(__lowerCamelCase					,							buff=0     )

								UpperCAmelCase_					:int      			=      Text('''Model'''					,							font_size=24     )
								UpperCAmelCase_					:str      			=      Group(__lowerCamelCase					,							__lowerCamelCase     ).arrange(__lowerCamelCase					,							buff=0.5					,							aligned_edge=__lowerCamelCase     )
								model.move_to([3, -1.0, 0]     )
								self.add(__lowerCamelCase     )

								UpperCAmelCase_					:Tuple      			=      []
								UpperCAmelCase_					:Any      			=      []

								for i, rect in enumerate(__lowerCamelCase     ):
												rect.set_stroke(__lowerCamelCase     )

												UpperCAmelCase_					:Any      			=      Rectangle(height=0.46 / 4					,							width=0.46 / 3     ).set_stroke(width=0.0     ).set_fill(__lowerCamelCase					,							opacity=0.7     )

												if i == 0:
																cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT     )					,							buff=0.02					,							direction=__lowerCamelCase     )
																cpu_target.set_x(cpu_target.get_x() + 0.1     )
												elif i == 3:
																cpu_target.next_to(model_cpu_arr[0]					,							direction=__lowerCamelCase					,							buff=0.0     )
												else:
																cpu_target.next_to(model_cpu_arr[i - 1]					,							direction=__lowerCamelCase					,							buff=0.0     )
												self.add(__lowerCamelCase     )
												model_cpu_arr.append(__lowerCamelCase     )

								self.add(*__lowerCamelCase					,							*__lowerCamelCase     )

								UpperCAmelCase_					:List[Any]      			=      [meta_mem.copy() for i in range(6     )]
								UpperCAmelCase_					:int      			=      [meta_mem.copy() for i in range(6     )]
								UpperCAmelCase_					:int      			=      VGroup(*__lowerCamelCase     ).arrange(__lowerCamelCase					,							buff=0     )
								UpperCAmelCase_					:Optional[int]      			=      VGroup(*__lowerCamelCase     ).arrange(__lowerCamelCase					,							buff=0     )
								UpperCAmelCase_					:Optional[int]      			=      VGroup(__lowerCamelCase					,							__lowerCamelCase     ).arrange(__lowerCamelCase					,							buff=0     )
								UpperCAmelCase_					:Dict      			=      Text('''Disk'''					,							font_size=24     )
								UpperCAmelCase_					:Optional[Any]      			=      Group(__lowerCamelCase					,							__lowerCamelCase     ).arrange(__lowerCamelCase					,							buff=0.5					,							aligned_edge=__lowerCamelCase     )
								disk.move_to([-4.0, -1.25, 0]     )
								self.add(__lowerCamelCase					,							__lowerCamelCase     )

								UpperCAmelCase_					:Union[str, Any]      			=      []

								for i in range(6     ):
												UpperCAmelCase_					:Optional[int]      			=      fill.copy().set_fill(__lowerCamelCase					,							opacity=0.8     )
												target.move_to(disk_left_col_base[i]     ).scale(0.5     )
												cpu_disk_arr.append(__lowerCamelCase     )

								self.add(*__lowerCamelCase     )

								UpperCAmelCase_					:Dict      			=      Square(side_length=2.2     )
								key.move_to([-5, 2, 0]     )

								UpperCAmelCase_					:Optional[int]      			=      MarkupText(
								    f'<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'					,							font_size=18					,							)

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

								self.add(__lowerCamelCase					,							__lowerCamelCase     )

								UpperCAmelCase_					:Tuple      			=      MarkupText(
								    f'<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'					,							font_size=18					,							)

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

								UpperCAmelCase_					:int      			=      MarkupText(f'The offloaded weights are all sent to the CPU.'					,							font_size=24     )
								step_a.move_to([2, 2, 0]     )

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

								for i in range(6     ):
												UpperCAmelCase_					:str      			=      cpu_disk_arr[i]
												UpperCAmelCase_					:Dict      			=      rect.copy().set_fill(__lowerCamelCase					,							opacity=0.8     ).scale(2.0     )
												cpa.generate_target()
												cpa.target.move_to(model_base[i]     )

												if i == 0:
																rect.set_fill(__lowerCamelCase					,							opacity=0.8     )
																rect.generate_target()
																rect.target.move_to(cpu_left_col_base[0]     ).scale(2.0     )

																self.remove(
																    *__lowerCamelCase					,							*__lowerCamelCase					,							)

												else:
																rect.generate_target()
																rect.target.move_to(cpu_left_col_base[i]     ).scale(2.0     )
												self.play(MoveToTarget(__lowerCamelCase     )					,							MoveToTarget(__lowerCamelCase     )					,							model_base[i].animate.set_stroke(__lowerCamelCase     )     )
								self.play(FadeOut(__lowerCamelCase     )     )

								UpperCAmelCase_					:Any      			=      MarkupText(
								    f'Finally, hooks are added to each weight in the model\nto transfer the weights from CPU to GPU\n\t\tand back when needed.'					,							font_size=24					,							)
								step_a.move_to([2, 2, 0]     )

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

								UpperCAmelCase_					:Any      			=      []
								UpperCAmelCase_					:str      			=      []
								for i in range(6     ):
												UpperCAmelCase_					:Optional[Any]      			=      Arrow(start=__lowerCamelCase					,							end=__lowerCamelCase					,							color=__lowerCamelCase					,							buff=0.5     )
												a.next_to(model_base[i].get_left()					,							__lowerCamelCase					,							buff=0.2     )
												arrows.append(__lowerCamelCase     )
												animations.append(Write(__lowerCamelCase     )     )
								self.play(*__lowerCamelCase     )
								self.wait()








"""simple docstring"""

from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__lowerCamelCase								=						{
    "configuration_bart": ["BART_PRETRAINED_CONFIG_ARCHIVE_MAP", "BartConfig", "BartOnnxConfig"],
    "tokenization_bart": ["BartTokenizer"],
}

try:
						if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__lowerCamelCase								=						["BartTokenizerFast"]

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__lowerCamelCase								=						[
						    "BART_PRETRAINED_MODEL_ARCHIVE_LIST",
						    "BartForCausalLM",
						    "BartForConditionalGeneration",
						    "BartForQuestionAnswering",
						    "BartForSequenceClassification",
						    "BartModel",
						    "BartPretrainedModel",
						    "PretrainedBartModel",
						]

try:
						if not is_tf_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__lowerCamelCase								=						[
						    "TFBartForConditionalGeneration",
						    "TFBartForSequenceClassification",
						    "TFBartModel",
						    "TFBartPretrainedModel",
						]

try:
						if not is_flax_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__lowerCamelCase								=						[
						    "FlaxBartDecoderPreTrainedModel",
						    "FlaxBartForCausalLM",
						    "FlaxBartForConditionalGeneration",
						    "FlaxBartForQuestionAnswering",
						    "FlaxBartForSequenceClassification",
						    "FlaxBartModel",
						    "FlaxBartPreTrainedModel",
						]

if TYPE_CHECKING:
						from .configuration_bart import BART_PRETRAINED_CONFIG_ARCHIVE_MAP, BartConfig, BartOnnxConfig
						from .tokenization_bart import BartTokenizer

						try:
												if not is_tokenizers_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_bart_fast import BartTokenizerFast

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_bart import (
												    BART_PRETRAINED_MODEL_ARCHIVE_LIST,
												    BartForCausalLM,
												    BartForConditionalGeneration,
												    BartForQuestionAnswering,
												    BartForSequenceClassification,
												    BartModel,
												    BartPretrainedModel,
												    PretrainedBartModel,
												)

						try:
												if not is_tf_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_tf_bart import (
												    TFBartForConditionalGeneration,
												    TFBartForSequenceClassification,
												    TFBartModel,
												    TFBartPretrainedModel,
												)

						try:
												if not is_flax_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_flax_bart import (
												    FlaxBartDecoderPreTrainedModel,
												    FlaxBartForCausalLM,
												    FlaxBartForConditionalGeneration,
												    FlaxBartForQuestionAnswering,
												    FlaxBartForSequenceClassification,
												    FlaxBartModel,
												    FlaxBartPreTrainedModel,
												)

else:
						import sys

						__lowerCamelCase								=						_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


