



import jax
import jax.numpy as jnp
from bigbird_flax import FlaxBigBirdForNaturalQuestions
from datasets import load_from_disk

from transformers import BigBirdTokenizerFast


_lowerCamelCase         =	{0: """null""", 1: """short""", 2: """long""", 3: """yes""", 4: """no"""}
_lowerCamelCase         =	set("""""".join(["""‘""", """’""", """´""", """`""", """.""", """,""", """-""", """\""""]))






def 						_lowerCAmelCase    (					__lowerCamelCase :  Union[str, Any]		,						__lowerCamelCase :  List[str]=0		,						__lowerCamelCase :  Optional[int]=None  ):



 """simple docstring"""





 return [" ".join(x.split(" "  )[begin:end]  ) for x in answers if len(x.split(" "  )  ) > 1]






def 						_lowerCAmelCase    (					__lowerCamelCase :  Optional[int]		,						__lowerCamelCase :  List[Any]=False  ):



 """simple docstring"""





 if make_sub_answers:
  # if answers are longer than one word, make sure a predictions is correct if it coresponds to the complete 1: or :-1 sub word
  # *e.g.* if the correct answer contains a prefix such as "the", or "a"
  __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						(
      given_answers + get_sub_answers(__lowerCamelCase		,						begin=1  ) + get_sub_answers(__lowerCamelCase		,						end=-1  )
  )
 __SCREAMING_SNAKE_CASE							:     Tuple							=						[]
 for answer in given_answers:
  __SCREAMING_SNAKE_CASE							:     Any							=						answer.replace("_"		,						" "  ).lower()
  __SCREAMING_SNAKE_CASE							:     str							=						"".join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else " " for c in alias  )
  answers.append(" ".join(alias.split()  ).strip()  )
 return set(__lowerCamelCase  )






def 						_lowerCAmelCase    (					__lowerCamelCase :  Optional[int]		,						__lowerCamelCase :  Dict		,						__lowerCamelCase :  List[str]=1		,						__lowerCamelCase :  int=100  ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							:     Any							=						jax.lax.top_k(__lowerCamelCase		,						__lowerCamelCase  )
 __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						jax.lax.top_k(__lowerCamelCase		,						__lowerCamelCase  )

 __SCREAMING_SNAKE_CASE							:     Dict							=						best_end_idx[:, None] - best_start_idx[None, :]
 __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						jnp.logical_or(widths < 0		,						widths > max_size  )
 __SCREAMING_SNAKE_CASE							:     List[Any]							=						(best_end_scores[:, None] + best_start_scores[None, :]) - (1E8 * mask)
 __SCREAMING_SNAKE_CASE							:     str							=						jnp.argmax(__lowerCamelCase  ).item()

 return best_start_idx[best_score % top_k], best_end_idx[best_score // top_k]






def 						_lowerCAmelCase    (					__lowerCamelCase :  Optional[int]  ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     List[str]							=						sample["question"]["text"]
 __SCREAMING_SNAKE_CASE							:     str							=						sample["document"]["tokens"]["token"]
 __SCREAMING_SNAKE_CASE							:     Dict							=						sample["document"]["tokens"]["is_html"]
 __SCREAMING_SNAKE_CASE							:     str							=						sample["annotations"]["long_answer"]
 __SCREAMING_SNAKE_CASE							:     List[Any]							=						sample["annotations"]["short_answers"]

 __SCREAMING_SNAKE_CASE							:     Optional[int]							=						" ".join([context[i] for i in range(len(__lowerCamelCase  )  ) if not is_html[i]]  )

 # 0 - No ; 1 - Yes
 for answer in sample["annotations"]["yes_no_answer"]:
  if answer == 0 or answer == 1:
   return {
       "question": question,
       "context": context_string,
       "short": [],
       "long": [],
       "category": "no" if answer == 0 else "yes",
   }

 __SCREAMING_SNAKE_CASE							:     Dict							=						[]
 for s in short_answers:
  short_targets.extend(s["text"]  )
 __SCREAMING_SNAKE_CASE							:     Tuple							=						list(set(__lowerCamelCase  )  )

 __SCREAMING_SNAKE_CASE							:     str							=						[]
 for s in long_answers:
  if s["start_token"] == -1:
   continue
  __SCREAMING_SNAKE_CASE							:     Any							=						context[s["start_token"] : s["end_token"]]
  __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						is_html[s["start_token"] : s["end_token"]]
  __SCREAMING_SNAKE_CASE							:     List[str]							=						" ".join([answer[i] for i in range(len(__lowerCamelCase  )  ) if not html[i]]  )
  if new_answer not in long_targets:
   long_targets.append(__lowerCamelCase  )

 __SCREAMING_SNAKE_CASE							:     Dict							=						"long_short" if len(short_targets + long_targets  ) > 0 else "null"

 return {
     "question": question,
     "context": context_string,
     "short": short_targets,
     "long": long_targets,
     "category": category,
 }






def 						_lowerCAmelCase    (					):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     List[Any]							=						load_from_disk("natural-questions-validation"  )
 __SCREAMING_SNAKE_CASE							:     List[str]							=						dataset.map(__lowerCamelCase  ).remove_columns(["annotations", "document", "id"]  )
 print(__lowerCamelCase  )

 __SCREAMING_SNAKE_CASE							:     str							=						dataset.filter(lambda __lowerCamelCase  : (len(x["question"]  ) + len(x["context"]  )) < 4 * 4096  )
 __SCREAMING_SNAKE_CASE							:     Any							=						short_validation_dataset.filter(lambda __lowerCamelCase  : x["category"] != "null"  )
 short_validation_dataset

 __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						"vasudevgupta/flax-bigbird-natural-questions"
 __SCREAMING_SNAKE_CASE							:     Tuple							=						FlaxBigBirdForNaturalQuestions.from_pretrained(__lowerCamelCase  )
 __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						BigBirdTokenizerFast.from_pretrained(__lowerCamelCase  )

 @jax.jit
 def forward(*__lowerCamelCase :  Optional[int]		,						**__lowerCamelCase :  Dict  ):
  __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							:     List[str]							=						model(*__lowerCamelCase		,						**__lowerCamelCase  )
  return start_logits, end_logits, jnp.argmax(__lowerCamelCase		,						axis=-1  )

 def evaluate(__lowerCamelCase :  Tuple  ):
  # encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
  __SCREAMING_SNAKE_CASE							:     str							=						tokenizer(
      example["question"]		,						example["context"]		,						return_tensors="np"		,						max_length=4096		,						padding="max_length"		,						truncation=__lowerCamelCase		,						)

  __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							:     Optional[Any]							=						forward(**__lowerCamelCase  )

  __SCREAMING_SNAKE_CASE							:     Tuple							=						CATEGORY_MAPPING[category.item()]

  __SCREAMING_SNAKE_CASE							:     Any							=						example["long"] + example["short"]
  if example["category"] in ["yes", "no", "null"]:
   __SCREAMING_SNAKE_CASE							:     List[str]							=						[example["category"]]
  __SCREAMING_SNAKE_CASE							:     Optional[int]							=						example["category"] != "null"
  # Now target can be: "yes", "no", "null", "list of long & short answers"

  if predicted_category in ["yes", "no", "null"]:
   __SCREAMING_SNAKE_CASE							:     Any							=						[predicted_category]
   __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						example["output"] == example["targets"]
   __SCREAMING_SNAKE_CASE							:     List[str]							=						predicted_category != "null"
   return example

  __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						38 if predicted_category == "short" else 1024
  __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							:     str							=						get_best_valid_start_end_idx(
      start_scores[0]		,						end_scores[0]		,						top_k=8		,						max_size=__lowerCamelCase  )

  __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						inputs["input_ids"][0].tolist()
  __SCREAMING_SNAKE_CASE							:     int							=						[tokenizer.decode(input_ids[start_score : end_score + 1]  )]

  __SCREAMING_SNAKE_CASE							:     Dict							=						expand_to_aliases(example["targets"]		,						make_sub_answers=__lowerCamelCase  )
  __SCREAMING_SNAKE_CASE							:     str							=						expand_to_aliases(example["output"]  )

  # some preprocessing to both prediction and answer
  __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						{"".join(a.split()  ) for a in answers}
  __SCREAMING_SNAKE_CASE							:     str							=						{"".join(p.split()  ) for p in predictions}
  __SCREAMING_SNAKE_CASE							:     List[Any]							=						{s for s in predictions if s not in ["``", "''", "`", "'"]}

  # if there is a common element, it's a exact match
  __SCREAMING_SNAKE_CASE							:     Any							=						len(list(answers & predictions  )  ) > 0
  __SCREAMING_SNAKE_CASE							:     Optional[int]							=						predicted_category != "null" and len(__lowerCamelCase  ) > 0

  return example

 __SCREAMING_SNAKE_CASE							:     Any							=						short_validation_dataset.map(__lowerCamelCase  )

 __SCREAMING_SNAKE_CASE							:     Tuple							=						len(__lowerCamelCase  )
 __SCREAMING_SNAKE_CASE							:     str							=						len(short_validation_dataset.filter(lambda __lowerCamelCase  : x["match"] == 1  )  )
 print("EM score:"		,						(matched / total) * 100		,						"%"  )


if __name__ == "__main__":
      main()







import os
from dataclasses import replace

import jax
import wandb
from bigbird_flax import Args, DataCollator, FlaxBigBirdForNaturalQuestions, Trainer, build_tx, train_step, val_step
from datasets import load_dataset
from flax import jax_utils

from transformers import BigBirdTokenizerFast


if __name__ == "__main__":
      print('''#################### AVAILABLE DEVICES ####################''')
      print(jax.devices())
      print('''###########################################################''')

      # setup for wandb sweep
      UpperCAmelCase_    		=				Args()
      UpperCAmelCase_    		=				wandb.init(project='''bigbird-natural-questions''', config=args.__dict__)
      UpperCAmelCase_    		=				dict(logger.config)
      del wandb_args["batch_size"]
      UpperCAmelCase_    		=				replace(args, **wandb_args)
      UpperCAmelCase_    		=				args.base_dir + '''-''' + wandb.run.id
      UpperCAmelCase_    		=				replace(args, base_dir=base_dir)
      print(args)

      UpperCAmelCase_    		=				load_dataset('''json''', data_files=args.tr_data_path)['''train''']
      UpperCAmelCase_    		=				load_dataset('''json''', data_files=args.val_data_path)['''train''']

      # drop extra batch for now
      UpperCAmelCase_    		=				range(len(tr_dataset) - len(tr_dataset) % args.batch_size)
      UpperCAmelCase_    		=				tr_dataset.shuffle().select(indices)
      UpperCAmelCase_    		=				range(len(val_dataset) - len(val_dataset) % args.batch_size)
      UpperCAmelCase_    		=				val_dataset.shuffle().select(indices)

      if os.environ.get('''TRAIN_ON_SMALL''', '''false''') == "true":
            UpperCAmelCase_    		=				tr_dataset.shuffle().select(range(80_000))
            UpperCAmelCase_    		=				val_dataset.shuffle().select(range(8_000))

      print(tr_dataset)
      print(val_dataset)

      UpperCAmelCase_    		=				FlaxBigBirdForNaturalQuestions.from_pretrained(
          args.model_id, block_size=args.block_size, num_random_blocks=args.num_random_blocks
      )
      UpperCAmelCase_    		=				BigBirdTokenizerFast.from_pretrained(args.model_id)
      UpperCAmelCase_    		=				DataCollator(pad_id=tokenizer.pad_token_id, max_length=4_096)

      UpperCAmelCase_    		=				{
          '''lr''': args.lr,
          '''init_lr''': args.init_lr,
          '''warmup_steps''': args.warmup_steps,
          '''num_train_steps''': args.max_epochs * (len(tr_dataset) // args.batch_size),
          '''weight_decay''': args.weight_decay,
      }
      UpperCAmelCase_		,     UpperCAmelCase_    		=				build_tx(**tx_args)

      UpperCAmelCase_    		=				Trainer(
          args=args,
          data_collator=data_collator,
          model_save_fn=model.save_pretrained,
          train_step_fn=train_step,
          val_step_fn=val_step,
          logger=logger,
          scheduler_fn=lr,
      )

      UpperCAmelCase_    		=				None
      UpperCAmelCase_    		=				trainer.create_state(model, tx, num_train_steps=tx_args['''num_train_steps'''], ckpt_dir=ckpt_dir)
      try:
            trainer.train(state, tr_dataset, val_dataset)
      except KeyboardInterrupt:
            print('''Oooops; TRAINING STOPPED UNFORTUNATELY''')

      print('''SAVING WEIGHTS IN `final-weights`''')
      UpperCAmelCase_    		=				jax_utils.unreplicate(state.params)
      model.save_pretrained(os.path.join(args.base_dir, '''final-weights'''), params=params)