Metadata-Version: 2.4
Name: rtd-gdes
Version: 0.1.4.1
Summary: ELECTRA-style RTD pretraining with Gradient-Disentangled Embedding Sharing (GDES)
Project-URL: Homepage, https://github.com/rsyue/rtd-gdes
Project-URL: Issues, https://github.com/rsyue/rtd-gdes/issues
Author-email: Richard Samuel Yue <richard.samuel.yue@outlook.com>
License-Expression: MIT
License-File: LICENSE
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Requires-Python: >=3.10
Requires-Dist: accelerate
Requires-Dist: datasets
Requires-Dist: protobuf
Requires-Dist: scikit-learn
Requires-Dist: sentencepiece
Requires-Dist: tiktoken
Requires-Dist: torch>=2.6.0
Requires-Dist: tqdm
Requires-Dist: transformers
Provides-Extra: amd-gfx1151
Provides-Extra: dev
Requires-Dist: black; extra == 'dev'
Requires-Dist: mypy; extra == 'dev'
Requires-Dist: pytest; extra == 'dev'
Requires-Dist: pytest-cov; extra == 'dev'
Requires-Dist: ruff; extra == 'dev'
Requires-Dist: types-tqdm; extra == 'dev'
Provides-Extra: jetson
Description-Content-Type: text/markdown

# GDES-DeBERTaV3

ELECTRA-style training with **Gradient-Disentangled Embedding Sharing (GDES)** in native PyTorch using HuggingFace Transformers.

## Overview

GDES-DeBERTaV3 implements the replaced token detection (RTD) pretraining objective from [ELECTRA](https://arxiv.org/abs/2003.10555) with the gradient-disentangled embedding sharing technique introduced in [DeBERTaV3](https://arxiv.org/abs/2111.09543). Rather than relying on a separate generator and discriminator network, GDES shares embeddings between the two while disentangling their gradient flows — enabling more parameter-efficient training without the instability of naive weight tying.

The training loop performs two forward passes per step:

1. **Generator pass** — predict masked tokens via MLM
2. **Discriminator pass** — classify each token as original or replaced, with embedding gradients frozen to disentangle the two objectives

## Installation

```bash
git clone https://github.com/rsyue/rtd-gdes
cd rtd-gdes
```

```bash
pip install .
```

If you intend to contribute or modify the source, an editable install is recommended so changes are reflected immediately without reinstalling:

```bash
pip install -e .
```

### Requirements

- Python ≥ 3.10
- PyTorch ≥ 2.0 (CUDA recommended, ROCm builds supported)
- Transformers
- Datasets
- Safetensors
- scikit-learn
- tqdm

## Quick Start

```bash
python -m rtd_gdes.train \
  --model microsoft/deberta-v3-base \
  --lambda_disc 0.5 \
  --batch_size 8 \
  --epochs 5 \
  --learning_rate 2e-5 \
  --weight_decay 0.01 \
  --gamma 0.9 \
  --output_dir ./checkpoints \
  --bf16
```

## Usage

### CLI Arguments

| Argument | Flag | Type | Default | Description |
|---|---|---|---|---|
| `--model` | `-m` | `str` | `microsoft/deberta-v3-base` | Pretrained model to train with RTD + GDES |
| `--lambda_disc` | `-ld` | `float` | `0.5` | Lambda coefficient scaling the discriminator loss |
| `--batch_size` | `-bs` | `int` | `8` | Batch size for training and evaluation |
| `--epochs` | `-ep` | `int` | `5` | Number of training epochs |
| `--learning_rate` | `-lr` | `float` | `2e-5` | Learning rate for AdamW |
| `--weight_decay` | `-wd` | `float` | `0.01` | Weight decay for AdamW |
| `--gamma` | `-g` | `float` | `0.9` | Gamma for exponential LR scheduler |
| `--dataset` | | `str` | `imdb` | HuggingFace dataset name |
| `--output_dir` | `-o` | `str` | `None` | Directory to save the model and tokenizer. The model is saved under `<output_dir>/<save_name>/`. Defaults to `./<save_name>/` |
| `--fp16` | | `flag` | `False` | Enable FP16 mixed precision |
| `--bf16` | | `flag` | `False` | Enable BF16 mixed precision |
| `--compile` | `-c` | `flag` | `False` | Run `torch.compile` with `max-autotune` mode |

### Training Details

The script trains on the [IMDB unsupervised split](https://huggingface.co/datasets/imdb) by default, with a configurable 90/10 train/eval split. The dataset can be changed via `--dataset`. The combined loss is computed as:

$$\mathcal{L} = \mathcal{L}_{\text{gen}} + \lambda \cdot \mathcal{L}_{\text{disc}}$$

where $\mathcal{L}_{\text{gen}}$ is the standard MLM cross-entropy loss and $\mathcal{L}_{\text{disc}}$ is binary cross-entropy over token-level replaced/original predictions.

Evaluation reports discriminator loss, accuracy, and F1 score on the held-out set.

### Saved Outputs

After training, the model and tokenizer are saved under `<output_dir>/<save_name>/` where `save_name` is derived from the model id (e.g. `deberta_v3_base_gdes`). If `--output_dir` is not specified, the model is saved to `./<save_name>/` in the current working directory.

## Development

Install with dev dependencies:

```bash
pip install -e ".[dev]"
```

Run the test suite:

```bash
pytest tests/ -v --cov=rtd_gdes
```

## Project Structure

```
rtd-gdes/
├── src/
│   └── rtd_gdes/
│       ├── config.py          # TrainConfig dataclass — all hyperparameter defaults
│       ├── train.py           # Entry point and CLI
│       └── gdes/
│           ├── data.py        # Dataset loading and DataLoader construction
│           ├── model.py       # DebertaV3GDES — generator + discriminator
│           ├── trainer.py     # train_one_epoch and evaluate loops
│           └── utils.py       # Shared exceptions
├── tests/
│   └── test_gdes.py           # Model, trainer, and config unit + integration tests
└── pyproject.toml
```

## Roadmap

- [ ] Distributed training (DDP / FSDP)
- [ ] Publish as PyPI package
- [ ] Support additional model architectures beyond DeBERTaV3

## Citation

If you use this code, please cite the original papers:

```bibtex
@article{he2021debertav3,
  title={DeBERTaV3: Improving DeBERTa using ELECTRA-Style Pre-Training with Gradient-Disentangled Embedding Sharing},
  author={He, Pengcheng and Liu, Jianfeng and Gao, Jianfeng and Chen, Weizhu},
  journal={arXiv preprint arXiv:2111.09543},
  year={2021}
}

@article{clark2020electra,
  title={ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators},
  author={Clark, Kevin and Luong, Minh-Thang and Le, Quoc V. and Manning, Christopher D.},
  journal={arXiv preprint arXiv:2003.10555},
  year={2020}
}
```

## License

MIT
