




'''simple docstring'''





import argparse
import collections
import json
import os
import re
import string
import sys

import numpy as np


_SCREAMING_SNAKE_CASE							:   List[str]      		= re.compile(r"\b(a|an|the)\b", re.UNICODE)

_SCREAMING_SNAKE_CASE							:   List[str]      		= None


def      UpperCamelCase_(	):

							'''simple docstring'''






							snake_case_						=      argparse.ArgumentParser("Official evaluation script for SQuAD version 2.0."     )
							parser.add_argument("data_file"							,				metavar="data.json"							,				help="Input data JSON file."     )
							parser.add_argument("pred_file"							,				metavar="pred.json"							,				help="Model predictions."     )
							parser.add_argument(
							    "--out-file"							,				"-o"							,				metavar="eval.json"							,				help="Write accuracy metrics to file (default is stdout)."     )
							parser.add_argument(
							    "--na-prob-file"							,				"-n"							,				metavar="na_prob.json"							,				help="Model estimates of probability of no answer."     )
							parser.add_argument(
							    "--na-prob-thresh"							,				"-t"							,				type=snake_case							,				default=1.0							,				help="Predict \"\" if no-answer probability exceeds this (default = 1.0)."							,				)
							parser.add_argument(
							    "--out-image-dir"							,				"-p"							,				metavar="out_images"							,				default=snake_case							,				help="Save precision-recall curves to directory."     )
							parser.add_argument("--verbose"							,				"-v"							,				action="store_true"     )
							if len(sys.argv     ) == 1:
														parser.print_help()
														sys.exit(1     )
							return parser.parse_args()


def      UpperCamelCase_(	snake_case							:       int     ):

							'''simple docstring'''






							snake_case_						=      {}
							for article in dataset:
														for p in article["paragraphs"]:
																					for qa in p["qas"]:
																												snake_case_						=      bool(qa["answers"]["text"]     )
							return qid_to_has_ans


def      UpperCamelCase_(	snake_case							:       Optional[int]     ):

							'''simple docstring'''






							def remove_articles(snake_case							:       Tuple     ):
														return ARTICLES_REGEX.sub(" "							,				snake_case     )

							def white_space_fix(snake_case							:       Dict     ):
														return " ".join(text.split()     )

							def remove_punc(snake_case							:       Optional[Any]     ):
														snake_case_						=      set(string.punctuation     )
														return "".join(ch for ch in text if ch not in exclude     )

							def lower(snake_case							:       Dict     ):
														return text.lower()

							return white_space_fix(remove_articles(remove_punc(lower(snake_case     )     )     )     )


def      UpperCamelCase_(	snake_case							:       str     ):

							'''simple docstring'''






							if not s:
														return []
							return normalize_answer(snake_case     ).split()


def      UpperCamelCase_(	snake_case							:       Optional[Any]							,				snake_case							:       int     ):

							'''simple docstring'''






							return int(normalize_answer(snake_case     ) == normalize_answer(snake_case     )     )


def      UpperCamelCase_(	snake_case							:       Union[str, Any]							,				snake_case							:       int     ):

							'''simple docstring'''






							snake_case_						=      get_tokens(snake_case     )
							snake_case_						=      get_tokens(snake_case     )
							snake_case_						=      collections.Counter(snake_case     ) & collections.Counter(snake_case     )
							snake_case_						=      sum(common.values()     )
							if len(snake_case     ) == 0 or len(snake_case     ) == 0:
														# If either is no-answer, then F1 is 1 if they agree, 0 otherwise
														return int(gold_toks == pred_toks     )
							if num_same == 0:
														return 0
							snake_case_						=      1.0 * num_same / len(snake_case     )
							snake_case_						=      1.0 * num_same / len(snake_case     )
							snake_case_						=      (2 * precision * recall) / (precision + recall)
							return fa


def      UpperCamelCase_(	snake_case							:       str							,				snake_case							:       str     ):

							'''simple docstring'''






							snake_case_						=      {}
							snake_case_						=      {}
							for article in dataset:
														for p in article["paragraphs"]:
																					for qa in p["qas"]:
																												snake_case_						=      qa["id"]
																												snake_case_						=      [t for t in qa["answers"]["text"] if normalize_answer(snake_case     )]
																												if not gold_answers:
																																			# For unanswerable questions, only correct answer is empty string
																																			snake_case_						=      [""]
																												if qid not in preds:
																																			print(f'Missing prediction for {qid}'     )
																																			continue
																												snake_case_						=      preds[qid]
																												# Take max over all gold answers
																												snake_case_						=      max(compute_exact(snake_case							,				snake_case     ) for a in gold_answers     )
																												snake_case_						=      max(compute_fa(snake_case							,				snake_case     ) for a in gold_answers     )
							return exact_scores, fa_scores


def      UpperCamelCase_(	snake_case							:       List[str]							,				snake_case							:       Union[str, Any]							,				snake_case							:       Union[str, Any]							,				snake_case							:       int     ):

							'''simple docstring'''






							snake_case_						=      {}
							for qid, s in scores.items():
														snake_case_						=      na_probs[qid] > na_prob_thresh
														if pred_na:
																					snake_case_						=      float(not qid_to_has_ans[qid]     )
														else:
																					snake_case_						=      s
							return new_scores


def      UpperCamelCase_(	snake_case							:       Tuple							,				snake_case							:       str							,				snake_case							:       Union[str, Any]=None     ):

							'''simple docstring'''






							if not qid_list:
														snake_case_						=      len(snake_case     )
														return collections.OrderedDict(
														    [
														        ("exact", 100.0 * sum(exact_scores.values()     ) / total),
														        ("f1", 100.0 * sum(fa_scores.values()     ) / total),
														        ("total", total),
														    ]     )
							else:
														snake_case_						=      len(snake_case     )
														return collections.OrderedDict(
														    [
														        ("exact", 100.0 * sum(exact_scores[k] for k in qid_list     ) / total),
														        ("f1", 100.0 * sum(fa_scores[k] for k in qid_list     ) / total),
														        ("total", total),
														    ]     )


def      UpperCamelCase_(	snake_case							:       Optional[int]							,				snake_case							:       Optional[Any]							,				snake_case							:       Optional[Any]     ):

							'''simple docstring'''






							for k in new_eval:
														snake_case_						=      new_eval[k]


def      UpperCamelCase_(	snake_case							:       Optional[int]							,				snake_case							:       List[Any]							,				snake_case							:       Optional[int]							,				snake_case							:       str     ):

							'''simple docstring'''






							plt.step(snake_case							,				snake_case							,				color="b"							,				alpha=0.2							,				where="post"     )
							plt.fill_between(snake_case							,				snake_case							,				step="post"							,				alpha=0.2							,				color="b"     )
							plt.xlabel("Recall"     )
							plt.ylabel("Precision"     )
							plt.xlim([0.0, 1.05]     )
							plt.ylim([0.0, 1.05]     )
							plt.title(snake_case     )
							plt.savefig(snake_case     )
							plt.clf()


def      UpperCamelCase_(	snake_case							:       Optional[Any]							,				snake_case							:       int							,				snake_case							:       Union[str, Any]							,				snake_case							:       int							,				snake_case							:       Optional[Any]=None							,				snake_case							:       Optional[Any]=None     ):

							'''simple docstring'''






							snake_case_						=      sorted(snake_case							,				key=lambda snake_case     : na_probs[k]     )
							snake_case_						=      0.0
							snake_case_						=      1.0
							snake_case_						=      0.0
							snake_case_						=      [1.0]
							snake_case_						=      [0.0]
							snake_case_						=      0.0
							for i, qid in enumerate(snake_case     ):
														if qid_to_has_ans[qid]:
																					true_pos += scores[qid]
														snake_case_						=      true_pos / float(i + 1     )
														snake_case_						=      true_pos / float(snake_case     )
														if i == len(snake_case     ) - 1 or na_probs[qid] != na_probs[qid_list[i + 1]]:
																					# i.e., if we can put a threshold after this point
																					avg_prec += cur_p * (cur_r - recalls[-1])
																					precisions.append(snake_case     )
																					recalls.append(snake_case     )
							if out_image:
														plot_pr_curve(snake_case							,				snake_case							,				snake_case							,				snake_case     )
							return {"ap": 100.0 * avg_prec}


def      UpperCamelCase_(	snake_case							:       Optional[int]							,				snake_case							:       Dict							,				snake_case							:       Optional[int]							,				snake_case							:       Dict							,				snake_case							:       List[str]							,				snake_case							:       Dict     ):

							'''simple docstring'''






							if out_image_dir and not os.path.exists(snake_case     ):
														os.makedirs(snake_case     )
							snake_case_						=      sum(1 for v in qid_to_has_ans.values() if v     )
							if num_true_pos == 0:
														return
							snake_case_						=      make_precision_recall_eval(
							    snake_case							,				snake_case							,				snake_case							,				snake_case							,				out_image=os.path.join(snake_case							,				"pr_exact.png"     )							,				title="Precision-Recall curve for Exact Match score"							,				)
							snake_case_						=      make_precision_recall_eval(
							    snake_case							,				snake_case							,				snake_case							,				snake_case							,				out_image=os.path.join(snake_case							,				"pr_f1.png"     )							,				title="Precision-Recall curve for F1 score"							,				)
							snake_case_						=      {k: float(snake_case     ) for k, v in qid_to_has_ans.items()}
							snake_case_						=      make_precision_recall_eval(
							    snake_case							,				snake_case							,				snake_case							,				snake_case							,				out_image=os.path.join(snake_case							,				"pr_oracle.png"     )							,				title="Oracle Precision-Recall curve (binary task of HasAns vs. NoAns)"							,				)
							merge_eval(snake_case							,				snake_case							,				"pr_exact"     )
							merge_eval(snake_case							,				snake_case							,				"pr_f1"     )
							merge_eval(snake_case							,				snake_case							,				"pr_oracle"     )


def      UpperCamelCase_(	snake_case							:       int							,				snake_case							:       Optional[int]							,				snake_case							:       str							,				snake_case							:       Any     ):

							'''simple docstring'''






							if not qid_list:
														return
							snake_case_						=      [na_probs[k] for k in qid_list]
							snake_case_						=      np.ones_like(snake_case     ) / float(len(snake_case     )     )
							plt.hist(snake_case							,				weights=snake_case							,				bins=2_0							,				range=(0.0, 1.0)     )
							plt.xlabel("Model probability of no-answer"     )
							plt.ylabel("Proportion of dataset"     )
							plt.title(f'Histogram of no-answer probability: {name}'     )
							plt.savefig(os.path.join(snake_case							,				f'na_prob_hist_{name}.png'     )     )
							plt.clf()


def      UpperCamelCase_(	snake_case							:       Dict							,				snake_case							:       int							,				snake_case							:       Optional[int]							,				snake_case							:       Optional[int]     ):

							'''simple docstring'''






							snake_case_						=      sum(1 for k in qid_to_has_ans if not qid_to_has_ans[k]     )
							snake_case_						=      num_no_ans
							snake_case_						=      cur_score
							snake_case_						=      0.0
							snake_case_						=      sorted(snake_case							,				key=lambda snake_case     : na_probs[k]     )
							for i, qid in enumerate(snake_case     ):
														if qid not in scores:
																					continue
														if qid_to_has_ans[qid]:
																					snake_case_						=      scores[qid]
														else:
																					if preds[qid]:
																												snake_case_						=      -1
																					else:
																												snake_case_						=      0
														cur_score += diff
														if cur_score > best_score:
																					snake_case_						=      cur_score
																					snake_case_						=      na_probs[qid]
							return 100.0 * best_score / len(snake_case     ), best_thresh


def      UpperCamelCase_(	snake_case							:       Dict							,				snake_case							:       List[Any]							,				snake_case							:       int							,				snake_case							:       str							,				snake_case							:       Dict							,				snake_case							:       Dict     ):

							'''simple docstring'''






							snake_case_					,      snake_case_						=      find_best_thresh(snake_case							,				snake_case							,				snake_case							,				snake_case     )
							snake_case_					,      snake_case_						=      find_best_thresh(snake_case							,				snake_case							,				snake_case							,				snake_case     )
							snake_case_						=      best_exact
							snake_case_						=      exact_thresh
							snake_case_						=      best_fa
							snake_case_						=      fa_thresh


def      UpperCamelCase_(	):

							'''simple docstring'''






							with open(OPTS.data_file     ) as f:
														snake_case_						=      json.load(snake_case     )
														snake_case_						=      dataset_json["data"]
							with open(OPTS.pred_file     ) as f:
														snake_case_						=      json.load(snake_case     )
							if OPTS.na_prob_file:
														with open(OPTS.na_prob_file     ) as f:
																					snake_case_						=      json.load(snake_case     )
							else:
														snake_case_						=      {k: 0.0 for k in preds}
							snake_case_						=      make_qid_to_has_ans(snake_case     )  # maps qid to True/False
							snake_case_						=      [k for k, v in qid_to_has_ans.items() if v]
							snake_case_						=      [k for k, v in qid_to_has_ans.items() if not v]
							snake_case_					,      snake_case_						=      get_raw_scores(snake_case							,				snake_case     )
							snake_case_						=      apply_no_ans_threshold(snake_case							,				snake_case							,				snake_case							,				OPTS.na_prob_thresh     )
							snake_case_						=      apply_no_ans_threshold(snake_case							,				snake_case							,				snake_case							,				OPTS.na_prob_thresh     )
							snake_case_						=      make_eval_dict(snake_case							,				snake_case     )
							if has_ans_qids:
														snake_case_						=      make_eval_dict(snake_case							,				snake_case							,				qid_list=snake_case     )
														merge_eval(snake_case							,				snake_case							,				"HasAns"     )
							if no_ans_qids:
														snake_case_						=      make_eval_dict(snake_case							,				snake_case							,				qid_list=snake_case     )
														merge_eval(snake_case							,				snake_case							,				"NoAns"     )
							if OPTS.na_prob_file:
														find_all_best_thresh(snake_case							,				snake_case							,				snake_case							,				snake_case							,				snake_case							,				snake_case     )
							if OPTS.na_prob_file and OPTS.out_image_dir:
														run_precision_recall_analysis(snake_case							,				snake_case							,				snake_case							,				snake_case							,				snake_case							,				OPTS.out_image_dir     )
														histogram_na_prob(snake_case							,				snake_case							,				OPTS.out_image_dir							,				"hasAns"     )
														histogram_na_prob(snake_case							,				snake_case							,				OPTS.out_image_dir							,				"noAns"     )
							if OPTS.out_file:
														with open(OPTS.out_file							,				"w"     ) as f:
																					json.dump(snake_case							,				snake_case     )
							else:
														print(json.dumps(snake_case							,				indent=2     )     )


if __name__ == "__main__":
		_SCREAMING_SNAKE_CASE							:   Any      		= parse_args()
		if OPTS.out_image_dir:
				import matplotlib

				matplotlib.use("Agg")
				import matplotlib.pyplot as plt
		main()





from unittest.mock import patch

import pyspark

from datasets.packaged_modules.spark.spark import (
    Spark,
    SparkExamplesIterable,
    _generate_iterable_examples,
)

from ..utils import (
    require_dill_gt_0_3_2,
    require_not_windows,
)






def  snake_case  (       snake_case__			:Union[str, Any]   ,	snake_case__			:Dict)						->   Any:
       _A        =  []
       for part_id in partition_order:
              _A        =  df.where(F'''SPARK_PARTITION_ID() = {part_id}''').collect()
              for row_idx, row in enumerate(snake_case__):
                     expected_row_ids_and_row_dicts.append((F'''{part_id}_{row_idx}''', row.asDict()))
       return expected_row_ids_and_row_dicts






@require_not_windows
@require_dill_gt_0_3_2
def  snake_case  (       )						->   Optional[Any]:
       _A        =  pyspark.sql.SparkSession.builder.master("""local[*]""").appName("""pyspark""").getOrCreate()
       _A        =  spark.range(100).repartition(1)
       _A        =  Spark(snake_case__)
       # The id ints will be converted to Pyarrow int64s, so each row will be 8 bytes. Setting a max_shard_size of 16 means
       # that each partition can hold 2 rows.
       spark_builder._repartition_df_if_needed(max_shard_size=16)
       # Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
       assert spark_builder.df.rdd.getNumPartitions() == 50






@require_not_windows
@require_dill_gt_0_3_2
def  snake_case  (       )						->   Union[str, Any]:
       _A        =  pyspark.sql.SparkSession.builder.master("""local[*]""").appName("""pyspark""").getOrCreate()
       _A        =  spark.range(10).repartition(2)
       _A        =  [1, 0]
       _A        =  _generate_iterable_examples(snake_case__   ,	snake_case__)  # Reverse the partitions.
       _A        =  _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__   ,	snake_case__)

       for i, (row_id, row_dict) in enumerate(generate_fn()):
              _A   ,  _A        =  expected_row_ids_and_row_dicts[i]
              assert row_id == expected_row_id
              assert row_dict == expected_row_dict






@require_not_windows
@require_dill_gt_0_3_2
def  snake_case  (       )						->   int:
       _A        =  pyspark.sql.SparkSession.builder.master("""local[*]""").appName("""pyspark""").getOrCreate()
       _A        =  spark.range(10).repartition(1)
       _A        =  SparkExamplesIterable(snake_case__)
       assert it.n_shards == 1
       for i, (row_id, row_dict) in enumerate(snake_case__):
              assert row_id == F'''0_{i}'''
              assert row_dict == {"id": i}






@require_not_windows
@require_dill_gt_0_3_2
def  snake_case  (       )						->   Union[str, Any]:
       _A        =  pyspark.sql.SparkSession.builder.master("""local[*]""").appName("""pyspark""").getOrCreate()
       _A        =  spark.range(30).repartition(3)
       # Mock the generator so that shuffle reverses the partition indices.
       with patch("""numpy.random.Generator""") as generator_mock:
              _A        =  lambda snake_case__: x.reverse()
              _A        =  _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__   ,	[2, 1, 0])

              _A        =  SparkExamplesIterable(snake_case__).shuffle_data_sources(snake_case__)
              assert shuffled_it.n_shards == 3
              for i, (row_id, row_dict) in enumerate(snake_case__):
                     _A   ,  _A        =  expected_row_ids_and_row_dicts[i]
                     assert row_id == expected_row_id
                     assert row_dict == expected_row_dict






@require_not_windows
@require_dill_gt_0_3_2
def  snake_case  (       )						->   List[str]:
       _A        =  pyspark.sql.SparkSession.builder.master("""local[*]""").appName("""pyspark""").getOrCreate()
       _A        =  spark.range(20).repartition(4)

       # Partitions 0 and 2
       _A        =  SparkExamplesIterable(snake_case__).shard_data_sources(worker_id=0   ,	num_workers=2)
       assert shard_it_a.n_shards == 2
       _A        =  _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__   ,	[0, 2])
       for i, (row_id, row_dict) in enumerate(snake_case__):
              _A   ,  _A        =  expected_row_ids_and_row_dicts_a[i]
              assert row_id == expected_row_id
              assert row_dict == expected_row_dict

       # Partitions 1 and 3
       _A        =  SparkExamplesIterable(snake_case__).shard_data_sources(worker_id=1   ,	num_workers=2)
       assert shard_it_a.n_shards == 2
       _A        =  _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__   ,	[1, 3])
       for i, (row_id, row_dict) in enumerate(snake_case__):
              _A   ,  _A        =  expected_row_ids_and_row_dicts_a[i]
              assert row_id == expected_row_id
              assert row_dict == expected_row_dict







@require_not_windows
@require_dill_gt_0_3_2
def  snake_case  (       )						->   Tuple:
       _A        =  pyspark.sql.SparkSession.builder.master("""local[*]""").appName("""pyspark""").getOrCreate()
       _A        =  spark.range(100).repartition(1)
       _A        =  Spark(snake_case__)
       # Choose a small max_shard_size for maximum partitioning.
       spark_builder._repartition_df_if_needed(max_shard_size=1)
       # The new number of partitions should not be greater than the number of rows.
       assert spark_builder.df.rdd.getNumPartitions() == 100




