



import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, LDMTextToImagePipeline, UNetaDConditionModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_numpy,
    nightly,
    require_torch_gpu,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()

class 	__SCREAMING_SNAKE_CASE     (				snake_case_		,       unittest.TestCase       ):


					'''simple docstring'''





					_a             =							LDMTextToImagePipeline
					_a             =							TEXT_TO_IMAGE_PARAMS - {
					    'negative_prompt',
					    'negative_prompt_embeds',
					    'cross_attention_kwargs',
					    'prompt_embeds',
					}
					_a             =							PipelineTesterMixin.required_optional_params - {
					    'num_images_per_prompt',
					    'callback',
					    'callback_steps',
					}
					_a             =							TEXT_TO_IMAGE_BATCH_PARAMS
					_a             =							False



					def 		snake_case	(	self							:  str					)-> int:
								torch.manual_seed(0					)
								lowerCamelCase__		:		Any		     =UNetaDConditionModel(
								    block_out_channels=(32, 64),  layers_per_block=2,  sample_size=32,  in_channels=4,  out_channels=4,  down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D'''),  up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D'''),  cross_attention_dim=32,  )
								lowerCamelCase__		:		Optional[int]		     =DDIMScheduler(
								    beta_start=0.00_085,  beta_end=0.012,  beta_schedule='''scaled_linear''',  clip_sample=lowerCamelCase,  set_alpha_to_one=lowerCamelCase,  )
								torch.manual_seed(0					)
								lowerCamelCase__		:		Tuple		     =AutoencoderKL(
								    block_out_channels=(32, 64),  in_channels=3,  out_channels=3,  down_block_types=('''DownEncoderBlock2D''', '''DownEncoderBlock2D'''),  up_block_types=('''UpDecoderBlock2D''', '''UpDecoderBlock2D'''),  latent_channels=4,  )
								torch.manual_seed(0					)
								lowerCamelCase__		:		Any		     =CLIPTextConfig(
								    bos_token_id=0,  eos_token_id=2,  hidden_size=32,  intermediate_size=37,  layer_norm_eps=1E-05,  num_attention_heads=4,  num_hidden_layers=5,  pad_token_id=1,  vocab_size=1000,  )
								lowerCamelCase__		:		Optional[int]		     =CLIPTextModel(lowerCamelCase					)
								lowerCamelCase__		:		Tuple		     =CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)

								lowerCamelCase__		:		Any		     ={
								    '''unet''': unet,
								    '''scheduler''': scheduler,
								    '''vqvae''': vae,
								    '''bert''': text_encoder,
								    '''tokenizer''': tokenizer,
								}
								return components



					def 		snake_case	(	self							:  Optional[Any],  lowerCamelCase							:  Optional[int],  lowerCamelCase							:  Optional[Any]=0					)-> Optional[int]:
								if str(lowerCamelCase					).startswith('''mps'''					):
											lowerCamelCase__		:		List[str]		     =torch.manual_seed(lowerCamelCase					)
								else:
											lowerCamelCase__		:		int		     =torch.Generator(device=lowerCamelCase					).manual_seed(lowerCamelCase					)
								lowerCamelCase__		:		int		     ={
								    '''prompt''': '''A painting of a squirrel eating a burger''',
								    '''generator''': generator,
								    '''num_inference_steps''': 2,
								    '''guidance_scale''': 6.0,
								    '''output_type''': '''numpy''',
								}
								return inputs




					def 		snake_case	(	self							:  Tuple					)-> Union[str, Any]:
								lowerCamelCase__		:		str		     ='''cpu'''  # ensure determinism for the device-dependent torch.Generator

								lowerCamelCase__		:		Union[str, Any]		     =self.get_dummy_components()
								lowerCamelCase__		:		Optional[Any]		     =LDMTextToImagePipeline(**lowerCamelCase					)
								pipe.to(lowerCamelCase					)
								pipe.set_progress_bar_config(disable=lowerCamelCase					)

								lowerCamelCase__		:		Any		     =self.get_dummy_inputs(lowerCamelCase					)
								lowerCamelCase__		:		Tuple		     =pipe(**lowerCamelCase					).images
								lowerCamelCase__		:		Tuple		     =image[0, -3:, -3:, -1]

								assert image.shape == (1, 16, 16, 3)
								lowerCamelCase__		:		List[Any]		     =np.array([0.6_101, 0.6_156, 0.5_622, 0.4_895, 0.6_661, 0.3_804, 0.5_748, 0.6_136, 0.5_014]					)

								assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-3





@slow
@require_torch_gpu
class 	__SCREAMING_SNAKE_CASE     (				unittest.TestCase       ):


					'''simple docstring'''





					def 		snake_case	(	self							:  Any					)-> Union[str, Any]:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()



					def 		snake_case	(	self							:  List[Any],  lowerCamelCase							:  List[str],  lowerCamelCase							:  Dict=torch.floataa,  lowerCamelCase							:  Tuple=0					)-> str:
								lowerCamelCase__		:		List[Any]		     =torch.manual_seed(lowerCamelCase					)
								lowerCamelCase__		:		str		     =np.random.RandomState(lowerCamelCase					).standard_normal((1, 4, 32, 32)					)
								lowerCamelCase__		:		Any		     =torch.from_numpy(lowerCamelCase					).to(device=lowerCamelCase,  dtype=lowerCamelCase					)
								lowerCamelCase__		:		Union[str, Any]		     ={
								    '''prompt''': '''A painting of a squirrel eating a burger''',
								    '''latents''': latents,
								    '''generator''': generator,
								    '''num_inference_steps''': 3,
								    '''guidance_scale''': 6.0,
								    '''output_type''': '''numpy''',
								}
								return inputs




					def 		snake_case	(	self							:  List[str]					)-> Tuple:
								lowerCamelCase__		:		Optional[Any]		     =LDMTextToImagePipeline.from_pretrained('''CompVis/ldm-text2im-large-256'''					).to(lowerCamelCase					)
								pipe.set_progress_bar_config(disable=lowerCamelCase					)

								lowerCamelCase__		:		Union[str, Any]		     =self.get_inputs(lowerCamelCase					)
								lowerCamelCase__		:		Optional[Any]		     =pipe(**lowerCamelCase					).images
								lowerCamelCase__		:		Optional[int]		     =image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 256, 256, 3)
								lowerCamelCase__		:		int		     =np.array([0.51_825, 0.52_850, 0.52_543, 0.54_258, 0.52_304, 0.52_569, 0.54_363, 0.55_276, 0.56_878]					)
								lowerCamelCase__		:		Tuple		     =np.abs(expected_slice - image_slice					).max()
								assert max_diff < 1E-3





@nightly
@require_torch_gpu
class 	__SCREAMING_SNAKE_CASE     (				unittest.TestCase       ):


					'''simple docstring'''





					def 		snake_case	(	self							:  List[str]					)-> Tuple:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()



					def 		snake_case	(	self							:  Any,  lowerCamelCase							:  Dict,  lowerCamelCase							:  List[str]=torch.floataa,  lowerCamelCase							:  Union[str, Any]=0					)-> int:
								lowerCamelCase__		:		Optional[int]		     =torch.manual_seed(lowerCamelCase					)
								lowerCamelCase__		:		Any		     =np.random.RandomState(lowerCamelCase					).standard_normal((1, 4, 32, 32)					)
								lowerCamelCase__		:		Tuple		     =torch.from_numpy(lowerCamelCase					).to(device=lowerCamelCase,  dtype=lowerCamelCase					)
								lowerCamelCase__		:		List[str]		     ={
								    '''prompt''': '''A painting of a squirrel eating a burger''',
								    '''latents''': latents,
								    '''generator''': generator,
								    '''num_inference_steps''': 50,
								    '''guidance_scale''': 6.0,
								    '''output_type''': '''numpy''',
								}
								return inputs




					def 		snake_case	(	self							:  List[str]					)-> Optional[int]:
								lowerCamelCase__		:		Union[str, Any]		     =LDMTextToImagePipeline.from_pretrained('''CompVis/ldm-text2im-large-256'''					).to(lowerCamelCase					)
								pipe.set_progress_bar_config(disable=lowerCamelCase					)

								lowerCamelCase__		:		Dict		     =self.get_inputs(lowerCamelCase					)
								lowerCamelCase__		:		Optional[Any]		     =pipe(**lowerCamelCase					).images[0]

								lowerCamelCase__		:		Any		     =load_numpy(
								    '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/ldm_text2img/ldm_large_256_ddim.npy'''					)
								lowerCamelCase__		:		str		     =np.abs(expected_image - image					).max()
								assert max_diff < 1E-3









"""simple docstring"""






# Usage:
# ./gen-card-allenai-wmt16.py

import os
from pathlib import Path






def 					snake_case__       (		__lowerCamelCase    :   Union[str, Any]							,		__lowerCamelCase    :   int							,		__lowerCamelCase    :   Tuple							,		__lowerCamelCase    :   Union[str, Any]     ):
			"""simple docstring"""
			lowerCamelCase__		:		Any		     ={
			    '''en''': '''Machine learning is great, isn\'t it?''',
			    '''ru''': '''Машинное обучение - это здорово, не так ли?''',
			    '''de''': '''Maschinelles Lernen ist großartig, nicht wahr?''',
			}

			# BLUE scores as follows:
			# "pair": [fairseq, transformers]
			lowerCamelCase__		:		Optional[Any]		     ={
			    '''wmt16-en-de-dist-12-1''': [28.3, 27.52],
			    '''wmt16-en-de-dist-6-1''': [27.4, 27.11],
			    '''wmt16-en-de-12-1''': [26.9, 25.75],
			}
			lowerCamelCase__		:		Any		     =f'''{src_lang}-{tgt_lang}'''

			lowerCamelCase__		:		Any		     =f'''
---
language:
- {src_lang}
- {tgt_lang}
thumbnail:
tags:
- translation
- wmt16
- allenai
license: apache-2.0
datasets:
- wmt16
metrics:
- bleu
---

# FSMT

## Model description

This is a ported version of fairseq-based [wmt16 transformer](https://github.com/jungokasai/deep-shallow/) for {src_lang}-{tgt_lang}.

For more details, please, see [Deep Encoder, Shallow Decoder: Reevaluating the Speed-Quality Tradeoff in Machine Translation](https://arxiv.org/abs/2006.10369).

All 3 models are available:

* [wmt16-en-de-dist-12-1](https://huggingface.co/allenai/wmt16-en-de-dist-12-1)
* [wmt16-en-de-dist-6-1](https://huggingface.co/allenai/wmt16-en-de-dist-6-1)
* [wmt16-en-de-12-1](https://huggingface.co/allenai/wmt16-en-de-12-1)


## Intended uses & limitations

#### How to use

```python
from transformers import FSMTForConditionalGeneration, FSMTTokenizer
mname = "allenai/{model_name}"
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


## Training data

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

## Eval results

Here are the BLEU scores:

model   | fairseq | transformers
-------|---------|----------
{model_name}  | {scores[model_name][0]} | {scores[model_name][1]}

The score is slightly below the score reported in the paper, as the researchers don\'t use `sacrebleu` and measure the score on tokenized outputs. `transformers` score was measured using `sacrebleu` on detokenized outputs.

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=5
mkdir -p $DATA_DIR
sacrebleu -t wmt16 -l $PAIR --echo src > $DATA_DIR/val.source
sacrebleu -t wmt16 -l $PAIR --echo ref > $DATA_DIR/val.target
echo $PAIR
PYTHONPATH="src:examples/seq2seq" python examples/seq2seq/run_eval.py allenai/{model_name} $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
```

## Data Sources

- [training, etc.](http://www.statmt.org/wmt16/)
- [test set](http://matrix.statmt.org/test_sets/newstest2016.tgz?1504722372)


### BibTeX entry and citation info

```
@misc{{kasai2020deep,
    title={{Deep Encoder, Shallow Decoder: Reevaluating the Speed-Quality Tradeoff in Machine Translation}},
    author={{Jungo Kasai and Nikolaos Pappas and Hao Peng and James Cross and Noah A. Smith}},
    year={{2020}},
    eprint={{2006.10369}},
    archivePrefix={{arXiv}},
    primaryClass={{cs.CL}}
}}
```

'''
			model_card_dir.mkdir(parents=__lowerCamelCase							,		exist_ok=__lowerCamelCase     )
			lowerCamelCase__		:		str		     =os.path.join(__lowerCamelCase							,		'''README.md'''     )
			print(f'''Generating {path}'''     )
			with open(__lowerCamelCase							,		'''w'''							,		encoding='''utf-8'''     ) as f:
						f.write(__lowerCamelCase     )


# make sure we are under the root of the project
_lowercase			:			List[str]										=						Path(__file__).resolve().parent.parent.parent
_lowercase			:			Dict										=						repo_dir / "model_cards"

for model_name in ["wmt16-en-de-dist-12-1", "wmt16-en-de-dist-6-1", "wmt16-en-de-12-1"]:
		_lowercase			:			int										=						model_cards_dir / "allenai" / model_name
		write_model_card(model_card_dir, src_lang="en", tgt_lang="de", model_name=model_name)




