

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

import numpy as np


SCREAMING_SNAKE_CASE_        =      re.compile(R'\b(a|an|the)\b', re.UNICODE)

SCREAMING_SNAKE_CASE_        =      None

def 							__SCREAMING_SNAKE_CASE							(				)   ->       Tuple:
				_UpperCAmelCase       :		List[Any] 		=    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=lowerCAmelCase_  ,      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=lowerCAmelCase_  ,      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 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Optional[int] )   ->       str:
				_UpperCAmelCase       :		Tuple 		=    {}
				for article in dataset:
								for p in article["paragraphs"]:
												for qa in p["qas"]:
																_UpperCAmelCase       :		Dict 		=    bool(qa["answers"]["text"] )
				return qid_to_has_ans

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Union[str, Any] )   ->       List[str]:


				def remove_articles(lowerCAmelCase:  Dict ):
								return ARTICLES_REGEX.sub(" "  ,      lowerCAmelCase_ )

				def white_space_fix(lowerCAmelCase:  Optional[Any] ):
								return " ".join(text.split() )

				def remove_punc(lowerCAmelCase:  str ):
								_UpperCAmelCase       :		List[Any] 		=    set(string.punctuation )
								return "".join(ch for ch in text if ch not in exclude )

				def lower(lowerCAmelCase:  List[str] ):
								return text.lower()

				return white_space_fix(remove_articles(remove_punc(lower(lowerCAmelCase_ ) ) ) )

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  List[str] )   ->       Dict:


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

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Dict  ,      lowerCAmelCase:  Optional[Any] )   ->       Any:
				return int(normalize_answer(lowerCAmelCase_ ) == normalize_answer(lowerCAmelCase_ ) )

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Dict  ,      lowerCAmelCase:  List[Any] )   ->       List[str]:
				_UpperCAmelCase       :		int 		=    get_tokens(lowerCAmelCase_ )
				_UpperCAmelCase       :		Tuple 		=    get_tokens(lowerCAmelCase_ )
				_UpperCAmelCase       :		Any 		=    collections.Counter(lowerCAmelCase_ ) & collections.Counter(lowerCAmelCase_ )
				_UpperCAmelCase       :		Union[str, Any] 		=    sum(common.values() )
				if len(lowerCAmelCase_ ) == 0 or len(lowerCAmelCase_ ) == 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
				_UpperCAmelCase       :		Optional[Any] 		=    1.0 * num_same / len(lowerCAmelCase_ )
				_UpperCAmelCase       :		Any 		=    1.0 * num_same / len(lowerCAmelCase_ )
				_UpperCAmelCase       :		Union[str, Any] 		=    (2 * precision * recall) / (precision + recall)
				return fa

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Union[str, Any]  ,      lowerCAmelCase:  List[Any] )   ->       List[str]:
				_UpperCAmelCase       :		str 		=    {}
				_UpperCAmelCase       :		Union[str, Any] 		=    {}
				for article in dataset:
								for p in article["paragraphs"]:
												for qa in p["qas"]:
																_UpperCAmelCase       :		Union[str, Any] 		=    qa["id"]
																_UpperCAmelCase       :		Tuple 		=    [t for t in qa["answers"]["text"] if normalize_answer(lowerCAmelCase_ )]
																if not gold_answers:
																				# For unanswerable questions, only correct answer is empty string
																				_UpperCAmelCase       :		Dict 		=    [""]
																if qid not in preds:
																				print(F'Missing prediction for {qid}' )
																				continue
																_UpperCAmelCase       :		str 		=    preds[qid]
																# Take max over all gold answers
																_UpperCAmelCase       :		int 		=    max(compute_exact(lowerCAmelCase_  ,      lowerCAmelCase_ ) for a in gold_answers )
																_UpperCAmelCase       :		List[str] 		=    max(compute_fa(lowerCAmelCase_  ,      lowerCAmelCase_ ) for a in gold_answers )
				return exact_scores, fa_scores

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  List[Any]  ,      lowerCAmelCase:  Any  ,      lowerCAmelCase:  List[Any]  ,      lowerCAmelCase:  int )   ->       List[Any]:
				_UpperCAmelCase       :		int 		=    {}
				for qid, s in scores.items():
								_UpperCAmelCase       :		Union[str, Any] 		=    na_probs[qid] > na_prob_thresh
								if pred_na:
												_UpperCAmelCase       :		Any 		=    float(not qid_to_has_ans[qid] )
								else:
												_UpperCAmelCase       :		Dict 		=    s
				return new_scores

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Optional[int]  ,      lowerCAmelCase:  Tuple  ,      lowerCAmelCase:  Any=None )   ->       int:


				if not qid_list:
								_UpperCAmelCase       :		List[Any] 		=    len(lowerCAmelCase_ )
								return collections.OrderedDict(
								    [
								        ("exact", 100.0 * sum(exact_scores.values() ) / total),
								        ("f1", 100.0 * sum(fa_scores.values() ) / total),
								        ("total", total),
								    ] )
				else:
								_UpperCAmelCase       :		Tuple 		=    len(lowerCAmelCase_ )
								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 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Any  ,      lowerCAmelCase:  str  ,      lowerCAmelCase:  Optional[int] )   ->       Tuple:


				for k in new_eval:
								_UpperCAmelCase       :		Tuple 		=    new_eval[k]

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Tuple  ,      lowerCAmelCase:  Dict  ,      lowerCAmelCase:  List[Any]  ,      lowerCAmelCase:  int )   ->       List[str]:
				plt.step(lowerCAmelCase_  ,      lowerCAmelCase_  ,      color="b"  ,      alpha=0.2  ,      where="post" )
				plt.fill_between(lowerCAmelCase_  ,      lowerCAmelCase_  ,      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(lowerCAmelCase_ )
				plt.savefig(lowerCAmelCase_ )
				plt.clf()

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Optional[int]  ,      lowerCAmelCase:  Union[str, Any]  ,      lowerCAmelCase:  Any  ,      lowerCAmelCase:  int  ,      lowerCAmelCase:  List[Any]=None  ,      lowerCAmelCase:  Union[str, Any]=None )   ->       int:
				_UpperCAmelCase       :		Optional[Any] 		=    sorted(lowerCAmelCase_  ,      key=lambda lowerCAmelCase : na_probs[k] )
				_UpperCAmelCase       :		List[Any] 		=    0.0
				_UpperCAmelCase       :		List[Any] 		=    1.0
				_UpperCAmelCase       :		Optional[Any] 		=    0.0
				_UpperCAmelCase       :		Union[str, Any] 		=    [1.0]
				_UpperCAmelCase       :		int 		=    [0.0]
				_UpperCAmelCase       :		Tuple 		=    0.0
				for i, qid in enumerate(lowerCAmelCase_ ):
								if qid_to_has_ans[qid]:
												true_pos += scores[qid]
								_UpperCAmelCase       :		List[Any] 		=    true_pos / float(i + 1 )
								_UpperCAmelCase       :		List[str] 		=    true_pos / float(lowerCAmelCase_ )
								if i == len(lowerCAmelCase_ ) - 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(lowerCAmelCase_ )
												recalls.append(lowerCAmelCase_ )
				if out_image:
								plot_pr_curve(lowerCAmelCase_  ,      lowerCAmelCase_  ,      lowerCAmelCase_  ,      lowerCAmelCase_ )
				return {"ap": 100.0 * avg_prec}

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  List[str]  ,      lowerCAmelCase:  List[Any]  ,      lowerCAmelCase:  List[Any]  ,      lowerCAmelCase:  Any  ,      lowerCAmelCase:  Optional[int]  ,      lowerCAmelCase:  int )   ->       Any:


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

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Any  ,      lowerCAmelCase:  Union[str, Any]  ,      lowerCAmelCase:  str  ,      lowerCAmelCase:  int )   ->       Optional[int]:


				if not qid_list:
								return
				_UpperCAmelCase       :		Tuple 		=    [na_probs[k] for k in qid_list]
				_UpperCAmelCase       :		str 		=    np.ones_like(lowerCAmelCase_ ) / float(len(lowerCAmelCase_ ) )
				plt.hist(lowerCAmelCase_  ,      weights=lowerCAmelCase_  ,      bins=20  ,      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(lowerCAmelCase_  ,      F'na_prob_hist_{name}.png' ) )
				plt.clf()

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Any  ,      lowerCAmelCase:  List[Any]  ,      lowerCAmelCase:  List[Any]  ,      lowerCAmelCase:  int )   ->       Union[str, Any]:
				_UpperCAmelCase       :		Dict 		=    sum(1 for k in qid_to_has_ans if not qid_to_has_ans[k] )
				_UpperCAmelCase       :		Optional[Any] 		=    num_no_ans
				_UpperCAmelCase       :		Union[str, Any] 		=    cur_score
				_UpperCAmelCase       :		List[str] 		=    0.0
				_UpperCAmelCase       :		str 		=    sorted(lowerCAmelCase_  ,      key=lambda lowerCAmelCase : na_probs[k] )
				for i, qid in enumerate(lowerCAmelCase_ ):
								if qid not in scores:
												continue
								if qid_to_has_ans[qid]:
												_UpperCAmelCase       :		Tuple 		=    scores[qid]
								else:
												if preds[qid]:
																_UpperCAmelCase       :		str 		=    -1
												else:
																_UpperCAmelCase       :		int 		=    0
								cur_score += diff
								if cur_score > best_score:
												_UpperCAmelCase       :		int 		=    cur_score
												_UpperCAmelCase       :		List[Any] 		=    na_probs[qid]
				return 100.0 * best_score / len(lowerCAmelCase_ ), best_thresh

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Optional[int]  ,      lowerCAmelCase:  Optional[Any]  ,      lowerCAmelCase:  int  ,      lowerCAmelCase:  str  ,      lowerCAmelCase:  Optional[Any]  ,      lowerCAmelCase:  Optional[int] )   ->       Union[str, Any]:
				_UpperCAmelCase       :		Optional[int] 		=    find_best_thresh(lowerCAmelCase_  ,      lowerCAmelCase_  ,      lowerCAmelCase_  ,      lowerCAmelCase_ )
				_UpperCAmelCase       :		List[Any] 		=    find_best_thresh(lowerCAmelCase_  ,      lowerCAmelCase_  ,      lowerCAmelCase_  ,      lowerCAmelCase_ )
				_UpperCAmelCase       :		Optional[int] 		=    best_exact
				_UpperCAmelCase       :		List[str] 		=    exact_thresh
				_UpperCAmelCase       :		Optional[Any] 		=    best_fa
				_UpperCAmelCase       :		List[str] 		=    fa_thresh

def 							__SCREAMING_SNAKE_CASE							(				)   ->       Tuple:


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


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

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





import argparse
import json
import os
from collections import OrderedDict

import numpy as np
import tensorflow as tf
import torch

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Any )   ->       Optional[int]:
				_UpperCAmelCase       :		Optional[Any] 		=    os.path.join(args.tf_model_dir  ,      "parameters.json" )
				_UpperCAmelCase       :		List[str] 		=    json.loads(open(lowerCAmelCase ).read() )
				if not params:
								raise ValueError(
								    F'It seems that the json file at {parameter_file} is empty. Make sure you have a correct json file.' )
				if not args.output.endswith(".pt" ):
								_UpperCAmelCase       :		Union[str, Any] 		=    args.output + ".pt"
				_UpperCAmelCase       :		Dict 		=    OrderedDict()
				with tf.device("/CPU:0" ):
								_UpperCAmelCase       :		int 		=    tf.train.load_checkpoint(args.tf_model_dir )
								_UpperCAmelCase       :		int 		=    reader.get_variable_to_shape_map()
								for key_name in shapes.keys():
												_UpperCAmelCase       :		Dict 		=    reader.get_tensor(lowerCAmelCase ).astype(np.floataa )
												if key_name.endswith("/adam_m" ) or key_name.endswith("/adam_v" ):
																continue
												if key_name.startswith("pasts/" ):
																if key_name.startswith("pasts/mlp" ):
																				_UpperCAmelCase       :		int 		=    int(key_name[9] )
																elif key_name.startswith("pasts/out" ):
																				_UpperCAmelCase       :		Any 		=    8
																_UpperCAmelCase       :		str 		=    "model.sqout.%d.weight" % (player * 2)  # enter to nn.Sequencial with Tanh, so 2 at a time
																_UpperCAmelCase       :		Optional[Any] 		=    vnp.transpose([1, 0] ).copy()  # Mesh-Tensorflow is a diagonal matrix
																_UpperCAmelCase       :		str 		=    torch.tensor(lowerCAmelCase )
												elif key_name.startswith("model/moe" ):
																_UpperCAmelCase       :		List[str] 		=    int(key_name[9:].split("/" )[0] )
																if key_name.endswith("/switch_gating/kernel" ):
																				_UpperCAmelCase       :		Union[str, Any] 		=    "model.blocks.%d.feed_forward.mlp.router.classifier.weight" % player
																				_UpperCAmelCase       :		List[str] 		=    vnp.transpose([1, 0] ).copy()  # Mesh-Tensorflow is a diagonal matrix
																				_UpperCAmelCase       :		Tuple 		=    torch.tensor(lowerCAmelCase )
																elif key_name.endswith("/softmlp/kernel" ):
																				_UpperCAmelCase       :		List[str] 		=    "model.blocks.%d.feed_forward.soft_bypass_mlp.weight" % player
																				_UpperCAmelCase       :		List[str] 		=    vnp.transpose([1, 0] ).copy()  # Mesh-Tensorflow is a diagonal matrix
																				_UpperCAmelCase       :		Union[str, Any] 		=    torch.tensor(lowerCAmelCase )
																elif key_name.endswith("/wo/kernel" ) or key_name.endswith("/wi/kernel" ):
																				_UpperCAmelCase       :		str 		=    key_name[-9:-7]
																				for i in range(16 ):
																								_UpperCAmelCase       :		Optional[int] 		=    "model.blocks.%d.feed_forward.mlp.experts.expert_%d.%s.weight" % (player, i, nlayer)
																								_UpperCAmelCase       :		Optional[int] 		=    (
																								    vnp[i].transpose([1, 0] ).copy()
																								)  # In Mesh-Tensorflow, it is one array, so it is divided
																								_UpperCAmelCase       :		Dict 		=    torch.tensor(lowerCAmelCase )
												elif key_name.startswith("model/mlp" ):
																_UpperCAmelCase       :		Tuple 		=    int(key_name[9:].split("/" )[0] )
																if key_name.endswith("/p1/kernel" ):
																				_UpperCAmelCase       :		Any 		=    "model.blocks.%d.feed_forward.mlp.wi.weight" % player
																				_UpperCAmelCase       :		int 		=    vnp.transpose([1, 0] ).copy()  # Mesh-Tensorflow is a diagonal matrix
																				_UpperCAmelCase       :		List[str] 		=    torch.tensor(lowerCAmelCase )
																elif key_name.endswith("/p1/bias" ):
																				_UpperCAmelCase       :		Optional[Any] 		=    "model.blocks.%d.feed_forward.mlp.wi.bias" % player
																				_UpperCAmelCase       :		Tuple 		=    vnp.copy()  # same because it is one dimensional
																				_UpperCAmelCase       :		Any 		=    torch.tensor(lowerCAmelCase )
																elif key_name.endswith("/p2/kernel" ):
																				_UpperCAmelCase       :		str 		=    "model.blocks.%d.feed_forward.mlp.wo.weight" % player
																				_UpperCAmelCase       :		Union[str, Any] 		=    vnp.transpose([1, 0] ).copy()  # Mesh-Tensorflow is a diagonal matrix
																				_UpperCAmelCase       :		List[Any] 		=    torch.tensor(lowerCAmelCase )
																elif key_name.endswith("/p2/bias" ):
																				_UpperCAmelCase       :		int 		=    "model.blocks.%d.feed_forward.mlp.wo.bias" % player
																				_UpperCAmelCase       :		Optional[int] 		=    vnp.copy()  # same because it is one dimensional
																				_UpperCAmelCase       :		str 		=    torch.tensor(lowerCAmelCase )
												elif key_name.startswith("model/ln" ):
																_UpperCAmelCase       :		Union[str, Any] 		=    int(key_name[8:].split("/" )[0] )
																if key_name.endswith("/b" ):
																				_UpperCAmelCase       :		Union[str, Any] 		=    "model.blocks.%d.feed_forward.norm.bias" % player
																				_UpperCAmelCase       :		int 		=    vnp.copy()  # same because it is one dimensional
																				_UpperCAmelCase       :		Union[str, Any] 		=    torch.tensor(lowerCAmelCase )
																elif key_name.endswith("/g" ):
																				_UpperCAmelCase       :		Optional[Any] 		=    "model.blocks.%d.feed_forward.norm.weight" % player
																				_UpperCAmelCase       :		Union[str, Any] 		=    vnp.copy()  # same because it is one dimensional
																				_UpperCAmelCase       :		List[Any] 		=    torch.tensor(lowerCAmelCase )
												elif key_name.startswith("model/att" ):
																_UpperCAmelCase       :		Optional[Any] 		=    int(key_name[9:].split("/" )[0] )
																if key_name.endswith("/qkv/kernel" ):
																				_UpperCAmelCase       :		List[Any] 		=    vnp.copy()  # Compute same dimension as Mesh-tensorflow using einsum
																				_UpperCAmelCase       :		Dict 		=    state[:, 0, :, :]
																				_UpperCAmelCase       :		List[Any] 		=    state[:, 1, :, :]
																				_UpperCAmelCase       :		Tuple 		=    state[:, 2, :, :]
																				_UpperCAmelCase       :		List[str] 		=    (
																				    state_q.reshape([state_q.shape[0], state_q.shape[1] * state_q.shape[2]] )
																				    .transpose([1, 0] )
																				    .copy()
																				)  # Mesh-Tensorflow is a diagonal matrix
																				_UpperCAmelCase       :		int 		=    (
																				    state_k.reshape([state_k.shape[0], state_k.shape[1] * state_k.shape[2]] )
																				    .transpose([1, 0] )
																				    .copy()
																				)  # Mesh-Tensorflow is a diagonal matrix
																				_UpperCAmelCase       :		Tuple 		=    (
																				    state_v.reshape([state_v.shape[0], state_v.shape[1] * state_v.shape[2]] )
																				    .transpose([1, 0] )
																				    .copy()
																				)  # Mesh-Tensorflow is a diagonal matrix
																				_UpperCAmelCase       :		int 		=    "model.blocks.%d.self_attn.self_attn.q_proj.weight" % player
																				_UpperCAmelCase       :		Any 		=    torch.tensor(lowerCAmelCase )
																				_UpperCAmelCase       :		List[Any] 		=    "model.blocks.%d.self_attn.self_attn.k_proj.weight" % player
																				_UpperCAmelCase       :		Any 		=    torch.tensor(lowerCAmelCase )
																				_UpperCAmelCase       :		Union[str, Any] 		=    "model.blocks.%d.self_attn.self_attn.v_proj.weight" % player
																				_UpperCAmelCase       :		str 		=    torch.tensor(lowerCAmelCase )
																elif key_name.endswith("/o/kernel" ):
																				_UpperCAmelCase       :		Dict 		=    "model.blocks.%d.self_attn.self_attn.out_proj.weight" % player
																				_UpperCAmelCase       :		Tuple 		=    (
																				    vnp.reshape([vnp.shape[0] * vnp.shape[1], vnp.shape[2]] ).transpose([1, 0] ).copy()
																				)  # Mesh-Tensorflow is a diagonal matrix
																				_UpperCAmelCase       :		List[str] 		=    torch.tensor(lowerCAmelCase )
												elif key_name.startswith("model/an" ):
																_UpperCAmelCase       :		Optional[int] 		=    int(key_name[8:].split("/" )[0] )
																if key_name.endswith("/b" ):
																				_UpperCAmelCase       :		Any 		=    "model.blocks.%d.self_attn.norm.bias" % player
																				_UpperCAmelCase       :		Any 		=    vnp.copy()  # same because it is one dimensional
																				_UpperCAmelCase       :		Dict 		=    torch.tensor(lowerCAmelCase )
																elif key_name.endswith("/g" ):
																				_UpperCAmelCase       :		Dict 		=    "model.blocks.%d.self_attn.norm.weight" % player
																				_UpperCAmelCase       :		Any 		=    vnp.copy()  # same because it is one dimensional
																				_UpperCAmelCase       :		Tuple 		=    torch.tensor(lowerCAmelCase )
												elif (
												    key_name.startswith("model/wte" )
												    or key_name.startswith("model/wpe" )
												    or key_name.startswith("model/ete" )
												):
																_UpperCAmelCase       :		Optional[Any] 		=    {"wte": "embed_tokens", "wpe": "position_embeddings", "ete": "extra_position_embeddings"}[
																    key_name[-3:]
																]
																_UpperCAmelCase       :		Dict 		=    "model.%s.weight" % nlayer
																_UpperCAmelCase       :		Optional[Any] 		=    vnp.copy()  # same in embedded
																_UpperCAmelCase       :		List[Any] 		=    torch.tensor(lowerCAmelCase )
																if key_name.startswith("model/wte" ):
																				_UpperCAmelCase       :		List[str] 		=    "lm_head.weight"
																				_UpperCAmelCase       :		Optional[int] 		=    vnp.copy()  # same in embedded
																				_UpperCAmelCase       :		List[str] 		=    torch.tensor(lowerCAmelCase )
												elif key_name.startswith("model/wob" ):
																_UpperCAmelCase       :		Dict 		=    "final_logits_bias"
																_UpperCAmelCase       :		Optional[Any] 		=    vnp.copy()  # same in embedded
																_UpperCAmelCase       :		str 		=    state.reshape((1, -1) )
																_UpperCAmelCase       :		Dict 		=    torch.tensor(lowerCAmelCase )
												elif key_name == "model/dense/kernel":
																_UpperCAmelCase       :		List[Any] 		=    "model.last_project.weight"
																_UpperCAmelCase       :		int 		=    vnp.transpose([1, 0] ).copy()  # Mesh-Tensorflow is a diagonal matrix
																_UpperCAmelCase       :		str 		=    torch.tensor(lowerCAmelCase )
												elif key_name == "model/dense_1/bias":
																_UpperCAmelCase       :		List[Any] 		=    "model.last_project.bias"
																_UpperCAmelCase       :		Any 		=    vnp.copy()  # same because it is one dimensional
																_UpperCAmelCase       :		Any 		=    torch.tensor(lowerCAmelCase )
				torch.save(lowerCAmelCase  ,      args.output )


if __name__ == "__main__":
					SCREAMING_SNAKE_CASE_        =      argparse.ArgumentParser(
					    description='model converter.', formatter_class=argparse.ArgumentDefaultsHelpFormatter
					)
					parser.add_argument('--tf_model_dir', metavar='PATH', type=str, required=True, help='import model')
					parser.add_argument('--output', metavar='PATH', type=str, required=True, help='output model')
					SCREAMING_SNAKE_CASE_        =      parser.parse_args()
					convert_tf_gptsan_to_pt(args)


