



import gc
import unittest

import torch
from parameterized import parameterized

from diffusers import AutoencoderKL
from diffusers.utils import floats_tensor, load_hf_numpy, require_torch_gpu, slow, torch_all_close, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()
class 				_snake_case			(					lowercase__    ,	lowercase__    ,	unittest.TestCase):
				UpperCamelCase__				:					Any     		=AutoencoderKL
				UpperCamelCase__				:					Any     		="""sample"""
				UpperCamelCase__				:					str     		=1E-2




				@property
				def 							A__      (						self					:						int						):
							lowercase__					  =				4
							lowercase__					  =				3
							lowercase__					  =				(32, 32)

							lowercase__					  =				floats_tensor((batch_size, num_channels) + sizes						).to(__lowercase						)

							return {"sample": image}




				@property
				def 							A__      (						self					:						int						):
							return (3, 32, 32)




				@property
				def 							A__      (						self					:						Optional[int]						):
							return (3, 32, 32)




				def 							A__      (						self					:						int						):
							lowercase__					  =				{
							    "block_out_channels": [32, 64],
							    "in_channels": 3,
							    "out_channels": 3,
							    "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D"],
							    "up_block_types": ["UpDecoderBlock2D", "UpDecoderBlock2D"],
							    "latent_channels": 4,
							}
							lowercase__					  =				self.dummy_input
							return init_dict, inputs_dict




				def 							A__      (						self					:						Union[str, Any]						):
							pass




				def 							A__      (						self					:						Dict						):
							pass




				@unittest.skipIf(torch_device == "mps",					"Gradient checkpointing skipped on MPS"						)
				def 							A__      (						self					:						List[Any]						):
							# enable deterministic behavior for gradient checkpointing
							lowercase__  ,       lowercase__					  =				self.prepare_init_args_and_inputs_for_common()
							lowercase__					  =				self.model_class(**__lowercase						)
							model.to(__lowercase						)

							assert not model.is_gradient_checkpointing and model.training

							lowercase__					  =				model(**__lowercase						).sample
							# run the backwards pass on the model. For backwards pass, for simplicity purpose,
							# we won't calculate the loss and rather backprop on out.sum()
							model.zero_grad()

							lowercase__					  =				torch.randn_like(__lowercase						)
							lowercase__					  =				(out - labels).mean()
							loss.backward()

							# re-instantiate the model now enabling gradient checkpointing
							lowercase__					  =				self.model_class(**__lowercase						)
							# clone model
							model_a.load_state_dict(model.state_dict()						)
							model_a.to(__lowercase						)
							model_a.enable_gradient_checkpointing()

							assert model_a.is_gradient_checkpointing and model_a.training

							lowercase__					  =				model_a(**__lowercase						).sample
							# run the backwards pass on the model. For backwards pass, for simplicity purpose,
							# we won't calculate the loss and rather backprop on out.sum()
							model_a.zero_grad()
							lowercase__					  =				(out_a - labels).mean()
							loss_a.backward()

							# compare the output and parameters gradients
							self.assertTrue((loss - loss_a).abs() < 1e-5						)
							lowercase__					  =				dict(model.named_parameters()						)
							lowercase__					  =				dict(model_a.named_parameters()						)
							for name, param in named_params.items():
										self.assertTrue(torch_all_close(param.grad.data,					named_params_a[name].grad.data,					atol=5e-5						)						)




				def 							A__      (						self					:						Union[str, Any]						):
							lowercase__  ,       lowercase__					  =				AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy",					output_loading_info=__lowercase						)
							self.assertIsNotNone(__lowercase						)
							self.assertEqual(len(loading_info["missing_keys"]						),					0						)

							model.to(__lowercase						)
							lowercase__					  =				model(**self.dummy_input						)

							assert image is not None, "Make sure output is not None"




				def 							A__      (						self					:						Any						):
							lowercase__					  =				AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy"						)
							lowercase__					  =				model.to(__lowercase						)
							model.eval()

							if torch_device == "mps":
										lowercase__					  =				torch.manual_seed(0						)
							else:
										lowercase__					  =				torch.Generator(device=__lowercase						).manual_seed(0						)

							lowercase__					  =				torch.randn(
							    1,					model.config.in_channels,					model.config.sample_size,					model.config.sample_size,					generator=torch.manual_seed(0						),					)
							lowercase__					  =				image.to(__lowercase						)
							with torch.no_grad():
										lowercase__					  =				model(__lowercase,					sample_posterior=__lowercase,					generator=__lowercase						).sample

							lowercase__					  =				output[0, -1, -3:, -3:].flatten().cpu()

							# Since the VAE Gaussian prior's generator is seeded on the appropriate device,
							# the expected output slices are not the same for CPU and GPU.
							if torch_device == "mps":
										lowercase__					  =				torch.tensor(
										    [
										        -4.0_0_7_8e-0_1,
										        -3.8_3_2_3e-0_4,
										        -1.2_6_8_1e-0_1,
										        -1.1_4_6_2e-0_1,
										        2.0_0_9_5e-0_1,
										        1.0_8_9_3e-0_1,
										        -8.8_2_4_7e-0_2,
										        -3.0_3_6_1e-0_1,
										        -9.8_6_4_4e-0_3,
										    ]						)
							elif torch_device == "cpu":
										lowercase__					  =				torch.tensor(
										    [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]						)
							else:
										lowercase__					  =				torch.tensor(
										    [-0.2421, 0.4642, 0.2507, -0.0438, 0.0682, 0.3160, -0.2018, -0.0727, 0.2485]						)

							self.assertTrue(torch_all_close(__lowercase,					__lowercase,					rtol=1e-2						)						)
@slow
class 				_snake_case			(					unittest.TestCase):




				def 							A__      (						self					:						Dict,					__lowercase					:						int,					__lowercase					:						str						):
							return F'''gaussian_noise_s={seed}_shape={"_".join([str(__lowercase						) for s in shape]						)}.npy'''




				def 							A__      (						self					:						Optional[Any]						):
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()




				def 							A__      (						self					:						Optional[Any],					__lowercase					:						Optional[Any]=0,					__lowercase					:						List[str]=(4, 3, 512, 512),					__lowercase					:						List[Any]=False						):
							lowercase__					  =				torch.floataa if fpaa else torch.floataa
							lowercase__					  =				torch.from_numpy(load_hf_numpy(self.get_file_format(__lowercase,					__lowercase						)						)						).to(__lowercase						).to(__lowercase						)
							return image




				def 							A__      (						self					:						Optional[int],					__lowercase					:						Any="CompVis/stable-diffusion-v1-4",					__lowercase					:						List[Any]=False						):
							lowercase__					  =				"fp16" if fpaa else None
							lowercase__					  =				torch.floataa if fpaa else torch.floataa

							lowercase__					  =				AutoencoderKL.from_pretrained(
							    __lowercase,					subfolder="vae",					torch_dtype=__lowercase,					revision=__lowercase,					)
							model.to(__lowercase						).eval()

							return model




				def 							A__      (						self					:						str,					__lowercase					:						Optional[int]=0						):
							if torch_device == "mps":
										return torch.manual_seed(__lowercase						)
							return torch.Generator(device=__lowercase						).manual_seed(__lowercase						)




				@parameterized.expand(
				    [
				        # fmt: off
				        [33, [-0.1603, 0.9878, -0.0495, -0.0790, -0.2709, 0.8375, -0.2060, -0.0824], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
				        [47, [-0.2376, 0.1168, 0.1332, -0.4840, -0.2508, -0.0791, -0.0493, -0.4089], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
				        # fmt: on
				    ]						)
				def 							A__      (						self					:						int,					__lowercase					:						str,					__lowercase					:						Union[str, Any],					__lowercase					:						List[Any]						):
							lowercase__					  =				self.get_sd_vae_model()
							lowercase__					  =				self.get_sd_image(__lowercase						)
							lowercase__					  =				self.get_generator(__lowercase						)

							with torch.no_grad():
										lowercase__					  =				model(__lowercase,					generator=__lowercase,					sample_posterior=__lowercase						).sample

							assert sample.shape == image.shape

							lowercase__					  =				sample[-1, -2:, -2:, :2].flatten().float().cpu()
							lowercase__					  =				torch.tensor(expected_slice_mps if torch_device == "mps" else expected_slice						)

							assert torch_all_close(__lowercase,					__lowercase,					atol=3e-3						)




				@parameterized.expand(
				    [
				        # fmt: off
				        [33, [-0.0513, 0.0289, 1.3799, 0.2166, -0.2573, -0.0871, 0.5103, -0.0999]],
				        [47, [-0.4128, -0.1320, -0.3704, 0.1965, -0.4116, -0.2332, -0.3340, 0.2247]],
				        # fmt: on
				    ]						)
				@require_torch_gpu
				def 							A__      (						self					:						int,					__lowercase					:						Optional[int],					__lowercase					:						Tuple						):
							lowercase__					  =				self.get_sd_vae_model(fpaa=__lowercase						)
							lowercase__					  =				self.get_sd_image(__lowercase,					fpaa=__lowercase						)
							lowercase__					  =				self.get_generator(__lowercase						)

							with torch.no_grad():
										lowercase__					  =				model(__lowercase,					generator=__lowercase,					sample_posterior=__lowercase						).sample

							assert sample.shape == image.shape

							lowercase__					  =				sample[-1, -2:, :2, -2:].flatten().float().cpu()
							lowercase__					  =				torch.tensor(__lowercase						)

							assert torch_all_close(__lowercase,					__lowercase,					atol=1e-2						)




				@parameterized.expand(
				    [
				        # fmt: off
				        [33, [-0.1609, 0.9866, -0.0487, -0.0777, -0.2716, 0.8368, -0.2055, -0.0814], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
				        [47, [-0.2377, 0.1147, 0.1333, -0.4841, -0.2506, -0.0805, -0.0491, -0.4085], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
				        # fmt: on
				    ]						)
				def 							A__      (						self					:						Any,					__lowercase					:						Optional[Any],					__lowercase					:						Tuple,					__lowercase					:						List[Any]						):
							lowercase__					  =				self.get_sd_vae_model()
							lowercase__					  =				self.get_sd_image(__lowercase						)

							with torch.no_grad():
										lowercase__					  =				model(__lowercase						).sample

							assert sample.shape == image.shape

							lowercase__					  =				sample[-1, -2:, -2:, :2].flatten().float().cpu()
							lowercase__					  =				torch.tensor(expected_slice_mps if torch_device == "mps" else expected_slice						)

							assert torch_all_close(__lowercase,					__lowercase,					atol=3e-3						)




				@parameterized.expand(
				    [
				        # fmt: off
				        [13, [-0.2051, -0.1803, -0.2311, -0.2114, -0.3292, -0.3574, -0.2953, -0.3323]],
				        [37, [-0.2632, -0.2625, -0.2199, -0.2741, -0.4539, -0.4990, -0.3720, -0.4925]],
				        # fmt: on
				    ]						)
				@require_torch_gpu
				def 							A__      (						self					:						Optional[Any],					__lowercase					:						Union[str, Any],					__lowercase					:						Tuple						):
							lowercase__					  =				self.get_sd_vae_model()
							lowercase__					  =				self.get_sd_image(__lowercase,					shape=(3, 4, 64, 64)						)

							with torch.no_grad():
										lowercase__					  =				model.decode(__lowercase						).sample

							assert list(sample.shape						) == [3, 3, 512, 512]

							lowercase__					  =				sample[-1, -2:, :2, -2:].flatten().cpu()
							lowercase__					  =				torch.tensor(__lowercase						)

							assert torch_all_close(__lowercase,					__lowercase,					atol=1e-3						)




				@parameterized.expand(
				    [
				        # fmt: off
				        [27, [-0.0369, 0.0207, -0.0776, -0.0682, -0.1747, -0.1930, -0.1465, -0.2039]],
				        [16, [-0.1628, -0.2134, -0.2747, -0.2642, -0.3774, -0.4404, -0.3687, -0.4277]],
				        # fmt: on
				    ]						)
				@require_torch_gpu
				def 							A__      (						self					:						Optional[int],					__lowercase					:						List[Any],					__lowercase					:						Any						):
							lowercase__					  =				self.get_sd_vae_model(fpaa=__lowercase						)
							lowercase__					  =				self.get_sd_image(__lowercase,					shape=(3, 4, 64, 64),					fpaa=__lowercase						)

							with torch.no_grad():
										lowercase__					  =				model.decode(__lowercase						).sample

							assert list(sample.shape						) == [3, 3, 512, 512]

							lowercase__					  =				sample[-1, -2:, :2, -2:].flatten().float().cpu()
							lowercase__					  =				torch.tensor(__lowercase						)

							assert torch_all_close(__lowercase,					__lowercase,					atol=5e-3						)




				@parameterized.expand([(13,), (16,), (27,)]						)
				@require_torch_gpu
				@unittest.skipIf(not is_xformers_available(),					reason="xformers is not required when using PyTorch 2.0."						)
				def 							A__      (						self					:						Tuple,					__lowercase					:						Union[str, Any]						):
							lowercase__					  =				self.get_sd_vae_model(fpaa=__lowercase						)
							lowercase__					  =				self.get_sd_image(__lowercase,					shape=(3, 4, 64, 64),					fpaa=__lowercase						)

							with torch.no_grad():
										lowercase__					  =				model.decode(__lowercase						).sample

							model.enable_xformers_memory_efficient_attention()
							with torch.no_grad():
										lowercase__					  =				model.decode(__lowercase						).sample

							assert list(sample.shape						) == [3, 3, 512, 512]

							assert torch_all_close(__lowercase,					__lowercase,					atol=1e-1						)




				@parameterized.expand([(13,), (16,), (37,)]						)
				@require_torch_gpu
				@unittest.skipIf(not is_xformers_available(),					reason="xformers is not required when using PyTorch 2.0."						)
				def 							A__      (						self					:						int,					__lowercase					:						Any						):
							lowercase__					  =				self.get_sd_vae_model()
							lowercase__					  =				self.get_sd_image(__lowercase,					shape=(3, 4, 64, 64)						)

							with torch.no_grad():
										lowercase__					  =				model.decode(__lowercase						).sample

							model.enable_xformers_memory_efficient_attention()
							with torch.no_grad():
										lowercase__					  =				model.decode(__lowercase						).sample

							assert list(sample.shape						) == [3, 3, 512, 512]

							assert torch_all_close(__lowercase,					__lowercase,					atol=1e-2						)




				@parameterized.expand(
				    [
				        # fmt: off
				        [33, [-0.3001, 0.0918, -2.6984, -3.9720, -3.2099, -5.0353, 1.7338, -0.2065, 3.4267]],
				        [47, [-1.5030, -4.3871, -6.0355, -9.1157, -1.6661, -2.7853, 2.1607, -5.0823, 2.5633]],
				        # fmt: on
				    ]						)
				def 							A__      (						self					:						Optional[Any],					__lowercase					:						Union[str, Any],					__lowercase					:						Optional[Any]						):
							lowercase__					  =				self.get_sd_vae_model()
							lowercase__					  =				self.get_sd_image(__lowercase						)
							lowercase__					  =				self.get_generator(__lowercase						)

							with torch.no_grad():
										lowercase__					  =				model.encode(__lowercase						).latent_dist
										lowercase__					  =				dist.sample(generator=__lowercase						)

							assert list(sample.shape						) == [image.shape[0], 4] + [i // 8 for i in image.shape[2:]]

							lowercase__					  =				sample[0, -1, -3:, -3:].flatten().cpu()
							lowercase__					  =				torch.tensor(__lowercase						)

							lowercase__					  =				3e-3 if torch_device != "mps" else 1e-2
							assert torch_all_close(__lowercase,					__lowercase,					atol=__lowercase						)







# Usage:
# ./gen-card-facebook-wmt19.py

import os
from pathlib import Path



def __lowerCAmelCase (     SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_						,  SCREAMING_SNAKE_CASE_  ):
			lowercase__					  =				{
			    "en": "Machine learning is great, isn't it?",
			    "ru": "Машинное обучение - это здорово, не так ли?",
			    "de": "Maschinelles Lernen ist großartig, oder?",
			}

			# BLUE scores as follows:
			# "pair": [fairseq, transformers]
			lowercase__					  =				{
			    "ru-en": ["[41.3](http://matrix.statmt.org/matrix/output/1907?run_id=6937)", "39.20"],
			    "en-ru": ["[36.4](http://matrix.statmt.org/matrix/output/1914?run_id=6724)", "33.47"],
			    "en-de": ["[43.1](http://matrix.statmt.org/matrix/output/1909?run_id=6862)", "42.83"],
			    "de-en": ["[42.3](http://matrix.statmt.org/matrix/output/1902?run_id=6750)", "41.35"],
			}
			lowercase__					  =				f'''{src_lang}-{tgt_lang}'''

			lowercase__					  =				f'''
---
language: 
- {src_lang}
- {tgt_lang}
thumbnail:
tags:
- translation
- wmt19
- facebook
license: apache-2.0
datasets:
- wmt19
metrics:
- bleu
---

# FSMT

## Model description

This is a ported version of [fairseq wmt19 transformer](https://github.com/pytorch/fairseq/blob/master/examples/wmt19/README.md) for {src_lang}-{tgt_lang}.

For more details, please see, [Facebook FAIR\'s WMT19 News Translation Task Submission](https://arxiv.org/abs/1907.06616).

The abbreviation FSMT stands for FairSeqMachineTranslation

All four models are available:

* [wmt19-en-ru](https://huggingface.co/facebook/wmt19-en-ru)
* [wmt19-ru-en](https://huggingface.co/facebook/wmt19-ru-en)
* [wmt19-en-de](https://huggingface.co/facebook/wmt19-en-de)
* [wmt19-de-en](https://huggingface.co/facebook/wmt19-de-en)

## Intended uses & limitations

#### How to use

```python
from transformers import FSMTForConditionalGeneration, FSMTTokenizer
mname = "facebook/wmt19-{src_lang}-{tgt_lang}"
tokenizer = FSMTTokenizer.from_pretrained(mname)
model = FSMTForConditionalGeneration.from_pretrained(mname)

input = "{texts[src_lang]}"
input_ids = tokenizer.encode(input, return_tensors="pt")
outputs = model.generate(input_ids)
decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(decoded) # {texts[tgt_lang]}

```

#### Limitations and bias

- The original (and this ported model) doesn\'t seem to handle well inputs with repeated sub-phrases, [content gets truncated](https://discuss.huggingface.co/t/issues-with-translating-inputs-containing-repeated-phrases/981)

## Training data

Pretrained weights were left identical to the original model released by fairseq. For more details, please, see the [paper](https://arxiv.org/abs/1907.06616).

## Eval results

pair   | fairseq | transformers
-------|---------|----------
{pair}  | {scores[pair][0]} | {scores[pair][1]}

The score is slightly below the score reported by `fairseq`, since `transformers`` currently doesn\'t support:
- model ensemble, therefore the best performing checkpoint was ported (``model4.pt``).
- re-ranking

The score was calculated using this code:

```bash
git clone https://github.com/huggingface/transformers
cd transformers
export PAIR={pair}
export DATA_DIR=data/$PAIR
export SAVE_DIR=data/$PAIR
export BS=8
export NUM_BEAMS=15
mkdir -p $DATA_DIR
sacrebleu -t wmt19 -l $PAIR --echo src > $DATA_DIR/val.source
sacrebleu -t wmt19 -l $PAIR --echo ref > $DATA_DIR/val.target
echo $PAIR
PYTHONPATH="src:examples/seq2seq" python examples/seq2seq/run_eval.py facebook/wmt19-$PAIR $DATA_DIR/val.source $SAVE_DIR/test_translations.txt --reference_path $DATA_DIR/val.target --score_path $SAVE_DIR/test_bleu.json --bs $BS --task translation --num_beams $NUM_BEAMS
```
note: fairseq reports using a beam of 50, so you should get a slightly higher score if re-run with `--num_beams 50`.

## Data Sources

- [training, etc.](http://www.statmt.org/wmt19/)
- [test set](http://matrix.statmt.org/test_sets/newstest2019.tgz?1556572561)


### BibTeX entry and citation info

```bibtex
@inproceedings{{...,
  year={{2020}},
  title={{Facebook FAIR\'s WMT19 News Translation Task Submission}},
  author={{Ng, Nathan and Yee, Kyra and Baevski, Alexei and Ott, Myle and Auli, Michael and Edunov, Sergey}},
  booktitle={{Proc. of WMT}},
}}
```


## TODO

- port model ensemble (fairseq uses 4 model checkpoints)

'''
			os.makedirs(SCREAMING_SNAKE_CASE_						,  exist_ok=SCREAMING_SNAKE_CASE_  )
			lowercase__					  =				os.path.join(SCREAMING_SNAKE_CASE_						,  "README.md"  )
			print(f'''Generating {path}'''  )
			with open(SCREAMING_SNAKE_CASE_						,  "w"						,  encoding="utf-8"  ) as f:
						f.write(SCREAMING_SNAKE_CASE_  )


# make sure we are under the root of the project
lowercase_								=     Path(__file__).resolve().parent.parent.parent
lowercase_								=     repo_dir / """model_cards"""

for model_name in ["wmt19-ru-en", "wmt19-en-ru", "wmt19-en-de", "wmt19-de-en"]:
							lowercase_				,    lowercase_				,    lowercase_								=     model_name.split("""-""")
							lowercase_								=     model_cards_dir / """facebook""" / model_name
							write_model_card(model_card_dir, src_lang=src_lang, tgt_lang=tgt_lang)




