"""simple docstring"""




import argparse
import gc
import os
import sys
from pathlib import Path
from typing import List  # noqa: F401

import pytorch_lightning as pl
import torch
from finetune import SummarizationModule, TranslationModule
from finetune import main as ft_main
from make_student import create_student_by_copying_alternating_layers, get_layers_to_supervise
from torch import nn

from transformers import AutoModelForSeqaSeqLM, MBartTokenizer, TaForConditionalGeneration
from transformers.models.bart.modeling_bart import shift_tokens_right
from utils import calculate_bleu, check_output_dir, freeze_params, label_smoothed_nll_loss, use_task_specific_params


# need the parent dir module
sys.path.insert(2, str(Path(__file__).resolve().parents[1]))
from lightning_base import generic_train  # noqa
class      UpperCamelCase__(      __A			):
						lowerCAmelCase__	:    List[str]    				=		['loss', 'ce_loss', 'mlm_loss', 'hid_loss_enc', 'hid_loss_dec']





						def __init__(							self		,__UpperCAmelCase  )       ->					str:
													assert Path(hparams.data_dir  ).exists()
													A__						  =				Path(hparams.output_dir  )
													self.output_dir.mkdir(exist_ok=__UpperCAmelCase  )

													A__						  =				self.output_dir.joinpath('student'  )

													A__						  =				str(__UpperCAmelCase  )  # Tell lightning we are training the student
													A__						  =				AutoModelForSeqaSeqLM.from_pretrained(hparams.teacher  ).eval()
													use_task_specific_params(__UpperCAmelCase		,hparams.task  )  # We copy good generation parameters to student by default
													if hparams.student is not None:
																				A__						  =				AutoModelForSeqaSeqLM.from_pretrained(hparams.student  )
																				use_task_specific_params(__UpperCAmelCase		,hparams.task  )
																				A__  ,	A__						  =				None, None
													else:
																				A__  ,	A__  ,	A__						  =				create_student_by_copying_alternating_layers(
																				    __UpperCAmelCase		,e=hparams.student_encoder_layers		,d=hparams.student_decoder_layers		,save_path=__UpperCAmelCase  )

													if hparams.length_penalty != -1:
																				A__						  =				hparams.length_penalty
													A__						  =				hparams.teacher  # Use teacher's tokenizer
													super().__init__(__UpperCAmelCase		,model=__UpperCAmelCase		,config=student.config  )
													assert student.config.model_type == teacher.config.model_type, (
													    f'''teacher, student model types should be the same, got {student.config.model_type} !='''
													    f''' {teacher.config.model_type}'''
													)

													if student.config.model_type == "t5":
																				A__						  =				len(student.get_encoder().block  )
																				A__						  =				len(student.get_decoder().block  )
																				A__						  =				len(teacher.get_encoder().block  )
																				A__						  =				len(teacher.get_decoder().block  )
													else:
																				A__						  =				student.config.encoder_layers
																				A__						  =				student.config.decoder_layers
																				A__						  =				teacher.config.encoder_layers
																				A__						  =				teacher.config.decoder_layers

													A__						  =				not (hparams.student is None or hparams.teacher == hparams.student)
													A__						  =				(not self.different_base_models) and hparams.alpha_hid > 0
													A__						  =				self.different_base_models or (student_encoder_layers != teacher_encoder_layers)
													# self.different_encoder determines whether we need to run the teacher encoder
													A__						  =				teacher
													freeze_params(self.teacher  )

													if not self.different_encoder:  # To save RAM, delete teacher encoder and freeze student encoder.
																				try:
																											del self.teacher.model.encoder
																				except AttributeError:  # T5
																											del self.teacher.encoder

													if e_layer_ids is None:
																				A__						  =				list(range(__UpperCAmelCase  )  )
													if d_layer_ids is None:
																				A__						  =				list(range(__UpperCAmelCase  )  )

													A__  ,	A__						  =				e_layer_ids, d_layer_ids  # type: List[int], List[int]

													if self.do_calc_hidden_loss:  # Intermediate supervision: Decide which layers to supervise
																				if hparams.supervise_forward:
																											A__						  =				get_layers_to_supervise(
																											    n_student=len(self.e_layer_ids  )		,n_teacher=__UpperCAmelCase  )
																											A__						  =				get_layers_to_supervise(
																											    n_student=len(self.d_layer_ids  )		,n_teacher=__UpperCAmelCase  )
																				else:  # student layer should emulate hidden states of the teacher layer it was copied from
																											A__						  =				self.e_layer_ids
																											A__						  =				self.d_layer_ids
													else:
																				A__						  =				None
																				A__						  =				None

													A__						  =				nn.KLDivLoss(reduction='batchmean'  )
													A__						  =				2.0
													A__						  =				hparams.alpha_mlm
													A__						  =				hparams.alpha_ce
													A__						  =				hparams.alpha_hid
													gc.collect()
													torch.cuda.empty_cache()





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )       ->					Tuple:
													A__						  =				mask[:, :, None].expand_as(__UpperCAmelCase  )
													A__						  =				s_logits.size(-1  )
													A__						  =				torch.masked_select(__UpperCAmelCase		,__UpperCAmelCase  )  # (bs * seq_length * voc_size) modulo the 1s in mask
													A__						  =				torch.masked_select(__UpperCAmelCase		,__UpperCAmelCase  )  # (bs * seq_length * voc_size) modulo the 1s in mask
													A__						  =				s_logits_slct.view(-1		,__UpperCAmelCase  )  # (bs * seq_length, voc_size) modulo the 1s in mask
													A__						  =				t_logits_slct.view(-1		,__UpperCAmelCase  )  # (bs * seq_length, voc_size) modulo the 1s in mask
													assert t_logits_slct.size() == s_logits_slct.size()
													A__						  =				(
													    self.ce_loss_fct(
													        nn.functional.log_softmax(s_logits_slct / self.temperature		,dim=-1  )		,nn.functional.softmax(t_logits_slct / self.temperature		,dim=-1  )		,)
													    * (self.temperature) ** 2
													)
													return loss_ce





						@staticmethod
						def 					snake_case__						(							__UpperCAmelCase		,__UpperCAmelCase  )       ->					Any:
													SummarizationModule.add_model_specific_args(__UpperCAmelCase		,__UpperCAmelCase  )
													add_distill_args(__UpperCAmelCase  )
													return parser





						def 					snake_case__						(							self		,__UpperCAmelCase  )       ->					tuple:
													A__						  =				self.tokenizer.pad_token_id
													A__  ,	A__  ,	A__						  =				batch['input_ids'], batch['attention_mask'], batch['labels']
													if isinstance(self.model		,__UpperCAmelCase  ):
																				A__						  =				self.model._shift_right(__UpperCAmelCase  )
													else:
																				A__						  =				shift_tokens_right(__UpperCAmelCase		,__UpperCAmelCase  )

													# noinspection PyCallingNonCallable
													A__						  =				self(
													    __UpperCAmelCase		,attention_mask=__UpperCAmelCase		,decoder_input_ids=__UpperCAmelCase		,output_hidden_states=self.do_calc_hidden_loss		,output_attentions=__UpperCAmelCase		,use_cache=__UpperCAmelCase		,)
													A__						  =				student_outputs['logits']

													# Same cross entropy vs. label smoothing logic as finetune.py
													assert lm_logits.shape[-1] == self.model.config.vocab_size
													if self.hparams.label_smoothing == 0:
																				# Same behavior as modeling_bart.py, besides ignoring pad_token_id
																				A__						  =				nn.CrossEntropyLoss(ignore_index=__UpperCAmelCase  )
																				A__						  =				loss_fct(lm_logits.view(-1		,lm_logits.shape[-1]  )		,labels.view(-1  )  )
													else:
																				A__						  =				nn.functional.log_softmax(__UpperCAmelCase		,dim=-1  )
																				A__  ,	A__						  =				label_smoothed_nll_loss(
																				    __UpperCAmelCase		,__UpperCAmelCase		,self.hparams.label_smoothing		,ignore_index=__UpperCAmelCase  )

													def zero_tensor():
																				return torch.tensor(0.0  ).type_as(__UpperCAmelCase  )

													A__						  =				student_outputs[
													    'encoder_last_hidden_state'
													]  # use this unless self.different_base_models
													A__  ,	A__						  =				zero_tensor(), zero_tensor()
													if self.different_encoder:  # compute encoder hidden state loss
																				A__						  =				self.teacher.get_encoder()(
																				    __UpperCAmelCase		,attention_mask=__UpperCAmelCase		,output_hidden_states=self.do_calc_hidden_loss		,)
																				if self.different_base_models:
																											A__						  =				all_teacher_encoder_outputs['last_hidden_state']
																				elif self.do_calc_hidden_loss:
																											A__						  =				self.calc_hidden_loss(
																											    __UpperCAmelCase		,student_outputs['encoder_hidden_states']		,all_teacher_encoder_outputs['hidden_states']		,self.e_matches		,normalize_hidden=self.hparams.normalize_hidden		,)

													A__						  =				self.teacher(
													    __UpperCAmelCase		,attention_mask=__UpperCAmelCase		,encoder_outputs=(teacher_enc_outputs,)		,decoder_input_ids=__UpperCAmelCase		,output_hidden_states=self.do_calc_hidden_loss		,use_cache=__UpperCAmelCase		,)
													A__						  =				decoder_input_ids.ne(__UpperCAmelCase  )
													A__						  =				self.calc_ce_loss(__UpperCAmelCase		,__UpperCAmelCase		,teacher_outputs['logits']  )
													if self.do_calc_hidden_loss:  # Intermediate supervision of decoder hidden states
																				A__						  =				self.calc_hidden_loss(
																				    __UpperCAmelCase		,student_outputs['decoder_hidden_states']		,teacher_outputs['decoder_hidden_states']		,self.d_matches		,normalize_hidden=self.hparams.normalize_hidden		,)

													A__						  =				(
													    self.alpha_ce * loss_ce
													    + self.alpha_mlm * student_lm_loss
													    + self.hparams.alpha_hid * (hid_loss_enc + hid_loss_dec)
													)
													return blended_loss, loss_ce, student_lm_loss, hid_loss_enc, hid_loss_dec





						@staticmethod
						def 					snake_case__						(							__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )       ->					int:
													A__						  =				'expected list or tuple for hidden_states, got tensor of shape: '
													assert not isinstance(__UpperCAmelCase		,torch.Tensor  ), f'''{msg}{hidden_states.shape}'''
													assert not isinstance(__UpperCAmelCase		,torch.Tensor  ), f'''{msg}{hidden_states_T.shape}'''
													A__						  =				attention_mask.to(hidden_states[0]  )
													A__						  =				mask.sum() * hidden_states[0].size(-1  )
													A__						  =				torch.stack([hidden_states[i] for i in range(len(__UpperCAmelCase  )  )]  )
													A__						  =				torch.stack([hidden_states_T[j] for j in matches]  )
													assert student_states.shape == teacher_states.shape, f'''{student_states.shape} != {teacher_states.shape}'''
													if normalize_hidden:
																				A__						  =				nn.functional.layer_norm(__UpperCAmelCase		,student_states.shape[1:]  )
																				A__						  =				nn.functional.layer_norm(__UpperCAmelCase		,teacher_states.shape[1:]  )
													A__						  =				nn.functional.mse_loss(__UpperCAmelCase		,__UpperCAmelCase		,reduction='none'  )
													A__						  =				(mse * mask.unsqueeze(0  ).unsqueeze(-1  )).sum() / valid_count
													return masked_mse






def 	UpperCAmelCase       (						UpperCamelCase__					):






							"""simple docstring"""





							parser.add_argument('--teacher'							,     type=UpperCamelCase__					)
							parser.add_argument('--alpha_ce'							,     default=0.8							,     type=UpperCamelCase__					)
							parser.add_argument('--alpha_mlm'							,     default=0.2							,     type=UpperCamelCase__					)
							parser.add_argument('--alpha_hid'							,     default=0.0							,     type=UpperCamelCase__							,     required=UpperCamelCase__					)
							parser.add_argument('--student'							,     type=UpperCamelCase__							,     required=UpperCamelCase__					)
							parser.add_argument('--student_decoder_layers'							,     default=12							,     type=UpperCamelCase__							,     required=UpperCamelCase__					)
							parser.add_argument('--student_encoder_layers'							,     default=12							,     type=UpperCamelCase__							,     required=UpperCamelCase__					)
							parser.add_argument('--no_teacher'							,     action='store_true'							,     default=UpperCamelCase__					)
							parser.add_argument('--length_penalty'							,     type=UpperCamelCase__							,     default=-1					)
							parser.add_argument('--supervise_forward'							,     action='store_true'							,     default=UpperCamelCase__					)
							parser.add_argument('--normalize_hidden'							,     action='store_true'							,     default=UpperCamelCase__					)
class      UpperCamelCase__(      __A			):
						lowerCAmelCase__	:    Tuple    				=		'translation'
						lowerCAmelCase__	:    str    				=		['bleu']
						lowerCAmelCase__	:    Optional[int]    				=		'bleu'





						def __init__(							self		,__UpperCAmelCase		,**__UpperCAmelCase  )       ->					str:
													super().__init__(__UpperCAmelCase		,**__UpperCAmelCase  )
													assert hparams.src_lang is not None
													assert hparams.tgt_lang is not None
													A__						  =				hparams.src_lang
													A__						  =				hparams.tgt_lang
													if self.model.config.decoder_start_token_id is None and isinstance(self.tokenizer		,__UpperCAmelCase  ):
																				A__						  =				self.tokenizer.lang_code_to_id[hparams.tgt_lang]





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase  )       ->					dict:
													return calculate_bleu(__UpperCAmelCase		,__UpperCAmelCase  )





						@staticmethod
						def 					snake_case__						(							__UpperCAmelCase		,__UpperCAmelCase  )       ->					List[Any]:
													TranslationModule.add_model_specific_args(__UpperCAmelCase		,__UpperCAmelCase  )
													add_distill_args(__UpperCAmelCase  )
													return parser






def 	UpperCAmelCase       (						UpperCamelCase__					):






							"""simple docstring"""





							if args.no_teacher:
														A__						  =				TranslationModule if 'translation' in args.task else SummarizationModule
							else:  # DISTILL WITH TEACHER
														A__						  =				TranslationDistiller if 'translation' in args.task else SummarizationDistiller
							A__						  =				module_cls.__name__
							print(F'''using module {args.setup_cls}'''					)
							A__						  =				module_cls(UpperCamelCase__					)
							return model






def 	UpperCAmelCase       (						UpperCamelCase__					):






							"""simple docstring"""





							Path(args.output_dir					).mkdir(exist_ok=UpperCamelCase__					)
							check_output_dir(UpperCamelCase__							,     expected_items=3					)

							A__						  =				create_module(UpperCamelCase__					)
							return ft_main(UpperCamelCase__							,     model=UpperCamelCase__					)


if __name__ == "__main__":
				__lowerCamelCase									=	argparse.ArgumentParser()
				__lowerCamelCase									=	pl.Trainer.add_argparse_args(parser)
				__lowerCamelCase									=	SummarizationDistiller.add_model_specific_args(parser, os.getcwd())
				__lowerCamelCase									=	parser.parse_args()

				distill_main(args)






"""simple docstring"""




import inspect
import unittest

from transformers import is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
				import torch

				from transformers import MODEL_MAPPING, PoolFormerConfig, PoolFormerForImageClassification, PoolFormerModel
				from transformers.models.poolformer.modeling_poolformer import POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
				from PIL import Image

				from transformers import PoolFormerImageProcessor
class      UpperCamelCase__(      __A			):





						def 					snake_case__						(							self  )       ->					int:
													A__						  =				self.config_class(**self.inputs_dict  )
													self.parent.assertTrue(hasattr(__UpperCAmelCase		,'hidden_sizes'  )  )
													self.parent.assertTrue(hasattr(__UpperCAmelCase		,'num_encoder_blocks'  )  )






class      UpperCamelCase__:





						def __init__(							self		,__UpperCAmelCase		,__UpperCAmelCase=13		,__UpperCAmelCase=64		,__UpperCAmelCase=3		,__UpperCAmelCase=4		,__UpperCAmelCase=[2, 2, 2, 2]		,__UpperCAmelCase=[8, 4, 2, 1]		,__UpperCAmelCase=[16, 32, 64, 1_28]		,__UpperCAmelCase=[1, 4, 8, 16]		,__UpperCAmelCase=False		,__UpperCAmelCase=True		,__UpperCAmelCase="gelu"		,__UpperCAmelCase=0.1		,__UpperCAmelCase=0.0_2		,__UpperCAmelCase=3		,__UpperCAmelCase=None		,)       ->					Tuple:
													A__						  =				parent
													A__						  =				batch_size
													A__						  =				image_size
													A__						  =				num_channels
													A__						  =				num_encoder_blocks
													A__						  =				sr_ratios
													A__						  =				depths
													A__						  =				hidden_sizes
													A__						  =				downsampling_rates
													A__						  =				is_training
													A__						  =				use_labels
													A__						  =				hidden_act
													A__						  =				hidden_dropout_prob
													A__						  =				initializer_range
													A__						  =				num_labels
													A__						  =				scope





						def 					snake_case__						(							self  )       ->					Tuple:
													A__						  =				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

													A__						  =				None
													if self.use_labels:
																				A__						  =				ids_tensor([self.batch_size, self.image_size, self.image_size]		,self.num_labels  )

													A__						  =				PoolFormerConfig(
													    image_size=self.image_size		,num_channels=self.num_channels		,num_encoder_blocks=self.num_encoder_blocks		,depths=self.depths		,hidden_sizes=self.hidden_sizes		,hidden_act=self.hidden_act		,hidden_dropout_prob=self.hidden_dropout_prob		,initializer_range=self.initializer_range		,)

													return config, pixel_values, labels





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )       ->					Union[str, Any]:
													A__						  =				PoolFormerModel(config=__UpperCAmelCase  )
													model.to(__UpperCAmelCase  )
													model.eval()
													A__						  =				model(__UpperCAmelCase  )
													A__						  =				A__						  =				self.image_size // 3_2.0
													self.parent.assertEqual(
													    result.last_hidden_state.shape		,(self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)  )





						def 					snake_case__						(							self  )       ->					Any:
													A__						  =				self.prepare_config_and_inputs()
													A__  ,	A__  ,	A__						  =				config_and_inputs
													A__						  =				{'pixel_values': pixel_values}
													return config, inputs_dict






@require_torch
class      UpperCamelCase__(      __A		,    __A		,    unittest.TestCase			):
						lowerCAmelCase__	:    Optional[int]    				=		(PoolFormerModel, PoolFormerForImageClassification) if is_torch_available() else ()
						lowerCAmelCase__	:    Optional[Any]    				=		(
						    {'feature-extraction': PoolFormerModel, 'image-classification': PoolFormerForImageClassification}
						    if is_torch_available()
						    else {}
						)

						lowerCAmelCase__	:    Any    				=		False
						lowerCAmelCase__	:    str    				=		False
						lowerCAmelCase__	:    Any    				=		False
						lowerCAmelCase__	:    str    				=		False
						lowerCAmelCase__	:    Optional[int]    				=		False





						def 					snake_case__						(							self  )       ->					Any:
													A__						  =				PoolFormerModelTester(self  )
													A__						  =				PoolFormerConfigTester(self		,config_class=__UpperCAmelCase  )





						def 					snake_case__						(							self  )       ->					int:
													self.config_tester.run_common_tests()





						def 					snake_case__						(							self  )       ->					Tuple:
													A__						  =				self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*__UpperCAmelCase  )





						@unittest.skip('PoolFormer does not use inputs_embeds'  )
						def 					snake_case__						(							self  )       ->					int:
													pass





						@unittest.skip('PoolFormer does not have get_input_embeddings method and get_output_embeddings methods'  )
						def 					snake_case__						(							self  )       ->					List[str]:
													pass





						def 					snake_case__						(							self  )       ->					List[Any]:
													def check_hidden_states_output(__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  ):
																				A__						  =				model_class(__UpperCAmelCase  )
																				model.to(__UpperCAmelCase  )
																				model.eval()

																				with torch.no_grad():
																											A__						  =				model(**self._prepare_for_class(__UpperCAmelCase		,__UpperCAmelCase  )  )

																				A__						  =				outputs.hidden_states

																				A__						  =				self.model_tester.num_encoder_blocks
																				self.assertEqual(len(__UpperCAmelCase  )		,__UpperCAmelCase  )

																				# verify the first hidden states (first block)
																				self.assertListEqual(
																				    list(hidden_states[0].shape[-3:]  )		,[
																				        self.model_tester.hidden_sizes[0],
																				        self.model_tester.image_size // 4,
																				        self.model_tester.image_size // 4,
																				    ]		,)

													A__  ,	A__						  =				self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				A__						  =				True
																				check_hidden_states_output(__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )

																				# check that output_hidden_states also work using config
																				del inputs_dict["output_hidden_states"]
																				A__						  =				True

																				check_hidden_states_output(__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )





						def 					snake_case__						(							self  )       ->					Union[str, Any]:
													if not self.model_tester.is_training:
																				return

													A__  ,	A__						  =				self.model_tester.prepare_config_and_inputs_for_common()
													A__						  =				True

													for model_class in self.all_model_classes:
																				if model_class in get_values(__UpperCAmelCase  ):
																											continue
																				A__						  =				model_class(__UpperCAmelCase  )
																				model.to(__UpperCAmelCase  )
																				model.train()
																				A__						  =				self._prepare_for_class(__UpperCAmelCase		,__UpperCAmelCase		,return_labels=__UpperCAmelCase  )
																				A__						  =				model(**__UpperCAmelCase  ).loss
																				loss.backward()





						def 					snake_case__						(							self  )       ->					Union[str, Any]:
													A__  ,	A__						  =				self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				A__						  =				model_class(__UpperCAmelCase  )
																				A__						  =				inspect.signature(model.forward  )
																				# signature.parameters is an OrderedDict => so arg_names order is deterministic
																				A__						  =				[*signature.parameters.keys()]

																				A__						  =				['pixel_values']
																				self.assertListEqual(arg_names[:1]		,__UpperCAmelCase  )





						@slow
						def 					snake_case__						(							self  )       ->					Union[str, Any]:
													for model_name in POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																				A__						  =				PoolFormerModel.from_pretrained(__UpperCAmelCase  )
																				self.assertIsNotNone(__UpperCAmelCase  )






def 	UpperCAmelCase       (						):






							"""simple docstring"""





							A__						  =				Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'					)
							return image


@require_torch
class      UpperCamelCase__(      unittest.TestCase			):





						@slow
						def 					snake_case__						(							self  )       ->					Union[str, Any]:
													A__						  =				PoolFormerImageProcessor()
													A__						  =				PoolFormerForImageClassification.from_pretrained('sail/poolformer_s12'  ).to(__UpperCAmelCase  )

													A__						  =				image_processor(images=prepare_img()		,return_tensors='pt'  ).to(__UpperCAmelCase  )

													# forward pass
													with torch.no_grad():
																				A__						  =				model(**__UpperCAmelCase  )

													# verify the logits
													A__						  =				torch.Size((1, 10_00)  )
													self.assertEqual(outputs.logits.shape		,__UpperCAmelCase  )

													A__						  =				torch.tensor([-0.6_1_1_3, 0.1_6_8_5, -0.0_4_9_2]  ).to(__UpperCAmelCase  )
													self.assertTrue(torch.allclose(outputs.logits[0, :3]		,__UpperCAmelCase		,atol=1e-4  )  )





