






'''simple docstring'''





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

import numpy as np


__A							:							Any  							=		re.compile(R"""\b(a|an|the)\b""", re.UNICODE)

__A							:							List[Any]  							=		None






def    lowerCamelCase_     (    ):
 lowerCamelCase__  =     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=lowercase__				,     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=lowercase__				,     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    lowerCamelCase_     (    lowercase__):
 lowerCamelCase__  =     {}
 for article in dataset:
  for p in article["paragraphs"]:
   for qa in p["qas"]:
    lowerCamelCase__  =     bool(qa["answers"]["text"])
 return qid_to_has_ans






def    lowerCamelCase_     (    lowercase__):

 def remove_articles(lowercase__):
  return ARTICLES_REGEX.sub(" "				,     lowercase__)

 def white_space_fix(lowercase__):
  return " ".join(text.split())

 def remove_punc(lowercase__):
  lowerCamelCase__  =     set(string.punctuation)
  return "".join(ch for ch in text if ch not in exclude)

 def lower(lowercase__):
  return text.lower()

 return white_space_fix(remove_articles(remove_punc(lower(lowercase__))))






def    lowerCamelCase_     (    lowercase__):
 if not s:
  return []
 return normalize_answer(lowercase__).split()






def    lowerCamelCase_     (    lowercase__				,     lowercase__):
 return int(normalize_answer(lowercase__) == normalize_answer(lowercase__))






def    lowerCamelCase_     (    lowercase__				,     lowercase__):
 lowerCamelCase__  =     get_tokens(lowercase__)
 lowerCamelCase__  =     get_tokens(lowercase__)
 lowerCamelCase__  =     collections.Counter(lowercase__) & collections.Counter(lowercase__)
 lowerCamelCase__  =     sum(common.values())
 if len(lowercase__) == 0 or len(lowercase__) == 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
 lowerCamelCase__  =     1.0 * num_same / len(lowercase__)
 lowerCamelCase__  =     1.0 * num_same / len(lowercase__)
 lowerCamelCase__  =     (2 * precision * recall) / (precision + recall)
 return fa






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






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






def    lowerCamelCase_     (    lowercase__				,     lowercase__				,     lowercase__=None):
 if not qid_list:
  lowerCamelCase__  =     len(lowercase__)
  return collections.OrderedDict(
      [
          ("exact", 100.0 * sum(exact_scores.values()) / total),
          ("f1", 100.0 * sum(fa_scores.values()) / total),
          ("total", total),
      ])
 else:
  lowerCamelCase__  =     len(lowercase__)
  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    lowerCamelCase_     (    lowercase__				,     lowercase__				,     lowercase__):
 for k in new_eval:
  lowerCamelCase__  =     new_eval[k]






def    lowerCamelCase_     (    lowercase__				,     lowercase__				,     lowercase__				,     lowercase__):
 plt.step(lowercase__				,     lowercase__				,     color="b"				,     alpha=0.2				,     where="post")
 plt.fill_between(lowercase__				,     lowercase__				,     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(lowercase__)
 plt.savefig(lowercase__)
 plt.clf()






def    lowerCamelCase_     (    lowercase__				,     lowercase__				,     lowercase__				,     lowercase__				,     lowercase__=None				,     lowercase__=None):
 lowerCamelCase__  =     sorted(lowercase__				,     key=lambda lowercase__: na_probs[k])
 lowerCamelCase__  =     0.0
 lowerCamelCase__  =     1.0
 lowerCamelCase__  =     0.0
 lowerCamelCase__  =     [1.0]
 lowerCamelCase__  =     [0.0]
 lowerCamelCase__  =     0.0
 for i, qid in enumerate(lowercase__):
  if qid_to_has_ans[qid]:
   true_pos += scores[qid]
  lowerCamelCase__  =     true_pos / float(i + 1)
  lowerCamelCase__  =     true_pos / float(lowercase__)
  if i == len(lowercase__) - 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(lowercase__)
   recalls.append(lowercase__)
 if out_image:
  plot_pr_curve(lowercase__				,     lowercase__				,     lowercase__				,     lowercase__)
 return {"ap": 100.0 * avg_prec}






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






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






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






def    lowerCamelCase_     (    lowercase__				,     lowercase__				,     lowercase__				,     lowercase__				,     lowercase__				,     lowercase__):
 lowerCamelCase__  ,	lowerCamelCase__  =     find_best_thresh(lowercase__				,     lowercase__				,     lowercase__				,     lowercase__)
 lowerCamelCase__  ,	lowerCamelCase__  =     find_best_thresh(lowercase__				,     lowercase__				,     lowercase__				,     lowercase__)
 lowerCamelCase__  =     best_exact
 lowerCamelCase__  =     exact_thresh
 lowerCamelCase__  =     best_fa
 lowerCamelCase__  =     fa_thresh






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


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

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









'''simple docstring'''





import math

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from .attention_processor import Attention
from .embeddings import get_timestep_embedding
from .modeling_utils import ModelMixin
class 				lowercase       (							_lowerCamelCase ,       _lowerCamelCase      ):


    '''simple docstring'''

    @register_to_config
    def __init__(					self     :     Optional[Any]				,     __lowerCamelCase     :     int = 128				,     __lowerCamelCase     :     int = 256				,     __lowerCamelCase     :     float = 2_0_0_0.0				,     __lowerCamelCase     :     int = 768				,     __lowerCamelCase     :     int = 12				,     __lowerCamelCase     :     int = 12				,     __lowerCamelCase     :     int = 64				,     __lowerCamelCase     :     int = 2048				,     __lowerCamelCase     :     float = 0.1				,     )							->							str:

     '''simple docstring'''




     super().__init__()

     lowerCamelCase__  =     nn.Sequential(
         nn.Linear(__lowerCamelCase				,     d_model * 4				,     bias=__lowerCamelCase     )				,     nn.SiLU()				,     nn.Linear(d_model * 4				,     d_model * 4				,     bias=__lowerCamelCase     )				,     nn.SiLU()				,     )

     lowerCamelCase__  =     nn.Embedding(__lowerCamelCase				,     __lowerCamelCase     )
     lowerCamelCase__  =     False

     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     __lowerCamelCase				,     bias=__lowerCamelCase     )

     lowerCamelCase__  =     nn.Dropout(p=__lowerCamelCase     )

     lowerCamelCase__  =     nn.ModuleList()
     for lyr_num in range(__lowerCamelCase     ):
      # FiLM conditional T5 decoder
      lowerCamelCase__  =     DecoderLayer(d_model=__lowerCamelCase				,     d_kv=__lowerCamelCase				,     num_heads=__lowerCamelCase				,     d_ff=__lowerCamelCase				,     dropout_rate=__lowerCamelCase     )
      self.decoders.append(__lowerCamelCase     )

     lowerCamelCase__  =     TaLayerNorm(__lowerCamelCase     )

     lowerCamelCase__  =     nn.Dropout(p=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     __lowerCamelCase				,     bias=__lowerCamelCase     )

    def    a__				(					self     :     Dict				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     Optional[int]     )							->							Any:

     '''simple docstring'''




     lowerCamelCase__  =     torch.mul(query_input.unsqueeze(-1     )				,     key_input.unsqueeze(-2     )     )
     return mask.unsqueeze(-3     )




    def    a__				(					self     :     Tuple				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     Union[str, Any]     )							->							Optional[int]:

     '''simple docstring'''




     lowerCamelCase__  ,	lowerCamelCase__  ,	lowerCamelCase__  =     decoder_input_tokens.shape
     assert decoder_noise_time.shape == (batch,)

     # decoder_noise_time is in [0, 1), so rescale to expected timing range.
     lowerCamelCase__  =     get_timestep_embedding(
         decoder_noise_time * self.config.max_decoder_noise_time				,     embedding_dim=self.config.d_model				,     max_period=self.config.max_decoder_noise_time				,     ).to(dtype=self.dtype     )

     lowerCamelCase__  =     self.conditioning_emb(__lowerCamelCase     ).unsqueeze(1     )

     assert conditioning_emb.shape == (batch, 1, self.config.d_model * 4)

     lowerCamelCase__  =     decoder_input_tokens.shape[1]

     # If we want to use relative positions for audio context, we can just offset
     # this sequence by the length of encodings_and_masks.
     lowerCamelCase__  =     torch.broadcast_to(
         torch.arange(__lowerCamelCase				,     device=decoder_input_tokens.device     )				,     (batch, seq_length)				,     )

     lowerCamelCase__  =     self.position_encoding(__lowerCamelCase     )

     lowerCamelCase__  =     self.continuous_inputs_projection(__lowerCamelCase     )
     inputs += position_encodings
     lowerCamelCase__  =     self.dropout(__lowerCamelCase     )

     # decoder: No padding present.
     lowerCamelCase__  =     torch.ones(
         decoder_input_tokens.shape[:2]				,     device=decoder_input_tokens.device				,     dtype=inputs.dtype     )

     # Translate encoding masks to encoder-decoder masks.
     lowerCamelCase__  =     [(x, self.encoder_decoder_mask(__lowerCamelCase				,     __lowerCamelCase     )) for x, y in encodings_and_masks]

     # cross attend style: concat encodings
     lowerCamelCase__  =     torch.cat([x[0] for x in encodings_and_encdec_masks]				,     dim=1     )
     lowerCamelCase__  =     torch.cat([x[1] for x in encodings_and_encdec_masks]				,     dim=-1     )

     for lyr in self.decoders:
      lowerCamelCase__  =     lyr(
          __lowerCamelCase				,     conditioning_emb=__lowerCamelCase				,     encoder_hidden_states=__lowerCamelCase				,     encoder_attention_mask=__lowerCamelCase				,     )[0]

     lowerCamelCase__  =     self.decoder_norm(__lowerCamelCase     )
     lowerCamelCase__  =     self.post_dropout(__lowerCamelCase     )

     lowerCamelCase__  =     self.spec_out(__lowerCamelCase     )
     return spec_out
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     Union[str, Any]				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     str=1E-6     )							->							int:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     nn.ModuleList()

     # cond self attention: layer 0
     self.layer.append(
         TaLayerSelfAttentionCond(d_model=__lowerCamelCase				,     d_kv=__lowerCamelCase				,     num_heads=__lowerCamelCase				,     dropout_rate=__lowerCamelCase     )     )

     # cross attention: layer 1
     self.layer.append(
         TaLayerCrossAttention(
             d_model=__lowerCamelCase				,     d_kv=__lowerCamelCase				,     num_heads=__lowerCamelCase				,     dropout_rate=__lowerCamelCase				,     layer_norm_epsilon=__lowerCamelCase				,     )     )

     # Film Cond MLP + dropout: last layer
     self.layer.append(
         TaLayerFFCond(d_model=__lowerCamelCase				,     d_ff=__lowerCamelCase				,     dropout_rate=__lowerCamelCase				,     layer_norm_epsilon=__lowerCamelCase     )     )




    def    a__				(					self     :     List[Any]				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     Optional[Any]=None				,     __lowerCamelCase     :     Optional[int]=None				,     __lowerCamelCase     :     int=None				,     __lowerCamelCase     :     List[Any]=None				,     __lowerCamelCase     :     List[str]=None				,     )							->							Dict:

     '''simple docstring'''




     lowerCamelCase__  =     self.layer[0](
         __lowerCamelCase				,     conditioning_emb=__lowerCamelCase				,     attention_mask=__lowerCamelCase				,     )

     if encoder_hidden_states is not None:
      lowerCamelCase__  =     torch.where(encoder_attention_mask > 0				,     0				,     -1E10     ).to(
          encoder_hidden_states.dtype     )

      lowerCamelCase__  =     self.layer[1](
          __lowerCamelCase				,     key_value_states=__lowerCamelCase				,     attention_mask=__lowerCamelCase				,     )

     # Apply Film Conditional Feed Forward layer
     lowerCamelCase__  =     self.layer[-1](__lowerCamelCase				,     __lowerCamelCase     )

     return (hidden_states,)
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     List[Any]				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     Tuple     )							->							Optional[int]:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     TaLayerNorm(__lowerCamelCase     )
     lowerCamelCase__  =     TaFiLMLayer(in_features=d_model * 4				,     out_features=__lowerCamelCase     )
     lowerCamelCase__  =     Attention(query_dim=__lowerCamelCase				,     heads=__lowerCamelCase				,     dim_head=__lowerCamelCase				,     out_bias=__lowerCamelCase				,     scale_qk=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Dropout(__lowerCamelCase     )




    def    a__				(					self     :     Dict				,     __lowerCamelCase     :     Any				,     __lowerCamelCase     :     Optional[Any]=None				,     __lowerCamelCase     :     str=None				,     )							->							List[Any]:

     '''simple docstring'''




     lowerCamelCase__  =     self.layer_norm(__lowerCamelCase     )

     if conditioning_emb is not None:
      lowerCamelCase__  =     self.FiLMLayer(__lowerCamelCase				,     __lowerCamelCase     )

     # Self-attention block
     lowerCamelCase__  =     self.attention(__lowerCamelCase     )

     lowerCamelCase__  =     hidden_states + self.dropout(__lowerCamelCase     )

     return hidden_states
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     Any				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     Any				,     __lowerCamelCase     :     Dict     )							->							Union[str, Any]:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     Attention(query_dim=__lowerCamelCase				,     heads=__lowerCamelCase				,     dim_head=__lowerCamelCase				,     out_bias=__lowerCamelCase				,     scale_qk=__lowerCamelCase     )
     lowerCamelCase__  =     TaLayerNorm(__lowerCamelCase				,     eps=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Dropout(__lowerCamelCase     )




    def    a__				(					self     :     List[str]				,     __lowerCamelCase     :     Any				,     __lowerCamelCase     :     Optional[Any]=None				,     __lowerCamelCase     :     Dict=None				,     )							->							Any:

     '''simple docstring'''




     lowerCamelCase__  =     self.layer_norm(__lowerCamelCase     )
     lowerCamelCase__  =     self.attention(
         __lowerCamelCase				,     encoder_hidden_states=__lowerCamelCase				,     attention_mask=attention_mask.squeeze(1     )				,     )
     lowerCamelCase__  =     hidden_states + self.dropout(__lowerCamelCase     )
     return layer_output
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     List[str]				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     List[str]     )							->							Optional[Any]:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     TaDenseGatedActDense(d_model=__lowerCamelCase				,     d_ff=__lowerCamelCase				,     dropout_rate=__lowerCamelCase     )
     lowerCamelCase__  =     TaFiLMLayer(in_features=d_model * 4				,     out_features=__lowerCamelCase     )
     lowerCamelCase__  =     TaLayerNorm(__lowerCamelCase				,     eps=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Dropout(__lowerCamelCase     )




    def    a__				(					self     :     Any				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     Union[str, Any]=None     )							->							Dict:

     '''simple docstring'''




     lowerCamelCase__  =     self.layer_norm(__lowerCamelCase     )
     if conditioning_emb is not None:
      lowerCamelCase__  =     self.film(__lowerCamelCase				,     __lowerCamelCase     )

     lowerCamelCase__  =     self.DenseReluDense(__lowerCamelCase     )
     lowerCamelCase__  =     hidden_states + self.dropout(__lowerCamelCase     )
     return hidden_states
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     List[str]				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     str     )							->							str:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     __lowerCamelCase				,     bias=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     __lowerCamelCase				,     bias=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     __lowerCamelCase				,     bias=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Dropout(__lowerCamelCase     )
     lowerCamelCase__  =     NewGELUActivation()




    def    a__				(					self     :     Dict				,     __lowerCamelCase     :     Union[str, Any]     )							->							Optional[Any]:

     '''simple docstring'''




     lowerCamelCase__  =     self.act(self.wi_a(__lowerCamelCase     )     )
     lowerCamelCase__  =     self.wi_a(__lowerCamelCase     )
     lowerCamelCase__  =     hidden_gelu * hidden_linear
     lowerCamelCase__  =     self.dropout(__lowerCamelCase     )

     lowerCamelCase__  =     self.wo(__lowerCamelCase     )
     return hidden_states
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     str				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     List[Any]=1E-6     )							->							Union[str, Any]:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     nn.Parameter(torch.ones(__lowerCamelCase     )     )
     lowerCamelCase__  =     eps




    def    a__				(					self     :     Optional[int]				,     __lowerCamelCase     :     str     )							->							int:

     '''simple docstring'''




     lowerCamelCase__  =     hidden_states.to(torch.floataa     ).pow(2     ).mean(-1				,     keepdim=__lowerCamelCase     )
     lowerCamelCase__  =     hidden_states * torch.rsqrt(variance + self.variance_epsilon     )

     # convert into half-precision if necessary
     if self.weight.dtype in [torch.floataa, torch.bfloataa]:
      lowerCamelCase__  =     hidden_states.to(self.weight.dtype     )

     return self.weight * hidden_states
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def    a__				(					self     :     Tuple				,     __lowerCamelCase     :     torch.Tensor     )							->							torch.Tensor:

     '''simple docstring'''




     return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi     ) * (input + 0.0_4_4_7_1_5 * torch.pow(__lowerCamelCase				,     3.0     ))     ))
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     Optional[int]				,     __lowerCamelCase     :     Union[str, Any]				,     __lowerCamelCase     :     List[Any]     )							->							Optional[int]:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     out_features * 2				,     bias=__lowerCamelCase     )




    def    a__				(					self     :     str				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     str     )							->							List[str]:

     '''simple docstring'''




     lowerCamelCase__  =     self.scale_bias(__lowerCamelCase     )
     lowerCamelCase__  ,	lowerCamelCase__  =     torch.chunk(__lowerCamelCase				,     2				,     -1     )
     lowerCamelCase__  =     x * (1 + scale) + shift
     return x

