


import os

import jsonlines
import numpy as np
from tqdm import tqdm


_snake_case             =     20_48
_snake_case             =     40_96
_snake_case             =     42
_snake_case             =     os.environ.pop('''PROCESS_TRAIN''', '''false''')
_snake_case             =     {"""null""": 0, """short""": 1, """long""": 2, """yes""": 3, """no""": 4}






def 		lowercase_(  SCREAMING_SNAKE_CASE_	):


  '''simple docstring'''

  def choose_first(SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_=False	):
    assert isinstance(_UpperCAmelCase      ,						_UpperCAmelCase	)
    if len(_UpperCAmelCase	) == 1:
      lowerCamelCase :					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:
        lowerCamelCase :					Optional[Any]    =				{k: [a[k]] for k in a}
      if len(a["start_token"]	) > 0:
        break
    return a

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

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

  lowerCamelCase :					Tuple    =				["start_token", "end_token", "start_byte", "end_byte", "text"]
  if not all(isinstance(answer[k]      ,						_UpperCAmelCase	) for k in cols	):
    raise ValueError("Issue in ID"      ,						example["id"]	)

  return answer






def 		lowercase_(  SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_=False	):


  '''simple docstring'''

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

  lowerCamelCase :					Optional[int]    =				["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

  lowerCamelCase :					List[str]    =				example["document"]["tokens"]
  lowerCamelCase :					Optional[int]    =				answer["start_token"]
  lowerCamelCase :					Union[str, Any]    =				answer["end_token"]

  lowerCamelCase :					Any    =				[]
  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
  lowerCamelCase :					Optional[Any]    =				" ".join(context[start_token:end_token]	)

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

  return {
      "context": " ".join(_UpperCAmelCase	),
      "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 		lowercase_(  SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_=2048      ,						SCREAMING_SNAKE_CASE_=4096      ,						SCREAMING_SNAKE_CASE_=True	):


  '''simple docstring'''

  lowerCamelCase :					Dict    =				get_context_and_ans(_UpperCAmelCase      ,						assertion=_UpperCAmelCase	)
  lowerCamelCase :					int    =				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"],
        },
    }

  lowerCamelCase :					Union[str, Any]    =				tokenizer(example["question"]["text"]      ,						out["context"]	).input_ids
  lowerCamelCase :					List[Any]    =				input_ids.index(tokenizer.sep_token_id	) + 1

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

  lowerCamelCase :					Dict    =				out["context"].split()
  lowerCamelCase :					Optional[int]    =				splitted_context[answer["end_token"]]
  lowerCamelCase :					Union[str, Any]    =				len(
      tokenizer(
          " ".join(splitted_context[: answer["start_token"]]	)      ,						add_special_tokens=_UpperCAmelCase      ,						).input_ids	)
  lowerCamelCase :					int    =				len(
      tokenizer(" ".join(splitted_context[: answer["end_token"]]	)      ,						add_special_tokens=_UpperCAmelCase	).input_ids	)

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

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

  lowerCamelCase :					List[str]    =				input_ids[answer["start_token"] : answer["end_token"] + 1]  # right & left are inclusive
  lowerCamelCase :					List[Any]    =				answer["start_token"]
  lowerCamelCase :					List[Any]    =				answer["end_token"]

  if assertion:
    lowerCamelCase :					Optional[Any]    =				tokenizer.decode(_UpperCAmelCase	)
    if answer["span"] != new:
      print("ISSUE IN TOKENIZATION"	)
      print("OLD:"      ,						answer["span"]	)
      print("NEW:"      ,						_UpperCAmelCase      ,						end="\n\n"	)

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

  lowerCamelCase :					Union[str, Any]    =				input_ids[:q_len]
  lowerCamelCase :					List[str]    =				range(_UpperCAmelCase      ,						len(_UpperCAmelCase	)      ,						max_length - doc_stride	)

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

    answers_start_token.append(_UpperCAmelCase	)
    answers_end_token.append(_UpperCAmelCase	)
    if assertion:
      if new != old and new != [tokenizer.cls_token_id]:
        print("ISSUE in strided for ID:"      ,						example["id"]	)
        print("New:"      ,						tokenizer.decode(_UpperCAmelCase	)	)
        print("Old:"      ,						tokenizer.decode(_UpperCAmelCase	)      ,						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 		lowercase_(  SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_=2048      ,						SCREAMING_SNAKE_CASE_=4096      ,						SCREAMING_SNAKE_CASE_=False	):


  '''simple docstring'''

  lowerCamelCase :					str    =				get_strided_contexts_and_ans(
      _UpperCAmelCase      ,						_UpperCAmelCase      ,						doc_stride=_UpperCAmelCase      ,						max_length=_UpperCAmelCase      ,						assertion=_UpperCAmelCase      ,						)

  return example






def 		lowercase_(  SCREAMING_SNAKE_CASE_      ,						SCREAMING_SNAKE_CASE_	):


  '''simple docstring'''

  with jsonlines.open(_UpperCAmelCase      ,						"a"	) as writer:
    for example in tqdm(_UpperCAmelCase      ,						total=len(_UpperCAmelCase	)      ,						desc="Saving samples ... "	):
      lowerCamelCase :					Tuple    =				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

   _snake_case             =     load_dataset('''natural_questions''')
   _snake_case             =     BigBirdTokenizer.from_pretrained('''google/bigbird-roberta-base''')

   _snake_case             =     data["""train""" if PROCESS_TRAIN == """true""" else """validation"""]

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

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






import argparse

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType


########################################################################
# This is a fully working simple example to use Accelerate
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


lowerCAmelCase							:  Optional[Any]   	=						16
lowerCAmelCase							:  List[str]   	=						32







def 			A_      (    _UpperCAmelCase				,	_UpperCAmelCase = 16    ):
						SCREAMING_SNAKE_CASE_: Tuple  		=       AutoTokenizer.from_pretrained("bert-base-cased"    )
						SCREAMING_SNAKE_CASE_: List[Any]  		=       load_dataset("glue"				,	"mrpc"    )

						def tokenize_function(_UpperCAmelCase    ):
												# max_length=None => use the model max length (it's actually the default)
												SCREAMING_SNAKE_CASE_: Any  		=       tokenizer(examples["sentence1"]				,	examples["sentence2"]				,	truncation=_UpperCAmelCase				,	max_length=_UpperCAmelCase    )
												return outputs

						# Apply the method we just defined to all the examples in all the splits of the dataset
						# starting with the main process first:
						with accelerator.main_process_first():
												SCREAMING_SNAKE_CASE_: Tuple  		=       datasets.map(
												    _UpperCAmelCase				,	batched=_UpperCAmelCase				,	remove_columns=["idx", "sentence1", "sentence2"]				,	)

						# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
						# transformers library
						SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       tokenized_datasets.rename_column("label"				,	"labels"    )

						def collate_fn(_UpperCAmelCase    ):
												# On TPU it's best to pad everything to the same length or training will be very slow.
												SCREAMING_SNAKE_CASE_: List[str]  		=       1_28 if accelerator.distributed_type == DistributedType.TPU else None
												# When using mixed precision we want round multiples of 8/16
												if accelerator.mixed_precision == "fp8":
																		SCREAMING_SNAKE_CASE_: List[Any]  		=       16
												elif accelerator.mixed_precision != "no":
																		SCREAMING_SNAKE_CASE_: Optional[Any]  		=       8
												else:
																		SCREAMING_SNAKE_CASE_: List[str]  		=       None

												return tokenizer.pad(
												    _UpperCAmelCase				,	padding="longest"				,	max_length=_UpperCAmelCase				,	pad_to_multiple_of=_UpperCAmelCase				,	return_tensors="pt"				,	)

						# Instantiate dataloaders.
						SCREAMING_SNAKE_CASE_: Dict  		=       DataLoader(
						    tokenized_datasets["train"]				,	shuffle=_UpperCAmelCase				,	collate_fn=_UpperCAmelCase				,	batch_size=_UpperCAmelCase				,	drop_last=_UpperCAmelCase    )
						SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       DataLoader(
						    tokenized_datasets["validation"]				,	shuffle=_UpperCAmelCase				,	collate_fn=_UpperCAmelCase				,	batch_size=_UpperCAmelCase				,	drop_last=(accelerator.mixed_precision == "fp8")				,	)

						return train_dataloader, eval_dataloader







def 			A_      (    _UpperCAmelCase				,	_UpperCAmelCase    ):
						# Initialize accelerator
						SCREAMING_SNAKE_CASE_: str  		=       Accelerator(cpu=args.cpu				,	mixed_precision=args.mixed_precision    )
						# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
						SCREAMING_SNAKE_CASE_: int  		=       config["lr"]
						SCREAMING_SNAKE_CASE_: Any  		=       int(config["num_epochs"]    )
						SCREAMING_SNAKE_CASE_: Optional[int]  		=       int(config["seed"]    )
						SCREAMING_SNAKE_CASE_: List[Any]  		=       int(config["batch_size"]    )

						SCREAMING_SNAKE_CASE_: List[str]  		=       evaluate.load("glue"				,	"mrpc"    )

						# If the batch size is too big we use gradient accumulation
						SCREAMING_SNAKE_CASE_: Optional[int]  		=       1
						if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
												SCREAMING_SNAKE_CASE_: Tuple  		=       batch_size // MAX_GPU_BATCH_SIZE
												SCREAMING_SNAKE_CASE_: Dict  		=       MAX_GPU_BATCH_SIZE

						set_seed(_UpperCAmelCase    )
						SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_: str  		=       get_dataloaders(_UpperCAmelCase				,	_UpperCAmelCase    )
						# Instantiate the model (we build the model here so that the seed also control new weights initialization)
						SCREAMING_SNAKE_CASE_: List[Any]  		=       AutoModelForSequenceClassification.from_pretrained("bert-base-cased"				,	return_dict=_UpperCAmelCase    )

						# We could avoid this line since the accelerator is set with `device_placement=True` (default value).
						# Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
						# creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
						SCREAMING_SNAKE_CASE_: Tuple  		=       model.to(accelerator.device    )
						# Instantiate optimizer
						SCREAMING_SNAKE_CASE_: Optional[int]  		=       AdamW(params=model.parameters()				,	lr=_UpperCAmelCase    )

						# Instantiate scheduler
						SCREAMING_SNAKE_CASE_: Optional[int]  		=       get_linear_schedule_with_warmup(
						    optimizer=_UpperCAmelCase				,	num_warmup_steps=1_00				,	num_training_steps=(len(_UpperCAmelCase    ) * num_epochs) // gradient_accumulation_steps				,	)

						# Prepare everything
						# There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
						# prepare method.

						SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_: Optional[int]  		=       accelerator.prepare(
						    _UpperCAmelCase				,	_UpperCAmelCase				,	_UpperCAmelCase				,	_UpperCAmelCase				,	_UpperCAmelCase    )

						# Now we train the model
						for epoch in range(_UpperCAmelCase    ):
												model.train()
												for step, batch in enumerate(_UpperCAmelCase    ):
																		# We could avoid this line since we set the accelerator with `device_placement=True`.
																		batch.to(accelerator.device    )
																		SCREAMING_SNAKE_CASE_: Tuple  		=       model(**_UpperCAmelCase    )
																		SCREAMING_SNAKE_CASE_: Optional[int]  		=       outputs.loss
																		SCREAMING_SNAKE_CASE_: Tuple  		=       loss / gradient_accumulation_steps
																		accelerator.backward(_UpperCAmelCase    )
																		if step % gradient_accumulation_steps == 0:
																								optimizer.step()
																								lr_scheduler.step()
																								optimizer.zero_grad()

												model.eval()
												for step, batch in enumerate(_UpperCAmelCase    ):
																		# We could avoid this line since we set the accelerator with `device_placement=True`.
																		batch.to(accelerator.device    )
																		with torch.no_grad():
																								SCREAMING_SNAKE_CASE_: Optional[int]  		=       model(**_UpperCAmelCase    )
																		SCREAMING_SNAKE_CASE_: int  		=       outputs.logits.argmax(dim=-1    )
																		SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_: int  		=       accelerator.gather_for_metrics((predictions, batch["labels"])    )
																		metric.add_batch(
																		    predictions=_UpperCAmelCase				,	references=_UpperCAmelCase				,	)

												SCREAMING_SNAKE_CASE_: List[Any]  		=       metric.compute()
												# Use accelerator.print to print only on the main process.
												accelerator.print(f"epoch {epoch}:"				,	_UpperCAmelCase    )







def 			A_      (    ):
						SCREAMING_SNAKE_CASE_: Any  		=       argparse.ArgumentParser(description="Simple example of training script."    )
						parser.add_argument(
						    "--mixed_precision"				,	type=_UpperCAmelCase				,	default=_UpperCAmelCase				,	choices=["no", "fp16", "bf16", "fp8"]				,	help="Whether to use mixed precision. Choose"
						    "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
						    "and an Nvidia Ampere GPU."				,	)
						parser.add_argument("--cpu"				,	action="store_true"				,	help="If passed, will train on the CPU."    )
						SCREAMING_SNAKE_CASE_: Optional[Any]  		=       parser.parse_args()
						SCREAMING_SNAKE_CASE_: Optional[int]  		=       {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
						training_function(_UpperCAmelCase				,	_UpperCAmelCase    )


if __name__ == "__main__":
						main()






