from typing import List, Optional, Tuple, Union

import torch
from torch import nn
from torch.nn import CrossEntropyLoss

from ... import AutoBackbone
from ...modeling_outputs import SemanticSegmenterOutput
from ...modeling_utils import PreTrainedModel
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
from ...utils.backbone_utils import BackboneMixin
from .configuration_upernet import UperNetConfig


__UpperCamelCase	: List[str]    		=  [
    'openmmlab/upernet-convnext-tiny',
    # See all UperNet models at https://huggingface.co/models?filter=upernet
]

# General docstring
__UpperCamelCase	: Optional[Any]    		=  'UperNetConfig'





class       lowercase__		(     nn.Module):


						def __init__( self			:       Any					,	UpperCamelCase__			:       int					,	UpperCamelCase__			:       int					,	UpperCamelCase__			:       Union[int, Tuple[int, int]]					,	UpperCamelCase__			:       Union[int, Tuple[int, int], str] = 0					,	UpperCamelCase__			:       bool = False					,	UpperCamelCase__			:       Union[int, Tuple[int, int]] = 1					,	):





												'''simple docstring'''





												super().__init__()
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							nn.Convad(
												    in_channels=UpperCamelCase__					,	out_channels=UpperCamelCase__					,	kernel_size=UpperCamelCase__					,	padding=UpperCamelCase__					,	bias=UpperCamelCase__					,	dilation=UpperCamelCase__					,	)
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							nn.BatchNormad(UpperCamelCase__		)
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							nn.ReLU()



						def  __A ( self			:       Tuple					,	UpperCamelCase__			:       torch.Tensor		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[Any]           =							self.conv(UpperCamelCase__		)
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							self.batch_norm(UpperCamelCase__		)
												SCREAMING_SNAKE_CASE							:					List[str]           =							self.activation(UpperCamelCase__		)

												return output






class       lowercase__		(     nn.Module):


						def __init__( self			:       Optional[int]					,	UpperCamelCase__			:       int					,	UpperCamelCase__			:       int					,	UpperCamelCase__			:       int		):





												'''simple docstring'''





												super().__init__()
												SCREAMING_SNAKE_CASE							:					Dict           =							[
												    nn.AdaptiveAvgPoolad(UpperCamelCase__		),
												    UperNetConvModule(UpperCamelCase__					,	UpperCamelCase__					,	kernel_size=1		),
												]
												for i, layer in enumerate(self.layers		):
																		self.add_module(str(UpperCamelCase__		)					,	UpperCamelCase__		)



						def  __A ( self			:       List[Any]					,	UpperCamelCase__			:       torch.Tensor		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[str]           =							input
												for layer in self.layers:
																		SCREAMING_SNAKE_CASE							:					int           =							layer(UpperCamelCase__		)
												return hidden_state






class       lowercase__		(     nn.Module):


						def __init__( self			:       str					,	UpperCamelCase__			:       Tuple[int, ...]					,	UpperCamelCase__			:       int					,	UpperCamelCase__			:       int					,	UpperCamelCase__			:       bool		):





												'''simple docstring'''





												super().__init__()
												SCREAMING_SNAKE_CASE							:					Any           =							pool_scales
												SCREAMING_SNAKE_CASE							:					List[Any]           =							align_corners
												SCREAMING_SNAKE_CASE							:					List[Any]           =							in_channels
												SCREAMING_SNAKE_CASE							:					str           =							channels
												SCREAMING_SNAKE_CASE							:					List[Any]           =							[]
												for i, pool_scale in enumerate(UpperCamelCase__		):
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							UperNetPyramidPoolingBlock(pool_scale=UpperCamelCase__					,	in_channels=UpperCamelCase__					,	channels=UpperCamelCase__		)
																		self.blocks.append(UpperCamelCase__		)
																		self.add_module(str(UpperCamelCase__		)					,	UpperCamelCase__		)



						def  __A ( self			:       Dict					,	UpperCamelCase__			:       torch.Tensor		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Optional[int]           =							[]
												for ppm in self.blocks:
																		SCREAMING_SNAKE_CASE							:					Tuple           =							ppm(UpperCamelCase__		)
																		SCREAMING_SNAKE_CASE							:					Any           =							nn.functional.interpolate(
																		    UpperCamelCase__					,	size=x.size()[2:]					,	mode='''bilinear'''					,	align_corners=self.align_corners		)
																		ppm_outs.append(UpperCamelCase__		)
												return ppm_outs






class       lowercase__		(     nn.Module):


						def __init__( self			:       Union[str, Any]					,	UpperCamelCase__			:       Dict					,	UpperCamelCase__			:       Tuple		):





												'''simple docstring'''





												super().__init__()

												SCREAMING_SNAKE_CASE							:					str           =							config
												SCREAMING_SNAKE_CASE							:					Dict           =							config.pool_scales  # e.g. (1, 2, 3, 6)
												SCREAMING_SNAKE_CASE							:					Any           =							in_channels
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							config.hidden_size
												SCREAMING_SNAKE_CASE							:					Dict           =							False
												SCREAMING_SNAKE_CASE							:					int           =							nn.Convad(self.channels					,	config.num_labels					,	kernel_size=1		)

												# PSP Module
												SCREAMING_SNAKE_CASE							:					Tuple           =							UperNetPyramidPoolingModule(
												    self.pool_scales					,	self.in_channels[-1]					,	self.channels					,	align_corners=self.align_corners					,	)
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							UperNetConvModule(
												    self.in_channels[-1] + len(self.pool_scales		) * self.channels					,	self.channels					,	kernel_size=3					,	padding=1					,	)
												# FPN Module
												SCREAMING_SNAKE_CASE							:					Tuple           =							nn.ModuleList()
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							nn.ModuleList()
												for in_channels in self.in_channels[:-1]:  # skip the top layer
																		SCREAMING_SNAKE_CASE							:					int           =							UperNetConvModule(UpperCamelCase__					,	self.channels					,	kernel_size=1		)
																		SCREAMING_SNAKE_CASE							:					List[str]           =							UperNetConvModule(self.channels					,	self.channels					,	kernel_size=3					,	padding=1		)
																		self.lateral_convs.append(UpperCamelCase__		)
																		self.fpn_convs.append(UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							UperNetConvModule(
												    len(self.in_channels		) * self.channels					,	self.channels					,	kernel_size=3					,	padding=1					,	)


						def  __A ( self			:       Optional[int]		):





												'''simple docstring'''





												self.apply(self._init_weights		)


						def  __A ( self			:       Union[str, Any]					,	UpperCamelCase__			:       Tuple		):





												'''simple docstring'''





												if isinstance(UpperCamelCase__					,	nn.Convad		):
																		module.weight.data.normal_(mean=0.0					,	std=self.config.initializer_range		)
																		if module.bias is not None:
																								module.bias.data.zero_()


						def  __A ( self			:       Any					,	UpperCamelCase__			:       Tuple		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Dict           =							inputs[-1]
												SCREAMING_SNAKE_CASE							:					List[Any]           =							[x]
												psp_outs.extend(self.psp_modules(UpperCamelCase__		)		)
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							torch.cat(UpperCamelCase__					,	dim=1		)
												SCREAMING_SNAKE_CASE							:					List[str]           =							self.bottleneck(UpperCamelCase__		)

												return output



						def  __A ( self			:       List[Any]					,	UpperCamelCase__			:       torch.Tensor		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							[lateral_conv(encoder_hidden_states[i]		) for i, lateral_conv in enumerate(self.lateral_convs		)]

												laterals.append(self.psp_forward(UpperCamelCase__		)		)

												# build top-down path
												SCREAMING_SNAKE_CASE							:					List[Any]           =							len(UpperCamelCase__		)
												for i in range(used_backbone_levels - 1					,	0					,	-1		):
																		SCREAMING_SNAKE_CASE							:					List[Any]           =							laterals[i - 1].shape[2:]
																		SCREAMING_SNAKE_CASE							:					str           =							laterals[i - 1] + nn.functional.interpolate(
																		    laterals[i]					,	size=UpperCamelCase__					,	mode='''bilinear'''					,	align_corners=self.align_corners		)

												# build outputs
												SCREAMING_SNAKE_CASE							:					List[str]           =							[self.fpn_convs[i](laterals[i]		) for i in range(used_backbone_levels - 1		)]
												# append psp feature
												fpn_outs.append(laterals[-1]		)

												for i in range(used_backbone_levels - 1					,	0					,	-1		):
																		SCREAMING_SNAKE_CASE							:					List[Any]           =							nn.functional.interpolate(
																		    fpn_outs[i]					,	size=fpn_outs[0].shape[2:]					,	mode='''bilinear'''					,	align_corners=self.align_corners		)
												SCREAMING_SNAKE_CASE							:					Any           =							torch.cat(UpperCamelCase__					,	dim=1		)
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							self.fpn_bottleneck(UpperCamelCase__		)
												SCREAMING_SNAKE_CASE							:					List[Any]           =							self.classifier(UpperCamelCase__		)

												return output






class       lowercase__		(     nn.Module):


						def __init__( self			:       List[Any]					,	UpperCamelCase__			:       Union[str, Any]					,	UpperCamelCase__			:       int = 2					,	UpperCamelCase__			:       int = 3					,	UpperCamelCase__			:       Union[int, Tuple[int, int]] = 1		):





												'''simple docstring'''





												super().__init__()

												SCREAMING_SNAKE_CASE							:					Dict           =							config
												SCREAMING_SNAKE_CASE							:					List[str]           =							config.auxiliary_in_channels
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							config.auxiliary_channels
												SCREAMING_SNAKE_CASE							:					int           =							config.auxiliary_num_convs
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							config.auxiliary_concat_input
												SCREAMING_SNAKE_CASE							:					List[str]           =							in_index

												SCREAMING_SNAKE_CASE							:					Tuple           =							(kernel_size // 2) * dilation
												SCREAMING_SNAKE_CASE							:					str           =							[]
												convs.append(
												    UperNetConvModule(
												        self.in_channels					,	self.channels					,	kernel_size=UpperCamelCase__					,	padding=UpperCamelCase__					,	dilation=UpperCamelCase__		)		)
												for i in range(self.num_convs - 1		):
																		convs.append(
																		    UperNetConvModule(
																		        self.channels					,	self.channels					,	kernel_size=UpperCamelCase__					,	padding=UpperCamelCase__					,	dilation=UpperCamelCase__		)		)
												if self.num_convs == 0:
																		SCREAMING_SNAKE_CASE							:					Optional[int]           =							nn.Identity()
												else:
																		SCREAMING_SNAKE_CASE							:					List[Any]           =							nn.Sequential(*UpperCamelCase__		)
												if self.concat_input:
																		SCREAMING_SNAKE_CASE							:					Any           =							UperNetConvModule(
																		    self.in_channels + self.channels					,	self.channels					,	kernel_size=UpperCamelCase__					,	padding=kernel_size // 2		)

												SCREAMING_SNAKE_CASE							:					List[Any]           =							nn.Convad(self.channels					,	config.num_labels					,	kernel_size=1		)


						def  __A ( self			:       List[Any]		):





												'''simple docstring'''





												self.apply(self._init_weights		)


						def  __A ( self			:       Optional[Any]					,	UpperCamelCase__			:       Tuple		):





												'''simple docstring'''





												if isinstance(UpperCamelCase__					,	nn.Convad		):
																		module.weight.data.normal_(mean=0.0					,	std=self.config.initializer_range		)
																		if module.bias is not None:
																								module.bias.data.zero_()



						def  __A ( self			:       Dict					,	UpperCamelCase__			:       torch.Tensor		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					int           =							encoder_hidden_states[self.in_index]
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							self.convs(UpperCamelCase__		)
												if self.concat_input:
																		SCREAMING_SNAKE_CASE							:					Optional[Any]           =							self.conv_cat(torch.cat([hidden_states, output]					,	dim=1		)		)
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							self.classifier(UpperCamelCase__		)
												return output






class       lowercase__		(     UpperCamelCase_):
						UpperCamelCase_  =	UperNetConfig
						UpperCamelCase_  =	"""pixel_values"""
						UpperCamelCase_  =	True


						def  __A ( self			:       Optional[int]					,	UpperCamelCase__			:       Dict		):





												'''simple docstring'''





												if isinstance(UpperCamelCase__					,	UpperCamelCase__		):
																		module.backbone.init_weights()
																		module.decode_head.init_weights()
																		module.auxiliary_head.init_weights()


						def  __A ( self			:       List[str]		):





												'''simple docstring'''





												self.backbone.init_weights()
												self.decode_head.init_weights()
												self.auxiliary_head.init_weights()



						def  __A ( self			:       List[Any]					,	UpperCamelCase__			:       Any					,	UpperCamelCase__			:       int=False		):





												'''simple docstring'''





												if isinstance(UpperCamelCase__					,	UpperCamelCase__		):
																		SCREAMING_SNAKE_CASE							:					Optional[Any]           =							value






__UpperCamelCase	: str    		=  R'\n    Parameters:\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use\n    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n        config ([`UperNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n'

__UpperCamelCase	: int    		=  R'\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using\n            [`AutoImageProcessor`]. See [`SegformerImageProcessor.__call__`] for details.\n        output_attentions (`bool`, *optional*):\n            Whether or not to return the attentions tensors of all attention layers in case the backbone has them. See\n            `attentions` under returned tensors for more detail.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers of the backbone. See `hidden_states` under\n            returned tensors for more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n'





@add_start_docstrings(
    """UperNet framework leveraging any vision backbone e.g. for ADE20k, CityScapes."""			,							UpperCamelCase_			,							)
class       lowercase__		(     UpperCamelCase_):


						def __init__( self			:       Any					,	UpperCamelCase__			:       List[str]		):





												'''simple docstring'''





												super().__init__(UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							AutoBackbone.from_config(config.backbone_config		)

												# Semantic segmentation head(s)
												SCREAMING_SNAKE_CASE							:					List[Any]           =							UperNetHead(UpperCamelCase__					,	in_channels=self.backbone.channels		)
												SCREAMING_SNAKE_CASE							:					List[Any]           =							UperNetFCNHead(UpperCamelCase__		) if config.use_auxiliary_head else None

												# Initialize weights and apply final processing
												self.post_init()



						@add_start_docstrings_to_model_forward(UPERNET_INPUTS_DOCSTRING.format('''batch_size, sequence_length'''		)		)
						@replace_return_docstrings(output_type=UpperCamelCase__					,	config_class=_CONFIG_FOR_DOC		)
						def  __A ( self			:       List[str]					,	UpperCamelCase__			:       Optional[torch.Tensor] = None					,	UpperCamelCase__			:       Optional[bool] = None					,	UpperCamelCase__			:       Optional[bool] = None					,	UpperCamelCase__			:       Optional[torch.Tensor] = None					,	UpperCamelCase__			:       Optional[bool] = None					,	):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					int           =							return_dict if return_dict is not None else self.config.use_return_dict
												SCREAMING_SNAKE_CASE							:					str           =							(
												    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
												)
												SCREAMING_SNAKE_CASE							:					Dict           =							output_attentions if output_attentions is not None else self.config.output_attentions

												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							self.backbone.forward_with_filtered_kwargs(
												    UpperCamelCase__					,	output_hidden_states=UpperCamelCase__					,	output_attentions=UpperCamelCase__		)
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							outputs.feature_maps

												SCREAMING_SNAKE_CASE							:					Optional[int]           =							self.decode_head(UpperCamelCase__		)
												SCREAMING_SNAKE_CASE							:					List[Any]           =							nn.functional.interpolate(UpperCamelCase__					,	size=pixel_values.shape[2:]					,	mode='''bilinear'''					,	align_corners=UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					Optional[int]           =							None
												if self.auxiliary_head is not None:
																		SCREAMING_SNAKE_CASE							:					Any           =							self.auxiliary_head(UpperCamelCase__		)
																		SCREAMING_SNAKE_CASE							:					Optional[Any]           =							nn.functional.interpolate(
																		    UpperCamelCase__					,	size=pixel_values.shape[2:]					,	mode='''bilinear'''					,	align_corners=UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							None
												if labels is not None:
																		if self.config.num_labels == 1:
																								raise ValueError('''The number of labels should be greater than one'''		)
																		else:
																								# compute weighted loss
																								SCREAMING_SNAKE_CASE							:					str           =							CrossEntropyLoss(ignore_index=self.config.loss_ignore_index		)
																								SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							loss_fct(UpperCamelCase__					,	UpperCamelCase__		)
																								SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							loss_fct(UpperCamelCase__					,	UpperCamelCase__		)
																								SCREAMING_SNAKE_CASE							:					List[str]           =							main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

												if not return_dict:
																		if output_hidden_states:
																								SCREAMING_SNAKE_CASE							:					Optional[Any]           =							(logits,) + outputs[1:]
																		else:
																								SCREAMING_SNAKE_CASE							:					Optional[Any]           =							(logits,) + outputs[2:]
																		return ((loss,) + output) if loss is not None else output

												return SemanticSegmenterOutput(
												    loss=UpperCamelCase__					,	logits=UpperCamelCase__					,	hidden_states=outputs.hidden_states					,	attentions=outputs.attentions					,	)






import argparse
import json
import logging
import os
import sys
from unittest.mock import patch

from transformers.testing_utils import TestCasePlus, get_gpu_count, slow


__UpperCamelCase	: str    		=  [
    os.path.join(os.path.dirname(__file__), dirname)
    for dirname in [
        'text-classification',
        'language-modeling',
        'summarization',
        'token-classification',
        'question-answering',
    ]
]
sys.path.extend(SRC_DIRS)


if SRC_DIRS is not None:
					import run_clm_flax
					import run_flax_glue
					import run_flax_ner
					import run_mlm_flax
					import run_qa
					import run_summarization_flax
					import run_ta_mlm_flax


logging.basicConfig(level=logging.DEBUG)

__UpperCamelCase	: int    		=  logging.getLogger()
def    A			(					):
						SCREAMING_SNAKE_CASE							:					str           =							argparse.ArgumentParser()
						parser.add_argument('''-f'''		)
						SCREAMING_SNAKE_CASE							:					List[str]           =							parser.parse_args()
						return args.f
def    A			(					_lowercase	,						_lowercase="eval"		):
						SCREAMING_SNAKE_CASE							:					Dict           =							os.path.join(_lowercase	,						f"""{split}_results.json"""		)
						if os.path.exists(_lowercase		):
												with open(_lowercase	,						'''r'''		) as f:
																		return json.load(_lowercase		)
						raise ValueError(f"""can't find {path}"""		)


__UpperCamelCase	: Optional[Any]    		=  logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)





class       lowercase__		(     UpperCamelCase_):


						def  __A ( self			:       Optional[int]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					int           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Tuple           =							f"""
            run_glue.py
            --model_name_or_path distilbert-base-uncased
            --output_dir {tmp_dir}
            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv
            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --learning_rate=1e-4
            --eval_steps=2
            --warmup_steps=2
            --seed=42
            --max_seq_length=128
            """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_flax_glue.main()
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							get_results(UpperCamelCase__		)
																		self.assertGreaterEqual(result['''eval_accuracy''']					,	0.75		)


						@slow
						def  __A ( self			:       Optional[int]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Dict           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					str           =							f"""
            run_clm_flax.py
            --model_name_or_path distilgpt2
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --do_train
            --do_eval
            --block_size 128
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --num_train_epochs 2
            --logging_steps 2 --eval_steps 2
            --output_dir {tmp_dir}
            --overwrite_output_dir
            """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_clm_flax.main()
																		SCREAMING_SNAKE_CASE							:					Dict           =							get_results(UpperCamelCase__		)
																		self.assertLess(result['''eval_perplexity''']					,	100		)


						@slow
						def  __A ( self			:       Any		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[str]           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							f"""
            run_summarization.py
            --model_name_or_path t5-small
            --train_file tests/fixtures/tests_samples/xsum/sample.json
            --validation_file tests/fixtures/tests_samples/xsum/sample.json
            --test_file tests/fixtures/tests_samples/xsum/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --num_train_epochs=3
            --warmup_steps=8
            --do_train
            --do_eval
            --do_predict
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --predict_with_generate
        """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_summarization_flax.main()
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							get_results(UpperCamelCase__					,	split='''test'''		)
																		self.assertGreaterEqual(result['''test_rouge1''']					,	10		)
																		self.assertGreaterEqual(result['''test_rouge2''']					,	2		)
																		self.assertGreaterEqual(result['''test_rougeL''']					,	7		)
																		self.assertGreaterEqual(result['''test_rougeLsum''']					,	7		)


						@slow
						def  __A ( self			:       Optional[int]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[Any]           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Dict           =							f"""
            run_mlm.py
            --model_name_or_path distilroberta-base
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --max_seq_length 128
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --logging_steps 2 --eval_steps 2
            --do_train
            --do_eval
            --num_train_epochs=1
        """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_mlm_flax.main()
																		SCREAMING_SNAKE_CASE							:					List[Any]           =							get_results(UpperCamelCase__		)
																		self.assertLess(result['''eval_perplexity''']					,	42		)


						@slow
						def  __A ( self			:       List[Any]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Dict           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							f"""
            run_t5_mlm_flax.py
            --model_name_or_path t5-small
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --do_train
            --do_eval
            --max_seq_length 128
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --num_train_epochs 2
            --logging_steps 2 --eval_steps 2
            --output_dir {tmp_dir}
            --overwrite_output_dir
            """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_ta_mlm_flax.main()
																		SCREAMING_SNAKE_CASE							:					Optional[int]           =							get_results(UpperCamelCase__		)
																		self.assertGreaterEqual(result['''eval_accuracy''']					,	0.42		)


						@slow
						def  __A ( self			:       Union[str, Any]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Tuple           =							7 if get_gpu_count() > 1 else 2

												SCREAMING_SNAKE_CASE							:					Dict           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Any           =							f"""
            run_flax_ner.py
            --model_name_or_path bert-base-uncased
            --train_file tests/fixtures/tests_samples/conll/sample.json
            --validation_file tests/fixtures/tests_samples/conll/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --do_train
            --do_eval
            --warmup_steps=2
            --learning_rate=2e-4
            --logging_steps 2 --eval_steps 2
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=2
            --num_train_epochs={epochs}
            --seed 7
        """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_flax_ner.main()
																		SCREAMING_SNAKE_CASE							:					List[str]           =							get_results(UpperCamelCase__		)
																		self.assertGreaterEqual(result['''eval_accuracy''']					,	0.75		)
																		self.assertGreaterEqual(result['''eval_f1''']					,	0.3		)



						@slow
						def  __A ( self			:       List[str]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[str]           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							f"""
            run_qa.py
            --model_name_or_path bert-base-uncased
            --version_2_with_negative
            --train_file tests/fixtures/tests_samples/SQUAD/sample.json
            --validation_file tests/fixtures/tests_samples/SQUAD/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --num_train_epochs=3
            --warmup_steps=2
            --do_train
            --do_eval
            --logging_steps 2 --eval_steps 2
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
        """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_qa.main()
																		SCREAMING_SNAKE_CASE							:					str           =							get_results(UpperCamelCase__		)
																		self.assertGreaterEqual(result['''eval_f1''']					,	30		)
																		self.assertGreaterEqual(result['''eval_exact''']					,	30		)





