## Structure
```
.
├── Dockerfile
├── README.md
├── bin
│     └── train.sh  # training_script
├── c-tokenizer # tokenizer for ddlm
├── ddlm
│     ├── __init__.py
│     ├── data  # dataset preprocessing
│     │     ├── __init__.py
│     │     ├── c4.py
│     │     ├── collate_fn.py
│     │     └── pile.py
│     ├── modeling  # modeles
│     │     ├── __init__.py
│     │     ├── diffusion.py  # <-- model we use
│     │     ├── diffusion_simplified.py
│     │     ├── simplified_kerple.py
│     │     └── variational_diffusion.py
│     ├── sampler  # sampling (generating) from diffusion model
│     │     ├── __init__.py
│     │     ├── early_exit.py
│     │     └── euler.py
│     ├── time  # time warping
│     │     ├── __init__.py
│     │     ├── low_discrepency_sampling.py
│     │     └── time_wrapping.py  # time warping model
│     ├── utils.py
│     └── validation
│         ├── __init__.py
│         ├── classification.py  # GLUE utils
│         ├── data.py
│         ├── dist.py  # dist metrics
│         ├── entropy.py
│         ├── generate.py  # generation functions
│         ├── load_model.py  # wandb utils
│         ├── logits_to_outputs.py
│         ├── peplexity.py  # AR-NLL metric
│         └── zipf.py  # Zipf's score metric
├── ee.py
├── generate_ddlm.py
├── generate_ssd.py
├── load_dataset.py  # load c4 chunk
├── requirements.txt
├── setup.py
├── ssd  # from the official ssd repo
│     ├── __init__.py
│     ├── configs.py
│     ├── decode.py
│     ├── early_exit.py
│     └── validate_ssd.py
└── validate_ssd.py
```
### Requirements

To reproduce our result, we recommend using our Docker container.

Use 
```bash
pip install -r requirements.txt
pip install -e .
```

to install requirements.

### Pre-Trained Weights

First of all, you can download our model weights through this [link](https://drive.google.com/drive/folders/1U8XDhU7U2WGFgqIjL8IjV5BXmw6RDRPn).
Or you can use `gdown`:
```bash
gdown 18j_R39rAcrRQ_uOhfcPfSi6Des9QGUOJ
gdown 18s5qheqsvjp0d1jJZBgvnHPPzJgM1Pfo
```


We are downloading a model trained on the span task above.

After downloading, you can use the code sample below to load model weights:

```python
import torch

from ddlm.modeling.diffusion import DiffusionTransformer, DiffusionConfig
from ddlm.time.time_wrapping import TimeWrapping

config = DiffusionConfig()
model = DiffusionTransformer(config)
model.load_state_dict(torch.load("pythorch_model.bin"))
tw_model = TimeWrapping(config=config)
tw_model.load_state_dict(torch.load("pytorch_model_1.bin"))
```

## Reproduce Paper Results

### Pre-Training

We obtained our model with `bin/train.sh` script.

### Generating & Early Exit

We used `generate.py` to obtain results in Table 2-5—also, this script logs data from all figures except time measurements.

In our work, we highly relate to [wandb](www.wandb.ai), so weights are loading from runs, but you can load it directly from your disk via [previous section](###pre-trained_weights).

## Early Exit

### SSD

To Generate samples from SSD model and validate metrics we use `generate_ssd.py` 

```commandline
python generate_ssd.py
```

After running this script we can validate samples with

```commandline
python validate_ssd.py
```

### DDLM

Generating samples and validation performed in `generate_ddlm.py`

