

import argparse
import csv
import json
import math
import time

import numpy as np
import torch
import torch.optim as optim
import torch.utils.data as data
from nltk.tokenize.treebank import TreebankWordDetokenizer
from pplm_classification_head import ClassificationHead
from torch import nn
from torchtext import data as torchtext_data
from torchtext import datasets
from tqdm import tqdm, trange

from transformers import GPTaLMHeadModel, GPTaTokenizer


torch.manual_seed(0)
np.random.seed(0)
lowerCamelCase 	=   1e-10
lowerCamelCase 	=   '''This is incredible! I love it, this is the best chicken I have ever had.'''
lowerCamelCase 	=   100







class 		__magic_name__     (				nn.Module				):

   '''simple docstring'''




   def __init__(							self,       lowercase_,       lowercase_="gpt2-medium",       lowercase_=False,       lowercase_="cpu"			)				->       Optional[Any]:






     """simple docstring"""






     super().__init__()
     a__    =GPTaTokenizer.from_pretrained(UpperCamelCase_			)
     a__    =GPTaLMHeadModel.from_pretrained(UpperCamelCase_			)
     a__    =self.encoder.transformer.config.hidden_size
     a__    =ClassificationHead(class_size=UpperCamelCase_,       embed_size=self.embed_size			)
     a__    =cached_mode
     a__    =device



   def 					_UpperCAmelCase					(							self			)				->       Union[str, Any]:






     """simple docstring"""






     return self.classifier_head



   def 					_UpperCAmelCase					(							self			)				->       Optional[Any]:






     """simple docstring"""






     for param in self.encoder.parameters():
       a__    =False
     self.classifier_head.train()



   def 					_UpperCAmelCase					(							self,       lowercase_			)				->       Dict:






     """simple docstring"""






     a__    =x.ne(0			).unsqueeze(2			).repeat(1,       1,       self.embed_size			).float().to(self.device			).detach()
     a__    =self.encoder.transformer(UpperCamelCase_			)['''last_hidden_state''']
     a__    =hidden * mask
     a__    =torch.sum(UpperCamelCase_,       dim=1			) / (torch.sum(UpperCamelCase_,       dim=1			).detach() + EPSILON)
     return avg_hidden



   def 					_UpperCAmelCase					(							self,       lowercase_			)				->       Any:






     """simple docstring"""






     if self.cached_mode:
       a__    =x.to(self.device			)
     else:
       a__    =self.avg_representation(x.to(self.device			)			)

     a__    =self.classifier_head(UpperCamelCase_			)
     a__    =nn.functional.log_softmax(UpperCamelCase_,       dim=-1			)

     return probs







class 		__magic_name__     (				data.Dataset				):

   '''simple docstring'''




   def __init__(							self,       lowercase_,       lowercase_			)				->       str:






     """simple docstring"""






     a__    =X
     a__    =y



   def __len__(							self			)				->       Any:






     """simple docstring"""






     return len(self.X			)



   def __getitem__(							self,       lowercase_			)				->       Any:






     """simple docstring"""






     a__    ={}
     a__    =self.X[index]
     a__    =self.y[index]
     return data





def      UpperCAmelCase__							(     _A :					Dict				):





  '''simple docstring'''







  def pad_sequences(_A :					Optional[Any]				):
    a__    =[len(_snake_case				) for seq in sequences]

    a__    =torch.zeros(len(_snake_case				)      , max(_snake_case				)				).long()  # padding value = 0

    for i, seq in enumerate(_snake_case				):
      a__    =lengths[i]
      a__    =seq[:end]

    return padded_sequences, lengths

  a__    ={}
  for key in data[0].keys():
    a__    =[d[key] for d in data]

  a__    =pad_sequences(item_info['''X''']				)
  a__    =torch.tensor(item_info['''y''']      , dtype=torch.long				)

  return x_batch, y_batch




def      UpperCAmelCase__							(     _A :					List[Any]				):





  '''simple docstring'''







  a__    ={}
  for key in data[0].keys():
    a__    =[d[key] for d in data]

  a__    =torch.cat(item_info['''X''']      , 0				)
  a__    =torch.tensor(item_info['''y''']      , dtype=torch.long				)

  return x_batch, y_batch




def      UpperCAmelCase__							(     _A :					List[str]      , _A :					str      , _A :					Any      , _A :					int=0      , _A :					Union[str, Any]=10      , _A :					Optional[int]="cpu"				):





  '''simple docstring'''







  a__    =0
  discriminator.train_custom()
  for batch_idx, (input_t, target_t) in enumerate(_snake_case				):
    a__    =input_t.to(_snake_case				), target_t.to(_snake_case				)

    optimizer.zero_grad()

    a__    =discriminator(_snake_case				)
    a__    =nn.functional.nll_loss(_snake_case      , _snake_case				)
    loss.backward(retain_graph=_snake_case				)
    optimizer.step()

    samples_so_far += len(_snake_case				)

    if batch_idx % log_interval == 0:
      print(
          '''Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'''.format(
              epoch + 1      , _snake_case      , len(data_loader.dataset				)      , 1_00 * samples_so_far / len(data_loader.dataset				)      , loss.item()      , )				)




def      UpperCAmelCase__							(     _A :					Optional[int]      , _A :					str      , _A :					Optional[Any]="cpu"				):





  '''simple docstring'''







  discriminator.eval()
  a__    =0
  a__    =0
  with torch.no_grad():
    for input_t, target_t in data_loader:
      a__    =input_t.to(_snake_case				), target_t.to(_snake_case				)
      a__    =discriminator(_snake_case				)
      # sum up batch loss
      test_loss += nn.functional.nll_loss(_snake_case      , _snake_case      , reduction='''sum'''				).item()
      # get the index of the max log-probability
      a__    =output_t.argmax(dim=1      , keepdim=_snake_case				)
      correct += pred_t.eq(target_t.view_as(_snake_case				)				).sum().item()

  test_loss /= len(data_loader.dataset				)

  print(
      '''Performance on test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'''.format(
          _snake_case      , _snake_case      , len(data_loader.dataset				)      , 100.0 * correct / len(data_loader.dataset				)				)				)




def      UpperCAmelCase__							(     _A :					List[str]      , _A :					Union[str, Any]      , _A :					Optional[int]      , _A :					List[str]=False      , _A :					Tuple="cpu"				):





  '''simple docstring'''







  a__    =model.tokenizer.encode(_snake_case				)
  a__    =torch.tensor([input_t]      , dtype=torch.long      , device=_snake_case				)
  if cached:
    a__    =model.avg_representation(_snake_case				)

  a__    =model(_snake_case				).data.cpu().numpy().flatten().tolist()
  print('''Input sentence:'''      , _snake_case				)
  print(
      '''Predictions:'''      , ''', '''.join('''{}: {:.4f}'''.format(_snake_case      , math.exp(_snake_case				)				) for c, log_prob in zip(_snake_case      , _snake_case				)				)      , )




def      UpperCAmelCase__							(     _A :					Any      , _A :					List[Any]      , _A :					Union[str, Any]      , _A :					Any=False      , _A :					List[Any]="cpu"				):





  '''simple docstring'''







  a__    =torch.utils.data.DataLoader(dataset=_snake_case      , batch_size=_snake_case      , collate_fn=_snake_case				)

  a__    =[]
  a__    =[]
  for batch_idx, (x, y) in enumerate(tqdm(_snake_case      , ascii=_snake_case				)				):
    with torch.no_grad():
      a__    =x.to(_snake_case				)
      a__    =discriminator.avg_representation(_snake_case				).cpu().detach()
      a__    =torch.unbind(avg_rep.unsqueeze(1				)				)
      xs += avg_rep_list
      ys += y.cpu().numpy().tolist()

  a__    =torch.utils.data.DataLoader(
      dataset=Dataset(_snake_case      , _snake_case				)      , batch_size=_snake_case      , shuffle=_snake_case      , collate_fn=_snake_case				)

  return data_loader




def      UpperCAmelCase__							(     _A :					Dict      , _A :					Optional[int]=None      , _A :					int="gpt2-medium"      , _A :					str=10      , _A :					int=64      , _A :					Optional[Any]=10      , _A :					str=False      , _A :					Optional[int]=False      , _A :					List[str]=False      , ):





  '''simple docstring'''







  a__    ='''cuda''' if torch.cuda.is_available() and not no_cuda else '''cpu'''

  print('''Preprocessing {} dataset...'''.format(_snake_case				)				)
  a__    =time.time()

  if dataset == "SST":
    a__    =['''positive''', '''negative''', '''very positive''', '''very negative''', '''neutral''']
    a__    ={c: i for i, c in enumerate(_snake_case				)}

    a__    =Discriminator(
        class_size=len(_snake_case				)      , pretrained_model=_snake_case      , cached_mode=_snake_case      , device=_snake_case				).to(_snake_case				)

    a__    =torchtext_data.Field()
    a__    =torchtext_data.Field(sequential=_snake_case				)
    a__    =datasets.SST.splits(
        _snake_case      , _snake_case      , fine_grained=_snake_case      , train_subtrees=_snake_case      , )

    a__    =[]
    a__    =[]
    for i in trange(len(_snake_case				)      , ascii=_snake_case				):
      a__    =TreebankWordDetokenizer().detokenize(vars(train_data[i]				)['''text''']				)
      a__    =discriminator.tokenizer.encode(_snake_case				)
      a__    =torch.tensor([5_02_56] + seq      , device=_snake_case      , dtype=torch.long				)
      x.append(_snake_case				)
      y.append(classaidx[vars(train_data[i]				)['''label''']]				)
    a__    =Dataset(_snake_case      , _snake_case				)

    a__    =[]
    a__    =[]
    for i in trange(len(_snake_case				)      , ascii=_snake_case				):
      a__    =TreebankWordDetokenizer().detokenize(vars(test_data[i]				)['''text''']				)
      a__    =discriminator.tokenizer.encode(_snake_case				)
      a__    =torch.tensor([5_02_56] + seq      , device=_snake_case      , dtype=torch.long				)
      test_x.append(_snake_case				)
      test_y.append(classaidx[vars(test_data[i]				)['''label''']]				)
    a__    =Dataset(_snake_case      , _snake_case				)

    a__    ={
        '''class_size''': len(_snake_case				),
        '''embed_size''': discriminator.embed_size,
        '''pretrained_model''': pretrained_model,
        '''class_vocab''': classaidx,
        '''default_class''': 2,
    }

  elif dataset == "clickbait":
    a__    =['''non_clickbait''', '''clickbait''']
    a__    ={c: i for i, c in enumerate(_snake_case				)}

    a__    =Discriminator(
        class_size=len(_snake_case				)      , pretrained_model=_snake_case      , cached_mode=_snake_case      , device=_snake_case				).to(_snake_case				)

    with open('''datasets/clickbait/clickbait_train_prefix.txt'''				) as f:
      a__    =[]
      for i, line in enumerate(_snake_case				):
        try:
          data.append(eval(_snake_case				)				)
        except Exception:
          print('''Error evaluating line {}: {}'''.format(_snake_case      , _snake_case				)				)
          continue
    a__    =[]
    a__    =[]
    with open('''datasets/clickbait/clickbait_train_prefix.txt'''				) as f:
      for i, line in enumerate(tqdm(_snake_case      , ascii=_snake_case				)				):
        try:
          a__    =eval(_snake_case				)
          a__    =discriminator.tokenizer.encode(d['''text''']				)

          if len(_snake_case				) < max_length_seq:
            a__    =torch.tensor([5_02_56] + seq      , device=_snake_case      , dtype=torch.long				)
          else:
            print('''Line {} is longer than maximum length {}'''.format(_snake_case      , _snake_case				)				)
            continue
          x.append(_snake_case				)
          y.append(d['''label''']				)
        except Exception:
          print('''Error evaluating / tokenizing line {}, skipping it'''.format(_snake_case				)				)
          pass

    a__    =Dataset(_snake_case      , _snake_case				)
    a__    =int(0.9 * len(_snake_case				)				)
    a__    =len(_snake_case				) - train_size
    a__    =torch.utils.data.random_split(_snake_case      , [train_size, test_size]				)

    a__    ={
        '''class_size''': len(_snake_case				),
        '''embed_size''': discriminator.embed_size,
        '''pretrained_model''': pretrained_model,
        '''class_vocab''': classaidx,
        '''default_class''': 1,
    }

  elif dataset == "toxic":
    a__    =['''non_toxic''', '''toxic''']
    a__    ={c: i for i, c in enumerate(_snake_case				)}

    a__    =Discriminator(
        class_size=len(_snake_case				)      , pretrained_model=_snake_case      , cached_mode=_snake_case      , device=_snake_case				).to(_snake_case				)

    a__    =[]
    a__    =[]
    with open('''datasets/toxic/toxic_train.txt'''				) as f:
      for i, line in enumerate(tqdm(_snake_case      , ascii=_snake_case				)				):
        try:
          a__    =eval(_snake_case				)
          a__    =discriminator.tokenizer.encode(d['''text''']				)

          if len(_snake_case				) < max_length_seq:
            a__    =torch.tensor([5_02_56] + seq      , device=_snake_case      , dtype=torch.long				)
          else:
            print('''Line {} is longer than maximum length {}'''.format(_snake_case      , _snake_case				)				)
            continue
          x.append(_snake_case				)
          y.append(int(np.sum(d['''label''']				) > 0				)				)
        except Exception:
          print('''Error evaluating / tokenizing line {}, skipping it'''.format(_snake_case				)				)
          pass

    a__    =Dataset(_snake_case      , _snake_case				)
    a__    =int(0.9 * len(_snake_case				)				)
    a__    =len(_snake_case				) - train_size
    a__    =torch.utils.data.random_split(_snake_case      , [train_size, test_size]				)

    a__    ={
        '''class_size''': len(_snake_case				),
        '''embed_size''': discriminator.embed_size,
        '''pretrained_model''': pretrained_model,
        '''class_vocab''': classaidx,
        '''default_class''': 0,
    }

  else:  # if dataset == "generic":
    # This assumes the input dataset is a TSV with the following structure:
    # class \t text

    if dataset_fp is None:
      raise ValueError('''When generic dataset is selected, dataset_fp needs to be specified aswell.'''				)

    a__    =set()
    with open(_snake_case				) as f:
      a__    =csv.reader(_snake_case      , delimiter='''\t'''				)
      for row in tqdm(_snake_case      , ascii=_snake_case				):
        if row:
          classes.add(row[0]				)

    a__    =sorted(_snake_case				)
    a__    ={c: i for i, c in enumerate(_snake_case				)}

    a__    =Discriminator(
        class_size=len(_snake_case				)      , pretrained_model=_snake_case      , cached_mode=_snake_case      , device=_snake_case				).to(_snake_case				)

    a__    =[]
    a__    =[]
    with open(_snake_case				) as f:
      a__    =csv.reader(_snake_case      , delimiter='''\t'''				)
      for i, row in enumerate(tqdm(_snake_case      , ascii=_snake_case				)				):
        if row:
          a__    =row[0]
          a__    =row[1]

          try:
            a__    =discriminator.tokenizer.encode(_snake_case				)
            if len(_snake_case				) < max_length_seq:
              a__    =torch.tensor([5_02_56] + seq      , device=_snake_case      , dtype=torch.long				)

            else:
              print('''Line {} is longer than maximum length {}'''.format(_snake_case      , _snake_case				)				)
              continue

            x.append(_snake_case				)
            y.append(classaidx[label]				)

          except Exception:
            print('''Error tokenizing line {}, skipping it'''.format(_snake_case				)				)
            pass

    a__    =Dataset(_snake_case      , _snake_case				)
    a__    =int(0.9 * len(_snake_case				)				)
    a__    =len(_snake_case				) - train_size
    a__    =torch.utils.data.random_split(_snake_case      , [train_size, test_size]				)

    a__    ={
        '''class_size''': len(_snake_case				),
        '''embed_size''': discriminator.embed_size,
        '''pretrained_model''': pretrained_model,
        '''class_vocab''': classaidx,
        '''default_class''': 0,
    }

  a__    =time.time()
  print('''Preprocessed {} data points'''.format(len(_snake_case				) + len(_snake_case				)				)				)
  print('''Data preprocessing took: {:.3f}s'''.format(end - start				)				)

  if cached:
    print('''Building representation cache...'''				)

    a__    =time.time()

    a__    =get_cached_data_loader(_snake_case      , _snake_case      , _snake_case      , shuffle=_snake_case      , device=_snake_case				)

    a__    =get_cached_data_loader(_snake_case      , _snake_case      , _snake_case      , device=_snake_case				)

    a__    =time.time()
    print('''Building representation cache took: {:.3f}s'''.format(end - start				)				)

  else:
    a__    =torch.utils.data.DataLoader(
        dataset=_snake_case      , batch_size=_snake_case      , shuffle=_snake_case      , collate_fn=_snake_case				)
    a__    =torch.utils.data.DataLoader(dataset=_snake_case      , batch_size=_snake_case      , collate_fn=_snake_case				)

  if save_model:
    with open('''{}_classifier_head_meta.json'''.format(_snake_case				)      , '''w'''				) as meta_file:
      json.dump(_snake_case      , _snake_case				)

  a__    =optim.Adam(discriminator.parameters()      , lr=0.0_001				)

  for epoch in range(_snake_case				):
    a__    =time.time()
    print('''\nEpoch'''      , epoch + 1				)

    train_epoch(
        discriminator=_snake_case      , data_loader=_snake_case      , optimizer=_snake_case      , epoch=_snake_case      , log_interval=_snake_case      , device=_snake_case      , )
    evaluate_performance(data_loader=_snake_case      , discriminator=_snake_case      , device=_snake_case				)

    a__    =time.time()
    print('''Epoch took: {:.3f}s'''.format(end - start				)				)

    print('''\nExample prediction'''				)
    predict(_snake_case      , _snake_case      , _snake_case      , cached=_snake_case      , device=_snake_case				)

    if save_model:
      # torch.save(discriminator.state_dict(),
      #           "{}_discriminator_{}.pt".format(
      #               args.dataset, epoch + 1
      #               ))
      torch.save(
          discriminator.get_classifier().state_dict()      , '''{}_classifier_head_epoch_{}.pt'''.format(_snake_case      , epoch + 1				)      , )


if __name__ == "__main__":
       lowerCamelCase 	=   argparse.ArgumentParser(description='''Train a discriminator on top of GPT-2 representations''')
       parser.add_argument(
           '''--dataset''',
           type=str,
           default='''SST''',
           choices=('''SST''', '''clickbait''', '''toxic''', '''generic'''),
           help=(
               '''dataset to train the discriminator on.'''
               '''In case of generic, the dataset is expected'''
               '''to be a TSBV file with structure: class \\t text'''
           ),
       )
       parser.add_argument(
           '''--dataset_fp''',
           type=str,
           default='''''',
           help='''File path of the dataset to use. Needed only in case of generic datadset''',
       )
       parser.add_argument(
           '''--pretrained_model''', type=str, default='''gpt2-medium''', help='''Pretrained model to use as encoder'''
       )
       parser.add_argument('''--epochs''', type=int, default=10, metavar='''N''', help='''Number of training epochs''')
       parser.add_argument(
           '''--batch_size''', type=int, default=64, metavar='''N''', help='''input batch size for training (default: 64)'''
       )
       parser.add_argument(
           '''--log_interval''',
           type=int,
           default=10,
           metavar='''N''',
           help='''how many batches to wait before logging training status''',
       )
       parser.add_argument('''--save_model''', action='''store_true''', help='''whether to save the model''')
       parser.add_argument('''--cached''', action='''store_true''', help='''whether to cache the input representations''')
       parser.add_argument('''--no_cuda''', action='''store_true''', help='''use to turn off cuda''')
       lowerCamelCase 	=   parser.parse_args()

       train_discriminator(**(vars(args)))




"""simple docstring"""





import unittest

from transformers import ErnieMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import (
     ErnieMForInformationExtraction,
     ErnieMForMultipleChoice,
     ErnieMForQuestionAnswering,
     ErnieMForSequenceClassification,
     ErnieMForTokenClassification,
     ErnieMModel,
 )
 from transformers.models.ernie_m.modeling_ernie_m import ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST



class 		snake_case_:






     def __init__(						self	:       Optional[int]       ,					UpperCamelCase_	:       Dict       ,					UpperCamelCase_	:       Optional[int]=1_3       ,					UpperCamelCase_	:       Dict=7       ,					UpperCamelCase_	:       Dict=True       ,					UpperCamelCase_	:       str=True       ,					UpperCamelCase_	:       int=True       ,					UpperCamelCase_	:       Optional[Any]=9_9       ,					UpperCamelCase_	:       Optional[int]=3_2       ,					UpperCamelCase_	:       Any=5       ,					UpperCamelCase_	:       Any=4       ,					UpperCamelCase_	:       Tuple=3_7       ,					UpperCamelCase_	:       Tuple="gelu"       ,					UpperCamelCase_	:       Union[str, Any]=0.1       ,					UpperCamelCase_	:       List[Any]=0.1       ,					UpperCamelCase_	:       int=5_1_2       ,					UpperCamelCase_	:       int=1_6       ,					UpperCamelCase_	:       List[Any]=2       ,					UpperCamelCase_	:       str=0.02       ,					UpperCamelCase_	:       int=3       ,					UpperCamelCase_	:       str=4       ,					UpperCamelCase_	:       Union[str, Any]=None       ,					):
           lowerCAmelCase       :						List[str]							= parent
           lowerCAmelCase       :						str							= batch_size
           lowerCAmelCase       :						Dict							= seq_length
           lowerCAmelCase       :						Dict							= is_training
           lowerCAmelCase       :						Tuple							= use_input_mask
           lowerCAmelCase       :						str							= use_labels
           lowerCAmelCase       :						Optional[Any]							= vocab_size
           lowerCAmelCase       :						Optional[Any]							= hidden_size
           lowerCAmelCase       :						Optional[int]							= num_hidden_layers
           lowerCAmelCase       :						List[Any]							= num_attention_heads
           lowerCAmelCase       :						Tuple							= intermediate_size
           lowerCAmelCase       :						Tuple							= hidden_act
           lowerCAmelCase       :						Optional[int]							= hidden_dropout_prob
           lowerCAmelCase       :						str							= attention_probs_dropout_prob
           lowerCAmelCase       :						str							= max_position_embeddings
           lowerCAmelCase       :						List[str]							= type_vocab_size
           lowerCAmelCase       :						str							= type_sequence_label_size
           lowerCAmelCase       :						int							= initializer_range
           lowerCAmelCase       :						Optional[Any]							= num_labels
           lowerCAmelCase       :						Any							= num_choices
           lowerCAmelCase       :						str							= scope






     def   lowerCamelCase__  (						self	:       Dict					):
           lowerCAmelCase       :						str							= ids_tensor([self.batch_size, self.seq_length]       ,					self.vocab_size					)

           lowerCAmelCase       :						List[Any]							= None
           if self.use_input_mask:
                 lowerCAmelCase       :						str							= random_attention_mask([self.batch_size, self.seq_length]					)

           lowerCAmelCase       :						List[Any]							= None
           lowerCAmelCase       :						str							= None
           lowerCAmelCase       :						str							= None
           if self.use_labels:
                 lowerCAmelCase       :						Tuple							= ids_tensor([self.batch_size]       ,					self.type_sequence_label_size					)
                 lowerCAmelCase       :						Dict							= ids_tensor([self.batch_size, self.seq_length]       ,					self.num_labels					)
                 lowerCAmelCase       :						List[Any]							= ids_tensor([self.batch_size]       ,					self.num_choices					)

           lowerCAmelCase       :						Any							= self.get_config()

           return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels






     def   lowerCamelCase__  (						self	:       str					):
           lowerCAmelCase       :						Tuple							= ids_tensor([self.batch_size, self.seq_length]       ,					self.vocab_size					)

           lowerCAmelCase       :						Optional[Any]							= None
           if self.use_input_mask:
                 lowerCAmelCase       :						Union[str, Any]							= random_attention_mask([self.batch_size, self.seq_length]					)
           lowerCAmelCase       :						Dict							= self.get_config()

           return config, input_ids, input_mask






     def   lowerCamelCase__  (						self	:       Union[str, Any]					):
           return ErnieMConfig(
               vocab_size=self.vocab_size       ,					hidden_size=self.hidden_size       ,					num_hidden_layers=self.num_hidden_layers       ,					num_attention_heads=self.num_attention_heads       ,					intermediate_size=self.intermediate_size       ,					hidden_act=self.hidden_act       ,					hidden_dropout_prob=self.hidden_dropout_prob       ,					attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,					max_position_embeddings=self.max_position_embeddings       ,					type_vocab_size=self.type_vocab_size       ,					initializer_range=self.initializer_range       ,					)






     def   lowerCamelCase__  (						self	:       Union[str, Any]       ,					UpperCamelCase_	:       Tuple       ,					UpperCamelCase_	:       Union[str, Any]       ,					UpperCamelCase_	:       Optional[Any]       ,					UpperCamelCase_	:       Dict       ,					UpperCamelCase_	:       Union[str, Any]       ,					UpperCamelCase_	:       Optional[Any]					):
           lowerCAmelCase       :						List[str]							= ErnieMModel(config=UpperCamelCase_					)
           model.to(UpperCamelCase_					)
           model.eval()
           lowerCAmelCase       :						List[Any]							= model(UpperCamelCase_       ,					return_dict=UpperCamelCase_					)
           self.parent.assertEqual(result.last_hidden_state.shape       ,					(self.batch_size, self.seq_length, self.hidden_size)					)






     def   lowerCamelCase__  (						self	:       Union[str, Any]       ,					UpperCamelCase_	:       Union[str, Any]       ,					UpperCamelCase_	:       Optional[Any]       ,					UpperCamelCase_	:       int       ,					UpperCamelCase_	:       List[Any]       ,					UpperCamelCase_	:       List[Any]       ,					UpperCamelCase_	:       Optional[Any]					):
           lowerCAmelCase       :						Optional[int]							= ErnieMForQuestionAnswering(config=UpperCamelCase_					)
           model.to(UpperCamelCase_					)
           model.eval()
           lowerCAmelCase       :						List[Any]							= model(
               UpperCamelCase_       ,					attention_mask=UpperCamelCase_       ,					start_positions=UpperCamelCase_       ,					end_positions=UpperCamelCase_       ,					)
           self.parent.assertEqual(result.start_logits.shape       ,					(self.batch_size, self.seq_length)					)
           self.parent.assertEqual(result.end_logits.shape       ,					(self.batch_size, self.seq_length)					)






     def   lowerCamelCase__  (						self	:       Dict       ,					UpperCamelCase_	:       Optional[int]       ,					UpperCamelCase_	:       Any       ,					UpperCamelCase_	:       List[str]       ,					UpperCamelCase_	:       Dict       ,					UpperCamelCase_	:       Optional[Any]       ,					UpperCamelCase_	:       Union[str, Any]					):
           lowerCAmelCase       :						Tuple							= ErnieMForInformationExtraction(config=UpperCamelCase_					)
           model.to(UpperCamelCase_					)
           model.eval()
           lowerCAmelCase       :						Tuple							= torch.ones_like(UpperCamelCase_       ,					dtype=torch.floataa					)
           lowerCAmelCase       :						str							= model(
               UpperCamelCase_       ,					attention_mask=UpperCamelCase_       ,					start_positions=UpperCamelCase_       ,					end_positions=UpperCamelCase_       ,					)
           self.parent.assertEqual(result.start_logits.shape       ,					(self.batch_size, self.seq_length)					)
           self.parent.assertEqual(result.end_logits.shape       ,					(self.batch_size, self.seq_length)					)






     def   lowerCamelCase__  (						self	:       Optional[int]       ,					UpperCamelCase_	:       Optional[int]       ,					UpperCamelCase_	:       Dict       ,					UpperCamelCase_	:       int       ,					UpperCamelCase_	:       int       ,					UpperCamelCase_	:       Union[str, Any]       ,					UpperCamelCase_	:       Any					):
           lowerCAmelCase       :						Union[str, Any]							= self.num_labels
           lowerCAmelCase       :						Tuple							= ErnieMForSequenceClassification(UpperCamelCase_					)
           model.to(UpperCamelCase_					)
           model.eval()
           lowerCAmelCase       :						Union[str, Any]							= model(UpperCamelCase_       ,					attention_mask=UpperCamelCase_       ,					labels=UpperCamelCase_					)
           self.parent.assertEqual(result.logits.shape       ,					(self.batch_size, self.num_labels)					)






     def   lowerCamelCase__  (						self	:       Optional[int]       ,					UpperCamelCase_	:       Any       ,					UpperCamelCase_	:       List[str]       ,					UpperCamelCase_	:       int       ,					UpperCamelCase_	:       Dict       ,					UpperCamelCase_	:       str       ,					UpperCamelCase_	:       List[str]					):
           lowerCAmelCase       :						Union[str, Any]							= self.num_labels
           lowerCAmelCase       :						List[Any]							= ErnieMForTokenClassification(config=UpperCamelCase_					)
           model.to(UpperCamelCase_					)
           model.eval()
           input_ids.to(UpperCamelCase_					)
           input_mask.to(UpperCamelCase_					)
           token_labels.to(UpperCamelCase_					)

           lowerCAmelCase       :						Optional[Any]							= model(UpperCamelCase_       ,					attention_mask=UpperCamelCase_       ,					labels=UpperCamelCase_					)

           self.parent.assertEqual(result.logits.shape       ,					(self.batch_size, self.seq_length, self.num_labels)					)






     def   lowerCamelCase__  (						self	:       Any       ,					UpperCamelCase_	:       str       ,					UpperCamelCase_	:       Optional[int]       ,					UpperCamelCase_	:       Tuple       ,					UpperCamelCase_	:       Dict       ,					UpperCamelCase_	:       Dict       ,					UpperCamelCase_	:       str					):
           lowerCAmelCase       :						Union[str, Any]							= self.num_choices
           lowerCAmelCase       :						Any							= ErnieMForMultipleChoice(config=UpperCamelCase_					)
           model.to(UpperCamelCase_					)
           model.eval()
           lowerCAmelCase       :						Any							= input_ids.unsqueeze(1					).expand(-1       ,					self.num_choices       ,					-1					).contiguous()
           lowerCAmelCase       :						int							= input_mask.unsqueeze(1					).expand(-1       ,					self.num_choices       ,					-1					).contiguous()
           lowerCAmelCase       :						Any							= model(
               UpperCamelCase_       ,					attention_mask=UpperCamelCase_       ,					labels=UpperCamelCase_       ,					)
           self.parent.assertEqual(result.logits.shape       ,					(self.batch_size, self.num_choices)					)






     def   lowerCamelCase__  (						self	:       Tuple					):
           lowerCAmelCase       :						List[Any]							= self.prepare_config_and_inputs()
           (
               (
               lowerCAmelCase
           ),   (
               lowerCAmelCase
           ),   (
               lowerCAmelCase
           ),   (
               lowerCAmelCase
           ),   (
               lowerCAmelCase
           ),   (
               lowerCAmelCase
           ),   
           )       :						str							= config_and_inputs
           lowerCAmelCase       :						Dict							= {'''input_ids''': input_ids, '''attention_mask''': input_mask}
           return config, inputs_dict



@require_torch
class 		snake_case_(							a__				,       a__				,       unittest.TestCase      ):
     __UpperCamelCase	=       (
         (
             ErnieMModel,
             ErnieMForMultipleChoice,
             ErnieMForQuestionAnswering,
             ErnieMForSequenceClassification,
             ErnieMForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     __UpperCamelCase	=       ()
     __UpperCamelCase	=       (
         {
             '''feature-extraction''': ErnieMModel,
             '''question-answering''': ErnieMForQuestionAnswering,
             '''text-classification''': ErnieMForSequenceClassification,
             '''token-classification''': ErnieMForTokenClassification,
             '''zero-shot''': ErnieMForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     __UpperCamelCase	=       False






     def   lowerCamelCase__  (						self	:       Optional[Any]       ,					UpperCamelCase_	:       str       ,					UpperCamelCase_	:       List[Any]       ,					UpperCamelCase_	:       Tuple       ,					UpperCamelCase_	:       int       ,					UpperCamelCase_	:       Any					):
           if pipeline_test_casse_name == "QAPipelineTests":
                 return True

           return False






     def   lowerCamelCase__  (						self	:       Dict					):
           lowerCAmelCase       :						str							= ErnieMModelTester(self					)
           lowerCAmelCase       :						Optional[Any]							= ConfigTester(self       ,					config_class=UpperCamelCase_       ,					hidden_size=3_7					)






     def   lowerCamelCase__  (						self	:       Tuple					):
           self.config_tester.run_common_tests()






     def   lowerCamelCase__  (						self	:       List[str]					):
           lowerCAmelCase       :						str							= self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCamelCase_					)






     def   lowerCamelCase__  (						self	:       Union[str, Any]					):
           lowerCAmelCase       :						Optional[Any]							= self.model_tester.prepare_config_and_inputs()
           for type in ["absolute", "relative_key", "relative_key_query"]:
                 lowerCAmelCase       :						Union[str, Any]							= type
                 self.model_tester.create_and_check_model(*UpperCamelCase_					)






     def   lowerCamelCase__  (						self	:       Tuple					):
           lowerCAmelCase       :						Dict							= self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_multiple_choice(*UpperCamelCase_					)






     def   lowerCamelCase__  (						self	:       Dict					):
           lowerCAmelCase       :						int							= self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(*UpperCamelCase_					)






     def   lowerCamelCase__  (						self	:       Optional[int]					):
           lowerCAmelCase       :						Tuple							= self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase_					)






     def   lowerCamelCase__  (						self	:       Tuple					):
           lowerCAmelCase       :						List[str]							= self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_information_extraction(*UpperCamelCase_					)






     def   lowerCamelCase__  (						self	:       List[Any]					):
           lowerCAmelCase       :						str							= self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*UpperCamelCase_					)






     @slow
     def   lowerCamelCase__  (						self	:       List[str]					):
           for model_name in ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCAmelCase       :						int							= ErnieMModel.from_pretrained(UpperCamelCase_					)
                 self.assertIsNotNone(UpperCamelCase_					)



@require_torch
class 		snake_case_(							unittest.TestCase      ):






     @slow
     def   lowerCamelCase__  (						self	:       Dict					):
           lowerCAmelCase       :						Optional[int]							= ErnieMModel.from_pretrained('''susnato/ernie-m-base_pytorch'''					)
           model.eval()
           lowerCAmelCase       :						Optional[int]							= torch.tensor([[0, 1, 2, 3, 4, 5]]					)
           lowerCAmelCase       :						str							= model(UpperCamelCase_					)[0]

           # TODO Replace vocab size
           lowerCAmelCase       :						Any							= 7_6_8

           lowerCAmelCase       :						int							= torch.Size((1, 6, hidden_size)					)
           self.assertEqual(output.shape       ,					UpperCamelCase_					)

           lowerCAmelCase       :						str							= torch.tensor(
               [[[-0.0_012, 0.1_245, -0.0_214], [-0.0_742, 0.0_244, -0.0_771], [-0.0_333, 0.1_164, -0.1_554]]]					)

           self.assertTrue(torch.allclose(output[:, :3, :3]       ,					UpperCamelCase_       ,					atol=1E-3					)					)


