






"""simple docstring"""
from manim import *







class 			UpperCamelCase_       (__A     ):





	def      _SCREAMING_SNAKE_CASE      (      self					: List[str]						)						->			List[str]:
				UpperCAmelCase_     : Any			     =					Rectangle(height=0.5	,	width=0.5						)
				UpperCAmelCase_     : List[Any]			     =					Rectangle(height=0.4_6	,	width=0.4_6						).set_stroke(width=0						)

				UpperCAmelCase_     : List[Any]			     =					Rectangle(height=0.2_5	,	width=0.2_5						)

				UpperCAmelCase_     : Tuple			     =					[mem.copy() for i in range(6						)]
				UpperCAmelCase_     : Union[str, Any]			     =					[mem.copy() for i in range(6						)]
				UpperCAmelCase_     : List[str]			     =					VGroup(*lowerCAmelCase_						).arrange(lowerCAmelCase_	,	buff=0						)
				UpperCAmelCase_     : int			     =					VGroup(*lowerCAmelCase_						).arrange(lowerCAmelCase_	,	buff=0						)
				UpperCAmelCase_     : int			     =					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_     : Union[str, Any]			     =					[mem.copy() for i in range(4						)]
				UpperCAmelCase_     : Tuple			     =					VGroup(*lowerCAmelCase_						).arrange(lowerCAmelCase_	,	buff=0						)
				UpperCAmelCase_     : List[Any]			     =					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_     : Any			     =					[mem.copy() for i in range(6						)]
				UpperCAmelCase_     : int			     =					VGroup(*lowerCAmelCase_						).arrange(lowerCAmelCase_	,	buff=0						)

				UpperCAmelCase_     : Tuple			     =					Text("Model"	,	font_size=24						)
				UpperCAmelCase_     : Optional[Any]			     =					Group(lowerCAmelCase_	,	lowerCAmelCase_						).arrange(lowerCAmelCase_	,	buff=0.5	,	aligned_edge=lowerCAmelCase_						)
				model.move_to([3, -1.0, 0]						)
				self.add(lowerCAmelCase_						)

				UpperCAmelCase_     : str			     =					[]
				UpperCAmelCase_     : int			     =					[]

				for i, rect in enumerate(lowerCAmelCase_						):
							UpperCAmelCase_     : Optional[int]			     =					fill.copy().set_fill(lowerCAmelCase_	,	opacity=0.8						)
							target.move_to(lowerCAmelCase_						)
							model_arr.append(lowerCAmelCase_						)

							UpperCAmelCase_     : str			     =					Rectangle(height=0.4_6	,	width=0.4_6						).set_stroke(width=0.0						).set_fill(lowerCAmelCase_	,	opacity=0.8						)
							cpu_target.move_to(cpu_left_col_base[i]						)
							model_cpu_arr.append(lowerCAmelCase_						)

				self.add(*lowerCAmelCase_	,	*lowerCAmelCase_						)

				UpperCAmelCase_     : Optional[Any]			     =					[meta_mem.copy() for i in range(6						)]
				UpperCAmelCase_     : Any			     =					[meta_mem.copy() for i in range(6						)]
				UpperCAmelCase_     : Dict			     =					VGroup(*lowerCAmelCase_						).arrange(lowerCAmelCase_	,	buff=0						)
				UpperCAmelCase_     : Optional[int]			     =					VGroup(*lowerCAmelCase_						).arrange(lowerCAmelCase_	,	buff=0						)
				UpperCAmelCase_     : Dict			     =					VGroup(lowerCAmelCase_	,	lowerCAmelCase_						).arrange(lowerCAmelCase_	,	buff=0						)
				UpperCAmelCase_     : Optional[int]			     =					Text("Disk"	,	font_size=24						)
				UpperCAmelCase_     : Union[str, Any]			     =					Group(lowerCAmelCase_	,	lowerCAmelCase_						).arrange(lowerCAmelCase_	,	buff=0.5	,	aligned_edge=lowerCAmelCase_						)
				disk.move_to([-4, -1.2_5, 0]						)
				self.add(lowerCAmelCase_	,	lowerCAmelCase_						)

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

				UpperCAmelCase_     : Dict			     =					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_     : int			     =					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_     : Dict			     =					MarkupText(
				    f"""Now watch as an input is passed through the model\nand how the memory is utilized and handled."""	,	font_size=24	,	)
				step_a.move_to([2, 2, 0]						)

				self.play(Write(lowerCAmelCase_						)						)

				UpperCAmelCase_     : str			     =					Square(0.3						)
				input.set_fill(lowerCAmelCase_	,	opacity=1.0						)
				input.set_stroke(width=0.0						)
				input.next_to(model_base[0]	,	lowerCAmelCase_	,	buff=0.5						)

				self.play(Write(lowerCAmelCase_						)						)

				input.generate_target()
				input.target.next_to(model_arr[0]	,	direction=lowerCAmelCase_	,	buff=0.0_2						)
				self.play(MoveToTarget(lowerCAmelCase_						)						)

				self.play(FadeOut(lowerCAmelCase_						)						)

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

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

				UpperCAmelCase_     : str			     =					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=24	,	)
				step_a.move_to([2, 2, 0]						)

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

				UpperCAmelCase_     : Optional[int]			     =					{"run_time": 1, "fade_in": True, "fade_out": True, "buff": 0.0_2}

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

				UpperCAmelCase_     : int			     =					a.copy()
				for i in range(6						):
							a_c.next_to(model_arr[i].get_right() + 0.0_2	,	lowerCAmelCase_	,	buff=0.2						)

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

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

							self.play(lowerCAmelCase_						)

							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]	,	**lowerCAmelCase_						)	,	Circumscribe(cpu_left_col_base[i]	,	**lowerCAmelCase_						)	,	Circumscribe(cpu_left_col_base[i + 1]	,	color=lowerCAmelCase_	,	**lowerCAmelCase_						)	,	Circumscribe(gpu_rect[0]	,	color=lowerCAmelCase_	,	**lowerCAmelCase_						)	,	Circumscribe(model_arr[i + 1]	,	color=lowerCAmelCase_	,	**lowerCAmelCase_						)	,	)
										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.0_2	,	buff=0.2						)

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

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

							UpperCAmelCase_     : List[str]			     =					a_c
							UpperCAmelCase_     : Tuple			     =					a_c.copy()

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

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

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

				self.wait()











"""simple docstring"""
from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase_					=      logging.get_logger(__name__)

lowerCamelCase_					=      {
    '''RWKV/rwkv-4-169m-pile''': '''https://huggingface.co/RWKV/rwkv-4-169m-pile/resolve/main/config.json''',
    '''RWKV/rwkv-4-430m-pile''': '''https://huggingface.co/RWKV/rwkv-4-430m-pile/resolve/main/config.json''',
    '''RWKV/rwkv-4-1b5-pile''': '''https://huggingface.co/RWKV/rwkv-4-1b5-pile/resolve/main/config.json''',
    '''RWKV/rwkv-4-3b-pile''': '''https://huggingface.co/RWKV/rwkv-4-3b-pile/resolve/main/config.json''',
    '''RWKV/rwkv-4-7b-pile''': '''https://huggingface.co/RWKV/rwkv-4-7b-pile/resolve/main/config.json''',
    '''RWKV/rwkv-4-14b-pile''': '''https://huggingface.co/RWKV/rwkv-4-14b-pile/resolve/main/config.json''',
    '''RWKV/rwkv-raven-1b5''': '''https://huggingface.co/RWKV/rwkv-raven-1b5/resolve/main/config.json''',
    '''RWKV/rwkv-raven-3b''': '''https://huggingface.co/RWKV/rwkv-raven-3b/resolve/main/config.json''',
    '''RWKV/rwkv-raven-7b''': '''https://huggingface.co/RWKV/rwkv-raven-7b/resolve/main/config.json''',
    '''RWKV/rwkv-raven-14b''': '''https://huggingface.co/RWKV/rwkv-raven-14b/resolve/main/config.json''',
}







class 			UpperCamelCase_       (__A     ):
	__magic_name__	      =				'''rwkv'''
	__magic_name__	      =				{'''max_position_embeddings''': '''context_length'''}





	def __init__(      self					: str	,	lowerCAmelCase_					: str=50_277	,	lowerCAmelCase_					: Optional[int]=1_024	,	lowerCAmelCase_					: Optional[int]=4_096	,	lowerCAmelCase_					: Union[str, Any]=32	,	lowerCAmelCase_					: List[Any]=None	,	lowerCAmelCase_					: int=None	,	lowerCAmelCase_					: List[Any]=1e-5	,	lowerCAmelCase_					: str=0	,	lowerCAmelCase_					: Tuple=0	,	lowerCAmelCase_					: int=6	,	lowerCAmelCase_					: Optional[Any]=False	,	lowerCAmelCase_					: Any=True	,	**lowerCAmelCase_					: List[Any]	,	)						->			List[str]:
				UpperCAmelCase_     : Tuple			     =					vocab_size
				UpperCAmelCase_     : List[str]			     =					context_length
				UpperCAmelCase_     : Dict			     =					hidden_size
				UpperCAmelCase_     : Optional[int]			     =					num_hidden_layers
				UpperCAmelCase_     : Optional[int]			     =					attention_hidden_size if attention_hidden_size is not None else hidden_size
				UpperCAmelCase_     : Dict			     =					intermediate_size if intermediate_size is not None else 4 * hidden_size
				UpperCAmelCase_     : Any			     =					layer_norm_epsilon
				UpperCAmelCase_     : List[Any]			     =					rescale_every
				UpperCAmelCase_     : List[str]			     =					use_cache

				UpperCAmelCase_     : List[str]			     =					bos_token_id
				UpperCAmelCase_     : Union[str, Any]			     =					eos_token_id

				super().__init__(
				    tie_word_embeddings=lowerCAmelCase_	,	bos_token_id=lowerCAmelCase_	,	eos_token_id=lowerCAmelCase_	,	**lowerCAmelCase_						)



