DOCSEARCHbyAlgolia
Submit Docs

DiffuSeq

repository·main·Indexed 21 days ago

https://github.com/shark-nlp/diffuseq

A framework for sequence-to-sequence text generation using diffusion models. It supports tasks such as dialogue, question generation, text simplification, and paraphrasing. The library provides standard and accelerated (v2) versions for training and sampling efficiency, integrates with PyTorch and HuggingFace transformers, and includes a customized DPM-Solver++ implementation to increase sampling speed.

Tokens
976
Snippets
5
Records
6
Agent score
24%

What's inside DiffuSeq

  2. Decode sequences with DiffuSeq

    main

    To perform decoding, you must first obtain the model_dir from the training stage. Navigate to the scripts directory and run the decoding script.

    For reproducing paper results, it is suggested to use a Majority Voting (MBR) candidate set size of 10 (running 10 times with different seeds). For diversity metrics, a candidate set size of 3 is used during computation.

    cd scripts
bash run_decode.sh

  3. Train a DiffuSeq model

    main

    To start training, navigate to the scripts directory and run the training script. You can use torch.distributed.launch for multi-GPU training.

    Training Arguments

    • --dataset: The name of the dataset (for notation).
    • --data_dir: Path to the dataset folder containing train.jsonl, test.jsonl, and valid.jsonl.
    • --seq_len: The maximum length of sequence $z$ ($x \oplus y$).
    • --resume_checkpoint: If provided, restores this checkpoint to continue training.
    • --vocab: Initialize the tokenizer using bert or provide a path to a preprocessed vocab dictionary (e.g., BPE).
    • --diff_steps: Number of diffusion steps.
    • --lr: Learning rate.
    • --learning_steps: Total number of training steps.
    • --save_interval: Interval for saving checkpoints.
    • --seed: Random seed.
    • --noise_schedule: The noise schedule to use (e.g., sqrt).
    • --hidden_dim: Dimension of the hidden state.
    • --bsz: Batch size.
    • --microbatch: Microbatch size (larger values can improve BLEU scores).
    • --schedule_sampler: The sampler to use (e.g., lossaware).
    • --notes: Notes for the run.

    DiffuSeq-v2 Arguments (Enhanced Version)

    • --learned_mean_embed: Whether to use the learned soft absorbing state.
    • --denoise: Whether to add discrete noise.
    • --use_fp16: Whether to use mixed precision training.
    • --denoise_rate: The denoise rate (default is 0.5).
    python -m torch.distributed.launch --nproc_per_node=4 --master_port=12233 --use_env run_train.py --diff_steps 2000 --lr 0.0001 --learning_steps 50000 --save_interval 10000 --seed 102 --noise_schedule sqrt --hidden_dim 128 --bsz 2048 --dataset qqp --data_dir {datasets/QQP} --vocab bert --seq_len 128 --schedule_sampler lossaware --notes qqp

  4. Prepare datasets for DiffuSeq

    main

    Datasets must be placed in the datasets folder. Each dataset should contain train.jsonl, test.jsonl, and valid.jsonl files.

    Supported tasks and datasets include:

    • Open-domain Dialogue: Commonsense Conversation (CCM)
    • Question Generation: Quasar-T
    • Text Simplification: Wiki-alignment
    • Paraphrase: QQP

  6. Evaluate model outputs and compute MBR

    main

    Use the evaluation script to compute metrics. You must provide a folder containing decoded text files.

    Important: If using output files from different models, ensure that the same line in every file refers to the same data point to ensure correct diversity scores.

    • To compute diversity scores only: Provide the folder path.
    • To perform Majority Voting (MBR) decoding: Add the --mbr flag.
    cd scripts
python eval_seq2seq.py --folder ../{your-path-to-outputs} --mbr