PhaGen / README.md

Upload PhaGen

264f262 17 days ago

5.77 kB

---
license: apache-2.0
language:
- en
tags:
- biology
- dna
- genomics
- bacteriophage
- diffusion
- masked-diffusion
- custom-code
library_name: transformers
inference: false
---
# PhaGen

PhaGen is a hierarchical masked diffusion model for bacteriophage DNA. It generates DNA
on a masked canvas instead of left-to-right, so the same model produces sequences,
representations and masked conditional scores.

| | |
| --- | --- |
| Parameters | 152,421,601 |
| Stages | 3, dimensions 512 / 256 / 196, depths 8 / 8 / 8 |
| Token canvas | 131,072 (`max_seq_len` 128 x 64 x 16) |
| Vocabulary | 9 tokens: `[BOS]` A T C G `[EOS]` `[MASK]` `[PAD]` `[UNK]` |
| Precision | BF16 safetensors |
| Model type / classes | `phagen` / `PhaGenModel`, `PhaGenConfig` |

## Load

Requires PyTorch 2.5, Transformers 4.54.1, einops, accelerate and safetensors. The model
code is custom, so `trust_remote_code=True` is required; review the supplied Python files
first. 
The public classes are `PhaGenModel` and `PhaGenConfig`, implemented in
`modeling_phagen.py` and `configuration_phagen.py`. 

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "YusenWithR/PhaGen"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(
    repo, trust_remote_code=True,
    attn_implementation="sdpa", torch_dtype=torch.float32,
).eval()
inputs = tokenizer("ACGTACGTAC", return_tensors="pt", add_special_tokens=False)
with torch.inference_mode():
    out = model(**inputs, is_causal=False)       # the tokenizer supplies attention_mask
print(out.logits.shape)                          # [1, 10, 9]
print([tuple(h.shape) for h in out.hidden_states])
# [(1, 2, 512), (1, 1, 65, 256), (1, 1, 64, 17, 196)] for this 10-base input.
# Inputs are padded to a multiple of 1024; the first stage grows for longer inputs.
```

`AutoModelForCausalLM` is the registered auto class for loading convenience. PhaGen is a
masked diffusion model, not an autoregressive LM, and it does not implement `generate()`
in the autoregressive sense.

### Representations

Use `model.embed` for the paper's 964-dimensional sequence representation:

```python
batch = tokenizer(
    ["ACGTACGTAC", "GATTACA"], padding=True,
    return_tensors="pt", add_special_tokens=False,
).to(model.device)
embeddings = model.embed(
    input_ids=batch["input_ids"], attention_mask=batch["attention_mask"],
)
print(embeddings.shape)  # torch.Size([2, 964]), float32, input order preserved
```

### Generation

Diffusion sampling runs through `diffusion_generate` with the `MDMGenerationConfig`
supplied in `generation_utils.py`:

```python
from transformers.dynamic_module_utils import get_class_from_dynamic_module
MDMGenerationConfig = get_class_from_dynamic_module(
    "generation_utils.MDMGenerationConfig", repo)
torch.manual_seed(42)
config = MDMGenerationConfig(
    max_length=256, steps=32, alg="entropy", alg_temp=0.1,
    temperature=1.1, top_p=0.90, top_k=4,
    mask_token_id=tokenizer.mask_token_id, pad_token_id=tokenizer.pad_token_id,
    return_dict_in_generate=True,
)
prefix = torch.tensor([[tokenizer.bos_token_id]])
with torch.inference_mode():
    result = model.diffusion_generate(
        input_ids=prefix, attention_mask=torch.ones_like(prefix), generation_config=config)
print(tokenizer.decode(result.sequences[0], skip_special_tokens=True).replace(" ", "")[:60])
```

That runs in a few seconds on CPU. The paper production setting keeps the same sampling
values and raises the budget to `max_length=131072` with `steps=256` and seed 42, which
needs a GPU. `alg` accepts `origin`, `entropy`, `topk_margin`, `p2` and `maskgit_plus`.

The canvas is a token budget, not an output DNA length. Extract complete records
from the token IDs before removing special tokens:

```python
# One token canvas at a time; returns ACGT strings in their original order.
contigs = model.extract_contigs(result.sequences[0], tokenizer, min_length=1)
for index, sequence in enumerate(contigs, start=1):
    print(f">contig_{index}\n{sequence}")
# The paper cohort retains complete records with at least 2000 DNA bases.
paper_contigs = model.extract_contigs(result.sequences[0], tokenizer, min_length=2000)
```

`min_length` counts DNA bases, excluding BOS/EOS. Only complete `[BOS] ACGT... [EOS]`
spans are returned. A new BOS restarts an unfinished span; an internal PAD, MASK,
UNK or other non-ACGT token invalidates the span. Empty, short and trailing
unclosed spans are dropped. Short canvases may return an empty list. Decoding the
entire canvas with `skip_special_tokens=True` is only a preview: it can join
separate records and discard evidence of invalid content.

For several generated canvases, call the helper for each row and retain the row
index alongside its records. The same extraction function is available without
loading the weights via `sequence_utils.extract_contigs(token_ids, tokenizer, min_length)`.

## Training data

98,665 phage genomes (4,800,631,119 bases) assembled from MGV (52,836), GPD (29,768) and
NCBI (16,061).

## Intended use and limitations

Intended for research on phage genome modelling: unconditional generation, sequence
representations for downstream probes, and masked conditional scoring.

- The model is trained on phage DNA; behaviour on other genomes is untested.
- Short runs can extract zero complete contigs; that is a property of the extraction
  protocol, not a failure.

## Acknowledgement

PhaGen builds on the ideas of several outstanding open-source projects.

- [MEGABYTE](https://github.com/lucidrains/MEGABYTE-pytorch)
- [Open-dLLM](https://github.com/pengzhangzhi/Open-dLLM)
- [VeOmni](https://github.com/ByteDance-Seed/VeOmni)

License details and upstream notices are in `NOTICE.md`.

	---
	license: apache-2.0
	language:
	- en
	tags:
	- biology
	- dna
	- genomics
	- bacteriophage
	- diffusion
	- masked-diffusion
	- custom-code
	library_name: transformers
	inference: false
	---

	# PhaGen

	PhaGen is a hierarchical masked diffusion model for bacteriophage DNA. It generates DNA
	on a masked canvas instead of left-to-right, so the same model produces sequences,
	representations and masked conditional scores.

	\| \| \|
	\| --- \| --- \|
	\| Parameters \| 152,421,601 \|
	\| Stages \| 3, dimensions 512 / 256 / 196, depths 8 / 8 / 8 \|
	\| Token canvas \| 131,072 (`max_seq_len` 128 x 64 x 16) \|
	\| Vocabulary \| 9 tokens: `[BOS]` A T C G `[EOS]` `[MASK]` `[PAD]` `[UNK]` \|
	\| Precision \| BF16 safetensors \|
	\| Model type / classes \| `phagen` / `PhaGenModel`, `PhaGenConfig` \|

	## Load

	Requires PyTorch 2.5, Transformers 4.54.1, einops, accelerate and safetensors. The model
	code is custom, so `trust_remote_code=True` is required; review the supplied Python files
	first.
	The public classes are `PhaGenModel` and `PhaGenConfig`, implemented in
	`modeling_phagen.py` and `configuration_phagen.py`.

	```python
	import torch
	from transformers import AutoModelForCausalLM, AutoTokenizer

	repo = "YusenWithR/PhaGen"
	tokenizer = AutoTokenizer.from_pretrained(repo)
	model = AutoModelForCausalLM.from_pretrained(
	repo, trust_remote_code=True,
	attn_implementation="sdpa", torch_dtype=torch.float32,
	).eval()

	inputs = tokenizer("ACGTACGTAC", return_tensors="pt", add_special_tokens=False)
	with torch.inference_mode():
	out = model(**inputs, is_causal=False) # the tokenizer supplies attention_mask
	print(out.logits.shape) # [1, 10, 9]
	print([tuple(h.shape) for h in out.hidden_states])
	# [(1, 2, 512), (1, 1, 65, 256), (1, 1, 64, 17, 196)] for this 10-base input.
	# Inputs are padded to a multiple of 1024; the first stage grows for longer inputs.
	```

	`AutoModelForCausalLM` is the registered auto class for loading convenience. PhaGen is a
	masked diffusion model, not an autoregressive LM, and it does not implement `generate()`
	in the autoregressive sense.

	### Representations

	Use `model.embed` for the paper's 964-dimensional sequence representation:

	```python
	batch = tokenizer(
	["ACGTACGTAC", "GATTACA"], padding=True,
	return_tensors="pt", add_special_tokens=False,
	).to(model.device)
	embeddings = model.embed(
	input_ids=batch["input_ids"], attention_mask=batch["attention_mask"],
	)
	print(embeddings.shape) # torch.Size([2, 964]), float32, input order preserved
	```

	### Generation

	Diffusion sampling runs through `diffusion_generate` with the `MDMGenerationConfig`
	supplied in `generation_utils.py`:

	```python
	from transformers.dynamic_module_utils import get_class_from_dynamic_module

	MDMGenerationConfig = get_class_from_dynamic_module(
	"generation_utils.MDMGenerationConfig", repo)

	torch.manual_seed(42)
	config = MDMGenerationConfig(
	max_length=256, steps=32, alg="entropy", alg_temp=0.1,
	temperature=1.1, top_p=0.90, top_k=4,
	mask_token_id=tokenizer.mask_token_id, pad_token_id=tokenizer.pad_token_id,
	return_dict_in_generate=True,
	)
	prefix = torch.tensor([[tokenizer.bos_token_id]])
	with torch.inference_mode():
	result = model.diffusion_generate(
	input_ids=prefix, attention_mask=torch.ones_like(prefix), generation_config=config)
	print(tokenizer.decode(result.sequences[0], skip_special_tokens=True).replace(" ", "")[:60])
	```

	That runs in a few seconds on CPU. The paper production setting keeps the same sampling
	values and raises the budget to `max_length=131072` with `steps=256` and seed 42, which
	needs a GPU. `alg` accepts `origin`, `entropy`, `topk_margin`, `p2` and `maskgit_plus`.

	The canvas is a token budget, not an output DNA length. Extract complete records
	from the token IDs before removing special tokens:

	```python
	# One token canvas at a time; returns ACGT strings in their original order.
	contigs = model.extract_contigs(result.sequences[0], tokenizer, min_length=1)
	for index, sequence in enumerate(contigs, start=1):
	print(f">contig_{index}\n{sequence}")

	# The paper cohort retains complete records with at least 2000 DNA bases.
	paper_contigs = model.extract_contigs(result.sequences[0], tokenizer, min_length=2000)
	```

	`min_length` counts DNA bases, excluding BOS/EOS. Only complete `[BOS] ACGT... [EOS]`
	spans are returned. A new BOS restarts an unfinished span; an internal PAD, MASK,
	UNK or other non-ACGT token invalidates the span. Empty, short and trailing
	unclosed spans are dropped. Short canvases may return an empty list. Decoding the
	entire canvas with `skip_special_tokens=True` is only a preview: it can join
	separate records and discard evidence of invalid content.

	For several generated canvases, call the helper for each row and retain the row
	index alongside its records. The same extraction function is available without
	loading the weights via `sequence_utils.extract_contigs(token_ids, tokenizer, min_length)`.

	## Training data

	98,665 phage genomes (4,800,631,119 bases) assembled from MGV (52,836), GPD (29,768) and
	NCBI (16,061).

	## Intended use and limitations

	Intended for research on phage genome modelling: unconditional generation, sequence
	representations for downstream probes, and masked conditional scoring.

	- The model is trained on phage DNA; behaviour on other genomes is untested.
	- Short runs can extract zero complete contigs; that is a property of the extraction
	protocol, not a failure.

	## Acknowledgement

	PhaGen builds on the ideas of several outstanding open-source projects.

	- [MEGABYTE](https://github.com/lucidrains/MEGABYTE-pytorch)
	- [Open-dLLM](https://github.com/pengzhangzhi/Open-dLLM)
	- [VeOmni](https://github.com/ByteDance-Seed/VeOmni)

	License details and upstream notices are in `NOTICE.md`.