



from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..utils import maybe_allow_in_graph
from .activations import get_activation
from .attention_processor import Attention
from .embeddings import CombinedTimestepLabelEmbeddings

@maybe_allow_in_graph
class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					def __init__(	self							:  str,  lowerCamelCase							:  Optional[Any],  lowerCamelCase							:  Union[str, Any],  lowerCamelCase							:  List[str],  lowerCamelCase							:  str=0.0,  lowerCamelCase							:  Optional[Any] = None,  lowerCamelCase							:  List[str] = "geglu",  lowerCamelCase							:  str = None,  lowerCamelCase							:  int = False,  lowerCamelCase							:  int = False,  lowerCamelCase							:  Optional[int] = False,  lowerCamelCase							:  List[str] = False,  lowerCamelCase							:  Optional[int] = True,  lowerCamelCase							:  Optional[Any] = "layer_norm",  lowerCamelCase							:  Optional[int] = False,  )-> Optional[int]:
								super().__init__()
								lowerCamelCase__		:		Optional[int]		     =only_cross_attention

								lowerCamelCase__		:		Tuple		     =(num_embeds_ada_norm is not None) and norm_type == 'ada_norm_zero'
								lowerCamelCase__		:		List[str]		     =(num_embeds_ada_norm is not None) and norm_type == 'ada_norm'

								if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
											raise ValueError(
											    F'''`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to'''
											    F''' define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}.'''					)

								# Define 3 blocks. Each block has its own normalization layer.
								# 1. Self-Attn
								if self.use_ada_layer_norm:
											lowerCamelCase__		:		List[Any]		     =AdaLayerNorm(snake_case__,  snake_case__					)
								elif self.use_ada_layer_norm_zero:
											lowerCamelCase__		:		List[Any]		     =AdaLayerNormZero(snake_case__,  snake_case__					)
								else:
											lowerCamelCase__		:		Optional[int]		     =nn.LayerNorm(snake_case__,  elementwise_affine=snake_case__					)
								lowerCamelCase__		:		Optional[Any]		     =Attention(
								    query_dim=snake_case__,  heads=snake_case__,  dim_head=snake_case__,  dropout=snake_case__,  bias=snake_case__,  cross_attention_dim=cross_attention_dim if only_cross_attention else None,  upcast_attention=snake_case__,  )

								# 2. Cross-Attn
								if cross_attention_dim is not None or double_self_attention:
											# We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
											# I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
											# the second cross attention block.
											lowerCamelCase__		:		Optional[Any]		     =(
											    AdaLayerNorm(snake_case__,  snake_case__					)
											    if self.use_ada_layer_norm
											    else nn.LayerNorm(snake_case__,  elementwise_affine=snake_case__					)
											)
											lowerCamelCase__		:		List[Any]		     =Attention(
											    query_dim=snake_case__,  cross_attention_dim=cross_attention_dim if not double_self_attention else None,  heads=snake_case__,  dim_head=snake_case__,  dropout=snake_case__,  bias=snake_case__,  upcast_attention=snake_case__,  )  # is self-attn if encoder_hidden_states is none
								else:
											lowerCamelCase__		:		Union[str, Any]		     =None
											lowerCamelCase__		:		int		     =None

								# 3. Feed-forward
								lowerCamelCase__		:		Optional[int]		     =nn.LayerNorm(snake_case__,  elementwise_affine=snake_case__					)
								lowerCamelCase__		:		Union[str, Any]		     =FeedForward(snake_case__,  dropout=snake_case__,  activation_fn=snake_case__,  final_dropout=snake_case__					)

								# let chunk size default to None
								lowerCamelCase__		:		Optional[int]		     =None
								lowerCamelCase__		:		Union[str, Any]		     =0



					def 		snake_case	(	self							:  List[Any],  lowerCamelCase							:  Tuple,  lowerCamelCase							:  List[Any]					)-> str:
								lowerCamelCase__		:		Optional[int]		     =chunk_size
								lowerCamelCase__		:		Tuple		     =dim




					def 		snake_case	(	self							:  Dict,  lowerCamelCase							:  Dict,  lowerCamelCase							:  List[str] = None,  lowerCamelCase							:  int = None,  lowerCamelCase							:  List[str] = None,  lowerCamelCase							:  Any = None,  lowerCamelCase							:  Optional[Any] = None,  lowerCamelCase							:  Tuple = None,  )-> Optional[Any]:

								if self.use_ada_layer_norm:
											lowerCamelCase__		:		Union[str, Any]		     =self.norma(snake_case__,  snake_case__					)
								elif self.use_ada_layer_norm_zero:
											lowerCamelCase__		:		str		     =self.norma(
											    snake_case__,  snake_case__,  snake_case__,  hidden_dtype=hidden_states.dtype					)
								else:
											lowerCamelCase__		:		Union[str, Any]		     =self.norma(snake_case__					)

								lowerCamelCase__		:		str		     =cross_attention_kwargs if cross_attention_kwargs is not None else {}

								lowerCamelCase__		:		Union[str, Any]		     =self.attna(
								    snake_case__,  encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,  attention_mask=snake_case__,  **snake_case__,  )
								if self.use_ada_layer_norm_zero:
											lowerCamelCase__		:		int		     =gate_msa.unsqueeze(1					) * attn_output
								lowerCamelCase__		:		str		     =attn_output + hidden_states

								# 2. Cross-Attention
								if self.attna is not None:
											lowerCamelCase__		:		List[str]		     =(
											    self.norma(snake_case__,  snake_case__					) if self.use_ada_layer_norm else self.norma(snake_case__					)
											)

											lowerCamelCase__		:		Optional[int]		     =self.attna(
											    snake_case__,  encoder_hidden_states=snake_case__,  attention_mask=snake_case__,  **snake_case__,  )
											lowerCamelCase__		:		Optional[int]		     =attn_output + hidden_states

								# 3. Feed-forward
								lowerCamelCase__		:		Optional[Any]		     =self.norma(snake_case__					)

								if self.use_ada_layer_norm_zero:
											lowerCamelCase__		:		Any		     =norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]

								if self._chunk_size is not None:
											# "feed_forward_chunk_size" can be used to save memory
											if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
														raise ValueError(
														    F'''`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`.'''					)

											lowerCamelCase__		:		Optional[int]		     =norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
											lowerCamelCase__		:		str		     =torch.cat(
											    [self.ff(snake_case__					) for hid_slice in norm_hidden_states.chunk(snake_case__,  dim=self._chunk_dim					)],  dim=self._chunk_dim,  )
								else:
											lowerCamelCase__		:		Tuple		     =self.ff(snake_case__					)

								if self.use_ada_layer_norm_zero:
											lowerCamelCase__		:		int		     =gate_mlp.unsqueeze(1					) * ff_output

								lowerCamelCase__		:		Dict		     =ff_output + hidden_states

								return hidden_states





class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					def __init__(	self							:  Union[str, Any],  lowerCamelCase							:  int,  lowerCamelCase							:  Union[str, Any] = None,  lowerCamelCase							:  Dict = 4,  lowerCamelCase							:  int = 0.0,  lowerCamelCase							:  Union[str, Any] = "geglu",  lowerCamelCase							:  List[Any] = False,  )-> Dict:
								super().__init__()
								lowerCamelCase__		:		Optional[Any]		     =int(dim * mult					)
								lowerCamelCase__		:		Optional[Any]		     =dim_out if dim_out is not None else dim

								if activation_fn == "gelu":
											lowerCamelCase__		:		Tuple		     =GELU(snake_case__,  snake_case__					)
								if activation_fn == "gelu-approximate":
											lowerCamelCase__		:		Any		     =GELU(snake_case__,  snake_case__,  approximate='''tanh'''					)
								elif activation_fn == "geglu":
											lowerCamelCase__		:		Tuple		     =GEGLU(snake_case__,  snake_case__					)
								elif activation_fn == "geglu-approximate":
											lowerCamelCase__		:		Any		     =ApproximateGELU(snake_case__,  snake_case__					)

								lowerCamelCase__		:		Union[str, Any]		     =nn.ModuleList([]					)
								# project in
								self.net.append(snake_case__					)
								# project dropout
								self.net.append(nn.Dropout(snake_case__					)					)
								# project out
								self.net.append(nn.Linear(snake_case__,  snake_case__					)					)
								# FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
								if final_dropout:
											self.net.append(nn.Dropout(snake_case__					)					)




					def 		snake_case	(	self							:  List[Any],  lowerCamelCase							:  int					)-> str:

								for module in self.net:
											lowerCamelCase__		:		List[Any]		     =module(snake_case__					)
								return hidden_states





class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					def __init__(	self							:  str,  lowerCamelCase							:  str,  lowerCamelCase							:  List[str],  lowerCamelCase							:  Any = "none"					)-> Optional[Any]:
								super().__init__()
								lowerCamelCase__		:		str		     =nn.Linear(snake_case__,  snake_case__					)
								lowerCamelCase__		:		List[Any]		     =approximate



					def 		snake_case	(	self							:  int,  lowerCamelCase							:  Any					)-> Any:

								if gate.device.type != "mps":
											return F.gelu(snake_case__,  approximate=self.approximate					)
								# mps: gelu is not implemented for float16
								return F.gelu(gate.to(dtype=torch.floataa					),  approximate=self.approximate					).to(dtype=gate.dtype					)




					def 		snake_case	(	self							:  Union[str, Any],  lowerCamelCase							:  List[Any]					)-> Dict:
								lowerCamelCase__		:		Union[str, Any]		     =self.proj(snake_case__					)
								lowerCamelCase__		:		str		     =self.gelu(snake_case__					)
								return hidden_states





class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					def __init__(	self							:  Dict,  lowerCamelCase							:  Union[str, Any],  lowerCamelCase							:  Dict					)-> str:
								super().__init__()
								lowerCamelCase__		:		Tuple		     =nn.Linear(snake_case__,  dim_out * 2					)



					def 		snake_case	(	self							:  Union[str, Any],  lowerCamelCase							:  List[Any]					)-> Optional[int]:

								if gate.device.type != "mps":
											return F.gelu(snake_case__					)
								# mps: gelu is not implemented for float16
								return F.gelu(gate.to(dtype=torch.floataa					)					).to(dtype=gate.dtype					)




					def 		snake_case	(	self							:  Optional[int],  lowerCamelCase							:  int					)-> str:
								lowerCamelCase__		:		Dict		     =self.proj(snake_case__					).chunk(2,  dim=-1					)
								return hidden_states * self.gelu(snake_case__					)





class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					def __init__(	self							:  Union[str, Any],  lowerCamelCase							:  Optional[int],  lowerCamelCase							:  List[str]					)-> Any:
								super().__init__()
								lowerCamelCase__		:		List[str]		     =nn.Linear(snake_case__,  snake_case__					)




					def 		snake_case	(	self							:  Dict,  lowerCamelCase							:  Union[str, Any]					)-> Union[str, Any]:
								lowerCamelCase__		:		Dict		     =self.proj(snake_case__					)
								return x * torch.sigmoid(1.702 * x					)





class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					def __init__(	self							:  Dict,  lowerCamelCase							:  str,  lowerCamelCase							:  int					)-> Any:
								super().__init__()
								lowerCamelCase__		:		Optional[Any]		     =nn.Embedding(snake_case__,  snake_case__					)
								lowerCamelCase__		:		List[Any]		     =nn.SiLU()
								lowerCamelCase__		:		Optional[int]		     =nn.Linear(snake_case__,  embedding_dim * 2					)
								lowerCamelCase__		:		Optional[int]		     =nn.LayerNorm(snake_case__,  elementwise_affine=snake_case__					)




					def 		snake_case	(	self							:  List[str],  lowerCamelCase							:  str,  lowerCamelCase							:  Any					)-> Any:
								lowerCamelCase__		:		Optional[Any]		     =self.linear(self.silu(self.emb(snake_case__					)					)					)
								lowerCamelCase__		:		List[Any]		     =torch.chunk(snake_case__,  2					)
								lowerCamelCase__		:		Optional[Any]		     =self.norm(snake_case__					) * (1 + scale) + shift
								return x





class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					def __init__(	self							:  int,  lowerCamelCase							:  Optional[Any],  lowerCamelCase							:  Union[str, Any]					)-> Optional[Any]:
								super().__init__()

								lowerCamelCase__		:		Union[str, Any]		     =CombinedTimestepLabelEmbeddings(snake_case__,  snake_case__					)

								lowerCamelCase__		:		Dict		     =nn.SiLU()
								lowerCamelCase__		:		Tuple		     =nn.Linear(snake_case__,  6 * embedding_dim,  bias=snake_case__					)
								lowerCamelCase__		:		Optional[Any]		     =nn.LayerNorm(snake_case__,  elementwise_affine=snake_case__,  eps=1E-6					)




					def 		snake_case	(	self							:  int,  lowerCamelCase							:  Optional[Any],  lowerCamelCase							:  Optional[int],  lowerCamelCase							:  Optional[int],  lowerCamelCase							:  Optional[Any]=None					)-> Optional[Any]:
								lowerCamelCase__		:		List[str]		     =self.linear(self.silu(self.emb(snake_case__,  snake_case__,  hidden_dtype=snake_case__					)					)					)
								lowerCamelCase__		:		Any		     =emb.chunk(6,  dim=1					)
								lowerCamelCase__		:		Any		     =self.norm(snake_case__					) * (1 + scale_msa[:, None]) + shift_msa[:, None]
								return x, gate_msa, shift_mlp, scale_mlp, gate_mlp





class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					def __init__(	self							:  str,  lowerCamelCase							:  Any,  lowerCamelCase							:  Union[str, Any],  lowerCamelCase							:  Optional[Any],  lowerCamelCase							:  int = None,  lowerCamelCase							:  Dict = 1E-5					)-> Tuple:
								super().__init__()
								lowerCamelCase__		:		Dict		     =num_groups
								lowerCamelCase__		:		Dict		     =eps

								if act_fn is None:
											lowerCamelCase__		:		Union[str, Any]		     =None
								else:
											lowerCamelCase__		:		List[Any]		     =get_activation(snake_case__					)

								lowerCamelCase__		:		int		     =nn.Linear(snake_case__,  out_dim * 2					)




					def 		snake_case	(	self							:  Tuple,  lowerCamelCase							:  Optional[int],  lowerCamelCase							:  Optional[Any]					)-> str:

								if self.act:
											lowerCamelCase__		:		List[Any]		     =self.act(snake_case__					)
								lowerCamelCase__		:		Optional[int]		     =self.linear(snake_case__					)
								lowerCamelCase__		:		str		     =emb[:, :, None, None]
								lowerCamelCase__		:		Any		     =emb.chunk(2,  dim=1					)

								lowerCamelCase__		:		int		     =F.group_norm(snake_case__,  self.num_groups,  eps=self.eps					)
								lowerCamelCase__		:		List[Any]		     =x * (1 + scale) + shift
								return x









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




