

"""simple docstring"""




import warnings

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from sklearn.metrics import pairwise_distances

warnings.filterwarnings('ignore')

__A	:       Optional[Any]        =						'K-MEANS-CLUST/ '
def   __SCREAMING_SNAKE_CASE		(		lowercase__     ,				lowercase__     ,				lowercase__=None ):
					"""simple docstring"""
					if seed is not None:  # useful for obtaining consistent results
										np.random.seed(lowercase__ )
					A     =   data.shape[0]  # number of data points

					# Pick K indices from range [0, N).
					A     =   np.random.randint(0     ,				lowercase__     ,				lowercase__ )

					# Keep centroids as dense format, as many entries will be nonzero due to averaging.
					# As long as at least one document in a cluster contains a word,
					# it will carry a nonzero weight in the TF-IDF vector of the centroid.
					A     =   data[rand_indices, :]

					return centroids
def   __SCREAMING_SNAKE_CASE		(		lowercase__     ,				lowercase__ ):
					"""simple docstring"""
					return pairwise_distances(lowercase__     ,				lowercase__     ,				metric="euclidean" )
def   __SCREAMING_SNAKE_CASE		(		lowercase__     ,				lowercase__ ):
					"""simple docstring"""
					# Compute distances between each data point and the set of centroids:
					# Fill in the blank (RHS only)
					A     =   centroid_pairwise_dist(lowercase__     ,				lowercase__ )

					# Compute cluster assignments for each data point:
					# Fill in the blank (RHS only)
					A     =   np.argmin(lowercase__     ,				axis=1 )

					return cluster_assignment
def   __SCREAMING_SNAKE_CASE		(		lowercase__     ,				lowercase__     ,				lowercase__ ):
					"""simple docstring"""
					A     =   []
					for i in range(lowercase__ ):
										# Select all data points that belong to cluster i. Fill in the blank (RHS only)
										A     =   data[cluster_assignment == i]
										# Compute the mean of the data points. Fill in the blank (RHS only)
										A     =   member_data_points.mean(axis=0 )
										new_centroids.append(lowercase__ )
					A     =   np.array(lowercase__ )

					return new_centroids
def   __SCREAMING_SNAKE_CASE		(		lowercase__     ,				lowercase__     ,				lowercase__     ,				lowercase__ ):
					"""simple docstring"""
					A     =   0.0
					for i in range(lowercase__ ):
										# Select all data points that belong to cluster i. Fill in the blank (RHS only)
										A     =   data[cluster_assignment == i, :]

										if member_data_points.shape[0] > 0:  # check if i-th cluster is non-empty
															# Compute distances from centroid to data points (RHS only)
															A     =   pairwise_distances(
															    lowercase__     ,				[centroids[i]]     ,				metric="euclidean" )
															A     =   distances**2
															heterogeneity += np.sum(lowercase__ )

					return heterogeneity
def   __SCREAMING_SNAKE_CASE		(		lowercase__     ,				lowercase__ ):
					"""simple docstring"""
					plt.figure(figsize=(7, 4) )
					plt.plot(lowercase__     ,				linewidth=4 )
					plt.xlabel("# Iterations" )
					plt.ylabel("Heterogeneity" )
					plt.title(F"""Heterogeneity of clustering over time, K={k:d}""" )
					plt.rcParams.update({"font.size": 16} )
					plt.show()
def   __SCREAMING_SNAKE_CASE		(		lowercase__     ,				lowercase__     ,				lowercase__     ,				lowercase__=500     ,				lowercase__=None     ,				lowercase__=False ):
					"""simple docstring"""
					A     =   initial_centroids[:]
					A     =   None

					for itr in range(lowercase__ ):
										if verbose:
															print(lowercase__     ,				end="" )

										# 1. Make cluster assignments using nearest centroids
										A     =   assign_clusters(lowercase__     ,				lowercase__ )

										# 2. Compute a new centroid for each of the k clusters, averaging all data
										#    points assigned to that cluster.
										A     =   revise_centroids(lowercase__     ,				lowercase__     ,				lowercase__ )

										# Check for convergence: if none of the assignments changed, stop
										if (
										    prev_cluster_assignment is not None
										    and (prev_cluster_assignment == cluster_assignment).all()
										):
															break

										# Print number of new assignments
										if prev_cluster_assignment is not None:
															A     =   np.sum(prev_cluster_assignment != cluster_assignment )
															if verbose:
																				print(
																				    F"""    {num_changed:5d} elements changed their cluster assignment.""" )

        # Record heterogeneity convergence metric
										if record_heterogeneity is not None:
															# YOUR CODE HERE
															A     =   compute_heterogeneity(lowercase__     ,				lowercase__     ,				lowercase__     ,				lowercase__ )
															record_heterogeneity.append(lowercase__ )

										A     =   cluster_assignment[:]

					return centroids, cluster_assignment


# Mock test below
if False:  # change to true to run this test case.
	from sklearn import datasets as ds

	__A	:       Union[str, Any]        =						ds.load_iris()
	__A	:       List[Any]        =						3
	__A	:       Optional[Any]        =						[]
	__A	:       Optional[int]        =						get_initial_centroids(dataset['data'], k, seed=0)
	__A    ,       __A	:       Dict        =						kmeans(
	    dataset['data'],
	    k,
	    initial_centroids,
	    maxiter=4_00,
	    record_heterogeneity=heterogeneity,
	    verbose=True,
	)
	plot_heterogeneity(heterogeneity, k)
def   __SCREAMING_SNAKE_CASE		(		lowercase__     ,				lowercase__     ,				lowercase__=None ):
					"""simple docstring"""
					# Fill missing values with given rules
					if fill_missing_report:
										A     =   df.fillna(value=lowercase__ )
					A     =   1
					A     =   df.select_dtypes(np.number ).columns
					A     =   (
					    df.groupby(["Cluster"] )[  # construct report dataframe
					        numeric_cols
					    ]  # group by cluster number
					    .agg(
					        [
					            ("sum", np.sum),
					            ("mean_with_zeros", lambda lowercase__ : np.mean(np.nan_to_num(lowercase__ ) )),
					            ("mean_without_zeros", lambda lowercase__ : x.replace(0     ,				np.NaN ).mean()),
					            (
					                "mean_25-75",
					                lambda lowercase__ : np.mean(
					                    np.nan_to_num(
					                        sorted(lowercase__ )[
					                            round(len(lowercase__ ) * 25 / 100 ) : round(len(lowercase__ ) * 75 / 100 )
					                        ] ) ),
					            ),
					            ("mean_with_na", np.mean),
					            ("min", lambda lowercase__ : x.min()),
					            ("5%", lambda lowercase__ : x.quantile(0.05 )),
					            ("25%", lambda lowercase__ : x.quantile(0.25 )),
					            ("50%", lambda lowercase__ : x.quantile(0.50 )),
					            ("75%", lambda lowercase__ : x.quantile(0.75 )),
					            ("95%", lambda lowercase__ : x.quantile(0.95 )),
					            ("max", lambda lowercase__ : x.max()),
					            ("count", lambda lowercase__ : x.count()),
					            ("stdev", lambda lowercase__ : x.std()),
					            ("mode", lambda lowercase__ : x.mode()[0]),
					            ("median", lambda lowercase__ : x.median()),
					            ("# > 0", lambda lowercase__ : (x > 0).sum()),
					        ] )
					    .T.reset_index()
					    .rename(index=lowercase__     ,				columns={"level_0": "Features", "level_1": "Type"} )
					)  # rename columns
					# calculate the size of cluster(count of clientID's)
					A     =   report[
					    (report["Features"] == "dummy") & (report["Type"] == "count")
					].copy()  # avoid SettingWithCopyWarning
					A     =   (
					    "ClusterSize"  # rename created cluster df to match report column names
					)
					A     =   "# of Customers"
					A     =   pd.DataFrame(
					    clustersize.iloc[:, 2:].values
					    / clustersize.iloc[:, 2:].values.sum() )
					A     =   "% of Customers"  # rename created cluster df to match report column names
					A     =   "ClusterProportion"
					A     =   clusterproportion.columns.tolist()
					A     =   cols[-2:] + cols[:-2]
					A     =   clusterproportion[cols]  # rearrange columns to match report
					A     =   report.columns
					A     =   pd.DataFrame(
					    abs(
					        report[report["Type"] == "count"].iloc[:, 2:].values
					        - clustersize.iloc[:, 2:].values ) )  # generating df with count of nan values
					A     =   0
					A     =   "# of nan"
					A     =   report[
					    report["Type"] == "count"
					].Features.tolist()  # filling values in order to match report
					A     =   a.columns.tolist()
					A     =   cols[-2:] + cols[:-2]
					A     =   a[cols]  # rearrange columns to match report
					A     =   report.columns  # rename columns to match report
					A     =   report.drop(
					    report[report.Type == "count"].index )  # drop count values except cluster size
					A     =   pd.concat(
					    [report, a, clustersize, clusterproportion]     ,				axis=0 )  # concat report with clustert size and nan values
					A     =   report["Features"].isin(lowercase__ )
					A     =   report.columns.tolist()
					A     =   cols[0:2] + cols[-1:] + cols[2:-1]
					A     =   report[cols]
					A     =   {
					    "ClusterSize": 9,
					    "ClusterProportion": 8,
					    "mean_with_zeros": 7,
					    "mean_with_na": 6,
					    "max": 5,
					    "50%": 4,
					    "min": 3,
					    "25%": 2,
					    "75%": 1,
					    "# of nan": 0,
					    "# > 0": -1,
					    "sum_with_na": -2,
					}
					A     =   (
					    report.assign(
					        Sortera=lambda lowercase__ : x.Type.map(lowercase__ )     ,				Sortera=lambda lowercase__ : list(reversed(range(len(lowercase__ ) ) ) )     ,				)
					    .sort_values(["Sorter1", "Mark", "Sorter2"]     ,				ascending=lowercase__ )
					    .drop(["Sorter1", "Sorter2"]     ,				axis=1 )
					)
					A     =   ""
					A     =   report.reset_index()
					A     =   report.drop(columns=["index"] )
					return report


if __name__ == "__main__":
	import doctest

	doctest.testmod()











'''simple docstring'''




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

from transformers import BigBirdTokenizerFast


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




def     snake_case__     (			lowerCamelCase__					: List[str]      ,		lowerCamelCase__					: Tuple=0      ,		lowerCamelCase__					: Tuple=None      ) -> Optional[Any]:
					return [" ".join(x.split(''' '''      )[begin:end]      ) for x in answers if len(x.split(''' '''      )      ) > 1]




def     snake_case__     (			lowerCamelCase__					: Any      ,		lowerCamelCase__					: Any=False      ) -> List[Any]:
					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"
										A_			:  Union[str, Any]			=				(
										    given_answers + get_sub_answers(lowerCamelCase__      ,		begin=1      ) + get_sub_answers(lowerCamelCase__      ,		end=-1      )
										)
					A_			:  List[Any]			=				[]
					for answer in given_answers:
										A_			:  Dict			=				answer.replace('''_'''      ,		''' '''      ).lower()
										A_			:  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     snake_case__     (			lowerCamelCase__					: Optional[Any]      ,		lowerCamelCase__					: Dict      ,		lowerCamelCase__					: Union[str, Any]=1      ,		lowerCamelCase__					: Any=1_0_0      ) -> Any:
					A_			,A_			:  str			=				jax.lax.top_k(lowerCamelCase__      ,		lowerCamelCase__      )
					A_			,A_			:  str			=				jax.lax.top_k(lowerCamelCase__      ,		lowerCamelCase__      )

					A_			:  List[Any]			=				best_end_idx[:, None] - best_start_idx[None, :]
					A_			:  Optional[Any]			=				jnp.logical_or(widths < 0      ,		widths > max_size      )
					A_			:  Optional[int]			=				(best_end_scores[:, None] + best_start_scores[None, :]) - (1e8 * mask)
					A_			:  str			=				jnp.argmax(lowerCamelCase__      ).item()

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




def     snake_case__     (			lowerCamelCase__					: List[str]      ) -> Union[str, Any]:
					A_			:  str			=				sample['''question''']['''text''']
					A_			:  Tuple			=				sample['''document''']['''tokens''']['''token''']
					A_			:  List[Any]			=				sample['''document''']['''tokens''']['''is_html''']
					A_			:  Tuple			=				sample['''annotations''']['''long_answer''']
					A_			:  List[Any]			=				sample['''annotations''']['''short_answers''']

					A_			:  Optional[Any]			=				''' '''.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",
															}

					A_			:  List[Any]			=				[]
					for s in short_answers:
										short_targets.extend(s['''text''']      )
					A_			:  Tuple			=				list(set(lowerCamelCase__      )      )

					A_			:  Tuple			=				[]
					for s in long_answers:
										if s["start_token"] == -1:
															continue
										A_			:  Dict			=				context[s['''start_token'''] : s['''end_token''']]
										A_			:  Tuple			=				is_html[s['''start_token'''] : s['''end_token''']]
										A_			:  Optional[Any]			=				''' '''.join([answer[i] for i in range(len(lowerCamelCase__      )      ) if not html[i]]      )
										if new_answer not in long_targets:
															long_targets.append(lowerCamelCase__      )

					A_			:  str			=				'''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__     (			) -> List[Any]:
					A_			:  int			=				load_from_disk('''natural-questions-validation'''      )
					A_			:  Optional[int]			=				dataset.map(lowerCamelCase__      ).remove_columns(['''annotations''', '''document''', '''id''']      )
					print(lowerCamelCase__      )

					A_			:  Tuple			=				dataset.filter(lambda lowerCamelCase__      : (len(x['''question''']      ) + len(x['''context''']      )) < 4 * 4_0_9_6      )
					A_			:  Optional[Any]			=				short_validation_dataset.filter(lambda lowerCamelCase__      : x["category"] != "null"      )
					short_validation_dataset

					A_			:  Union[str, Any]			=				'''vasudevgupta/flax-bigbird-natural-questions'''
					A_			:  List[str]			=				FlaxBigBirdForNaturalQuestions.from_pretrained(lowerCamelCase__      )
					A_			:  int			=				BigBirdTokenizerFast.from_pretrained(lowerCamelCase__      )

					@jax.jit
					def forward(*lowerCamelCase__					: Any      ,		**lowerCamelCase__					: Optional[Any]      ):
										A_			,A_			,A_			:  Tuple			=				model(*lowerCamelCase__      ,		**lowerCamelCase__      )
										return start_logits, end_logits, jnp.argmax(lowerCamelCase__      ,		axis=-1      )

					def evaluate(lowerCamelCase__					: List[str]      ):
										# encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
										A_			:  List[Any]			=				tokenizer(
										    example['''question''']      ,		example['''context''']      ,		return_tensors='''np'''      ,		max_length=4_0_9_6      ,		padding='''max_length'''      ,		truncation=lowerCamelCase__      ,		)

										A_			,A_			,A_			:  Optional[Any]			=				forward(**lowerCamelCase__      )

										A_			:  Optional[Any]			=				CATEGORY_MAPPING[category.item()]

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

										if predicted_category in ["yes", "no", "null"]:
															A_			:  int			=				[predicted_category]
															A_			:  Optional[int]			=				example['''output'''] == example['''targets''']
															A_			:  Dict			=				predicted_category != '''null'''
															return example

										A_			:  List[Any]			=				3_8 if predicted_category == '''short''' else 1_0_2_4
										A_			,A_			:  str			=				get_best_valid_start_end_idx(
										    start_scores[0]      ,		end_scores[0]      ,		top_k=8      ,		max_size=lowerCamelCase__      )

										A_			:  List[Any]			=				inputs['''input_ids'''][0].tolist()
										A_			:  Dict			=				[tokenizer.decode(input_ids[start_score : end_score + 1]      )]

										A_			:  Optional[int]			=				expand_to_aliases(example['''targets''']      ,		make_sub_answers=lowerCamelCase__      )
										A_			:  str			=				expand_to_aliases(example['''output''']      )

										# some preprocessing to both prediction and answer
										A_			:  Tuple			=				{''''''.join(a.split()      ) for a in answers}
										A_			:  Optional[int]			=				{''''''.join(p.split()      ) for p in predictions}
										A_			:  List[str]			=				{s for s in predictions if s not in ['''``''', '''\'\'''', '''`''', '''\'''']}

										# if there is a common element, it's a exact match
										A_			:  List[Any]			=				len(list(answers & predictions      )      ) > 0
										A_			:  Dict			=				predicted_category != '''null''' and len(lowerCamelCase__      ) > 0

										return example

					A_			:  Dict			=				short_validation_dataset.map(lowerCamelCase__      )

					A_			:  List[str]			=				len(lowerCamelCase__      )
					A_			:  Any			=				len(short_validation_dataset.filter(lambda lowerCamelCase__      : x["match"] == 1      )      )
					print('''EM score:'''      ,		(matched / total) * 1_0_0      ,		'''%'''      )


if __name__ == "__main__":
							main()



