


import os

import jsonlines
import numpy as np
from tqdm import tqdm


lowerCamelCase__							:						Optional[Any]							     =   2_0_4_8
lowerCamelCase__							:						Tuple							     =   4_0_9_6
lowerCamelCase__							:						Any							     =   4_2
lowerCamelCase__							:						str							     =   os.environ.pop("""PROCESS_TRAIN""", """false""")
lowerCamelCase__							:						Any							     =   {"""null""": 0, """short""": 1, """long""": 2, """yes""": 3, """no""": 4}







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase				)    ->				Union[str, Any]:
  def choose_first(__lowerCAmelCase							,   __lowerCAmelCase=False				):
    assert isinstance(__lowerCAmelCase							,   __lowerCAmelCase				)
    if len(__lowerCAmelCase				) == 1:
      snake_case__						= answer[0]
      return {k: [answer[k]] for k in answer} if is_long_answer else answer
    for a in answer:
      if is_long_answer:
        snake_case__						= {k: [a[k]] for k in a}
      if len(a['''start_token''']				) > 0:
        break
    return a

  snake_case__						= {'''id''': example['''id''']}
  snake_case__						= example['''annotations''']
  snake_case__						= annotation['''yes_no_answer''']
  if 0 in yes_no_answer or 1 in yes_no_answer:
    snake_case__						= ['''yes'''] if 1 in yes_no_answer else ['''no''']
    snake_case__						= snake_case__						= []
    snake_case__						= snake_case__						= []
    snake_case__						= ['''<cls>''']
  else:
    snake_case__						= ['''short''']
    snake_case__						= choose_first(annotation['''short_answers''']				)
    if len(out['''start_token''']				) == 0:
      # answer will be long if short is not available
      snake_case__						= ['''long''']
      snake_case__						= choose_first(annotation['''long_answer''']							,   is_long_answer=__lowerCAmelCase				)
      snake_case__						= []
    answer.update(__lowerCAmelCase				)

  # disregard some samples
  if len(answer['''start_token''']				) > 1 or answer["start_token"] == answer["end_token"]:
    snake_case__						= True
  else:
    snake_case__						= False

  snake_case__						= ['''start_token''', '''end_token''', '''start_byte''', '''end_byte''', '''text''']
  if not all(isinstance(answer[k]							,   __lowerCAmelCase				) for k in cols				):
    raise ValueError('''Issue in ID'''							,   example['''id''']				)

  return answer







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase=False				)    ->				Dict:
  snake_case__						= _get_single_answer(__lowerCAmelCase				)
  # bytes are of no use
  del answer["start_byte"]
  del answer["end_byte"]

  # handle yes_no answers explicitly
  if answer["category"][0] in ["yes", "no"]:  # category is list with one element
    snake_case__						= example['''document''']['''tokens''']
    snake_case__						= []
    for i in range(len(doc['''token''']				)				):
      if not doc["is_html"][i]:
        context.append(doc['''token'''][i]				)
    return {
        "context": " ".join(__lowerCAmelCase				),
        "answer": {
            "start_token": -100,  # ignore index in cross-entropy
            "end_token": -100,  # ignore index in cross-entropy
            "category": answer["category"],
            "span": answer["category"],  # extra
        },
    }

  # later, help in removing all no answers
  if answer["start_token"] == [-1]:
    return {
        "context": "None",
        "answer": {
            "start_token": -1,
            "end_token": -1,
            "category": "null",
            "span": "None",  # extra
        },
    }

  # handling normal samples

  snake_case__						= ['''start_token''', '''end_token''']
  answer.update({k: answer[k][0] if len(answer[k]				) > 0 else answer[k] for k in cols}				)  # e.g. [10] == 10

  snake_case__						= example['''document''']['''tokens''']
  snake_case__						= answer['''start_token''']
  snake_case__						= answer['''end_token''']

  snake_case__						= []
  for i in range(len(doc['''token''']				)				):
    if not doc["is_html"][i]:
      context.append(doc['''token'''][i]				)
    else:
      if answer["start_token"] > i:
        start_token -= 1
      if answer["end_token"] > i:
        end_token -= 1
  snake_case__						= ''' '''.join(context[start_token:end_token]				)

  # checking above code
  if assertion:
    snake_case__						= doc['''is_html'''][answer['''start_token'''] : answer['''end_token''']]
    snake_case__						= doc['''token'''][answer['''start_token'''] : answer['''end_token''']]
    snake_case__						= ''' '''.join([old[i] for i in range(len(__lowerCAmelCase				)				) if not is_html[i]]				)
    if new != old:
      print('''ID:'''							,   example['''id''']				)
      print('''New:'''							,   __lowerCAmelCase							,   end='''\n'''				)
      print('''Old:'''							,   __lowerCAmelCase							,   end='''\n\n'''				)

  return {
      "context": " ".join(__lowerCAmelCase				),
      "answer": {
          "start_token": start_token,
          "end_token": end_token - 1,  # this makes it inclusive
          "category": answer["category"],  # either long or short
          "span": new,  # extra
      },
  }







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase							,   __lowerCAmelCase=2048							,   __lowerCAmelCase=4096							,   __lowerCAmelCase=True				)    ->				Dict:
  # overlap will be of doc_stride - q_len

  snake_case__						= get_context_and_ans(__lowerCAmelCase							,   assertion=__lowerCAmelCase				)
  snake_case__						= out['''answer''']

  # later, removing these samples
  if answer["start_token"] == -1:
    return {
        "example_id": example["id"],
        "input_ids": [[-1]],
        "labels": {
            "start_token": [-1],
            "end_token": [-1],
            "category": ["null"],
        },
    }

  snake_case__						= tokenizer(example['''question''']['''text''']							,   out['''context''']				).input_ids
  snake_case__						= input_ids.index(tokenizer.sep_token_id				) + 1

  # return yes/no
  if answer["category"][0] in ["yes", "no"]:  # category is list with one element
    snake_case__						= []
    snake_case__						= []
    snake_case__						= input_ids[:q_len]
    snake_case__						= range(__lowerCAmelCase							,   len(__lowerCAmelCase				)							,   max_length - doc_stride				)
    for i in doc_start_indices:
      snake_case__						= i + max_length - q_len
      snake_case__						= input_ids[i:end_index]
      inputs.append(q_indices + slice				)
      category.append(answer['''category'''][0]				)
      if slice[-1] == tokenizer.sep_token_id:
        break

    return {
        "example_id": example["id"],
        "input_ids": inputs,
        "labels": {
            "start_token": [-100] * len(__lowerCAmelCase				),
            "end_token": [-100] * len(__lowerCAmelCase				),
            "category": category,
        },
    }

  snake_case__						= out['''context'''].split()
  snake_case__						= splitted_context[answer['''end_token''']]
  snake_case__						= len(
      tokenizer(
          ''' '''.join(splitted_context[: answer['''start_token''']]				)							,   add_special_tokens=__lowerCAmelCase							,   ).input_ids				)
  snake_case__						= len(
      tokenizer(''' '''.join(splitted_context[: answer['''end_token''']]				)							,   add_special_tokens=__lowerCAmelCase				).input_ids				)

  answer["start_token"] += q_len
  answer["end_token"] += q_len

  # fixing end token
  snake_case__						= len(tokenizer(__lowerCAmelCase							,   add_special_tokens=__lowerCAmelCase				).input_ids				)
  if num_sub_tokens > 1:
    answer["end_token"] += num_sub_tokens - 1

  snake_case__						= input_ids[answer['''start_token'''] : answer['''end_token'''] + 1]  # right & left are inclusive
  snake_case__						= answer['''start_token''']
  snake_case__						= answer['''end_token''']

  if assertion:
    snake_case__						= tokenizer.decode(__lowerCAmelCase				)
    if answer["span"] != new:
      print('''ISSUE IN TOKENIZATION'''				)
      print('''OLD:'''							,   answer['''span''']				)
      print('''NEW:'''							,   __lowerCAmelCase							,   end='''\n\n'''				)

  if len(__lowerCAmelCase				) <= max_length:
    return {
        "example_id": example["id"],
        "input_ids": [input_ids],
        "labels": {
            "start_token": [answer["start_token"]],
            "end_token": [answer["end_token"]],
            "category": answer["category"],
        },
    }

  snake_case__						= input_ids[:q_len]
  snake_case__						= range(__lowerCAmelCase							,   len(__lowerCAmelCase				)							,   max_length - doc_stride				)

  snake_case__						= []
  snake_case__						= []
  snake_case__						= []
  snake_case__						= []  # null, yes, no, long, short
  for i in doc_start_indices:
    snake_case__						= i + max_length - q_len
    snake_case__						= input_ids[i:end_index]
    inputs.append(q_indices + slice				)
    assert len(inputs[-1]				) <= max_length, "Issue in truncating length"

    if start_token >= i and end_token <= end_index - 1:
      snake_case__						= start_token - i + q_len
      snake_case__						= end_token - i + q_len
      answers_category.append(answer['''category'''][0]				)  # ["short"] -> "short"
    else:
      snake_case__						= -100
      snake_case__						= -100
      answers_category.append('''null'''				)
    snake_case__						= inputs[-1][start_token : end_token + 1]

    answers_start_token.append(__lowerCAmelCase				)
    answers_end_token.append(__lowerCAmelCase				)
    if assertion:
      if new != old and new != [tokenizer.cls_token_id]:
        print('''ISSUE in strided for ID:'''							,   example['''id''']				)
        print('''New:'''							,   tokenizer.decode(__lowerCAmelCase				)				)
        print('''Old:'''							,   tokenizer.decode(__lowerCAmelCase				)							,   end='''\n\n'''				)
    if slice[-1] == tokenizer.sep_token_id:
      break

  return {
      "example_id": example["id"],
      "input_ids": inputs,
      "labels": {
          "start_token": answers_start_token,
          "end_token": answers_end_token,
          "category": answers_category,
      },
  }







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase							,   __lowerCAmelCase=2048							,   __lowerCAmelCase=4096							,   __lowerCAmelCase=False				)    ->				List[Any]:
  snake_case__						= get_strided_contexts_and_ans(
      __lowerCAmelCase							,   __lowerCAmelCase							,   doc_stride=__lowerCAmelCase							,   max_length=__lowerCAmelCase							,   assertion=__lowerCAmelCase							,   )

  return example







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase				)    ->				List[str]:
  with jsonlines.open(__lowerCAmelCase							,   '''a'''				) as writer:
    for example in tqdm(__lowerCAmelCase							,   total=len(__lowerCAmelCase				)							,   desc='''Saving samples ... '''				):
      snake_case__						= example['''labels''']
      for ids, start, end, cat in zip(
          example['''input_ids''']							,   labels['''start_token''']							,   labels['''end_token''']							,   labels['''category''']							,   ):
        if start == -1 and end == -1:
          continue  # leave waste samples with no answer
        if cat == "null" and np.random.rand() < 0.6:
          continue  # removing 50 % samples
        writer.write(
            {
                '''input_ids''': ids,
                '''start_token''': start,
                '''end_token''': end,
                '''category''': CATEGORY_MAPPING[cat],
            }				)


if __name__ == "__main__":
   from datasets import load_dataset

   from transformers import BigBirdTokenizer

   lowerCamelCase__							:						Union[str, Any]							     =   load_dataset("""natural_questions""")
   lowerCamelCase__							:						str							     =   BigBirdTokenizer.from_pretrained("""google/bigbird-roberta-base""")

   lowerCamelCase__							:						List[str]							     =   data["""train""" if PROCESS_TRAIN == """true""" else """validation"""]

   lowerCamelCase__							:						List[str]							     =   {
       """tokenizer""": tokenizer,
       """doc_stride""": DOC_STRIDE,
       """max_length""": MAX_LENGTH,
       """assertion""": False,
   }
   lowerCamelCase__							:						List[Any]							     =   data.map(prepare_inputs, fn_kwargs=fn_kwargs)
   lowerCamelCase__							:						Any							     =   data.remove_columns(["""annotations""", """document""", """id""", """question"""])
   print(data)

   np.random.seed(SEED)
   lowerCamelCase__							:						Tuple							     =   """nq-training.jsonl""" if PROCESS_TRAIN == """true""" else """nq-validation.jsonl"""
   save_to_disk(data, file_name=cache_file_name)





import argparse
import json
from pathlib import Path

import requests
import timm
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import DeiTImageProcessor, ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
lowerCamelCase__							:						int							     =   logging.get_logger(__name__)







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase=False				)    ->				int:
  snake_case__						= []
  for i in range(config.num_hidden_layers				):
    # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
    rename_keys.append((F"""blocks.{i}.norm1.weight""", F"""vit.encoder.layer.{i}.layernorm_before.weight""")				)
    rename_keys.append((F"""blocks.{i}.norm1.bias""", F"""vit.encoder.layer.{i}.layernorm_before.bias""")				)
    rename_keys.append((F"""blocks.{i}.attn.proj.weight""", F"""vit.encoder.layer.{i}.attention.output.dense.weight""")				)
    rename_keys.append((F"""blocks.{i}.attn.proj.bias""", F"""vit.encoder.layer.{i}.attention.output.dense.bias""")				)
    rename_keys.append((F"""blocks.{i}.norm2.weight""", F"""vit.encoder.layer.{i}.layernorm_after.weight""")				)
    rename_keys.append((F"""blocks.{i}.norm2.bias""", F"""vit.encoder.layer.{i}.layernorm_after.bias""")				)
    rename_keys.append((F"""blocks.{i}.mlp.fc1.weight""", F"""vit.encoder.layer.{i}.intermediate.dense.weight""")				)
    rename_keys.append((F"""blocks.{i}.mlp.fc1.bias""", F"""vit.encoder.layer.{i}.intermediate.dense.bias""")				)
    rename_keys.append((F"""blocks.{i}.mlp.fc2.weight""", F"""vit.encoder.layer.{i}.output.dense.weight""")				)
    rename_keys.append((F"""blocks.{i}.mlp.fc2.bias""", F"""vit.encoder.layer.{i}.output.dense.bias""")				)

  # projection layer + position embeddings
  rename_keys.extend(
      [
          ('''cls_token''', '''vit.embeddings.cls_token'''),
          ('''patch_embed.proj.weight''', '''vit.embeddings.patch_embeddings.projection.weight'''),
          ('''patch_embed.proj.bias''', '''vit.embeddings.patch_embeddings.projection.bias'''),
          ('''pos_embed''', '''vit.embeddings.position_embeddings'''),
      ]				)

  if base_model:
    # layernorm + pooler
    rename_keys.extend(
        [
            ('''norm.weight''', '''layernorm.weight'''),
            ('''norm.bias''', '''layernorm.bias'''),
            ('''pre_logits.fc.weight''', '''pooler.dense.weight'''),
            ('''pre_logits.fc.bias''', '''pooler.dense.bias'''),
        ]				)

    # if just the base model, we should remove "vit" from all keys that start with "vit"
    snake_case__						= [(pair[0], pair[1][4:]) if pair[1].startswith('''vit'''				) else pair for pair in rename_keys]
  else:
    # layernorm + classification head
    rename_keys.extend(
        [
            ('''norm.weight''', '''vit.layernorm.weight'''),
            ('''norm.bias''', '''vit.layernorm.bias'''),
            ('''head.weight''', '''classifier.weight'''),
            ('''head.bias''', '''classifier.bias'''),
        ]				)

  return rename_keys







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase							,   __lowerCAmelCase=False				)    ->				Dict:
  for i in range(config.num_hidden_layers				):
    if base_model:
      snake_case__						= ''''''
    else:
      snake_case__						= '''vit.'''
    # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
    snake_case__						= state_dict.pop(F"""blocks.{i}.attn.qkv.weight"""				)
    snake_case__						= state_dict.pop(F"""blocks.{i}.attn.qkv.bias"""				)
    # next, add query, keys and values (in that order) to the state dict
    snake_case__						= in_proj_weight[
        : config.hidden_size, :
    ]
    snake_case__						= in_proj_bias[: config.hidden_size]
    snake_case__						= in_proj_weight[
        config.hidden_size : config.hidden_size * 2, :
    ]
    snake_case__						= in_proj_bias[
        config.hidden_size : config.hidden_size * 2
    ]
    snake_case__						= in_proj_weight[
        -config.hidden_size :, :
    ]
    snake_case__						= in_proj_bias[-config.hidden_size :]







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase				)    ->				Optional[Any]:
  snake_case__						= ['''head.weight''', '''head.bias''']
  for k in ignore_keys:
    state_dict.pop(__lowerCAmelCase							,   __lowerCAmelCase				)







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase							,   __lowerCAmelCase				)    ->				Union[str, Any]:
  snake_case__						= dct.pop(__lowerCAmelCase				)
  snake_case__						= val







def SCREAMING_SNAKE_CASE	(			)    ->				str:
  snake_case__						= '''http://images.cocodataset.org/val2017/000000039769.jpg'''
  snake_case__						= Image.open(requests.get(__lowerCAmelCase							,   stream=__lowerCAmelCase				).raw				)
  return im







@torch.no_grad()
def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase				)    ->				Dict:
  snake_case__						= ViTConfig()
  snake_case__						= False
  # dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
  if vit_name[-5:] == "in21k":
    snake_case__						= True
    snake_case__						= int(vit_name[-12:-10]				)
    snake_case__						= int(vit_name[-9:-6]				)
  else:
    snake_case__						= 1000
    snake_case__						= '''huggingface/label-files'''
    snake_case__						= '''imagenet-1k-id2label.json'''
    snake_case__						= json.load(open(hf_hub_download(__lowerCAmelCase							,   __lowerCAmelCase							,   repo_type='''dataset'''				)							,   '''r'''				)				)
    snake_case__						= {int(__lowerCAmelCase				): v for k, v in idalabel.items()}
    snake_case__						= idalabel
    snake_case__						= {v: k for k, v in idalabel.items()}
    snake_case__						= int(vit_name[-6:-4]				)
    snake_case__						= int(vit_name[-3:]				)
  # size of the architecture
  if "deit" in vit_name:
    if vit_name[9:].startswith('''tiny'''				):
      snake_case__						= 192
      snake_case__						= 768
      snake_case__						= 12
      snake_case__						= 3
    elif vit_name[9:].startswith('''small'''				):
      snake_case__						= 384
      snake_case__						= 1536
      snake_case__						= 12
      snake_case__						= 6
    else:
      pass
  else:
    if vit_name[4:].startswith('''small'''				):
      snake_case__						= 768
      snake_case__						= 2304
      snake_case__						= 8
      snake_case__						= 8
    elif vit_name[4:].startswith('''base'''				):
      pass
    elif vit_name[4:].startswith('''large'''				):
      snake_case__						= 1024
      snake_case__						= 4096
      snake_case__						= 24
      snake_case__						= 16
    elif vit_name[4:].startswith('''huge'''				):
      snake_case__						= 1280
      snake_case__						= 5120
      snake_case__						= 32
      snake_case__						= 16

    # load original model from timm
  snake_case__						= timm.create_model(__lowerCAmelCase							,   pretrained=__lowerCAmelCase				)
  timm_model.eval()

  # load state_dict of original model, remove and rename some keys
  snake_case__						= timm_model.state_dict()
  if base_model:
    remove_classification_head_(__lowerCAmelCase				)
  snake_case__						= create_rename_keys(__lowerCAmelCase							,   __lowerCAmelCase				)
  for src, dest in rename_keys:
    rename_key(__lowerCAmelCase							,   __lowerCAmelCase							,   __lowerCAmelCase				)
  read_in_q_k_v(__lowerCAmelCase							,   __lowerCAmelCase							,   __lowerCAmelCase				)

  # load HuggingFace model
  if vit_name[-5:] == "in21k":
    snake_case__						= ViTModel(__lowerCAmelCase				).eval()
  else:
    snake_case__						= ViTForImageClassification(__lowerCAmelCase				).eval()
  model.load_state_dict(__lowerCAmelCase				)

  # Check outputs on an image, prepared by ViTImageProcessor/DeiTImageProcessor
  if "deit" in vit_name:
    snake_case__						= DeiTImageProcessor(size=config.image_size				)
  else:
    snake_case__						= ViTImageProcessor(size=config.image_size				)
  snake_case__						= image_processor(images=prepare_img()							,   return_tensors='''pt'''				)
  snake_case__						= encoding['''pixel_values''']
  snake_case__						= model(__lowerCAmelCase				)

  if base_model:
    snake_case__						= timm_model.forward_features(__lowerCAmelCase				)
    assert timm_pooled_output.shape == outputs.pooler_output.shape
    assert torch.allclose(__lowerCAmelCase							,   outputs.pooler_output							,   atol=1e-3				)
  else:
    snake_case__						= timm_model(__lowerCAmelCase				)
    assert timm_logits.shape == outputs.logits.shape
    assert torch.allclose(__lowerCAmelCase							,   outputs.logits							,   atol=1e-3				)

  Path(__lowerCAmelCase				).mkdir(exist_ok=__lowerCAmelCase				)
  print(F"""Saving model {vit_name} to {pytorch_dump_folder_path}"""				)
  model.save_pretrained(__lowerCAmelCase				)
  print(F"""Saving image processor to {pytorch_dump_folder_path}"""				)
  image_processor.save_pretrained(__lowerCAmelCase				)


if __name__ == "__main__":
   lowerCamelCase__							:						Optional[Any]							     =   argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       """--vit_name""",
       default="""vit_base_patch16_224""",
       type=str,
       help="""Name of the ViT timm model you'd like to convert.""",
   )
   parser.add_argument(
       """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
   )

   lowerCamelCase__							:						str							     =   parser.parse_args()
   convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path)

