| import torch | |
| CONFIG = { | |
| "model_type": "d3pm_cross_attention", | |
| "data": { | |
| "include_negative_examples": True, | |
| "dataset_size": 60000, | |
| }, | |
| "diffusion": { | |
| "mask_token_id": 0, | |
| }, | |
| "model": { | |
| "src_vocab_size": 16000, | |
| "tgt_vocab_size": 16000, | |
| "d_model": 384, | |
| "n_heads": 8, | |
| "d_ff": 1536, | |
| "n_layers": 6, | |
| "dropout": 0.1, | |
| "max_seq_len": 80, | |
| "diffusion_steps": 64, | |
| }, | |
| "training": { | |
| "device": "cuda" if torch.cuda.is_available() else "cpu", | |
| }, | |
| "inference": { | |
| "num_steps": 64, | |
| "temperature": 0.7, | |
| "top_k": 40, | |
| "repetition_penalty": 1.2, | |
| "diversity_penalty": 0.0, | |
| }, | |
| } | |