code: 
  '''simple docstring'''
  
  
  
  
  
  import logging
  
  import torch
  from torch import nn
  from torch.nn import CrossEntropyLoss, MSELoss
  
  from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
  from transformers.models.albert.modeling_albert import (
      ALBERT_INPUTS_DOCSTRING,
      ALBERT_START_DOCSTRING,
      AlbertModel,
      AlbertPreTrainedModel,
      AlbertTransformer,
  )
  
  
  A         =   logging.getLogger(__name__)
  
  
  
  class  lowerCAmelCase__   (       lowerCAmelCase_       ):
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    def 		_lowerCamelCase	(self  : List[str]	, lowerCAmelCase_  : Any	, lowerCAmelCase_  : Union[str, Any]	, lowerCAmelCase_  : List[Any]=None	, lowerCAmelCase_  : int=None    )   ->       List[str]:
  
  
  
  
  
          if current_layer == 0:
                snake_case_						       =							self.embedding_hidden_mapping_in(__a    )
          else:
                snake_case_						       =							hidden_states[0]
  
          snake_case_						       =							int(self.config.num_hidden_layers / self.config.num_hidden_groups    )
  
          # Index of the hidden group
          snake_case_						       =							int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)    )
  
          snake_case_						       =							self.albert_layer_groups[group_idx](
              __a	, __a	, head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]	, )
          snake_case_						       =							layer_group_output[0]
  
          return (hidden_states,)
  
  
  
  @add_start_docstrings(
      'The bare ALBERT Model transformer with PABEE outputting raw hidden-states without any specific head on top.'					,							lowerCAmelCase_					,							)
  class  lowerCAmelCase__   (       lowerCAmelCase_       ):
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    def __init__(self  : List[Any]	, lowerCAmelCase_  : str    )   ->       Optional[int]:
          super().__init__(__a    )
  
          snake_case_						       =							AlbertTransformerWithPabee(__a    )
  
          self.init_weights()
          snake_case_						       =							0
          snake_case_						       =							0
          snake_case_						       =							0
  
          snake_case_						       =							0
  
    def 		_lowerCamelCase	(self  : int	, lowerCAmelCase_  : Optional[Any]    )   ->       int:
          snake_case_						       =							threshold
  
    def 		_lowerCamelCase	(self  : Any	, lowerCAmelCase_  : Any    )   ->       str:
          snake_case_						       =							patience
  
    def 		_lowerCamelCase	(self  : List[Any]    )   ->       str:
          snake_case_						       =							0
          snake_case_						       =							0
  
    def 		_lowerCamelCase	(self  : Dict    )   ->       List[Any]:
          snake_case_						       =							self.inference_layers_num / self.inference_instances_num
          snake_case_						       =							(
              f'''*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ='''
              f''' {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***'''
          )
          print(__a    )
  
  
    @add_start_docstrings_to_model_forward(__a    )
    def 		_lowerCamelCase	(self  : Optional[int]	, lowerCAmelCase_  : List[Any]=None	, lowerCAmelCase_  : Union[str, Any]=None	, lowerCAmelCase_  : Any=None	, lowerCAmelCase_  : List[Any]=None	, lowerCAmelCase_  : Optional[int]=None	, lowerCAmelCase_  : Optional[int]=None	, lowerCAmelCase_  : List[Any]=None	, lowerCAmelCase_  : Union[str, Any]=None	, lowerCAmelCase_  : Optional[Any]=False	, )   ->       List[str]:
  
  
  
  
  
          if input_ids is not None and inputs_embeds is not None:
                raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time"    )
          elif input_ids is not None:
                snake_case_						       =							input_ids.size()
          elif inputs_embeds is not None:
                snake_case_						       =							inputs_embeds.size()[:-1]
          else:
                raise ValueError("You have to specify either input_ids or inputs_embeds"    )
  
          snake_case_						       =							input_ids.device if input_ids is not None else inputs_embeds.device
  
          if attention_mask is None:
                snake_case_						       =							torch.ones(__a	, device=__a    )
          if token_type_ids is None:
                snake_case_						       =							torch.zeros(__a	, dtype=torch.long	, device=__a    )
  
          snake_case_						       =							attention_mask.unsqueeze(1    ).unsqueeze(2    )
          snake_case_						       =							extended_attention_mask.to(dtype=self.dtype    )  # fp16 compatibility
          snake_case_						       =							(1.0 - extended_attention_mask) * -10_000.0
          snake_case_						       =							self.get_head_mask(__a	, self.config.num_hidden_layers    )
  
          snake_case_						       =							self.embeddings(
              __a	, position_ids=__a	, token_type_ids=__a	, inputs_embeds=__a    )
          snake_case_						       =							embedding_output
  
          if self.training:
                snake_case_						       =							[]
                for i in range(self.config.num_hidden_layers    ):
                      snake_case_						       =							self.encoder.adaptive_forward(
                          __a	, current_layer=__a	, attention_mask=__a	, head_mask=__a	, )
  
                      snake_case_						       =							self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]    )    )
                      snake_case_						       =							output_layers[i](output_dropout(__a    )    )
                      res.append(__a    )
          elif self.patience == 0:  # Use all layers for inference
                snake_case_						       =							self.encoder(__a	, __a	, head_mask=__a    )
                snake_case_						       =							self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]    )    )
                snake_case_						       =							[output_layers[self.config.num_hidden_layers - 1](__a    )]
          else:
                snake_case_						       =							0
                snake_case_						       =							None
                snake_case_						       =							0
                for i in range(self.config.num_hidden_layers    ):
                      calculated_layer_num += 1
                      snake_case_						       =							self.encoder.adaptive_forward(
                          __a	, current_layer=__a	, attention_mask=__a	, head_mask=__a	, )
  
                      snake_case_						       =							self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]    )    )
                      snake_case_						       =							output_layers[i](__a    )
                      if regression:
                            snake_case_						       =							logits.detach()
                            if patient_result is not None:
                                  snake_case_						       =							patient_result.detach()
                            if (patient_result is not None) and torch.abs(patient_result - labels    ) < self.regression_threshold:
                                  patient_counter += 1
                            else:
                                  snake_case_						       =							0
                      else:
                            snake_case_						       =							logits.detach().argmax(dim=1    )
                            if patient_result is not None:
                                  snake_case_						       =							patient_result.detach().argmax(dim=1    )
                            if (patient_result is not None) and torch.all(labels.eq(__a    )    ):
                                  patient_counter += 1
                            else:
                                  snake_case_						       =							0
  
                      snake_case_						       =							logits
                      if patient_counter == self.patience:
                            break
                snake_case_						       =							[patient_result]
                self.inference_layers_num += calculated_layer_num
                self.inference_instances_num += 1
  
          return res
  
  
  
  @add_start_docstrings(
      'Albert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of\n    the pooled output) e.g. for GLUE tasks. '					,							lowerCAmelCase_					,							)
  class  lowerCAmelCase__   (       lowerCAmelCase_       ):
  
  
  
    '''simple docstring'''
  
  
  
  
  
  
    def __init__(self  : List[str]	, lowerCAmelCase_  : str    )   ->       Optional[Any]:
          super().__init__(__a    )
          snake_case_						       =							config.num_labels
  
          snake_case_						       =							AlbertModelWithPabee(__a    )
          snake_case_						       =							nn.Dropout(config.classifier_dropout_prob    )
          snake_case_						       =							nn.ModuleList(
              [nn.Linear(config.hidden_size	, self.config.num_labels    ) for _ in range(config.num_hidden_layers    )]    )
  
          self.init_weights()
  
  
    @add_start_docstrings_to_model_forward(__a    )
    def 		_lowerCamelCase	(self  : Optional[Any]	, lowerCAmelCase_  : Union[str, Any]=None	, lowerCAmelCase_  : Union[str, Any]=None	, lowerCAmelCase_  : int=None	, lowerCAmelCase_  : Union[str, Any]=None	, lowerCAmelCase_  : Optional[int]=None	, lowerCAmelCase_  : int=None	, lowerCAmelCase_  : Union[str, Any]=None	, )   ->       Optional[Any]:
          snake_case_						       =							self.albert(
              input_ids=__a	, attention_mask=__a	, token_type_ids=__a	, position_ids=__a	, head_mask=__a	, inputs_embeds=__a	, output_dropout=self.dropout	, output_layers=self.classifiers	, regression=self.num_labels == 1	, )
  
          snake_case_						       =							(logits[-1],)
  
          if labels is not None:
                snake_case_						       =							None
                snake_case_						       =							0
                for ix, logits_item in enumerate(__a    ):
                      if self.num_labels == 1:
                            #  We are doing regression
                            snake_case_						       =							MSELoss()
                            snake_case_						       =							loss_fct(logits_item.view(-1    )	, labels.view(-1    )    )
                      else:
                            snake_case_						       =							CrossEntropyLoss()
                            snake_case_						       =							loss_fct(logits_item.view(-1	, self.num_labels    )	, labels.view(-1    )    )
                      if total_loss is None:
                            snake_case_						       =							loss
                      else:
                            total_loss += loss * (ix + 1)
                      total_weights += ix + 1
                snake_case_						       =							(total_loss / total_weights,) + outputs
  
          return outputs
  
  
style_context: 
  '''simple docstring'''
  
  
  
  
  
  import jax
  import jax.numpy as jnp
  from bigbird_flax import FlaxBigBirdForNaturalQuestions
  from datasets import load_from_disk
  
  from transformers import BigBirdTokenizerFast
  
  
  A           =   {0: "null", 1: "short", 2: "long", 3: "yes", 4: "no"}
  A           =   set("".join(["‘", "’", "´", "`", ".", ",", "-", "\""]))
  def   snake_case_       (							_a     :							Any ,					_a     :							Union[str, Any]=0 ,					_a     :							Optional[int]=None			):
        return [" ".join(x.split(" "			)[begin:end]			) for x in answers if len(x.split(" "			)			) > 1]
  def   snake_case_       (							_a     :							Optional[int] ,					_a     :							str=False			):
        if make_sub_answers:
              # if answers are longer than one word, make sure a predictions is correct if it coresponds to the complete 1: or :-1 sub word
              # *e.g.* if the correct answer contains a prefix such as "the", or "a"
              snake_case_						       =							(
                  given_answers + get_sub_answers(_a ,					begin=1			) + get_sub_answers(_a ,					end=-1			)
              )
        snake_case_						       =							[]
        for answer in given_answers:
              snake_case_						       =							answer.replace("_" ,					" "			).lower()
              snake_case_						       =							"".join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else " " for c in alias			)
              answers.append(" ".join(alias.split()			).strip()			)
        return set(_a			)
  def   snake_case_       (							_a     :							Optional[Any] ,					_a     :							int ,					_a     :							Optional[Any]=1 ,					_a     :							List[Any]=100			):
        snake_case_       ,	snake_case_						       =							jax.lax.top_k(_a ,					_a			)
        snake_case_       ,	snake_case_						       =							jax.lax.top_k(_a ,					_a			)
  
        snake_case_						       =							best_end_idx[:, None] - best_start_idx[None, :]
        snake_case_						       =							jnp.logical_or(widths < 0 ,					widths > max_size			)
        snake_case_						       =							(best_end_scores[:, None] + best_start_scores[None, :]) - (1E8 * mask)
        snake_case_						       =							jnp.argmax(_a			).item()
  
        return best_start_idx[best_score % top_k], best_end_idx[best_score // top_k]
  def   snake_case_       (							_a     :							Any			):
        snake_case_						       =							sample["question"]["text"]
        snake_case_						       =							sample["document"]["tokens"]["token"]
        snake_case_						       =							sample["document"]["tokens"]["is_html"]
        snake_case_						       =							sample["annotations"]["long_answer"]
        snake_case_						       =							sample["annotations"]["short_answers"]
  
        snake_case_						       =							" ".join([context[i] for i in range(len(_a			)			) if not is_html[i]]			)
  
        # 0 - No ; 1 - Yes
        for answer in sample["annotations"]["yes_no_answer"]:
              if answer == 0 or answer == 1:
                    return {
                        "question": question,
                        "context": context_string,
                        "short": [],
                        "long": [],
                        "category": "no" if answer == 0 else "yes",
                    }
  
        snake_case_						       =							[]
        for s in short_answers:
              short_targets.extend(s["text"]			)
        snake_case_						       =							list(set(_a			)			)
  
        snake_case_						       =							[]
        for s in long_answers:
              if s["start_token"] == -1:
                    continue
              snake_case_						       =							context[s["start_token"] : s["end_token"]]
              snake_case_						       =							is_html[s["start_token"] : s["end_token"]]
              snake_case_						       =							" ".join([answer[i] for i in range(len(_a			)			) if not html[i]]			)
              if new_answer not in long_targets:
                    long_targets.append(_a			)
  
        snake_case_						       =							"long_short" if len(short_targets + long_targets			) > 0 else "null"
  
        return {
            "question": question,
            "context": context_string,
            "short": short_targets,
            "long": long_targets,
            "category": category,
        }
  def   snake_case_       (							):
        snake_case_						       =							load_from_disk("natural-questions-validation"			)
        snake_case_						       =							dataset.map(_a			).remove_columns(["annotations", "document", "id"]			)
        print(_a			)
  
        snake_case_						       =							dataset.filter(lambda _a			: (len(x["question"]			) + len(x["context"]			)) < 4 * 4_096			)
        snake_case_						       =							short_validation_dataset.filter(lambda _a			: x["category"] != "null"			)
        short_validation_dataset
  
        snake_case_						       =							"vasudevgupta/flax-bigbird-natural-questions"
        snake_case_						       =							FlaxBigBirdForNaturalQuestions.from_pretrained(_a			)
        snake_case_						       =							BigBirdTokenizerFast.from_pretrained(_a			)
  
        @jax.jit
        def forward(*_a     :							List[Any] ,					**_a     :							int			):
              snake_case_       ,	snake_case_       ,	snake_case_						       =							model(*_a ,					**_a			)
              return start_logits, end_logits, jnp.argmax(_a ,					axis=-1			)
  
        def evaluate(_a     :							List[Any]			):
              # encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
              snake_case_						       =							tokenizer(
                  example["question"] ,					example["context"] ,					return_tensors="np" ,					max_length=4_096 ,					padding="max_length" ,					truncation=_a ,					)
  
              snake_case_       ,	snake_case_       ,	snake_case_						       =							forward(**_a			)
  
              snake_case_						       =							CATEGORY_MAPPING[category.item()]
  
              snake_case_						       =							example["long"] + example["short"]
              if example["category"] in ["yes", "no", "null"]:
                    snake_case_						       =							[example["category"]]
              snake_case_						       =							example["category"] != "null"
              # Now target can be: "yes", "no", "null", "list of long & short answers"
  
              if predicted_category in ["yes", "no", "null"]:
                    snake_case_						       =							[predicted_category]
                    snake_case_						       =							example["output"] == example["targets"]
                    snake_case_						       =							predicted_category != "null"
                    return example
  
              snake_case_						       =							38 if predicted_category == "short" else 1_024
              snake_case_       ,	snake_case_						       =							get_best_valid_start_end_idx(
                  start_scores[0] ,					end_scores[0] ,					top_k=8 ,					max_size=_a			)
  
              snake_case_						       =							inputs["input_ids"][0].tolist()
              snake_case_						       =							[tokenizer.decode(input_ids[start_score : end_score + 1]			)]
  
              snake_case_						       =							expand_to_aliases(example["targets"] ,					make_sub_answers=_a			)
              snake_case_						       =							expand_to_aliases(example["output"]			)
  
              # some preprocessing to both prediction and answer
              snake_case_						       =							{"".join(a.split()			) for a in answers}
              snake_case_						       =							{"".join(p.split()			) for p in predictions}
              snake_case_						       =							{s for s in predictions if s not in ["``", "''", "`", "'"]}
  
              # if there is a common element, it's a exact match
              snake_case_						       =							len(list(answers & predictions			)			) > 0
              snake_case_						       =							predicted_category != "null" and len(_a			) > 0
  
              return example
  
        snake_case_						       =							short_validation_dataset.map(_a			)
  
        snake_case_						       =							len(_a			)
        snake_case_						       =							len(short_validation_dataset.filter(lambda _a			: x["match"] == 1			)			)
        print("EM score:" ,					(matched / total) * 100 ,					"%"			)
  
  
  if __name__ == "__main__":
       main()
  
  