Hengchang-Liu/D3LM-scratch

pip install transformers torch tqdm

import sys
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

model_name = "Hengchang-Liu/D3LM-scratch"
model = AutoModelForMaskedLM.from_pretrained(model_name, trust_remote_code=True).eval()
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

# Import MDMGenerationConfig from the model's auto-downloaded module
MDMGenerationConfig = getattr(sys.modules[type(model).__module__], "MDMGenerationConfig")

# Unconditional generation: create a fully-masked prompt of desired length
length = 200
input_ids = torch.full((1, length), tokenizer.mask_token_id, dtype=torch.long)

config = MDMGenerationConfig(
    mask_token_id=tokenizer.mask_token_id,
    max_length=length,
    steps=50,
    temperature=1.0,
    top_p=0.9,
    alg="random",
    num_return_sequences=4,
    return_dict_in_generate=True,
)

with torch.no_grad():
    outputs = model.diffusion_generate(inputs=input_ids, generation_config=config)

for i, seq in enumerate(outputs.sequences):
    print(f">{i}
{tokenizer.decode(seq, skip_special_tokens=True).replace(' ', '')}")

@misc{yang2026d3lmdiscretednadiffusion,
      title={D3LM: A Discrete DNA Diffusion Language Model for Bidirectional DNA Understanding and Generation}, 
      author={Zhao Yang and Hengchang Liu and Chuan Cao and Bing Su},
      year={2026},
      eprint={2603.01780},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2603.01780}, 
}