

'''simple docstring'''






import argparse

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    CLIPTokenizer,
    CLIPTokenizerFast,
    VideoMAEImageProcessor,
    XCLIPConfig,
    XCLIPModel,
    XCLIPProcessor,
    XCLIPTextConfig,
    XCLIPVisionConfig,
)


def    __lowerCamelCase     ( SCREAMING_SNAKE_CASE_	:      Optional[Any]    ,    SCREAMING_SNAKE_CASE_	:      Optional[int]       )				->  str:
      """simple docstring"""

      SCREAMING_SNAKE_CASE_ :     List[str]	=			XCLIPTextConfig()

      # derive patch size from model name
      SCREAMING_SNAKE_CASE_ :     Dict	=			model_name.find("patch"       )
      SCREAMING_SNAKE_CASE_ :     Optional[Any]	=			int(model_name[start_idx + len("patch"       ) : start_idx + len("patch"       ) + 2]       )
      SCREAMING_SNAKE_CASE_ :     List[str]	=			XCLIPVisionConfig(patch_size=A__    ,    num_frames=A__       )

      if "large" in model_name:
            SCREAMING_SNAKE_CASE_ :     Dict	=			7_6_8
            SCREAMING_SNAKE_CASE_ :     int	=			3_0_7_2
            SCREAMING_SNAKE_CASE_ :     List[str]	=			1_2

            SCREAMING_SNAKE_CASE_ :     Optional[Any]	=			1_0_2_4
            SCREAMING_SNAKE_CASE_ :     Any	=			4_0_9_6
            SCREAMING_SNAKE_CASE_ :     List[str]	=			1_6
            SCREAMING_SNAKE_CASE_ :     Tuple	=			2_4
            SCREAMING_SNAKE_CASE_ :     str	=			7_6_8
            SCREAMING_SNAKE_CASE_ :     Dict	=			3_0_7_2

      if model_name == "xclip-large-patch14-16-frames":
            SCREAMING_SNAKE_CASE_ :     Optional[Any]	=			3_3_6

      SCREAMING_SNAKE_CASE_ :     Optional[Any]	=			XCLIPConfig.from_text_vision_configs(A__    ,    A__       )

      if "large" in model_name:
            SCREAMING_SNAKE_CASE_ :     int	=			7_6_8

      return config


def    __lowerCamelCase     ( SCREAMING_SNAKE_CASE_	:      Dict       )				->  List[Any]:
      """simple docstring"""

      if name == "token_embedding.weight":
            SCREAMING_SNAKE_CASE_ :     Optional[Any]	=			name.replace("token_embedding.weight"    ,    "text_model.embeddings.token_embedding.weight"       )
      if name == "positional_embedding":
            SCREAMING_SNAKE_CASE_ :     Union[str, Any]	=			name.replace("positional_embedding"    ,    "text_model.embeddings.position_embedding.weight"       )
      if "ln_1" in name:
            SCREAMING_SNAKE_CASE_ :     Dict	=			name.replace("ln_1"    ,    "layer_norm1"       )
      if "ln_2" in name:
            SCREAMING_SNAKE_CASE_ :     List[str]	=			name.replace("ln_2"    ,    "layer_norm2"       )
      if "c_fc" in name:
            SCREAMING_SNAKE_CASE_ :     Optional[Any]	=			name.replace("c_fc"    ,    "fc1"       )
      if "c_proj" in name:
            SCREAMING_SNAKE_CASE_ :     Tuple	=			name.replace("c_proj"    ,    "fc2"       )
      if name.startswith("transformer.resblocks"       ):
            SCREAMING_SNAKE_CASE_ :     List[str]	=			name.replace("transformer.resblocks"    ,    "text_model.encoder.layers"       )
      if "attn.out_proj" in name and "message" not in name:
            SCREAMING_SNAKE_CASE_ :     str	=			name.replace("attn.out_proj"    ,    "self_attn.out_proj"       )
      if "ln_final" in name:
            SCREAMING_SNAKE_CASE_ :     Optional[Any]	=			name.replace("ln_final"    ,    "text_model.final_layer_norm"       )
      # visual encoder
      if name == "visual.class_embedding":
            SCREAMING_SNAKE_CASE_ :     Tuple	=			name.replace("visual.class_embedding"    ,    "vision_model.embeddings.class_embedding"       )
      if name == "visual.positional_embedding":
            SCREAMING_SNAKE_CASE_ :     int	=			name.replace("visual.positional_embedding"    ,    "vision_model.embeddings.position_embedding.weight"       )
      if name.startswith("visual.transformer.resblocks"       ):
            SCREAMING_SNAKE_CASE_ :     Any	=			name.replace("visual.transformer.resblocks"    ,    "vision_model.encoder.layers"       )
      if "visual.conv1" in name:
            SCREAMING_SNAKE_CASE_ :     List[Any]	=			name.replace("visual.conv1"    ,    "vision_model.embeddings.patch_embedding"       )
      if "visual.ln_pre" in name:
            SCREAMING_SNAKE_CASE_ :     Dict	=			name.replace("visual.ln_pre"    ,    "vision_model.pre_layernorm"       )
      if "visual.ln_post" in name:
            SCREAMING_SNAKE_CASE_ :     Optional[int]	=			name.replace("visual.ln_post"    ,    "vision_model.post_layernorm"       )
      if "visual.proj" in name:
            SCREAMING_SNAKE_CASE_ :     Optional[int]	=			name.replace("visual.proj"    ,    "visual_projection.weight"       )
      if "text_projection" in name:
            SCREAMING_SNAKE_CASE_ :     str	=			name.replace("text_projection"    ,    "text_projection.weight"       )
      # things on top
      if "prompts_visual_proj" in name:
            SCREAMING_SNAKE_CASE_ :     Union[str, Any]	=			name.replace("prompts_visual_proj"    ,    "prompts_visual_projection"       )
      if "prompts_visual_ln" in name:
            SCREAMING_SNAKE_CASE_ :     List[str]	=			name.replace("prompts_visual_ln"    ,    "prompts_visual_layernorm"       )
      # mit
      if name == "mit.positional_embedding":
            SCREAMING_SNAKE_CASE_ :     List[Any]	=			name.replace("positional"    ,    "position"       )
      if name.startswith("mit.resblocks"       ):
            SCREAMING_SNAKE_CASE_ :     Any	=			name.replace("mit.resblocks"    ,    "mit.encoder.layers"       )
      # prompts generator
      if name.startswith("prompts_generator.norm"       ):
            SCREAMING_SNAKE_CASE_ :     List[Any]	=			name.replace("prompts_generator.norm"    ,    "prompts_generator.layernorm"       )

      return name


def    __lowerCamelCase     ( SCREAMING_SNAKE_CASE_	:      Optional[Any]    ,    SCREAMING_SNAKE_CASE_	:      Any       )				->  Tuple:
      """simple docstring"""

      for key in orig_state_dict.copy().keys():
            SCREAMING_SNAKE_CASE_ :     str	=			orig_state_dict.pop(A__       )

            if "attn.in_proj" in key:
                  SCREAMING_SNAKE_CASE_ :     int	=			key.split("."       )
                  if key.startswith("visual"       ):
                        SCREAMING_SNAKE_CASE_ :     str	=			key_split[3]
                        SCREAMING_SNAKE_CASE_ :     str	=			config.vision_config.hidden_size
                        if "message_attn" in key:
                              if "weight" in key:
                                    SCREAMING_SNAKE_CASE_ :     Optional[int]	=			val[
                                        :dim, :
                                    ]
                                    SCREAMING_SNAKE_CASE_ :     int	=			val[
                                        dim : dim * 2, :
                                    ]
                                    SCREAMING_SNAKE_CASE_ :     List[str]	=			val[
                                        -dim:, :
                                    ]
                              else:
                                    SCREAMING_SNAKE_CASE_ :     str	=			val[
                                        :dim
                                    ]
                                    SCREAMING_SNAKE_CASE_ :     Tuple	=			val[
                                        dim : dim * 2
                                    ]
                                    SCREAMING_SNAKE_CASE_ :     Union[str, Any]	=			val[
                                        -dim:
                                    ]
                        else:
                              if "weight" in key:
                                    SCREAMING_SNAKE_CASE_ :     Union[str, Any]	=			val[
                                        :dim, :
                                    ]
                                    SCREAMING_SNAKE_CASE_ :     List[Any]	=			val[
                                        dim : dim * 2, :
                                    ]
                                    SCREAMING_SNAKE_CASE_ :     Optional[int]	=			val[
                                        -dim:, :
                                    ]
                              else:
                                    SCREAMING_SNAKE_CASE_ :     List[Any]	=			val[:dim]
                                    SCREAMING_SNAKE_CASE_ :     List[str]	=			val[
                                        dim : dim * 2
                                    ]
                                    SCREAMING_SNAKE_CASE_ :     List[str]	=			val[-dim:]
                  elif key.startswith("mit"       ):
                        SCREAMING_SNAKE_CASE_ :     Optional[Any]	=			key_split[2]
                        SCREAMING_SNAKE_CASE_ :     List[Any]	=			config.vision_config.mit_hidden_size
                        if "weight" in key:
                              SCREAMING_SNAKE_CASE_ :     Optional[Any]	=			val[:dim, :]
                              SCREAMING_SNAKE_CASE_ :     Any	=			val[dim : dim * 2, :]
                              SCREAMING_SNAKE_CASE_ :     List[str]	=			val[-dim:, :]
                        else:
                              SCREAMING_SNAKE_CASE_ :     str	=			val[:dim]
                              SCREAMING_SNAKE_CASE_ :     Tuple	=			val[dim : dim * 2]
                              SCREAMING_SNAKE_CASE_ :     Optional[Any]	=			val[-dim:]
                  else:
                        SCREAMING_SNAKE_CASE_ :     Union[str, Any]	=			key_split[2]
                        SCREAMING_SNAKE_CASE_ :     Optional[Any]	=			config.text_config.hidden_size
                        if "weight" in key:
                              SCREAMING_SNAKE_CASE_ :     Dict	=			val[:dim, :]
                              SCREAMING_SNAKE_CASE_ :     List[str]	=			val[
                                  dim : dim * 2, :
                              ]
                              SCREAMING_SNAKE_CASE_ :     Any	=			val[-dim:, :]
                        else:
                              SCREAMING_SNAKE_CASE_ :     List[Any]	=			val[:dim]
                              SCREAMING_SNAKE_CASE_ :     Dict	=			val[
                                  dim : dim * 2
                              ]
                              SCREAMING_SNAKE_CASE_ :     List[str]	=			val[-dim:]
            else:
                  SCREAMING_SNAKE_CASE_ :     Optional[int]	=			rename_key(A__       )
                  if new_key_name in ["visual_projection.weight", "text_projection.weight"]:
                        SCREAMING_SNAKE_CASE_ :     Tuple	=			val.T
                  SCREAMING_SNAKE_CASE_ :     Optional[Any]	=			val

      return orig_state_dict


def    __lowerCamelCase     ( SCREAMING_SNAKE_CASE_	:      Optional[int]       )				->  Dict:
      """simple docstring"""

      if num_frames == 8:
            SCREAMING_SNAKE_CASE_ :     Dict	=			"""eating_spaghetti_8_frames.npy"""
      elif num_frames == 1_6:
            SCREAMING_SNAKE_CASE_ :     Any	=			"""eating_spaghetti.npy"""
      elif num_frames == 3_2:
            SCREAMING_SNAKE_CASE_ :     Tuple	=			"""eating_spaghetti_32_frames.npy"""
      SCREAMING_SNAKE_CASE_ :     Tuple	=			hf_hub_download(
          repo_id="hf-internal-testing/spaghetti-video"    ,    filename=A__    ,    repo_type="dataset"    ,    )
      SCREAMING_SNAKE_CASE_ :     Any	=			np.load(A__       )
      return list(A__       )


def    __lowerCamelCase     ( SCREAMING_SNAKE_CASE_	:      Optional[int]    ,    SCREAMING_SNAKE_CASE_	:      Tuple=None    ,    SCREAMING_SNAKE_CASE_	:      Optional[Any]=False       )				->  Dict:
      """simple docstring"""

      SCREAMING_SNAKE_CASE_ :     int	=			{
          # fully supervised kinetics-400 checkpoints
          """xclip-base-patch32""": """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_8.pth""",
          """xclip-base-patch32-16-frames""": (
              """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_16.pth"""
          ),
          """xclip-base-patch16""": """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_8.pth""",
          """xclip-base-patch16-16-frames""": (
              """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_16.pth"""
          ),
          """xclip-large-patch14""": """https://drive.google.com/u/0/uc?id=1NUOImq0o5DlQTST17iIP3vG7DgmHQuCx&amp;export=download&amp;confirm=t&amp;uuid=b26caedc-88e2-473e-830a-9d158b653cdb""",
          """xclip-large-patch14-16-frames""": """https://drive.google.com/u/0/uc?id=1FOYgnJc097OJ4lGwtRCCydQyVPJEOH7d&amp;export=download&amp;confirm=t&amp;uuid=538fa810-e671-4050-b385-9a623f89804f""",
          # fully supervised kinetics-600 checkpoints
          """xclip-base-patch16-kinetics-600""": (
              """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_8.pth"""
          ),
          """xclip-base-patch16-kinetics-600-16-frames""": (
              """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_16.pth"""
          ),
          """xclip-large-patch14-kinetics-600""": """https://drive.google.com/u/0/uc?id=1FV8C1INuM91sLAN4ImjzePLIlpMSihwV&amp;export=download&amp;confirm=t&amp;uuid=141d4977-4a65-44ae-864f-4b0c19f838be""",
          # few shot
          """xclip-base-patch16-hmdb-2-shot""": (
              """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_2.pth"""
          ),
          """xclip-base-patch16-hmdb-4-shot""": (
              """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_4.pth"""
          ),
          """xclip-base-patch16-hmdb-8-shot""": (
              """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_8.pth"""
          ),
          """xclip-base-patch16-hmdb-16-shot""": (
              """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_16.pth"""
          ),
          """xclip-base-patch16-ucf-2-shot""": (
              """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_2.pth"""
          ),
          """xclip-base-patch16-ucf-4-shot""": (
              """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_4.pth"""
          ),
          """xclip-base-patch16-ucf-8-shot""": (
              """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_8.pth"""
          ),
          """xclip-base-patch16-ucf-16-shot""": (
              """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_16.pth"""
          ),
          # zero shot
          """xclip-base-patch16-zero-shot""": """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/zero.pth""",
      }

      SCREAMING_SNAKE_CASE_ :     Any	=			model_to_url[model_name]
      SCREAMING_SNAKE_CASE_ :     Optional[int]	=			8
      if "16-frames" in model_name:
            SCREAMING_SNAKE_CASE_ :     Dict	=			1_6
      elif "shot" in model_name:
            SCREAMING_SNAKE_CASE_ :     Tuple	=			3_2

      SCREAMING_SNAKE_CASE_ :     Tuple	=			get_xclip_config(A__    ,    A__       )
      SCREAMING_SNAKE_CASE_ :     Union[str, Any]	=			XCLIPModel(A__       )
      model.eval()

      if "drive" in checkpoint_url:
            SCREAMING_SNAKE_CASE_ :     Tuple	=			"""pytorch_model.bin"""
            gdown.cached_download(A__    ,    A__    ,    quiet=A__       )
            SCREAMING_SNAKE_CASE_ :     Any	=			torch.load(A__    ,    map_location="cpu"       )["""model"""]
      else:
            SCREAMING_SNAKE_CASE_ :     List[str]	=			torch.hub.load_state_dict_from_url(A__       )["""model"""]

      SCREAMING_SNAKE_CASE_ :     List[Any]	=			convert_state_dict(A__    ,    A__       )

      SCREAMING_SNAKE_CASE_ :     Optional[Any]	=			XCLIPModel(A__       )
      SCREAMING_SNAKE_CASE_ :     str	=			model.load_state_dict(A__    ,    strict=A__       )
      assert missing_keys == ["text_model.embeddings.position_ids", "vision_model.embeddings.position_ids"]
      model.eval()

      SCREAMING_SNAKE_CASE_ :     str	=			3_3_6 if model_name == """xclip-large-patch14-16-frames""" else 2_2_4
      SCREAMING_SNAKE_CASE_ :     Dict	=			VideoMAEImageProcessor(size=A__       )
      SCREAMING_SNAKE_CASE_ :     Union[str, Any]	=			CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32"       )
      SCREAMING_SNAKE_CASE_ :     Any	=			CLIPTokenizerFast.from_pretrained("openai/clip-vit-base-patch32"       )
      SCREAMING_SNAKE_CASE_ :     Any	=			XCLIPProcessor(image_processor=A__    ,    tokenizer=A__       )

      SCREAMING_SNAKE_CASE_ :     Dict	=			prepare_video(A__       )
      SCREAMING_SNAKE_CASE_ :     Union[str, Any]	=			processor(
          text=["playing sports", "eating spaghetti", "go shopping"]    ,    videos=A__    ,    return_tensors="pt"    ,    padding=A__       )

      print("Shape of pixel values:"    ,    inputs.pixel_values.shape       )

      with torch.no_grad():
            SCREAMING_SNAKE_CASE_ :     Union[str, Any]	=			model(**A__       )

      # Verify outputs
      SCREAMING_SNAKE_CASE_ :     Union[str, Any]	=			outputs.logits_per_video
      SCREAMING_SNAKE_CASE_ :     Union[str, Any]	=			logits_per_video.softmax(dim=1       )
      print("Probs:"    ,    A__       )
      # kinetics-400
      if model_name == "xclip-base-patch32":
            SCREAMING_SNAKE_CASE_ :     Union[str, Any]	=			torch.tensor([[0.0019, 0.9951, 0.0030]]       )
      elif model_name == "xclip-base-patch32-16-frames":
            SCREAMING_SNAKE_CASE_ :     Tuple	=			torch.tensor([[7.0_9_9_9E-0_4, 9.9_8_8_3E-0_1, 4.5_5_8_0E-0_4]]       )
      elif model_name == "xclip-base-patch16":
            SCREAMING_SNAKE_CASE_ :     List[str]	=			torch.tensor([[0.0083, 0.9681, 0.0236]]       )
      elif model_name == "xclip-base-patch16-16-frames":
            SCREAMING_SNAKE_CASE_ :     Optional[Any]	=			torch.tensor([[7.6_9_3_7E-0_4, 9.9_7_2_8E-0_1, 1.9_4_7_3E-0_3]]       )
      elif model_name == "xclip-large-patch14":
            SCREAMING_SNAKE_CASE_ :     List[str]	=			torch.tensor([[0.0062, 0.9864, 0.0075]]       )
      elif model_name == "xclip-large-patch14-16-frames":
            SCREAMING_SNAKE_CASE_ :     Dict	=			torch.tensor([[3.3_8_7_7E-0_4, 9.9_9_3_7E-0_1, 2.8_8_8_8E-0_4]]       )
      # kinetics-600
      elif model_name == "xclip-base-patch16-kinetics-600":
            SCREAMING_SNAKE_CASE_ :     Optional[Any]	=			torch.tensor([[0.0555, 0.8914, 0.0531]]       )
      elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
            SCREAMING_SNAKE_CASE_ :     Optional[int]	=			torch.tensor([[3.8_5_5_4E-0_4, 9.9_9_2_9E-0_1, 3.2_7_5_4E-0_4]]       )
      elif model_name == "xclip-large-patch14-kinetics-600":
            SCREAMING_SNAKE_CASE_ :     List[str]	=			torch.tensor([[0.0036, 0.9920, 0.0045]]       )
      # few shot
      elif model_name == "xclip-base-patch16-hmdb-2-shot":
            SCREAMING_SNAKE_CASE_ :     Optional[int]	=			torch.tensor([[7.1_8_9_0E-0_6, 9.9_9_9_4E-0_1, 5.6_5_5_9E-0_5]]       )
      elif model_name == "xclip-base-patch16-hmdb-4-shot":
            SCREAMING_SNAKE_CASE_ :     Union[str, Any]	=			torch.tensor([[1.0_3_2_0E-0_5, 9.9_9_9_3E-0_1, 6.2_4_3_5E-0_5]]       )
      elif model_name == "xclip-base-patch16-hmdb-8-shot":
            SCREAMING_SNAKE_CASE_ :     Any	=			torch.tensor([[4.1_3_7_7E-0_6, 9.9_9_9_0E-0_1, 9.8_3_8_6E-0_5]]       )
      elif model_name == "xclip-base-patch16-hmdb-16-shot":
            SCREAMING_SNAKE_CASE_ :     Tuple	=			torch.tensor([[4.1_3_4_7E-0_5, 9.9_9_6_2E-0_1, 3.3_4_1_1E-0_4]]       )
      elif model_name == "xclip-base-patch16-ucf-2-shot":
            SCREAMING_SNAKE_CASE_ :     str	=			torch.tensor([[8.5_8_5_7E-0_5, 9.9_9_2_8E-0_1, 6.3_2_9_1E-0_4]]       )
      elif model_name == "xclip-base-patch16-ucf-4-shot":
            SCREAMING_SNAKE_CASE_ :     Optional[Any]	=			torch.tensor([[8.5_8_5_7E-0_5, 9.9_9_2_8E-0_1, 6.3_2_9_1E-0_4]]       )
      elif model_name == "xclip-base-patch16-ucf-8-shot":
            SCREAMING_SNAKE_CASE_ :     Union[str, Any]	=			torch.tensor([[0.0027, 0.9904, 0.0070]]       )
      elif model_name == "xclip-base-patch16-ucf-16-shot":
            SCREAMING_SNAKE_CASE_ :     Optional[int]	=			torch.tensor([[9.8_2_1_9E-0_4, 9.9_5_9_3E-0_1, 3.0_8_6_3E-0_3]]       )
      # zero shot
      elif model_name == "xclip-base-patch16-zero-shot":
            SCREAMING_SNAKE_CASE_ :     Tuple	=			torch.tensor([[3.5_0_8_2E-0_4, 9.9_7_8_5E-0_1, 1.7_9_6_6E-0_3]]       )
      else:
            raise ValueError(F"Model name {model_name} not supported"       )
      assert torch.allclose(A__    ,    A__    ,    atol=1E-3       )
      print("Looks ok!"       )

      if pytorch_dump_folder_path is not None:
            print(F"Saving model {model_name} to {pytorch_dump_folder_path}"       )
            model.save_pretrained(A__       )

      if push_to_hub:
            print("Pushing model, processor and slow tokenizer files to the hub..."       )
            model.push_to_hub(A__    ,    organization="nielsr"       )
            processor.push_to_hub(A__    ,    organization="nielsr"       )
            slow_tokenizer.push_to_hub(A__    ,    organization="nielsr"       )


if __name__ == "__main__":
      snake_case_	    =			argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          '--model_name',
          default='xclip-base-patch32',
          type=str,
          help='Name of the model.',
      )
      parser.add_argument(
          '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
      )
      parser.add_argument(
          '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
      )

      snake_case_	    =			parser.parse_args()
      convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


snake_case__				:	Dict					       =					logging.getLogger(__name__)


def    __lowerCamelCase   (      )     ->    Any:
   lowerCamelCase_						:   str		= argparse.ArgumentParser(
       description="""Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset."""					)
   parser.add_argument(
       """--dataset_name"""						,						type=A__						,						default="""wikitext"""						,						help="""Name of the training. Explore datasets at: hf.co/datasets."""						,						)
   parser.add_argument(
       """--dataset_config"""						,						type=A__						,						default="""wikitext-103-raw-v1"""						,						help="""Configuration name of the dataset."""					)
   parser.add_argument(
       """--tokenizer_name_or_path"""						,						type=A__						,						default="""sayakpaul/unigram-tokenizer-wikitext"""						,						help="""Tokenizer identifier. Can be a local filepath or a Hub identifier."""						,						)
   parser.add_argument(
       """--shard_size"""						,						type=A__						,						default=1000						,						help="""Number of entries to go in a single shard."""						,						)
   parser.add_argument("""--split"""						,						type=A__						,						default="""train"""						,						choices=["""train""", """test""", """validation"""]					)
   parser.add_argument(
       """--limit"""						,						default=A__						,						type=A__						,						help="""Limit the number of shards (used for debugging)."""						,						)
   parser.add_argument(
       """--max_length"""						,						type=A__						,						default=512						,						help="""Maximum sequence length. For training on TPUs, it helps to have a maximum"""
       """ sequence length that is a multiple of 8."""						,						)
   parser.add_argument(
       """--output_dir"""						,						default="""tf-tpu"""						,						type=A__						,						help="""Output directory where the TFRecord shards will be saved. If the"""
       """ path is appended with `gs://` ('gs://tf-tpu', for example) then the TFRecord"""
       """ shards will be directly saved to a Google Cloud Storage bucket."""						,						)

   lowerCamelCase_						:   Optional[int]		= parser.parse_args()
   return args


def    __lowerCamelCase   (      A__						: Tuple					)     ->    str:
   def fn(A__						: Dict					):
      return tokenizer(examples["""text"""]					)

   return fn


def    __lowerCamelCase   (      A__						: Union[str, Any]					)     ->    Union[str, Any]:
   lowerCamelCase_						:   Dict		= []
   for i in range(len(tokenized_data["""input_ids"""]					)					):
      lowerCamelCase_						:   List[str]		= {
          """input_ids""": tf.train.Feature(intaa_list=tf.train.IntaaList(value=tokenized_data["""input_ids"""][i]					)					),
          """attention_mask""": tf.train.Feature(
              intaa_list=tf.train.IntaaList(value=tokenized_data["""attention_mask"""][i]					)					),
      }
      lowerCamelCase_						:   Optional[int]		= tf.train.Features(feature=A__					)
      lowerCamelCase_						:   Tuple		= tf.train.Example(features=A__					)
      lowerCamelCase_						:   Optional[Any]		= example.SerializeToString()
      records.append(A__					)
   return records


def    __lowerCamelCase   (      A__						: Tuple					)     ->    Union[str, Any]:
   lowerCamelCase_						:   int		= datasets.load_dataset(args.dataset_name						,						args.dataset_config						,						split=args.split					)

   if args.limit is not None:
      lowerCamelCase_						:   str		= min(len(A__					)						,						args.limit					)
      lowerCamelCase_						:   Dict		= dataset.select(range(A__					)					)
      print(f'''Limiting the dataset to {args.limit} entries.'''					)

   lowerCamelCase_						:   str		= AutoTokenizer.from_pretrained(args.tokenizer_name_or_path					)

   # Handle output directory creation.
   # For serializing into a Google Cloud Storage Bucket, one needs to first
   # create a bucket.
   if "gs" not in args.output_dir:
      if not os.path.exists(args.output_dir					):
         os.makedirs(args.output_dir					)
      lowerCamelCase_						:   Optional[int]		= os.path.join(args.output_dir						,						args.split					)
      if not os.path.exists(A__					):
         os.makedirs(A__					)
   else:
      lowerCamelCase_						:   Any		= os.path.join(args.output_dir						,						args.split					)

   # Tokenize the whole dataset at once.
   lowerCamelCase_						:   List[str]		= tokenize_function(A__					)
   lowerCamelCase_						:   List[Any]		= dataset.map(A__						,						batched=A__						,						num_proc=4						,						remove_columns=["""text"""]					)

   # We need to concatenate all our texts together, and then split the result
   # into chunks of a fixed size, which we will call block_size. To do this, we
   # will use the map method again, with the option batched=True. When we use batched=True,
   # the function we pass to map() will be passed multiple inputs at once, allowing us
   # to group them into more or fewer examples than we had in the input.
   # This allows us to create our new fixed-length samples. The advantage of this
   # method is that we don't lose a whole lot of content from the dataset compared to the
   # case where we simply tokenize with a pre-defined max_length.

   def group_texts(A__						: Union[str, Any]					):
      # Concatenate all texts.
      lowerCamelCase_						:   Optional[int]		= {k: sum(examples[k]						,						[]					) for k in examples.keys()}
      lowerCamelCase_						:   Union[str, Any]		= len(concatenated_examples[list(examples.keys()					)[0]]					)
      # We drop the small remainder, though you could add padding instead if the model supports it
      # In this, as in all things, we advise you to follow your heart 🫀
      lowerCamelCase_						:   Union[str, Any]		= (total_length // args.max_length) * args.max_length
      # Split by chunks of max_len.
      lowerCamelCase_						:   List[str]		= {
          k: [t[i : i + args.max_length] for i in range(0						,						A__						,						args.max_length					)]
          for k, t in concatenated_examples.items()
      }
      return result

   lowerCamelCase_						:   Optional[int]		= dataset_tokenized.map(A__						,						batched=A__						,						batch_size=1000						,						num_proc=4					)

   lowerCamelCase_						:   List[str]		= 0
   lowerCamelCase_						:   Tuple		= 0
   for shard in range(0						,						len(A__					)						,						args.shard_size					):
      lowerCamelCase_						:   Union[str, Any]		= grouped_dataset[shard : shard + args.shard_size]
      lowerCamelCase_						:   Optional[int]		= len(dataset_snapshot["""input_ids"""]					)
      lowerCamelCase_						:   List[str]		= os.path.join(A__						,						f'''dataset-{shard_count}-{records_containing}.tfrecord'''					)
      lowerCamelCase_						:   Optional[int]		= get_serialized_examples(A__					)

      with tf.io.TFRecordWriter(A__					) as out_file:
         for i in range(len(A__					)					):
            lowerCamelCase_						:   Dict		= serialized_examples[i]
            out_file.write(A__					)
         print("""Wrote file {} containing {} records""".format(A__						,						A__					)					)

      shard_count += 1
      total_records += records_containing

   with open(f'''split-{args.split}-records-count.txt'''						,						"""w"""					) as f:
      print(f'''Total {args.split} records: {total_records}'''						,						file=A__					)


if __name__ == "__main__":
      snake_case__				:	str					       =					parse_args()
      main(args)





