
"""simple docstring"""

from manim import *

class __lowerCAmelCase		(				__SCREAMING_SNAKE_CASE      ):





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase      =	Rectangle(height=0.5     ,					width=0.5				)
      __UpperCamelCase      =	Rectangle(height=0.4_6     ,					width=0.4_6				).set_stroke(width=0				)

      __UpperCamelCase      =	Rectangle(height=0.2_5     ,					width=0.2_5				)

      __UpperCamelCase      =	[mem.copy() for i in range(6				)]
      __UpperCamelCase      =	[mem.copy() for i in range(6				)]
      __UpperCamelCase      =	VGroup(*_SCREAMING_SNAKE_CASE				).arrange(_SCREAMING_SNAKE_CASE     ,					buff=0				)
      __UpperCamelCase      =	VGroup(*_SCREAMING_SNAKE_CASE				).arrange(_SCREAMING_SNAKE_CASE     ,					buff=0				)
      __UpperCamelCase      =	VGroup(_SCREAMING_SNAKE_CASE     ,					_SCREAMING_SNAKE_CASE				).arrange(_SCREAMING_SNAKE_CASE     ,					buff=0				)
      __UpperCamelCase      =	Text('CPU'     ,					font_size=24				)
      __UpperCamelCase      =	Group(_SCREAMING_SNAKE_CASE     ,					_SCREAMING_SNAKE_CASE				).arrange(_SCREAMING_SNAKE_CASE     ,					buff=0.5     ,					aligned_edge=_SCREAMING_SNAKE_CASE				)
      cpu.move_to([-2.5, -0.5, 0]				)
      self.add(_SCREAMING_SNAKE_CASE				)

      __UpperCamelCase      =	[mem.copy() for i in range(4				)]
      __UpperCamelCase      =	VGroup(*_SCREAMING_SNAKE_CASE				).arrange(_SCREAMING_SNAKE_CASE     ,					buff=0				)
      __UpperCamelCase      =	Text('GPU'     ,					font_size=24				)
      __UpperCamelCase      =	Group(_SCREAMING_SNAKE_CASE     ,					_SCREAMING_SNAKE_CASE				).arrange(_SCREAMING_SNAKE_CASE     ,					buff=0.5     ,					aligned_edge=_SCREAMING_SNAKE_CASE				)
      gpu.move_to([-1, -1, 0]				)
      self.add(_SCREAMING_SNAKE_CASE				)

      __UpperCamelCase      =	[mem.copy() for i in range(6				)]
      __UpperCamelCase      =	VGroup(*_SCREAMING_SNAKE_CASE				).arrange(_SCREAMING_SNAKE_CASE     ,					buff=0				)

      __UpperCamelCase      =	Text('Model'     ,					font_size=24				)
      __UpperCamelCase      =	Group(_SCREAMING_SNAKE_CASE     ,					_SCREAMING_SNAKE_CASE				).arrange(_SCREAMING_SNAKE_CASE     ,					buff=0.5     ,					aligned_edge=_SCREAMING_SNAKE_CASE				)
      model.move_to([3, -1.0, 0]				)
      self.add(_SCREAMING_SNAKE_CASE				)

      __UpperCamelCase      =	[]
      __UpperCamelCase      =	[]

      for i, rect in enumerate(_SCREAMING_SNAKE_CASE				):
       __UpperCamelCase      =	fill.copy().set_fill(_SCREAMING_SNAKE_CASE     ,					opacity=0.8				)
       target.move_to(_SCREAMING_SNAKE_CASE				)
       model_arr.append(_SCREAMING_SNAKE_CASE				)

       __UpperCamelCase      =	Rectangle(height=0.4_6     ,					width=0.4_6				).set_stroke(width=0.0				).set_fill(_SCREAMING_SNAKE_CASE     ,					opacity=0.8				)
       cpu_target.move_to(cpu_left_col_base[i]				)
       model_cpu_arr.append(_SCREAMING_SNAKE_CASE				)

      self.add(*_SCREAMING_SNAKE_CASE     ,					*_SCREAMING_SNAKE_CASE				)

      __UpperCamelCase      =	[meta_mem.copy() for i in range(6				)]
      __UpperCamelCase      =	[meta_mem.copy() for i in range(6				)]
      __UpperCamelCase      =	VGroup(*_SCREAMING_SNAKE_CASE				).arrange(_SCREAMING_SNAKE_CASE     ,					buff=0				)
      __UpperCamelCase      =	VGroup(*_SCREAMING_SNAKE_CASE				).arrange(_SCREAMING_SNAKE_CASE     ,					buff=0				)
      __UpperCamelCase      =	VGroup(_SCREAMING_SNAKE_CASE     ,					_SCREAMING_SNAKE_CASE				).arrange(_SCREAMING_SNAKE_CASE     ,					buff=0				)
      __UpperCamelCase      =	Text('Disk'     ,					font_size=24				)
      __UpperCamelCase      =	Group(_SCREAMING_SNAKE_CASE     ,					_SCREAMING_SNAKE_CASE				).arrange(_SCREAMING_SNAKE_CASE     ,					buff=0.5     ,					aligned_edge=_SCREAMING_SNAKE_CASE				)
      disk.move_to([-4, -1.2_5, 0]				)
      self.add(_SCREAMING_SNAKE_CASE     ,					_SCREAMING_SNAKE_CASE				)

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

      __UpperCamelCase      =	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(_SCREAMING_SNAKE_CASE     ,					_SCREAMING_SNAKE_CASE				)

      __UpperCamelCase      =	MarkupText(
          F'<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'     ,					font_size=18     ,					)

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

      __UpperCamelCase      =	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(_SCREAMING_SNAKE_CASE				)				)

      __UpperCamelCase      =	Square(0.3				)
      input.set_fill(_SCREAMING_SNAKE_CASE     ,					opacity=1.0				)
      input.set_stroke(width=0.0				)
      input.next_to(model_base[0]     ,					_SCREAMING_SNAKE_CASE     ,					buff=0.5				)

      self.play(Write(_SCREAMING_SNAKE_CASE				)				)

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

      self.play(FadeOut(_SCREAMING_SNAKE_CASE				)				)

      __UpperCamelCase      =	Arrow(start=_SCREAMING_SNAKE_CASE     ,					end=_SCREAMING_SNAKE_CASE     ,					color=_SCREAMING_SNAKE_CASE     ,					buff=0.5				)
      a.next_to(model_arr[0].get_left()     ,					_SCREAMING_SNAKE_CASE     ,					buff=0.2				)

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

      __UpperCamelCase      =	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(_SCREAMING_SNAKE_CASE     ,					run_time=3				)				)

      __UpperCamelCase      =	{'''run_time''': 1, '''fade_in''': True, '''fade_out''': True, '''buff''': 0.0_2}

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

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

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

       __UpperCamelCase      =	AnimationGroup(
           FadeOut(_SCREAMING_SNAKE_CASE     ,					run_time=0.5				)     ,					MoveToTarget(_SCREAMING_SNAKE_CASE     ,					run_time=0.5				)     ,					FadeIn(_SCREAMING_SNAKE_CASE     ,					run_time=0.5				)     ,					lag_ratio=0.2				)

       self.play(_SCREAMING_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      =	0.7

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

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

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

       __UpperCamelCase      =	a_c
       __UpperCamelCase      =	a_c.copy()

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

      __UpperCamelCase      =	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(_SCREAMING_SNAKE_CASE     ,					run_time=3				)     ,					MoveToTarget(_SCREAMING_SNAKE_CASE				)				)

      self.wait()





"""simple docstring"""

import multiprocessing
import os
from typing import BinaryIO, Optional, Union

import fsspec

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.json.json import Json
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader

class __lowerCAmelCase		(				__SCREAMING_SNAKE_CASE      ):





     def __init__(     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase = None     ,					__UpperCAmelCase = None     ,					__UpperCAmelCase = None     ,					__UpperCAmelCase = False     ,					__UpperCAmelCase = False     ,					__UpperCAmelCase = None     ,					__UpperCAmelCase = None     ,					**__UpperCAmelCase     ,					):

      '''simple docstring'''




      super().__init__(
          __UpperCAmelCase     ,					split=__UpperCAmelCase     ,					features=__UpperCAmelCase     ,					cache_dir=__UpperCAmelCase     ,					keep_in_memory=__UpperCAmelCase     ,					streaming=__UpperCAmelCase     ,					num_proc=__UpperCAmelCase     ,					**__UpperCAmelCase     ,					)
      __UpperCamelCase      =	field
      __UpperCamelCase      =	path_or_paths if isinstance(__UpperCAmelCase     ,					__UpperCAmelCase				) else {self.split: path_or_paths}
      __UpperCamelCase      =	Json(
          cache_dir=__UpperCAmelCase     ,					data_files=__UpperCAmelCase     ,					features=__UpperCAmelCase     ,					field=__UpperCAmelCase     ,					**__UpperCAmelCase     ,					)







     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      if self.streaming:
       __UpperCamelCase      =	self.builder.as_streaming_dataset(split=self.split				)
      # Build regular (map-style) dataset
      else:
       __UpperCamelCase      =	None
       __UpperCamelCase      =	None
       __UpperCamelCase      =	None
       __UpperCamelCase      =	None

       self.builder.download_and_prepare(
           download_config=__UpperCAmelCase     ,					download_mode=__UpperCAmelCase     ,					verification_mode=__UpperCAmelCase     ,					base_path=__UpperCAmelCase     ,					num_proc=self.num_proc     ,					)
       __UpperCamelCase      =	self.builder.as_dataset(
           split=self.split     ,					verification_mode=__UpperCAmelCase     ,					in_memory=self.keep_in_memory				)
      return dataset


class __lowerCAmelCase		:





     def __init__(     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase = None     ,					__UpperCAmelCase = None     ,					**__UpperCAmelCase     ,					):

      '''simple docstring'''




      if num_proc is not None and num_proc <= 0:
       raise ValueError(F'num_proc {num_proc} must be an integer > 0.'				)

      __UpperCamelCase      =	dataset
      __UpperCamelCase      =	path_or_buf
      __UpperCamelCase      =	batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
      __UpperCamelCase      =	num_proc
      __UpperCamelCase      =	'utf-8'
      __UpperCamelCase      =	to_json_kwargs





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase      =	self.to_json_kwargs.pop('path_or_buf'     ,					__UpperCAmelCase				)
      __UpperCamelCase      =	self.to_json_kwargs.pop('orient'     ,					'records'				)
      __UpperCamelCase      =	self.to_json_kwargs.pop('lines'     ,					True if orient == 'records' else False				)
      __UpperCamelCase      =	self.to_json_kwargs.pop('index'     ,					False if orient in ['split', 'table'] else True				)
      __UpperCamelCase      =	self.to_json_kwargs.pop('compression'     ,					__UpperCAmelCase				)

      if compression not in [None, "infer", "gzip", "bz2", "xz"]:
       raise NotImplementedError(F'`datasets` currently does not support {compression} compression'				)

      if isinstance(self.path_or_buf     ,					(str, bytes, os.PathLike)				):
       with fsspec.open(self.path_or_buf     ,					'wb'     ,					compression=__UpperCAmelCase				) as buffer:
        __UpperCamelCase      =	self._write(file_obj=__UpperCAmelCase     ,					orient=__UpperCAmelCase     ,					lines=__UpperCAmelCase     ,					index=__UpperCAmelCase     ,					**self.to_json_kwargs				)
      else:
       if compression:
        raise NotImplementedError(
            F'The compression parameter is not supported when writing to a buffer, but compression={compression}'
            ' was passed. Please provide a local path instead.'				)
       __UpperCamelCase      =	self._write(
           file_obj=self.path_or_buf     ,					orient=__UpperCAmelCase     ,					lines=__UpperCAmelCase     ,					index=__UpperCAmelCase     ,					**self.to_json_kwargs				)
      return written





     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase				):

      '''simple docstring'''




      __UpperCamelCase  ,     __UpperCamelCase  ,     __UpperCamelCase  ,     __UpperCamelCase  ,     __UpperCamelCase      =	args

      __UpperCamelCase      =	query_table(
          table=self.dataset.data     ,					key=slice(__UpperCAmelCase     ,					offset + self.batch_size				)     ,					indices=self.dataset._indices     ,					)
      __UpperCamelCase      =	batch.to_pandas().to_json(
          path_or_buf=__UpperCAmelCase     ,					orient=__UpperCAmelCase     ,					lines=__UpperCAmelCase     ,					index=__UpperCAmelCase     ,					**__UpperCAmelCase				)
      if not json_str.endswith('\n'				):
       json_str += "\n"
      return json_str.encode(self.encoding				)







     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					**__UpperCAmelCase     ,					):

      '''simple docstring'''




      __UpperCamelCase      =	0

      if self.num_proc is None or self.num_proc == 1:
       for offset in logging.tqdm(
           range(0     ,					len(self.dataset				)     ,					self.batch_size				)     ,					unit='ba'     ,					disable=not logging.is_progress_bar_enabled()     ,					desc='Creating json from Arrow format'     ,					):
        __UpperCamelCase      =	self._batch_json((offset, orient, lines, index, to_json_kwargs)				)
        written += file_obj.write(__UpperCAmelCase				)
      else:
       __UpperCamelCase  ,     __UpperCamelCase      =	len(self.dataset				), self.batch_size
       with multiprocessing.Pool(self.num_proc				) as pool:
        for json_str in logging.tqdm(
            pool.imap(
                self._batch_json     ,					[(offset, orient, lines, index, to_json_kwargs) for offset in range(0     ,					__UpperCAmelCase     ,					__UpperCAmelCase				)]     ,					)     ,					total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size     ,					unit='ba'     ,					disable=not logging.is_progress_bar_enabled()     ,					desc='Creating json from Arrow format'     ,					):
         written += file_obj.write(__UpperCAmelCase				)

      return written



