




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

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

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

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

					return answer







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Dict   ,   SCREAMING_SNAKE_CASE_							:       List[str]=False  ):
					__lowerCAmelCase            = _get_single_answer(SCREAMING_SNAKE_CASE_  )
					# 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            = example["document"]["tokens"]
										__lowerCAmelCase            = []
										for i in range(len(doc["token"]  )  ):
															if not doc["is_html"][i]:
																				context.append(doc["token"][i]  )
										return {
										    "context": " ".join(SCREAMING_SNAKE_CASE_  ),
										    "answer": {
										        "start_token": -1_00,  # ignore index in cross-entropy
										        "end_token": -1_00,  # 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            = ["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            = example["document"]["tokens"]
					__lowerCAmelCase            = answer["start_token"]
					__lowerCAmelCase            = answer["end_token"]

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

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

					return {
					    "context": " ".join(SCREAMING_SNAKE_CASE_  ),
					    "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   (		SCREAMING_SNAKE_CASE_							:       int   ,   SCREAMING_SNAKE_CASE_							:       int   ,   SCREAMING_SNAKE_CASE_							:       Tuple=20_48   ,   SCREAMING_SNAKE_CASE_							:       Tuple=40_96   ,   SCREAMING_SNAKE_CASE_							:       List[str]=True  ):
					# overlap will be of doc_stride - q_len

					__lowerCAmelCase            = get_context_and_ans(SCREAMING_SNAKE_CASE_   ,   assertion=SCREAMING_SNAKE_CASE_  )
					__lowerCAmelCase            = 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            = tokenizer(example["question"]["text"]   ,   out["context"]  ).input_ids
					__lowerCAmelCase            = 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            = []
										__lowerCAmelCase            = []
										__lowerCAmelCase            = input_ids[:q_len]
										__lowerCAmelCase            = range(SCREAMING_SNAKE_CASE_   ,   len(SCREAMING_SNAKE_CASE_  )   ,   max_length - doc_stride  )
										for i in doc_start_indices:
															__lowerCAmelCase            = i + max_length - q_len
															__lowerCAmelCase            = 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": [-1_00] * len(SCREAMING_SNAKE_CASE_  ),
										        "end_token": [-1_00] * len(SCREAMING_SNAKE_CASE_  ),
										        "category": category,
										    },
										}

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

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

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

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

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

					if len(SCREAMING_SNAKE_CASE_  ) <= 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            = input_ids[:q_len]
					__lowerCAmelCase            = range(SCREAMING_SNAKE_CASE_   ,   len(SCREAMING_SNAKE_CASE_  )   ,   max_length - doc_stride  )

					__lowerCAmelCase            = []
					__lowerCAmelCase            = []
					__lowerCAmelCase            = []
					__lowerCAmelCase            = []  # null, yes, no, long, short
					for i in doc_start_indices:
										__lowerCAmelCase            = i + max_length - q_len
										__lowerCAmelCase            = 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            = start_token - i + q_len
															__lowerCAmelCase            = end_token - i + q_len
															answers_category.append(answer["category"][0]  )  # ["short"] -> "short"
										else:
															__lowerCAmelCase            = -1_00
															__lowerCAmelCase            = -1_00
															answers_category.append("null"  )
										__lowerCAmelCase            = inputs[-1][start_token : end_token + 1]

										answers_start_token.append(SCREAMING_SNAKE_CASE_  )
										answers_end_token.append(SCREAMING_SNAKE_CASE_  )
										if assertion:
															if new != old and new != [tokenizer.cls_token_id]:
																				print("ISSUE in strided for ID:"   ,   example["id"]  )
																				print("New:"   ,   tokenizer.decode(SCREAMING_SNAKE_CASE_  )  )
																				print("Old:"   ,   tokenizer.decode(SCREAMING_SNAKE_CASE_  )   ,   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   (		SCREAMING_SNAKE_CASE_							:       Optional[Any]   ,   SCREAMING_SNAKE_CASE_							:       int   ,   SCREAMING_SNAKE_CASE_							:       Optional[Any]=20_48   ,   SCREAMING_SNAKE_CASE_							:       Any=40_96   ,   SCREAMING_SNAKE_CASE_							:       Optional[Any]=False  ):
					__lowerCAmelCase            = get_strided_contexts_and_ans(
					    SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_   ,   doc_stride=SCREAMING_SNAKE_CASE_   ,   max_length=SCREAMING_SNAKE_CASE_   ,   assertion=SCREAMING_SNAKE_CASE_   ,   )

					return example







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Union[str, Any]   ,   SCREAMING_SNAKE_CASE_							:       List[Any]  ):
					with jsonlines.open(SCREAMING_SNAKE_CASE_   ,   "a"  ) as writer:
										for example in tqdm(SCREAMING_SNAKE_CASE_   ,   total=len(SCREAMING_SNAKE_CASE_  )   ,   desc="Saving samples ... "  ):
															__lowerCAmelCase            = 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)










from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
					from PIL import Image

					from ..image_utils import load_image

if is_tf_available():
					import tensorflow as tf

					from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
					from ..tf_utils import stable_softmax

if is_torch_available():
					from ..models.auto.modeling_auto import MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING

UpperCamelCase__					    =    logging.get_logger(__name__)
@add_end_docstrings(snake_case__   )
class     a__					(      snake_case__   ):






						def __init__(     self	,    *_A	,    **_A    ):




											"""simple docstring"""



											super().__init__(*_A	,    **_A    )
											requires_backends(self	,    "vision"    )
											self.check_model_type(
											    TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
											    if self.framework == "tf"
											    else MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING    )






						def        __SCREAMING_SNAKE_CASE(     self	,    _A=None    ):




											"""simple docstring"""



											__lowerCAmelCase            = {}
											if top_k is not None:
																__lowerCAmelCase            = top_k
											return {}, {}, postprocess_params






						def __call__(     self	,    _A	,    **_A    ):




											"""simple docstring"""



											return super().__call__(_A	,    **_A    )






						def        __SCREAMING_SNAKE_CASE(     self	,    _A    ):




											"""simple docstring"""



											__lowerCAmelCase            = load_image(_A    )
											__lowerCAmelCase            = self.image_processor(images=_A	,    return_tensors=self.framework    )
											return model_inputs






						def        __SCREAMING_SNAKE_CASE(     self	,    _A    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.model(**_A    )
											return model_outputs






						def        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A=5    ):




											"""simple docstring"""



											if top_k > self.model.config.num_labels:
																__lowerCAmelCase            = self.model.config.num_labels

											if self.framework == "pt":
																__lowerCAmelCase            = model_outputs.logits.softmax(-1    )[0]
																__lowerCAmelCase     ,							__lowerCAmelCase            = probs.topk(_A    )
											elif self.framework == "tf":
																__lowerCAmelCase            = stable_softmax(model_outputs.logits	,    axis=-1    )[0]
																__lowerCAmelCase            = tf.math.top_k(_A	,    k=_A    )
																__lowerCAmelCase     ,							__lowerCAmelCase            = topk.values.numpy(), topk.indices.numpy()
											else:
																raise ValueError(f"""Unsupported framework: {self.framework}"""    )

											__lowerCAmelCase            = scores.tolist()
											__lowerCAmelCase            = ids.tolist()
											return [{"score": score, "label": self.model.config.idalabel[_id]} for score, _id in zip(_A	,    _A    )]




