<div  align="center">
    <h1>WeFT: Weighted Entropy-Driven Fine-tuning
for dLLMs</h1>
</div>

We propose WeFT, an entropy-driven weighted fine-tuning method for diffusion language models that enhances reasoning ability compared to SFT while maintaining efficiency.

## Environment Setup

To setup the environment, run;

```
pip install -r requirements.txt
```

## Training

```bash
# First go to the SFT directory
cd SFT
(SFT) CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch --config_file ddp_config.yaml --main_process_port 29502 --num_processes 4 sft_train.py
(WeFT) CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch --config_file ddp_config.yaml --main_process_port 29502 --num_processes 4 sft_train.py --weighted
```

## Evaluation

The evaluation code is inside the `eval` directory.

- Run with `bash eval/****.sh`
- Make sure to point to the correct checkpoint.
- The evaluation file will only save the generations; use the parser to calculate accuracy
- For example, baseline generations are in the `eval_baselines` directory. Use `python parse_and_get_acc.py` to print the accuracy.

# Acknowledgement

The implementation is adapted from [wd1](https://github.com/xiaohangt/wd1). We appreciate the clear repository!
