



"""simple docstring"""
import copy
import logging
import random

import joblib
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm

from transformers import AdamW, GPTaLMHeadModel, get_linear_schedule_with_warmup


_lowerCAmelCase					:Tuple  =      logging.getLogger(__name__)





def 				lowerCamelCase_		(UpperCamelCase__	:	Dict							):
    random.seed(UpperCamelCase__							)
    np.random.seed(UpperCamelCase__							)
    torch.manual_seed(UpperCamelCase__							)
    torch.cuda.manual_seed_all(UpperCamelCase__							)





def 				lowerCamelCase_		(UpperCamelCase__	:	Dict  , UpperCamelCase__	:	List[Any]  , UpperCamelCase__	:	List[Any]							):
    model.eval()
    _UpperCAmelCase							:						Tuple       							=					next(model.parameters()							).device
    _UpperCAmelCase							:						List[Any]       							=					1
    _UpperCAmelCase							:						Dict       							=					torch.zeros((eval_batch_size, context_len)  , dtype=torch.long  , device=UpperCamelCase__							)
    _UpperCAmelCase							:						Dict       							=					DataLoader(UpperCamelCase__  , shuffle=UpperCamelCase__  , batch_size=UpperCamelCase__							)
    _UpperCAmelCase							:						List[Any]       							=					torch.zeros(1  , device=UpperCamelCase__							)
    _UpperCAmelCase							:						Dict       							=					0
    for batch in eval_dataloader:
        batch.to(UpperCamelCase__							)
        # pad
        context.zero_()
        for i in range(UpperCamelCase__							):
            _UpperCAmelCase							:						Union[str, Any]       							=					batch[i]
        _UpperCAmelCase							:						Tuple       							=					model(UpperCamelCase__  , labels=UpperCamelCase__							)
        eval_loss += outputs[0].sum().item()
        nb_eval_examples += batch.size(0							)
    _UpperCAmelCase							:						Tuple       							=					eval_loss / nb_eval_examples
    _UpperCAmelCase							:						List[Any]       							=					torch.exp(UpperCamelCase__							)
    model.train()
    return perplexity





def 				lowerCamelCase_		(UpperCamelCase__	:	List[Any]="gpt2"							):
    _UpperCAmelCase							:						Optional[Any]       							=					GPTaLMHeadModel.from_pretrained(UpperCamelCase__  , output_hidden_states=UpperCamelCase__							)
    torch.save(model.state_dict()  , model_name + '''local.pt'''							)
    return model





def 				lowerCamelCase_		(UpperCamelCase__	:	List[str]  , UpperCamelCase__	:	Any  , UpperCamelCase__	:	Union[str, Any]							):
    _UpperCAmelCase							:						Optional[int]       							=					copy.deepcopy(UpperCamelCase__							)
    model.to(UpperCamelCase__							)

    _UpperCAmelCase							:						str       							=					['''bias''', '''LayerNorm.weight''']
    _UpperCAmelCase							:						Dict       							=					[
        {
            '''params''': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay							)],
            '''weight_decay''': 0.0,
        },
        {'''params''': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay							)], '''weight_decay''': 0.0},
    ]
    _UpperCAmelCase							:						Optional[int]       							=					AdamW(UpperCamelCase__  , lr=5E-5  , eps=1E-8							)
    _UpperCAmelCase							:						Union[str, Any]       							=					get_linear_schedule_with_warmup(UpperCamelCase__  , 0  , UpperCamelCase__							)
    torch.cuda.empty_cache()
    return model, lm_optimizer, lm_scheduler





def 				lowerCamelCase_		(UpperCamelCase__	:	int  , UpperCamelCase__	:	str  , UpperCamelCase__	:	List[Any]  , UpperCamelCase__	:	Union[str, Any]							):
    _UpperCAmelCase							:						Tuple       							=					np.array(UpperCamelCase__							).mean()
    _UpperCAmelCase							:						Tuple       							=					np.array(UpperCamelCase__							).std()
    _UpperCAmelCase							:						Dict       							=					(real_perps - avg) / std
    _UpperCAmelCase							:						List[str]       							=					list(zip(UpperCamelCase__  , UpperCamelCase__  , UpperCamelCase__							)							)
    joblib.dump(UpperCamelCase__  , UpperCamelCase__							)





def 				lowerCamelCase_		(UpperCamelCase__	:	List[Any]  , UpperCamelCase__	:	Any  , UpperCamelCase__	:	str  , UpperCamelCase__	:	int  , UpperCamelCase__	:	Union[str, Any]  , UpperCamelCase__	:	List[Any]  , UpperCamelCase__	:	Tuple  , UpperCamelCase__	:	Optional[Any]="dev.jbl"  , UpperCamelCase__	:	int=recopy_gpta  , ):
    _UpperCAmelCase							:						Union[str, Any]       							=					[]
    _UpperCAmelCase							:						Optional[int]       							=					[]
    _UpperCAmelCase							:						int       							=					[]

    # Initialize the transformer model
    _UpperCAmelCase							:						Union[str, Any]       							=					copy.deepcopy(UpperCamelCase__							)
    orig_model.to(device='''cpu'''							)
    torch.cuda.empty_cache()

    # Compute perplexity of initial transformer model for comparison
    model.train()
    _UpperCAmelCase							,       _UpperCAmelCase							,       _UpperCAmelCase							:						str       							=					recopy_model(UpperCamelCase__  , UpperCamelCase__  , UpperCamelCase__							)

    for step in tqdm(range(UpperCamelCase__							)							):
        _UpperCAmelCase							:						int       							=					torch.zeros((1, context_len)  , dtype=torch.long  , device=UpperCamelCase__							)
        _UpperCAmelCase							:						List[str]       							=					random.choice(UpperCamelCase__							)
        _UpperCAmelCase							:						Any       							=					random.randint(0  , len(story[0]							) - context_len - 1							)
        _UpperCAmelCase							:						List[Any]       							=					story[0][start : start + context_len]
        lm_optimizer.zero_grad()
        _UpperCAmelCase							:						Union[str, Any]       							=					model(UpperCamelCase__  , labels=UpperCamelCase__							)
        _UpperCAmelCase							:						Union[str, Any]       							=					outputs[0]
        _UpperCAmelCase							:						int       							=					compute_perplexity(UpperCamelCase__  , UpperCamelCase__  , UpperCamelCase__							)
        model.train()
        lm_loss.backward()
        # Do LM backprop
        torch.nn.utils.clip_grad_norm_(model.parameters()  , 3.0							)
        lm_optimizer.step()
        lm_scheduler.step()  # Update learning rate schedule

        # Compute perplexity after back-propagating on the selected context
        _UpperCAmelCase							:						List[Any]       							=					compute_perplexity(UpperCamelCase__  , UpperCamelCase__  , UpperCamelCase__							)

        # Periodically save the stored (X, IG(X)) pairs
        if step % 1000 == 0 and step > 1:
            intermittent_save(UpperCamelCase__  , UpperCamelCase__  , UpperCamelCase__  , UpperCamelCase__							)

        # Reset the pretrained model to the original pretrained GPT-2 weights after each iteration
        _UpperCAmelCase							,       _UpperCAmelCase							,       _UpperCAmelCase							:						Optional[Any]       							=					recopy_model(UpperCamelCase__  , UpperCamelCase__  , UpperCamelCase__							)

        past_perps.append(past_perp.item()							)
        real_perps.append(orig_perp - real_perp.item()							)
        contexts.append(np.array(context.cpu()							)							)

    intermittent_save(UpperCamelCase__  , UpperCamelCase__  , UpperCamelCase__  , UpperCamelCase__							)





def 				lowerCamelCase_		(UpperCamelCase__	:	List[str]  , UpperCamelCase__	:	Any="data/tokenized_stories_train_wikitext103.jbl"  , UpperCamelCase__	:	Union[str, Any]=100  , UpperCamelCase__	:	str=1026  , UpperCamelCase__	:	Union[str, Any]=True							):
    _UpperCAmelCase							:						str       							=					joblib.load(UpperCamelCase__							)
    print('''data loaded'''							)
    _UpperCAmelCase							:						Tuple       							=					[]
    if trim:
        for i, example in enumerate(UpperCamelCase__							):
            if len(example[0]							) > min_len:
                _UpperCAmelCase							:						List[str]       							=					random.randint(0  , len(example[0]							) - context_len - 1							)
                objective_set.append(example[0, start : start + context_len]							)
            if len(UpperCamelCase__							) >= number:
                break
        _UpperCAmelCase							:						List[str]       							=					[]
        for j in range(i + 1  , len(UpperCamelCase__							)							):
            if len(data[j][0]							) > min_len:
                train_data.append(data[j]							)
    else:
        _UpperCAmelCase							:						Optional[Any]       							=					data[0:number]
        _UpperCAmelCase							:						List[str]       							=					data[number:]

    joblib.dump(UpperCamelCase__  , '''objective_set.jbl'''							)
    print('''objective set saved'''							)
    return train_data, objective_set





def 				lowerCamelCase_		(UpperCamelCase__	:	List[str]  , UpperCamelCase__	:	Optional[int]  , UpperCamelCase__	:	List[str]  , UpperCamelCase__	:	List[str]  , UpperCamelCase__	:	List[Any]=50  , UpperCamelCase__	:	Optional[int]="secondary_learner.pt"							):
    _UpperCAmelCase							:						Any       							=					torch.device('''cuda:0''' if torch.cuda.is_available() else '''cpu'''							)
    # We will use the first 512 pairs from our dataset as a test set for
    # our secondary learner and the rest to train
    _UpperCAmelCase							:						int       							=					train_dataset[:512]
    _UpperCAmelCase							:						Any       							=					train_dataset[512:]
    _UpperCAmelCase							:						Dict       							=					DataLoader(UpperCamelCase__  , shuffle=UpperCamelCase__  , batch_size=UpperCamelCase__							)
    _UpperCAmelCase							:						List[str]       							=					DataLoader(UpperCamelCase__  , shuffle=UpperCamelCase__  , batch_size=UpperCamelCase__							)

    # secondary learner model set up
    _UpperCAmelCase							:						Tuple       							=					nn.MSELoss()
    _UpperCAmelCase							:						Any       							=					nn.MSELoss(reduction='''sum'''							)
    secondary_learner.to(UpperCamelCase__							)
    _UpperCAmelCase							:						List[Any]       							=					torch.optim.Adam(secondary_learner.parameters()  , lr=0.0_0001							)
    secondary_learner.train()

    # TODO in original code this is written as number of actual batches seen
    # not number of items seen but other places it is number of items instead.
    # improve consistency! changed this to epochs for clarity
    _UpperCAmelCase							:						Dict       							=					float('''inf'''							)
    # Iterate through batches until we've used max_steps batches
    for epoch in range(int(UpperCamelCase__							)							):
        _UpperCAmelCase							:						Union[str, Any]       							=					0.0
        secondary_learner.train()
        for step, batch in enumerate(UpperCamelCase__							):
            _UpperCAmelCase							:						Optional[Any]       							=					batch[0].to(UpperCamelCase__							)
            _UpperCAmelCase							:						Dict       							=					batch[1].to(UpperCamelCase__							)
            _UpperCAmelCase							:						int       							=					secondary_learner(UpperCamelCase__							)
            q_optimizer.zero_grad()
            _UpperCAmelCase							:						List[Any]       							=					loss(UpperCamelCase__  , real_q.float()							)
            q_loss.backward()
            q_optimizer.step()
            tr_q_loss += q_loss.item()

            # model trains fairly quickly so we won't wait for a full epoch
            # eval is triggered at eval_freq and end of epochs
            if (step % eval_freq == 0 and step > 0) or ((step + 1) == len(UpperCamelCase__							)):
                _UpperCAmelCase							:						int       							=					tr_q_loss / (step + 1)

                secondary_learner.eval()
                _UpperCAmelCase							:						Optional[Any]       							=					0.0
                _UpperCAmelCase							:						Dict       							=					0.0
                _UpperCAmelCase							:						Union[str, Any]       							=					[]
                _UpperCAmelCase							:						int       							=					[]
                # Compute performance of the secondary learner after this batch
                for stepa, batcha in enumerate(UpperCamelCase__							):
                    _UpperCAmelCase							:						List[str]       							=					batcha[0].to(UpperCamelCase__							)
                    _UpperCAmelCase							:						str       							=					batcha[1].to(UpperCamelCase__							)
                    _UpperCAmelCase							:						Optional[int]       							=					secondary_learner(UpperCamelCase__							)
                    q_lossa += test_loss(UpperCamelCase__  , UpperCamelCase__							).item()
                    sum_qa += torch.sum(UpperCamelCase__							).item()
                    for ei, i in enumerate(predicted_qa.cpu().detach().numpy()							):
                        predicted.append(i.item()							)
                    for ei, i in enumerate(real_qa.cpu().detach().numpy()							):
                        actual.append(i.item()							)

                q_lossa /= len(UpperCamelCase__							)
                print(
                    '''Epoch: '''  , UpperCamelCase__  , '''step: '''  , UpperCamelCase__  , '''Avg. q:'''  , sum_qa / len(UpperCamelCase__							)  , '''Train Loss: '''  , UpperCamelCase__  , '''Test Loss: '''  , UpperCamelCase__  , )
                if q_lossa < best_test_loss:
                    joblib.dump((predicted, actual)  , '''pred_vs_actual.jbl'''							)
                    torch.save(secondary_learner.state_dict()  , UpperCamelCase__							)
                    _UpperCAmelCase							:						Tuple       							=					q_lossa

            secondary_learner.train()
    return secondary_learner


class 							_UpperCAmelCase							(				nn.Module   ):





     '''simple docstring'''



     def __init__(      self ,			A	)							->      Tuple:
         super(UpperCamelCase__ ,			self	).__init__()
         _UpperCAmelCase							:						Tuple       							=					model.transformer.wte
         _UpperCAmelCase							:						Union[str, Any]       							=					copy.deepcopy(model.transformer.wte.weight	)
         _UpperCAmelCase							:						Any       							=					nn.Convad(self.embeddings.weight.size(1	) ,			2_5_6 ,			3 ,			padding=1	)
         _UpperCAmelCase							:						List[str]       							=					nn.Sequential(nn.Linear(2_5_6 ,			3_2	) ,			nn.Dropout(p=0.1	) ,			nn.Linear(3_2 ,			3_2	) ,			nn.Linear(3_2 ,			1	)	)





     def 				__lowerCAmelCase			(      self ,			A	)							->      Optional[Any]:
         _UpperCAmelCase							:						Optional[int]       							=					torch.max(self.conv(self.embeddings(UpperCamelCase__	).squeeze(1	).transpose(1 ,			2	)	) ,			2	)[0]
         _UpperCAmelCase							:						Tuple       							=					self.fc(UpperCamelCase__	)
         return qs.squeeze(1	)





     @classmethod
     def 				__lowerCAmelCase			(      cls ,			A ,			A	)							->      List[str]:
         _UpperCAmelCase							:						Dict       							=					cls(UpperCamelCase__	)  # this calls __init__
         _UpperCAmelCase							:						List[str]       							=					torch.load(UpperCamelCase__	)
         secondary_learner.load_state_dict(UpperCamelCase__	)
         _UpperCAmelCase							:						Optional[int]       							=					model.transformer.wte
         _UpperCAmelCase							:						int       							=					copy.deepcopy(model.transformer.wte.weight	)
         return secondary_learner







"""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,
)


_lowerCAmelCase					:Union[str, Any]  =      logging.getLogger(__name__)
class 							_UpperCAmelCase							(				a   ):





     '''simple docstring'''



     def 				__lowerCAmelCase			(      self ,			A ,			A ,			A=None ,			A=None	)							->      Tuple:
         if current_layer == 0:
             _UpperCAmelCase							:						Optional[Any]       							=					self.embedding_hidden_mapping_in(A	)
         else:
             _UpperCAmelCase							:						Union[str, Any]       							=					hidden_states[0]

         _UpperCAmelCase							:						Dict       							=					int(self.config.num_hidden_layers / self.config.num_hidden_groups	)

         # Index of the hidden group
         _UpperCAmelCase							:						Dict       							=					int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)	)

         _UpperCAmelCase							:						List[Any]       							=					self.albert_layer_groups[group_idx](
             A ,			A ,			head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group] ,			)
         _UpperCAmelCase							:						str       							=					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.'''	,a	,)
class 							_UpperCAmelCase							(				a   ):





     '''simple docstring'''



     def __init__(      self ,			A	)							->      str:
         super().__init__(A	)

         _UpperCAmelCase							:						Dict       							=					AlbertTransformerWithPabee(A	)

         self.init_weights()
         _UpperCAmelCase							:						str       							=					0
         _UpperCAmelCase							:						Union[str, Any]       							=					0
         _UpperCAmelCase							:						Dict       							=					0

         _UpperCAmelCase							:						str       							=					0





     def 				__lowerCAmelCase			(      self ,			A	)							->      Optional[int]:
         _UpperCAmelCase							:						List[Any]       							=					threshold





     def 				__lowerCAmelCase			(      self ,			A	)							->      Optional[Any]:
         _UpperCAmelCase							:						int       							=					patience





     def 				__lowerCAmelCase			(      self	)							->      Any:
         _UpperCAmelCase							:						Dict       							=					0
         _UpperCAmelCase							:						List[str]       							=					0





     def 				__lowerCAmelCase			(      self	)							->      List[Any]:
         _UpperCAmelCase							:						Union[str, Any]       							=					self.inference_layers_num / self.inference_instances_num
         _UpperCAmelCase							:						str       							=					(
             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 ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=False ,			)							->      int:

         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:
             _UpperCAmelCase							:						Optional[Any]       							=					input_ids.size()
         elif inputs_embeds is not None:
             _UpperCAmelCase							:						Optional[Any]       							=					inputs_embeds.size()[:-1]
         else:
             raise ValueError('''You have to specify either input_ids or inputs_embeds'''	)

         _UpperCAmelCase							:						Tuple       							=					input_ids.device if input_ids is not None else inputs_embeds.device

         if attention_mask is None:
             _UpperCAmelCase							:						Union[str, Any]       							=					torch.ones(A ,			device=A	)
         if token_type_ids is None:
             _UpperCAmelCase							:						Tuple       							=					torch.zeros(A ,			dtype=torch.long ,			device=A	)

         _UpperCAmelCase							:						Optional[Any]       							=					attention_mask.unsqueeze(1	).unsqueeze(2	)
         _UpperCAmelCase							:						Union[str, Any]       							=					extended_attention_mask.to(dtype=self.dtype	)  # fp16 compatibility
         _UpperCAmelCase							:						Dict       							=					(1.0 - extended_attention_mask) * -10_000.0
         _UpperCAmelCase							:						Any       							=					self.get_head_mask(A ,			self.config.num_hidden_layers	)

         _UpperCAmelCase							:						int       							=					self.embeddings(
             A ,			position_ids=A ,			token_type_ids=A ,			inputs_embeds=A	)
         _UpperCAmelCase							:						Any       							=					embedding_output

         if self.training:
             _UpperCAmelCase							:						Optional[Any]       							=					[]
             for i in range(self.config.num_hidden_layers	):
                 _UpperCAmelCase							:						Union[str, Any]       							=					self.encoder.adaptive_forward(
                     A ,			current_layer=A ,			attention_mask=A ,			head_mask=A ,			)

                 _UpperCAmelCase							:						List[str]       							=					self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]	)	)
                 _UpperCAmelCase							:						Dict       							=					output_layers[i](output_dropout(A	)	)
                 res.append(A	)
         elif self.patience == 0:  # Use all layers for inference
             _UpperCAmelCase							:						Dict       							=					self.encoder(A ,			A ,			head_mask=A	)
             _UpperCAmelCase							:						Tuple       							=					self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]	)	)
             _UpperCAmelCase							:						Optional[int]       							=					[output_layers[self.config.num_hidden_layers - 1](A	)]
         else:
             _UpperCAmelCase							:						List[Any]       							=					0
             _UpperCAmelCase							:						int       							=					None
             _UpperCAmelCase							:						int       							=					0
             for i in range(self.config.num_hidden_layers	):
                 calculated_layer_num += 1
                 _UpperCAmelCase							:						Tuple       							=					self.encoder.adaptive_forward(
                     A ,			current_layer=A ,			attention_mask=A ,			head_mask=A ,			)

                 _UpperCAmelCase							:						Any       							=					self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]	)	)
                 _UpperCAmelCase							:						Optional[Any]       							=					output_layers[i](A	)
                 if regression:
                     _UpperCAmelCase							:						str       							=					logits.detach()
                     if patient_result is not None:
                         _UpperCAmelCase							:						str       							=					patient_result.detach()
                     if (patient_result is not None) and torch.abs(patient_result - labels	) < self.regression_threshold:
                         patient_counter += 1
                     else:
                         _UpperCAmelCase							:						List[Any]       							=					0
                 else:
                     _UpperCAmelCase							:						List[str]       							=					logits.detach().argmax(dim=1	)
                     if patient_result is not None:
                         _UpperCAmelCase							:						Optional[Any]       							=					patient_result.detach().argmax(dim=1	)
                     if (patient_result is not None) and torch.all(labels.eq(A	)	):
                         patient_counter += 1
                     else:
                         _UpperCAmelCase							:						Tuple       							=					0

                 _UpperCAmelCase							:						Any       							=					logits
                 if patient_counter == self.patience:
                     break
             _UpperCAmelCase							:						List[str]       							=					[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
    the pooled output) e.g. for GLUE tasks. '''	,a	,)
class 							_UpperCAmelCase							(				a   ):





     '''simple docstring'''



     def __init__(      self ,			A	)							->      Optional[int]:
         super().__init__(A	)
         _UpperCAmelCase							:						List[Any]       							=					config.num_labels

         _UpperCAmelCase							:						Optional[int]       							=					AlbertModelWithPabee(A	)
         _UpperCAmelCase							:						List[str]       							=					nn.Dropout(config.classifier_dropout_prob	)
         _UpperCAmelCase							:						Optional[Any]       							=					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 ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			)							->      List[str]:
         _UpperCAmelCase							:						Optional[Any]       							=					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 ,			)

         _UpperCAmelCase							:						int       							=					(logits[-1],)

         if labels is not None:
             _UpperCAmelCase							:						int       							=					None
             _UpperCAmelCase							:						Tuple       							=					0
             for ix, logits_item in enumerate(A	):
                 if self.num_labels == 1:
                     #  We are doing regression
                     _UpperCAmelCase							:						Dict       							=					MSELoss()
                     _UpperCAmelCase							:						Optional[int]       							=					loss_fct(logits_item.view(-1	) ,			labels.view(-1	)	)
                 else:
                     _UpperCAmelCase							:						List[str]       							=					CrossEntropyLoss()
                     _UpperCAmelCase							:						Any       							=					loss_fct(logits_item.view(-1 ,			self.num_labels	) ,			labels.view(-1	)	)
                 if total_loss is None:
                     _UpperCAmelCase							:						Optional[Any]       							=					loss
                 else:
                     total_loss += loss * (ix + 1)
                 total_weights += ix + 1
             _UpperCAmelCase							:						Any       							=					(total_loss / total_weights,) + outputs

         return outputs


