


"""simple docstring"""

import os

import jsonlines
import numpy as np
from tqdm import tqdm


UpperCAmelCase__  							=2048
UpperCAmelCase__  							=4096
UpperCAmelCase__  							=42
UpperCAmelCase__  							=os.environ.pop("PROCESS_TRAIN", "false")
UpperCAmelCase__  							={"null": 0, "short": 1, "long": 2, "yes": 3, "no": 4}


def lowerCAmelCase_      (				UpperCamelCase__	: Optional[int] ):






 """simple docstring"""







 def choose_first(UpperCamelCase__	: List[str]				,   UpperCamelCase__	: Union[str, Any]=False ):
  assert isinstance(UpperCamelCase__				,   UpperCamelCase__ )
  if len(UpperCamelCase__ ) == 1:
   __lowercase									=	answer[0]
   return {k: [answer[k]] for k in answer} if is_long_answer else answer
  for a in answer:
   if is_long_answer:
    __lowercase									=	{k: [a[k]] for k in a}
   if len(a["""start_token"""] ) > 0:
    break
  return a

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

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

 __lowercase									=	["""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 lowerCAmelCase_      (				UpperCamelCase__	: Any				,   UpperCamelCase__	: Optional[Any]=False ):






 """simple docstring"""







 __lowercase									=	_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
  __lowercase									=	example["""document"""]["""tokens"""]
  __lowercase									=	[]
  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

 __lowercase									=	["""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

 __lowercase									=	example["""document"""]["""tokens"""]
 __lowercase									=	answer["""start_token"""]
 __lowercase									=	answer["""end_token"""]

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

 # checking above code
 if assertion:
  __lowercase									=	doc["""is_html"""][answer["""start_token"""] : answer["""end_token"""]]
  __lowercase									=	doc["""token"""][answer["""start_token"""] : answer["""end_token"""]]
  __lowercase									=	""" """.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 lowerCAmelCase_      (				UpperCamelCase__	: str				,   UpperCamelCase__	: List[Any]				,   UpperCamelCase__	: Any=2048				,   UpperCamelCase__	: Optional[int]=4096				,   UpperCamelCase__	: List[Any]=True ):






 """simple docstring"""







 __lowercase									=	get_context_and_ans(UpperCamelCase__				,   assertion=UpperCamelCase__ )
 __lowercase									=	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"],
      },
  }

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

 # return yes/no
 if answer["category"][0] in ["yes", "no"]:  # category is list with one element
  __lowercase									=	[]
  __lowercase									=	[]
  __lowercase									=	input_ids[:q_len]
  __lowercase									=	range(UpperCamelCase__				,   len(UpperCamelCase__ )				,   max_length - doc_stride )
  for i in doc_start_indices:
   __lowercase									=	i + max_length - q_len
   __lowercase									=	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,
      },
  }

 __lowercase									=	out["""context"""].split()
 __lowercase									=	splitted_context[answer["""end_token"""]]
 __lowercase									=	len(
     tokenizer(
         """ """.join(splitted_context[: answer["""start_token"""]] )				,   add_special_tokens=UpperCamelCase__				,   ).input_ids )
 __lowercase									=	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
 __lowercase									=	len(tokenizer(UpperCamelCase__				,   add_special_tokens=UpperCamelCase__ ).input_ids )
 if num_sub_tokens > 1:
  answer["end_token"] += num_sub_tokens - 1

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

 if assertion:
  __lowercase									=	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"],
      },
  }

 __lowercase									=	input_ids[:q_len]
 __lowercase									=	range(UpperCamelCase__				,   len(UpperCamelCase__ )				,   max_length - doc_stride )

 __lowercase									=	[]
 __lowercase									=	[]
 __lowercase									=	[]
 __lowercase									=	[]  # null, yes, no, long, short
 for i in doc_start_indices:
  __lowercase									=	i + max_length - q_len
  __lowercase									=	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:
   __lowercase									=	start_token - i + q_len
   __lowercase									=	end_token - i + q_len
   answers_category.append(answer["""category"""][0] )  # ["short"] -> "short"
  else:
   __lowercase									=	-100
   __lowercase									=	-100
   answers_category.append("""null""" )
  __lowercase									=	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 lowerCAmelCase_      (				UpperCamelCase__	: Optional[Any]				,   UpperCamelCase__	: str				,   UpperCamelCase__	: Union[str, Any]=2048				,   UpperCamelCase__	: Any=4096				,   UpperCamelCase__	: Dict=False ):






 """simple docstring"""







 __lowercase									=	get_strided_contexts_and_ans(
     UpperCamelCase__				,   UpperCamelCase__				,   doc_stride=UpperCamelCase__				,   max_length=UpperCamelCase__				,   assertion=UpperCamelCase__				,   )

 return example


def lowerCAmelCase_      (				UpperCamelCase__	: List[str]				,   UpperCamelCase__	: Tuple ):






 """simple docstring"""







 with jsonlines.open(UpperCamelCase__				,   """a""" ) as writer:
  for example in tqdm(UpperCamelCase__				,   total=len(UpperCamelCase__ )				,   desc="""Saving samples ... """ ):
   __lowercase									=	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__  							=load_dataset("natural_questions")
       UpperCAmelCase__  							=BigBirdTokenizer.from_pretrained("google/bigbird-roberta-base")

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

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

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




"""simple docstring"""

import os
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers.models.realm.configuration_realm import RealmConfig
from transformers.models.realm.retrieval_realm import _REALM_BLOCK_RECORDS_FILENAME, RealmRetriever
from transformers.models.realm.tokenization_realm import VOCAB_FILES_NAMES, RealmTokenizer





class    lowerCamelCase__						(		_a    ):
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


       __lowercase									=	tempfile.mkdtemp()
       __lowercase									=	5

       # Realm tok
       __lowercase									=	[
           """[UNK]""",
           """[CLS]""",
           """[SEP]""",
           """[PAD]""",
           """[MASK]""",
           """test""",
           """question""",
           """this""",
           """is""",
           """the""",
           """first""",
           """second""",
           """third""",
           """fourth""",
           """fifth""",
           """record""",
           """want""",
           """##want""",
           """##ed""",
           """wa""",
           """un""",
           """runn""",
           """##ing""",
           """,""",
           """low""",
           """lowest""",
       ]
       __lowercase									=	os.path.join(self.tmpdirname      ,       """realm_tokenizer"""							)
       os.makedirs(A_      ,       exist_ok=A_							)
       __lowercase									=	os.path.join(A_      ,       VOCAB_FILES_NAMES["""vocab_file"""]							)
       with open(self.vocab_file      ,       """w"""      ,       encoding="""utf-8"""							) as vocab_writer:
        vocab_writer.write("""""".join([x + """\n""" for x in vocab_tokens]							)							)

       __lowercase									=	os.path.join(self.tmpdirname      ,       """realm_block_records"""							)
       os.makedirs(A_      ,       exist_ok=A_							)
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[int]							):






       '''simple docstring'''


       return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname      ,       """realm_tokenizer"""							)							)
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[int]							):






       '''simple docstring'''


       shutil.rmtree(self.tmpdirname							)
      def SCREAMING_SNAKE_CASE_	(		self	:	int							):






       '''simple docstring'''


       __lowercase									=	RealmConfig(num_block_records=self.num_block_records							)
       return config
      def SCREAMING_SNAKE_CASE_	(		self	:	Any							):






       '''simple docstring'''


       __lowercase									=	Dataset.from_dict(
           {
               """id""": ["""0""", """1"""],
               """question""": ["""foo""", """bar"""],
               """answers""": [["""Foo""", """Bar"""], ["""Bar"""]],
           }							)
       return dataset
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


       __lowercase									=	np.array(
           [
               B"""This is the first record""",
               B"""This is the second record""",
               B"""This is the third record""",
               B"""This is the fourth record""",
               B"""This is the fifth record""",
               B"""This is a longer longer longer record""",
           ]      ,       dtype=A_      ,       )
       return block_records
      def SCREAMING_SNAKE_CASE_	(		self	:	Union[str, Any]							):






       '''simple docstring'''


       __lowercase									=	RealmRetriever(
           block_records=self.get_dummy_block_records()      ,       tokenizer=self.get_tokenizer()      ,       )
       return retriever
      def SCREAMING_SNAKE_CASE_	(		self	:	Dict							):






       '''simple docstring'''


       __lowercase									=	self.get_config()
       __lowercase									=	self.get_dummy_retriever()
       __lowercase									=	retriever.tokenizer

       __lowercase									=	np.array([0, 3]      ,       dtype="""long"""							)
       __lowercase									=	tokenizer(["""Test question"""]							).input_ids
       __lowercase									=	tokenizer(
           ["""the fourth"""]      ,       add_special_tokens=A_      ,       return_token_type_ids=A_      ,       return_attention_mask=A_      ,       ).input_ids
       __lowercase									=	config.reader_seq_len

       __lowercase							,		__lowercase							,		__lowercase							,		__lowercase									=	retriever(
           A_      ,       A_      ,       answer_ids=A_      ,       max_length=A_      ,       return_tensors="""np"""							)

       self.assertEqual(len(A_							)      ,       2							)
       self.assertEqual(len(A_							)      ,       2							)
       self.assertEqual(len(A_							)      ,       2							)
       self.assertEqual(concat_inputs.input_ids.shape      ,       (2, 1_0)							)
       self.assertEqual(concat_inputs.attention_mask.shape      ,       (2, 1_0)							)
       self.assertEqual(concat_inputs.token_type_ids.shape      ,       (2, 1_0)							)
       self.assertEqual(concat_inputs.special_tokens_mask.shape      ,       (2, 1_0)							)
       self.assertEqual(
           tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[0]							)      ,       ["""[CLS]""", """test""", """question""", """[SEP]""", """this""", """is""", """the""", """first""", """record""", """[SEP]"""]      ,       )
       self.assertEqual(
           tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[1]							)      ,       ["""[CLS]""", """test""", """question""", """[SEP]""", """this""", """is""", """the""", """fourth""", """record""", """[SEP]"""]      ,       )
      def SCREAMING_SNAKE_CASE_	(		self	:	str							):






       '''simple docstring'''


       __lowercase									=	self.get_config()
       __lowercase									=	self.get_dummy_retriever()
       __lowercase									=	retriever.tokenizer

       __lowercase									=	np.array([0, 3, 5]      ,       dtype="""long"""							)
       __lowercase									=	tokenizer(["""Test question"""]							).input_ids
       __lowercase									=	tokenizer(
           ["""the fourth""", """longer longer"""]      ,       add_special_tokens=A_      ,       return_token_type_ids=A_      ,       return_attention_mask=A_      ,       ).input_ids
       __lowercase									=	config.reader_seq_len

       __lowercase							,		__lowercase							,		__lowercase							,		__lowercase									=	retriever(
           A_      ,       A_      ,       answer_ids=A_      ,       max_length=A_      ,       return_tensors="""np"""							)

       self.assertEqual([False, True, True]      ,       A_							)
       self.assertEqual([[-1, -1, -1], [6, -1, -1], [6, 7, 8]]      ,       A_							)
       self.assertEqual([[-1, -1, -1], [7, -1, -1], [7, 8, 9]]      ,       A_							)






      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


       __lowercase									=	self.get_dummy_retriever()
       retriever.save_pretrained(os.path.join(self.tmpdirname      ,       """realm_block_records"""							)							)

       # Test local path
       __lowercase									=	retriever.from_pretrained(os.path.join(self.tmpdirname      ,       """realm_block_records"""							)							)
       self.assertEqual(retriever.block_records[0]      ,       B"""This is the first record"""							)

       # Test mocked remote path
       with patch("""transformers.models.realm.retrieval_realm.hf_hub_download"""							) as mock_hf_hub_download:
        __lowercase									=	os.path.join(
            os.path.join(self.tmpdirname      ,       """realm_block_records"""							)      ,       _REALM_BLOCK_RECORDS_FILENAME							)
        __lowercase									=	RealmRetriever.from_pretrained("""google/realm-cc-news-pretrained-openqa"""							)

       self.assertEqual(retriever.block_records[0]      ,       B"""This is the first record"""							)
