
'''simple docstring'''





import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device


if is_torch_available():
     import torch

     from transformers import XLMProphetNetForConditionalGeneration, XLMProphetNetTokenizer



@require_torch
class  lowerCAmelCase__   (       unittest.TestCase       ):



  '''simple docstring'''






  @slow
  def 		_lowerCamelCase	(self  : Dict    )   ->       Union[str, Any]:
        snake_case_						       =							XLMProphetNetForConditionalGeneration.from_pretrained("microsoft/xprophetnet-large-wiki100-cased"    )
        model.to(lowerCAmelCase_    )

        # encoder-decoder outputs
        snake_case_						       =							torch.tensor([[17, 9_62_08, 10_34_71, 2]]    ).to(lowerCAmelCase_    )
        snake_case_						       =							torch.tensor(
            [[2, 2_50, 99_53, 34, 6_94_89, 16_20, 32, 11_84_24, 6_24, 2_10, 1_05, 29_13, 10_32, 3_51]]    ).to(lowerCAmelCase_    )
        snake_case_						       =							model(
            input_ids=lowerCAmelCase_	, attention_mask=lowerCAmelCase_	, encoder_outputs=lowerCAmelCase_	, decoder_input_ids=lowerCAmelCase_    )
        snake_case_						       =							output[0]
        snake_case_						       =							torch.Size((1, 14, 25_00_12)    )
        self.assertEqual(output_predited_logis.shape	, lowerCAmelCase_    )
        snake_case_						       =							torch.tensor(
            [[[-6.3_986, -8.2_391, 12.5_189], [-6.3_289, -8.0_864, 12.6_211], [-6.2_418, -8.0_445, 12.7_968]]]    ).to(lowerCAmelCase_    )
        self.assertTrue(torch.allclose(output_predited_logis[:, :3, :3]	, lowerCAmelCase_	, atol=1E-4    )    )

        # encoder outputs
        snake_case_						       =							model.prophetnet.encoder(lowerCAmelCase_    )[0]
        snake_case_						       =							torch.tensor(
            [[[-1.4_260, -0.7_628, 0.8_453], [-1.4_719, -0.1_391, 0.7_807], [-1.7_678, 0.0_114, 0.4_646]]]    ).to(lowerCAmelCase_    )
        snake_case_						       =							torch.Size((1, 4, 10_24)    )
        self.assertEqual(encoder_outputs.shape	, lowerCAmelCase_    )
        self.assertTrue(torch.allclose(encoder_outputs[:, :3, :3]	, lowerCAmelCase_	, atol=1E-4    )    )

        # decoder outputs
        snake_case_						       =							model.prophetnet.decoder(
            lowerCAmelCase_	, encoder_hidden_states=lowerCAmelCase_	, )
        snake_case_						       =							decoder_outputs[1].view(1	, model.config.ngram	, 14	, -1    )
        snake_case_						       =							model.lm_head(lowerCAmelCase_    )
        snake_case_						       =							predicting_streams_logits[:, 0]
        self.assertTrue(torch.allclose(next_first_stream_logits[:, :3, :3]	, lowerCAmelCase_	, atol=1E-4    )    )

  @slow
  def 		_lowerCamelCase	(self  : Optional[int]    )   ->       str:
        snake_case_						       =							XLMProphetNetForConditionalGeneration.from_pretrained(
            "microsoft/xprophetnet-large-wiki100-cased-xglue-ntg"    )
        model.to(lowerCAmelCase_    )

        snake_case_						       =							torch.tensor([[17, 9_62_08, 10_34_71, 2]]    ).to(lowerCAmelCase_    )
        snake_case_						       =							torch.tensor(
            [[2, 2_50, 99_53, 34, 6_94_89, 16_20, 32, 11_84_24, 6_24, 2_10, 1_05, 29_13, 10_32, 3_51]]    ).to(lowerCAmelCase_    )
        snake_case_						       =							model(
            input_ids=lowerCAmelCase_	, attention_mask=lowerCAmelCase_	, encoder_outputs=lowerCAmelCase_	, decoder_input_ids=lowerCAmelCase_    )
        snake_case_						       =							output[0]
        snake_case_						       =							torch.Size((1, 14, 25_00_12)    )
        self.assertEqual(output_predited_logis.shape	, lowerCAmelCase_    )
        # compare the actual values for a slice.
        snake_case_						       =							torch.tensor(
            [[[-9.2_253, -9.7_173, -6.3_529], [-7.6_701, -9.0_145, -1.9_382], [-8.0_195, -7.0_004, -0.1_523]]]    ).to(lowerCAmelCase_    )

        self.assertTrue(torch.allclose(output_predited_logis[:, :3, :3]	, lowerCAmelCase_	, atol=1E-4    )    )


  @slow
  def 		_lowerCamelCase	(self  : Optional[int]    )   ->       Union[str, Any]:
        snake_case_						       =							XLMProphetNetForConditionalGeneration.from_pretrained(
            "microsoft/xprophetnet-large-wiki100-cased-xglue-ntg"    )
        model.to(lowerCAmelCase_    )
        snake_case_						       =							5_12

        snake_case_						       =							XLMProphetNetTokenizer.from_pretrained("microsoft/xprophetnet-large-wiki100-cased-xglue-ntg"    )

        snake_case_						       =							(
            "Microsoft Corporation intends to officially end free support for the Windows 7 operating system after"
            " January 14, 2020, according to the official portal of the organization. From that day, users of this"
            " system will not be able to receive security updates, which could make their computers vulnerable to"
            " cyber attacks."
        )
        snake_case_						       =							(
            "орпорация Microsoft намерена официально прекратить бесплатную поддержку операционной системы Windows 7"
            " после 14 января 2020 года, сообщается на официальном портале организации . С указанного дня пользователи"
            " этой системы не смогут получать обновления безопасности, из-за чего их компьютеры могут стать уязвимыми"
            " к кибератакам."
        )
        snake_case_						       =							(
            "根据该组织的官方门户网站，微软公司打算在2020年1月14日之后正式终止对Windows 7操作系统的免费支持。从那时起，该系统的用户将无法接收安全更新，这可能会使他们的计算机容易受到网络攻击。"
        )

        snake_case_						       =							tokenizer(
            [EN_SENTENCE, RU_SENTENCE, ZH_SENTENCE]	, padding=lowerCAmelCase_	, max_length=2_55	, return_tensors="pt"    ).input_ids
        snake_case_						       =							input_ids.to(lowerCAmelCase_    )

        snake_case_						       =							model.generate(
            lowerCAmelCase_	, num_beams=10	, length_penalty=1.0	, no_repeat_ngram_size=3	, early_stopping=lowerCAmelCase_    )
        snake_case_						       =							[tokenizer.decode(lowerCAmelCase_	, skip_special_tokens=lowerCAmelCase_    ) for g in summary_ids]
        snake_case_						       =							"Microsoft to end Windows 7 free support after January 14, 2020"
        snake_case_						       =							"Microsoft намерена прекратить бесплатную поддержку Windows 7 после 14 января 2020 года"
        snake_case_						       =							"微软打算终止对Windows 7操作系统的免费支持"
        self.assertListEqual(
            [EXPECTED_TITLE_EN, EXPECTED_TITLE_RU, EXPECTED_TITLE_ZH]	, lowerCAmelCase_	, )

        snake_case_						       =							model.generate(
            lowerCAmelCase_	, num_beams=1	, length_penalty=1.0	, no_repeat_ngram_size=3	, early_stopping=lowerCAmelCase_    )
        snake_case_						       =							[
            tokenizer.convert_ids_to_tokens(lowerCAmelCase_	, skip_special_tokens=lowerCAmelCase_    ) for g in summary_ids_beama
        ]
        snake_case_						       =							"▁Microsoft ▁to ▁end ▁free ▁support ▁for ▁Windows ▁7".split(" "    )
        snake_case_						       =							"▁Microsoft ▁намерен а ▁прекрати ть ▁бес плат ную ▁поддержку ▁Windows ▁7 ▁после ▁14 ▁января ▁2020 ▁года".split(
            " "    )
        snake_case_						       =							"微软 公司 打算 终止 对 Windows ▁7 操作 系统的 免费 支持".split(" "    )
        self.assertListEqual(
            [EXPECTED_TITLE_EN_BEAM1_TOK, EXPECTED_TITLE_RU_BEAM1_TOK, EXPECTED_TITLE_ZH_BEAM1_TOK]	, lowerCAmelCase_	, )



'''simple docstring'''





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

from transformers import BigBirdTokenizerFast


A           =   {0: "null", 1: "short", 2: "long", 3: "yes", 4: "no"}
A           =   set("".join(["‘", "’", "´", "`", ".", ",", "-", "\""]))
def   snake_case_       (							_a     :							Any ,					_a     :							Union[str, Any]=0 ,					_a     :							Optional[int]=None			):
      return [" ".join(x.split(" "			)[begin:end]			) for x in answers if len(x.split(" "			)			) > 1]
def   snake_case_       (							_a     :							Optional[int] ,					_a     :							str=False			):
      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"
            snake_case_						       =							(
                given_answers + get_sub_answers(_a ,					begin=1			) + get_sub_answers(_a ,					end=-1			)
            )
      snake_case_						       =							[]
      for answer in given_answers:
            snake_case_						       =							answer.replace("_" ,					" "			).lower()
            snake_case_						       =							"".join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else " " for c in alias			)
            answers.append(" ".join(alias.split()			).strip()			)
      return set(_a			)
def   snake_case_       (							_a     :							Optional[Any] ,					_a     :							int ,					_a     :							Optional[Any]=1 ,					_a     :							List[Any]=100			):
      snake_case_       ,	snake_case_						       =							jax.lax.top_k(_a ,					_a			)
      snake_case_       ,	snake_case_						       =							jax.lax.top_k(_a ,					_a			)

      snake_case_						       =							best_end_idx[:, None] - best_start_idx[None, :]
      snake_case_						       =							jnp.logical_or(widths < 0 ,					widths > max_size			)
      snake_case_						       =							(best_end_scores[:, None] + best_start_scores[None, :]) - (1E8 * mask)
      snake_case_						       =							jnp.argmax(_a			).item()

      return best_start_idx[best_score % top_k], best_end_idx[best_score // top_k]
def   snake_case_       (							_a     :							Any			):
      snake_case_						       =							sample["question"]["text"]
      snake_case_						       =							sample["document"]["tokens"]["token"]
      snake_case_						       =							sample["document"]["tokens"]["is_html"]
      snake_case_						       =							sample["annotations"]["long_answer"]
      snake_case_						       =							sample["annotations"]["short_answers"]

      snake_case_						       =							" ".join([context[i] for i in range(len(_a			)			) 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",
                  }

      snake_case_						       =							[]
      for s in short_answers:
            short_targets.extend(s["text"]			)
      snake_case_						       =							list(set(_a			)			)

      snake_case_						       =							[]
      for s in long_answers:
            if s["start_token"] == -1:
                  continue
            snake_case_						       =							context[s["start_token"] : s["end_token"]]
            snake_case_						       =							is_html[s["start_token"] : s["end_token"]]
            snake_case_						       =							" ".join([answer[i] for i in range(len(_a			)			) if not html[i]]			)
            if new_answer not in long_targets:
                  long_targets.append(_a			)

      snake_case_						       =							"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   snake_case_       (							):
      snake_case_						       =							load_from_disk("natural-questions-validation"			)
      snake_case_						       =							dataset.map(_a			).remove_columns(["annotations", "document", "id"]			)
      print(_a			)

      snake_case_						       =							dataset.filter(lambda _a			: (len(x["question"]			) + len(x["context"]			)) < 4 * 4_096			)
      snake_case_						       =							short_validation_dataset.filter(lambda _a			: x["category"] != "null"			)
      short_validation_dataset

      snake_case_						       =							"vasudevgupta/flax-bigbird-natural-questions"
      snake_case_						       =							FlaxBigBirdForNaturalQuestions.from_pretrained(_a			)
      snake_case_						       =							BigBirdTokenizerFast.from_pretrained(_a			)

      @jax.jit
      def forward(*_a     :							List[Any] ,					**_a     :							int			):
            snake_case_       ,	snake_case_       ,	snake_case_						       =							model(*_a ,					**_a			)
            return start_logits, end_logits, jnp.argmax(_a ,					axis=-1			)

      def evaluate(_a     :							List[Any]			):
            # encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
            snake_case_						       =							tokenizer(
                example["question"] ,					example["context"] ,					return_tensors="np" ,					max_length=4_096 ,					padding="max_length" ,					truncation=_a ,					)

            snake_case_       ,	snake_case_       ,	snake_case_						       =							forward(**_a			)

            snake_case_						       =							CATEGORY_MAPPING[category.item()]

            snake_case_						       =							example["long"] + example["short"]
            if example["category"] in ["yes", "no", "null"]:
                  snake_case_						       =							[example["category"]]
            snake_case_						       =							example["category"] != "null"
            # Now target can be: "yes", "no", "null", "list of long & short answers"

            if predicted_category in ["yes", "no", "null"]:
                  snake_case_						       =							[predicted_category]
                  snake_case_						       =							example["output"] == example["targets"]
                  snake_case_						       =							predicted_category != "null"
                  return example

            snake_case_						       =							38 if predicted_category == "short" else 1_024
            snake_case_       ,	snake_case_						       =							get_best_valid_start_end_idx(
                start_scores[0] ,					end_scores[0] ,					top_k=8 ,					max_size=_a			)

            snake_case_						       =							inputs["input_ids"][0].tolist()
            snake_case_						       =							[tokenizer.decode(input_ids[start_score : end_score + 1]			)]

            snake_case_						       =							expand_to_aliases(example["targets"] ,					make_sub_answers=_a			)
            snake_case_						       =							expand_to_aliases(example["output"]			)

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

            # if there is a common element, it's a exact match
            snake_case_						       =							len(list(answers & predictions			)			) > 0
            snake_case_						       =							predicted_category != "null" and len(_a			) > 0

            return example

      snake_case_						       =							short_validation_dataset.map(_a			)

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


if __name__ == "__main__":
     main()

