


'''simple docstring'''



import os

import jsonlines
import numpy as np
from tqdm import tqdm


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




def      _a     (						_lowercase :     Any  ):





     '''simple docstring'''






     def choose_first(_lowercase :     str      , _lowercase :     Tuple=False  ):
          assert isinstance(_lowercase      , _lowercase  )
          if len(_lowercase  ) == 1:
               __UpperCAmelCase     :			Dict  							=      answer[0]
               return {k: [answer[k]] for k in answer} if is_long_answer else answer
          for a in answer:
               if is_long_answer:
                    __UpperCAmelCase     :			Dict  							=      {k: [a[k]] for k in a}
               if len(a['''start_token''']  ) > 0:
                    break
          return a

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

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

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

     return answer




def      _a     (						_lowercase :     List[str]      , _lowercase :     Union[str, Any]=False  ):





     '''simple docstring'''






     __UpperCAmelCase     :			Dict  							=      _get_single_answer(_lowercase  )
     # 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
          __UpperCAmelCase     :			List[Any]  							=      example['''document''']['''tokens''']
          __UpperCAmelCase     :			Optional[int]  							=      []
          for i in range(len(doc['''token''']  )  ):
               if not doc["is_html"][i]:
                    context.append(doc['''token'''][i]  )
          return {
              "context": " ".join(_lowercase  ),
              "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

     __UpperCAmelCase     :			Optional[Any]  							=      ['''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

     __UpperCAmelCase     :			int  							=      example['''document''']['''tokens''']
     __UpperCAmelCase     :			List[Any]  							=      answer['''start_token''']
     __UpperCAmelCase     :			str  							=      answer['''end_token''']

     __UpperCAmelCase     :			List[str]  							=      []
     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
     __UpperCAmelCase     :			str  							=      ''' '''.join(context[start_token:end_token]  )

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

     return {
         "context": " ".join(_lowercase  ),
         "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      _a     (						_lowercase :     Optional[Any]      , _lowercase :     str      , _lowercase :     Any=2048      , _lowercase :     str=4096      , _lowercase :     int=True  ):





     '''simple docstring'''






     __UpperCAmelCase     :			Tuple  							=      get_context_and_ans(_lowercase      , assertion=_lowercase  )
     __UpperCAmelCase     :			Tuple  							=      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"],
              },
          }

     __UpperCAmelCase     :			Tuple  							=      tokenizer(example['''question''']['''text''']      , out['''context''']  ).input_ids
     __UpperCAmelCase     :			Dict  							=      input_ids.index(tokenizer.sep_token_id  ) + 1

     # return yes/no
     if answer["category"][0] in ["yes", "no"]:  # category is list with one element
          __UpperCAmelCase     :			Optional[int]  							=      []
          __UpperCAmelCase     :			List[Any]  							=      []
          __UpperCAmelCase     :			List[Any]  							=      input_ids[:q_len]
          __UpperCAmelCase     :			Any  							=      range(_lowercase      , len(_lowercase  )      , max_length - doc_stride  )
          for i in doc_start_indices:
               __UpperCAmelCase     :			Dict  							=      i + max_length - q_len
               __UpperCAmelCase     :			int  							=      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(_lowercase  ),
                  "end_token": [-100] * len(_lowercase  ),
                  "category": category,
              },
          }

     __UpperCAmelCase     :			Union[str, Any]  							=      out['''context'''].split()
     __UpperCAmelCase     :			List[str]  							=      splitted_context[answer['''end_token''']]
     __UpperCAmelCase     :			Tuple  							=      len(
         tokenizer(
             ''' '''.join(splitted_context[: answer['''start_token''']]  )      , add_special_tokens=_lowercase      , ).input_ids  )
     __UpperCAmelCase     :			Optional[int]  							=      len(
         tokenizer(''' '''.join(splitted_context[: answer['''end_token''']]  )      , add_special_tokens=_lowercase  ).input_ids  )

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

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

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

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

     if len(_lowercase  ) <= 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"],
              },
          }

     __UpperCAmelCase     :			List[Any]  							=      input_ids[:q_len]
     __UpperCAmelCase     :			int  							=      range(_lowercase      , len(_lowercase  )      , max_length - doc_stride  )

     __UpperCAmelCase     :			List[str]  							=      []
     __UpperCAmelCase     :			Union[str, Any]  							=      []
     __UpperCAmelCase     :			Union[str, Any]  							=      []
     __UpperCAmelCase     :			Optional[int]  							=      []  # null, yes, no, long, short
     for i in doc_start_indices:
          __UpperCAmelCase     :			str  							=      i + max_length - q_len
          __UpperCAmelCase     :			Dict  							=      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:
               __UpperCAmelCase     :			List[Any]  							=      start_token - i + q_len
               __UpperCAmelCase     :			List[str]  							=      end_token - i + q_len
               answers_category.append(answer['''category'''][0]  )  # ["short"] -> "short"
          else:
               __UpperCAmelCase     :			Tuple  							=      -100
               __UpperCAmelCase     :			List[Any]  							=      -100
               answers_category.append('''null'''  )
          __UpperCAmelCase     :			Any  							=      inputs[-1][start_token : end_token + 1]

          answers_start_token.append(_lowercase  )
          answers_end_token.append(_lowercase  )
          if assertion:
               if new != old and new != [tokenizer.cls_token_id]:
                    print('''ISSUE in strided for ID:'''      , example['''id''']  )
                    print('''New:'''      , tokenizer.decode(_lowercase  )  )
                    print('''Old:'''      , tokenizer.decode(_lowercase  )      , 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      _a     (						_lowercase :     Dict      , _lowercase :     Optional[int]      , _lowercase :     Dict=2048      , _lowercase :     str=4096      , _lowercase :     Optional[int]=False  ):





     '''simple docstring'''






     __UpperCAmelCase     :			Optional[int]  							=      get_strided_contexts_and_ans(
         _lowercase      , _lowercase      , doc_stride=_lowercase      , max_length=_lowercase      , assertion=_lowercase      , )

     return example




def      _a     (						_lowercase :     Optional[int]      , _lowercase :     Dict  ):





     '''simple docstring'''






     with jsonlines.open(_lowercase      , '''a'''  ) as writer:
          for example in tqdm(_lowercase      , total=len(_lowercase  )      , desc='''Saving samples ... '''  ):
               __UpperCAmelCase     :			Optional[int]  							=      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

      __UpperCAmelCase      :Optional[int]			    =       load_dataset("natural_questions")
      __UpperCAmelCase      :str			    =       BigBirdTokenizer.from_pretrained("google/bigbird-roberta-base")

      __UpperCAmelCase      :Dict			    =       data["train" if PROCESS_TRAIN == "true" else "validation"]

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

      np.random.seed(SEED)
      __UpperCAmelCase      :List[Any]			    =       "nq-training.jsonl" if PROCESS_TRAIN == "true" else "nq-validation.jsonl"
      save_to_disk(data, file_name=cache_file_name)



'''simple docstring'''



import argparse

import torch

from transformers import BlenderbotConfig, BlenderbotForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()
__UpperCAmelCase      :Any			    =       logging.get_logger(__name__)

__UpperCAmelCase      :Dict			    =       [
    ["attention", "attn"],
    ["encoder_attention", "encoder_attn"],
    ["q_lin", "q_proj"],
    ["k_lin", "k_proj"],
    ["v_lin", "v_proj"],
    ["out_lin", "out_proj"],
    ["norm_embeddings", "layernorm_embedding"],
    ["position_embeddings", "embed_positions"],
    ["embeddings", "embed_tokens"],
    ["ffn.lin", "fc"],
]




def      _a     (						_lowercase :     Tuple  ):





     '''simple docstring'''






     if k == "embeddings.weight":
          return "shared.weight"

     for parlai_name, hf_name in PATTERNS:
          __UpperCAmelCase     :			Any  							=      k.replace(_lowercase      , _lowercase  )

     if k.startswith('''encoder'''  ):
          __UpperCAmelCase     :			str  							=      k.replace('''.attn'''      , '''.self_attn'''  )
          __UpperCAmelCase     :			Any  							=      k.replace('''norm1'''      , '''self_attn_layer_norm'''  )
          __UpperCAmelCase     :			List[str]  							=      k.replace('''norm2'''      , '''final_layer_norm'''  )
     elif k.startswith('''decoder'''  ):
          __UpperCAmelCase     :			int  							=      k.replace('''norm1'''      , '''self_attn_layer_norm'''  )
          __UpperCAmelCase     :			Union[str, Any]  							=      k.replace('''norm2'''      , '''encoder_attn_layer_norm'''  )
          __UpperCAmelCase     :			List[Any]  							=      k.replace('''norm3'''      , '''final_layer_norm'''  )
     return k




def      _a     (						_lowercase :     Union[str, Any]  ):





     '''simple docstring'''






     __UpperCAmelCase     :			int  							=      [
         '''model.encoder.layernorm_embedding.weight''',
         '''model.encoder.layernorm_embedding.bias''',
         '''model.decoder.layernorm_embedding.weight''',
         '''model.decoder.layernorm_embedding.bias''',
     ]
     for k in keys:
          __UpperCAmelCase     :			Any  							=      sd.pop(_lowercase  )
          __UpperCAmelCase     :			Optional[int]  							=      k.replace('''layernorm_embedding'''      , '''layer_norm'''  )
          assert new_k not in sd
          __UpperCAmelCase     :			List[str]  							=      v


__UpperCAmelCase      :str			    =       ["START"]




@torch.no_grad()
def      _a     (						_lowercase :     Optional[int]      , _lowercase :     Optional[int]      , _lowercase :     str  ):





     '''simple docstring'''






     __UpperCAmelCase     :			Any  							=      torch.load(_lowercase      , map_location='''cpu'''  )
     __UpperCAmelCase     :			List[str]  							=      model['''model''']
     __UpperCAmelCase     :			Optional[Any]  							=      BlenderbotConfig.from_json_file(_lowercase  )
     __UpperCAmelCase     :			Optional[Any]  							=      BlenderbotForConditionalGeneration(_lowercase  )
     __UpperCAmelCase     :			Optional[Any]  							=      m.model.state_dict().keys()
     __UpperCAmelCase     :			int  							=      []
     __UpperCAmelCase     :			List[str]  							=      {}
     for k, v in sd.items():
          if k in IGNORE_KEYS:
               continue

          __UpperCAmelCase     :			int  							=      rename_state_dict_key(_lowercase  )
          if new_k not in valid_keys:
               failures.append([k, new_k]  )
          else:
               __UpperCAmelCase     :			Union[str, Any]  							=      v
     if cfg.normalize_before:  # Blenderbot-3B checkpoints. Rename layernorm_embedding -> layer_norm
          rename_layernorm_keys(_lowercase  )
     m.model.load_state_dict(_lowercase      , strict=_lowercase  )
     m.half()
     m.save_pretrained(_lowercase  )


if __name__ == "__main__":
      __UpperCAmelCase      :Optional[int]			    =       argparse.ArgumentParser()
      # Required parameters
      parser.add_argument("--src_path", type=str, help="like blenderbot-model.bin")
      parser.add_argument("--save_dir", default="hf_blenderbot", type=str, help="Where to save converted model.")
      parser.add_argument(
          "--hf_config_json", default="blenderbot-3b-config.json", type=str, help="Path to config to use"
      )
      __UpperCAmelCase      :Tuple			    =       parser.parse_args()
      convert_parlai_checkpoint(args.src_path, args.save_dir, args.hf_config_json)