diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..db520b63741d228eef06b911efe3ce3275d14504 --- /dev/null +++ b/README.md @@ -0,0 +1,213 @@ +# pCoMole + +**pCoMole** (Pareto-Constrained Molecule editing) steers a pretrained Edit Flow toward user-specified objectives while enforcing hard terminal constraints. This repository is a minimal, clone-and-run package for: + +1. Training, evaluating, and sampling **Edit Flows** +2. Running **pCoMole** multi-objective editing on **GFP** and **peptidomimetics** + +Paper: *pCoMol: Pareto-Constrained Molecule Editing with Discrete Flows* + +The original working tree is left untouched. This folder is the public-facing copy. + +## Repository layout + +```text +pCoMole/ + train.py / evaluate.py # Edit Flow train / val-loss eval + configs/ # GFP (protein) and SELFIES configs + model/ logic/ flow_matching/ # Edit Flow implementation + smiles_tokenizer/ # SMILES SPE + SELFIES vocab + data/selfies/28k_mimetics/ # shipped SELFIES training split + gfp/ # GFP generation + pCoMole + pcomole.py # official GFP editor + generate.py + ckpt/last_2.ckpt + classifier_ckpt/best.pt + FPredX/ # excitation / brightness / emission models + peptidomimetics/ # peptidomimetic generation + pCoMole + pcomole.py + generate.py + objectives.py # PeptiVerse + Admetica/DeepDTAGen mix + ckpt/SELFIES_EditFlows.ckpt + ckpt/SMILES_BindEvaluator.ckpt + ckpt/admetica/ # LD50, solubility, Caco-2, half-life + ckpt/deepdtagen/ +``` + +## 1. Environment + +A CUDA GPU is strongly recommended. + +```bash +git clone pCoMole +cd pCoMole +python -m venv .venv +source .venv/bin/activate +pip install -r requirements.txt +``` + +HuggingFace model downloads happen on first use (`facebook/esm2_t33_650M_UR50D`, `aaronfeller/PeptideCLM-23M-all`, ChemBERTa). + +### PeptiVerse (required for peptidomimetic pCoMole) + +Peptidomimetic property oracles use the latest [PeptiVerse](https://huggingface.co/ChatterjeeLab/PeptiVerse) SMILES predictors. Clone it **next to** this repo, or set `PEPTIVERSE_ROOT`: + +```bash +# recommended: sibling of pCoMole +cd .. +git clone https://huggingface.co/ChatterjeeLab/PeptiVerse +cd PeptiVerse +pip install -r requirements.txt +# follow PeptiVerse README to download model weights +``` + +pCoMole looks for PeptiVerse in this order: + +1. `$PEPTIVERSE_ROOT` +2. `../PeptiVerse` relative to this repository + +PeptiVerse scores the peptide-like side of each candidate. Admetica + DeepDTAGen still score the small-molecule side, and the original peptide-likeness mix logic is unchanged. + +### MAFFT (required for GFP pCoMole) + +GFP excitation, brightness, and emission oracles run through FPredX and need [MAFFT](https://mafft.cbrc.jp/alignment/software/) on your `PATH`. + +```bash +# conda +conda install -c bioconda mafft + +# or point to a local binary +export MAFFT_PATH=/path/to/mafft +``` + +Confirm with `mafft --version` before running GFP editing. + +## 2. Train, evaluate, and sample Edit Flows + +Run every command from the **repository root**. + +### SELFIES peptidomimetic Edit Flow + +Training data is shipped at `data/selfies/28k_mimetics`. + +```bash +# train +python train.py --config configs/config_selfies.yaml +# optional: python train.py --config configs/config_selfies.yaml --wandb + +# evaluate a checkpoint on the validation split +python evaluate.py \ + --config configs/config_selfies.yaml \ + --ckpt peptidomimetics/ckpt/SELFIES_EditFlows.ckpt + +# unconditional / seeded generation +python peptidomimetics/generate.py \ + --config configs/config_selfies.yaml \ + --ckpt peptidomimetics/ckpt/SELFIES_EditFlows.ckpt \ + --input 'CSCC[C@H](NC(=O)[C@@H]1CCCN1C(=O)[C@H](Cc1ccc(O)cc1)NC(=O)[C@H](CCCNC(=N)N)NC(=O)[C@H](CO)NC(=O)[C@H](Cc1ccc(O)cc1)NC(=O)[C@@H]1CCCN1C(=O)[C@@H]1CCCN1C(=O)[C@@H]1CCCN1C(=O)[C@@H]1CCCN1C(=O)[C@@H](N)CO)C(=O)N[C@@H](CC(=O)O)C(=O)O' \ + --num_steps 30 \ + --num_samples 2 \ + --output_csv outputs/peptidomimetic_unconditional.csv +``` + +Or: `bash peptidomimetics/scripts/train.sh` and `bash peptidomimetics/scripts/generate.sh`. + +### GFP protein Edit Flow + +A trained GFP checkpoint is shipped at `gfp/ckpt/last_2.ckpt`. The original GFP training arrows are **not** included. To retrain, put a HuggingFace `load_from_disk` dataset under `data/gfp/{train,validation}` (or edit `configs/config_gfp.yaml`). + +```bash +python evaluate.py --config configs/config_gfp.yaml --ckpt gfp/ckpt/last_2.ckpt + +python gfp/generate.py \ + --config configs/config_gfp.yaml \ + --ckpt gfp/ckpt/last_2.ckpt \ + --input 'MSSGALLFHGKIPYVVEMEGNVDGHTFSIRGKGYGDASVGKVDAQFICTTGDVPVPWSTLVTTLTYGAQCFAKYGPELKDFYKSCMPDGYVQERTITFEGDGNFKTRAEVTFENGSVYNRVKLNGQGFKKDGHVLGKNLEFNFTPHCLYIWGDQANHGLKSAFKICHEITGSKGDFIVADHTQMNTPIGGGPVHVPEYHHMSYHVKLSKDVTDHRDNMSLKETVRAVDCRKTYDFDAGSGDTS' \ + --num_steps 10 \ + --num_samples 8 \ + --output_csv outputs/gfp_unconditional.csv +``` + +## 3. pCoMole multi-objective editing + +### GFP + +Official editor: `gfp/pcomole.py` (length + excitation + brightness, with GFP-classifier and emission constraints). + +```bash +bash gfp/scripts/pcomole.sh +``` + +Equivalent command: + +```bash +python gfp/pcomole.py \ + --root_dir gfp/FPredX \ + --config configs/config_gfp.yaml \ + --ckpt gfp/ckpt/last_2.ckpt \ + --num_steps 10 \ + --num_candidates 50 \ + --num_rollouts 10 \ + --objective_weights 3 1 1 \ + --output_file outputs/gfp_length_excitation_brightness.csv \ + --input 'MSSGALLFHGKIPYVVEMEGNVDGHTFSIRGKGYGDASVGKVDAQFICTTGDVPVPWSTLVTTLTYGAQCFAKYGPELKDFYKSCMPDGYVQERTITFEGDGNFKTRAEVTFENGSVYNRVKLNGQGFKKDGHVLGKNLEFNFTPHCLYIWGDQANHGLKSAFKICHEITGSKGDFIVADHTQMNTPIGGGPVHVPEYHHMSYHVKLSKDVTDHRDNMSLKETVRAVDCRKTYDFDAGSGDTS' +``` + +If the output filename contains `length_excitation_brightness`, `length_brightness`, `length_excitation`, or `length`, the matching objective subset is used. Any other name uses all three objectives. + +### Peptidomimetics + +Official editor: `peptidomimetics/pcomole.py`. + +Objectives (7 scores when `--specificity` is on): + +1. non-toxicity +2. solubility +3. permeability +4. half-life +5. affinity +6. motif +7. specificity + +PeptiVerse provides the peptide-side scores. Admetica + DeepDTAGen provide the small-molecule side. The two are mixed by peptide-likeness, same as the paper code. Motif / specificity still use the shipped BindEvaluator checkpoint. + +```bash +# after PeptiVerse is cloned and its weights are downloaded +bash peptidomimetics/scripts/pcomole.sh +``` + +`--target` and `--motifs` are required for affinity and motif oracles. `--specificity` adds the seventh objective. + +## 4. Included checkpoints + +| Path | Role | +|---|---| +| `gfp/ckpt/last_2.ckpt` | GFP Edit Flow | +| `gfp/classifier_ckpt/best.pt` | GFP hard constraint | +| `gfp/FPredX/{ex,bright,em}_model` | FPredX property models | +| `peptidomimetics/ckpt/SELFIES_EditFlows.ckpt` | peptidomimetic Edit Flow | +| `peptidomimetics/ckpt/SMILES_BindEvaluator.ckpt` | motif / specificity | +| `peptidomimetics/ckpt/admetica/*.ckpt` | Admetica blending oracles | +| `peptidomimetics/ckpt/deepdtagen/` | DeepDTAGen affinity | + +Together these are about 10 GB. HuggingFace uploads should use Git LFS for the `.ckpt` / `.pt` / `.pth` files. + +## 5. Environment variables + +| Variable | Purpose | +|---|---| +| `PEPTIVERSE_ROOT` | PeptiVerse repo root (manifest + `training_classifiers/`) | +| `MAFFT_PATH` | MAFFT binary, if it is not on `PATH` | +| `GFP_CLASSIFIER_CKPT` | optional override for the GFP classifier | + +## 6. Notes + +- Run scripts from the repo root, or use the wrappers under `*/scripts/`. +- GFP FPredX will fail immediately if `mafft` is missing; that is expected. +- Peptidomimetic pCoMole will fail at oracle init if PeptiVerse is missing or its weights were not downloaded. +- Cas9 is not part of this release. +- Training writes to `outputs/` by default. + +## Citation + +If you use this code, please cite the pCoMole paper and [PeptiVerse](https://www.biorxiv.org/content/10.64898/2025.12.31.697180). diff --git a/configs/config_gfp.yaml b/configs/config_gfp.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c5412e37dd66e95769dca027c6d0fc5496514d76 --- /dev/null +++ b/configs/config_gfp.yaml @@ -0,0 +1,75 @@ +task: protein + +compute: + ngpus: 1 + nodes: 1 + +logging: + enable_wandb: False + project: pCoMole + +data: + # GFP training arrows are not shipped. Point these at your own + # HuggingFace load_from_disk() dataset if you want to retrain. + train_path: data/gfp/train + val_path: data/gfp/validation + num_workers: 1 + +training: + batch_size: None + snapshot: 100 + eval_freq: 1000 + perplexity_freq: 100 + seed: 42 + grad_accumulation_steps: 2 + loc_prop_path: False + lam_prop: 1 + reparameterize: True + gamma_rate: 1 + gamma_edit: 1 + use_aux_ce: False + aux_ce_weight: 1.0 + no_training: False + +eval: + batch_size: None + sample_batch_size: 4 + perplexity: False + perplexity_batch_size: 4 + +optim: + weight_decay: 0.05 + optimizer: AdamW + lr: 1.0e-6 + beta1: 0.9 + beta2: 0.95 + eps: 1.0e-8 + grad_clip: 1.0 + eta_min_ratio: 0.1 + fused: false + n_epochs: 1000 + warmup_ratio: 0.1 + +flow: + source_distribution: uniform + loss_function: editflows + exponent: 2.0 + scheduler_type: polynomial + sampling_steps: 256 + eps_id: -1 + +model: + d_model: 768 + max_len: 2048 + n_layers: 12 + n_heads: 12 + dropout: 0.1 + attn_dropout: 0.0 + proj_dropout: 0.0 + compile: False + esm_model_name: facebook/esm2_t33_650M_UR50D + freeze_esm: true + p_optimal: 0 + scale_size: 0.5 + +work_dir: outputs/gfp diff --git a/configs/config_selfies.yaml b/configs/config_selfies.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cbbb1fb555e033a07830caf76449054363c2b008 --- /dev/null +++ b/configs/config_selfies.yaml @@ -0,0 +1,57 @@ +task: selfies + +compute: + ngpus: 1 + nodes: 1 + +data: + train_path: data/selfies/28k_mimetics/train + val_path: data/selfies/28k_mimetics/validation + num_workers: 4 + +training: + batch_size: 4 + snapshot: 100 + eval_freq: 1000 + perplexity_freq: 100 + seed: 42 + +eval: + batch_size: 4 + sample_batch_size: 4 + perplexity: False + perplexity_batch_size: 4 + +optim: + weight_decay: 0.03 + optimizer: AdamW + lr: 3.0e-4 + beta1: 0.9 + beta2: 0.95 + eps: 1.0e-8 + grad_clip: 1.0 + eta_min_ratio: 0.1 + fused: false + n_epochs: 100 + warmup_ratio: 0.1 + +flow: + source_distribution: uniform + loss_function: editflows_reparam + exponent: 2.0 + scheduler_type: polynomial + sampling_steps: 256 + eps_id: -1 + +model: + hidden_size: 768 + cond_dim: 128 + length: 2048 + n_blocks: 8 + n_heads: 12 + dropout: 0.1 + compile: False + p_optimal: 0 + scale_size: 1.0 + +work_dir: outputs/selfies diff --git a/data/.DS_Store b/data/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..6026ac55ee86b203c8e164fc65518d5f2febde53 Binary files /dev/null and b/data/.DS_Store differ diff --git a/data/selfies/.DS_Store b/data/selfies/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..59eab37fe16c35aba279ed57701f69e3cf519b68 Binary files /dev/null and b/data/selfies/.DS_Store differ diff --git a/data/selfies/28k_mimetics/dataset_dict.json b/data/selfies/28k_mimetics/dataset_dict.json new file mode 100644 index 0000000000000000000000000000000000000000..9195703312e22d2b9b9fe14951aa733949480a2d --- /dev/null +++ b/data/selfies/28k_mimetics/dataset_dict.json @@ -0,0 +1 @@ +{"splits": ["train", "validation", "test"]} \ No newline at end of file diff --git a/data/selfies/28k_mimetics/test/data-00000-of-00001.arrow b/data/selfies/28k_mimetics/test/data-00000-of-00001.arrow new file mode 100644 index 0000000000000000000000000000000000000000..0f6f48f3f764052ee57485a73bdf3beed28cecb2 --- /dev/null +++ b/data/selfies/28k_mimetics/test/data-00000-of-00001.arrow @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d709f3386113ebc6abe9400950274e2e41aaf7063b31e67d6ba15f3239462251 +size 786928 diff --git a/data/selfies/28k_mimetics/test/dataset_info.json b/data/selfies/28k_mimetics/test/dataset_info.json new file mode 100644 index 0000000000000000000000000000000000000000..b1e3cb417099ff98a2390bf74ca1f7db85b08a44 --- /dev/null +++ b/data/selfies/28k_mimetics/test/dataset_info.json @@ -0,0 +1,36 @@ +{ + "citation": "", + "description": "", + "features": { + "input_ids": { + "feature": { + "feature": { + "dtype": "int32", + "_type": "Value" + }, + "_type": "List" + }, + "_type": "List" + }, + "attention_mask": { + "feature": { + "feature": { + "dtype": "int8", + "_type": "Value" + }, + "_type": "List" + }, + "_type": "List" + }, + "seq_length": { + "dtype": "int64", + "_type": "Value" + }, + "batch_size": { + "dtype": "int64", + "_type": "Value" + } + }, + "homepage": "", + "license": "" +} \ No newline at end of file diff --git a/data/selfies/28k_mimetics/test/state.json b/data/selfies/28k_mimetics/test/state.json new file mode 100644 index 0000000000000000000000000000000000000000..f03e8fb946c06eed75319f6af1ebbbf19ef39e95 --- /dev/null +++ b/data/selfies/28k_mimetics/test/state.json @@ -0,0 +1,13 @@ +{ + "_data_files": [ + { + "filename": "data-00000-of-00001.arrow" + } + ], + "_fingerprint": "98db6912c1e59254", + "_format_columns": null, + "_format_kwargs": {}, + "_format_type": null, + "_output_all_columns": false, + "_split": null +} \ No newline at end of file diff --git a/data/selfies/28k_mimetics/tokenizer/tokenizer_config.json b/data/selfies/28k_mimetics/tokenizer/tokenizer_config.json new file mode 100644 index 0000000000000000000000000000000000000000..467ad8ef1476bd0862480ef0671e6e1a5cb7f3da --- /dev/null +++ b/data/selfies/28k_mimetics/tokenizer/tokenizer_config.json @@ -0,0 +1,10 @@ +{ + "type": "selfies-wordlevel", + "version": "2.1.1", + "special_tokens": { + "pad": "", + "bos": "", + "eos": "" + }, + "semantic_constraints": null +} \ No newline at end of file diff --git a/data/selfies/28k_mimetics/tokenizer/vocab.json b/data/selfies/28k_mimetics/tokenizer/vocab.json new file mode 100644 index 0000000000000000000000000000000000000000..a52c28786347a35bb445e64b2c2afca65a9b529e --- /dev/null +++ b/data/selfies/28k_mimetics/tokenizer/vocab.json @@ -0,0 +1,46 @@ +{ + "": 1, + "": 2, + "": 0, + "[#Branch1]": 3, + "[#Branch2]": 4, + "[#C]": 5, + "[#N]": 6, + "[/-Ring1]": 7, + "[/C]": 8, + "[/N]": 9, + "[=Branch1]": 10, + "[=Branch2]": 11, + "[=C]": 12, + "[=N]": 13, + "[=O]": 14, + "[=Ring1]": 15, + "[=Ring2]": 16, + "[=S]": 17, + "[Br]": 18, + "[Branch1]": 19, + "[Branch2]": 20, + "[C@@H1]": 21, + "[C@@]": 22, + "[C@H1]": 23, + "[C@]": 24, + "[CH1-1]": 25, + "[C]": 26, + "[Cl]": 27, + "[F]": 28, + "[I]": 29, + "[N+1]": 30, + "[N-1]": 31, + "[NH1]": 32, + "[N]": 33, + "[O-1]": 34, + "[O]": 35, + "[P]": 36, + "[Ring1]": 37, + "[Ring2]": 38, + "[S]": 39, + "[\\-Ring1]": 40, + "[\\-Ring2]": 41, + "[\\C]": 42, + "[\\N]": 43 +} \ No newline at end of file diff --git a/data/selfies/28k_mimetics/train/data-00000-of-00001.arrow b/data/selfies/28k_mimetics/train/data-00000-of-00001.arrow new file mode 100644 index 0000000000000000000000000000000000000000..cf392abd22ce1fd541d76237ea615168dc7260d0 --- /dev/null +++ b/data/selfies/28k_mimetics/train/data-00000-of-00001.arrow @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4d7e3d408b36d7c8a75b0d36f20d09b744c1c845961401c72a4023a9a6ffa3f6 +size 6084456 diff --git a/data/selfies/28k_mimetics/train/dataset_info.json b/data/selfies/28k_mimetics/train/dataset_info.json new file mode 100644 index 0000000000000000000000000000000000000000..b1e3cb417099ff98a2390bf74ca1f7db85b08a44 --- /dev/null +++ b/data/selfies/28k_mimetics/train/dataset_info.json @@ -0,0 +1,36 @@ +{ + "citation": "", + "description": "", + "features": { + "input_ids": { + "feature": { + "feature": { + "dtype": "int32", + "_type": "Value" + }, + "_type": "List" + }, + "_type": "List" + }, + "attention_mask": { + "feature": { + "feature": { + "dtype": "int8", + "_type": "Value" + }, + "_type": "List" + }, + "_type": "List" + }, + "seq_length": { + "dtype": "int64", + "_type": "Value" + }, + "batch_size": { + "dtype": "int64", + "_type": "Value" + } + }, + "homepage": "", + "license": "" +} \ No newline at end of file diff --git a/data/selfies/28k_mimetics/train/state.json b/data/selfies/28k_mimetics/train/state.json new file mode 100644 index 0000000000000000000000000000000000000000..1f99fc92d2c750ceb11e9f0adf48f1d0ccfa48c0 --- /dev/null +++ b/data/selfies/28k_mimetics/train/state.json @@ -0,0 +1,13 @@ +{ + "_data_files": [ + { + "filename": "data-00000-of-00001.arrow" + } + ], + "_fingerprint": "58b5e20848e4941d", + "_format_columns": null, + "_format_kwargs": {}, + "_format_type": null, + "_output_all_columns": false, + "_split": null +} \ No newline at end of file diff --git a/data/selfies/28k_mimetics/validation/data-00000-of-00001.arrow b/data/selfies/28k_mimetics/validation/data-00000-of-00001.arrow new file mode 100644 index 0000000000000000000000000000000000000000..44215994e31ab1591140d7715f4316234897d8c3 --- /dev/null +++ b/data/selfies/28k_mimetics/validation/data-00000-of-00001.arrow @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6667b7a12a440b50d996c8523813df0965a718e1a467ef2162fdea65d9b972cb +size 778640 diff --git a/data/selfies/28k_mimetics/validation/dataset_info.json b/data/selfies/28k_mimetics/validation/dataset_info.json new file mode 100644 index 0000000000000000000000000000000000000000..b1e3cb417099ff98a2390bf74ca1f7db85b08a44 --- /dev/null +++ b/data/selfies/28k_mimetics/validation/dataset_info.json @@ -0,0 +1,36 @@ +{ + "citation": "", + "description": "", + "features": { + "input_ids": { + "feature": { + "feature": { + "dtype": "int32", + "_type": "Value" + }, + "_type": "List" + }, + "_type": "List" + }, + "attention_mask": { + "feature": { + "feature": { + "dtype": "int8", + "_type": "Value" + }, + "_type": "List" + }, + "_type": "List" + }, + "seq_length": { + "dtype": "int64", + "_type": "Value" + }, + "batch_size": { + "dtype": "int64", + "_type": "Value" + } + }, + "homepage": "", + "license": "" +} \ No newline at end of file diff --git a/data/selfies/28k_mimetics/validation/state.json b/data/selfies/28k_mimetics/validation/state.json new file mode 100644 index 0000000000000000000000000000000000000000..822c4abbd8fafc919b82b55392f0ac2350f700ba --- /dev/null +++ b/data/selfies/28k_mimetics/validation/state.json @@ -0,0 +1,13 @@ +{ + "_data_files": [ + { + "filename": "data-00000-of-00001.arrow" + } + ], + "_fingerprint": "8dc499c569440f6b", + "_format_columns": null, + "_format_kwargs": {}, + "_format_type": null, + "_output_all_columns": false, + "_split": null +} \ No newline at end of file diff --git a/evaluate.py b/evaluate.py new file mode 100644 index 0000000000000000000000000000000000000000..1430013d1f630dc0eadb2fd1760fd7218c4f39f6 --- /dev/null +++ b/evaluate.py @@ -0,0 +1,39 @@ +import argparse + +import torch +import pytorch_lightning as pl + +from paths import add_repo_to_sys_path, resolve_path +from train import build_dataloaders, build_editflow, load_config + +add_repo_to_sys_path() + + +def main(): + parser = argparse.ArgumentParser(description="Evaluate an Edit Flow checkpoint on the validation split") + parser.add_argument("--config", type=str, required=True) + parser.add_argument("--ckpt", type=str, required=True) + args = parser.parse_args() + + cfg = load_config(args.config) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + editflow, _, _, _, _, _, _ = build_editflow(cfg, device=device) + + ckpt = torch.load(resolve_path(args.ckpt), map_location=device, weights_only=False) + state = ckpt["state_dict"] if isinstance(ckpt, dict) and "state_dict" in ckpt else ckpt + editflow.load_state_dict(state, strict=False) + editflow.eval() + + _, val_dataloader = build_dataloaders(cfg) + trainer = pl.Trainer( + accelerator="gpu" if torch.cuda.is_available() else "cpu", + devices=1, + logger=False, + enable_checkpointing=False, + ) + metrics = trainer.validate(editflow, val_dataloader, verbose=True) + print(metrics) + + +if __name__ == "__main__": + main() diff --git a/flow_matching/.DS_Store b/flow_matching/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..9ad827f175dd64dd0e57ec44fb48269281f893ec Binary files /dev/null and b/flow_matching/.DS_Store differ diff --git a/flow_matching/__init__.py b/flow_matching/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..719ae055c6a13d749beca75c7a496f3cd7cd8948 --- /dev/null +++ b/flow_matching/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +__version__ = "1.0.10" diff --git a/flow_matching/loss/.DS_Store b/flow_matching/loss/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..5008ddfcf53c02e82d7eee2e57c38e5672ef89f6 Binary files /dev/null and b/flow_matching/loss/.DS_Store differ diff --git a/flow_matching/loss/__init__.py b/flow_matching/loss/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f80c1bacab8818ccf5071459e2d9e7603a6b64f4 --- /dev/null +++ b/flow_matching/loss/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from .generalized_loss import MixturePathGeneralizedKL +from .editflows_loss import EditFlowsLoss, EditFlowsLossReParam + +__all__ = [ + "MixturePathGeneralizedKL", + "EditFlowsLoss", + "EditFlowsLossReParam" +] diff --git a/flow_matching/loss/editflows_loss.py b/flow_matching/loss/editflows_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..a33608060931275cdd08a10e77544e57cc91d374 --- /dev/null +++ b/flow_matching/loss/editflows_loss.py @@ -0,0 +1,320 @@ +# Copyright (c) Meta Platforms, Inc. +# All rights reserved. + +from typing import List, Optional, Sequence +import torch +import torch.nn.functional as F +from torch import Tensor +from torch.nn.modules.loss import _Loss + +# If you kept the adapter around only for scheduler access you can drop it entirely, +# because we pass precomputed_weight from training.step. Keeping it is harmless though. +class EditFlowsLoss(_Loss): + """ + Edit Flows loss (Eq. 23), ragged version: + L_i = (sum_j λ_ins[i][j] + sum_j λ_del[i][j] + sum_j λ_sub[i][j]) + - w_i * ( sum_{ins events e} [log λ_ins[i][slot_e] + log Q_ins[i](y_e)] + + sum_{del j} [log λ_del[i][j]] + + sum_{sub j} [log λ_sub[i][j] + log Q_sub[i](y_j)] ) + """ + + def __init__(self, reduction: str = "mean") -> None: + super().__init__(None, None, reduction) + + def forward( + self, + lam_ins: torch.Tensor, # (B, L) + logits_ins: torch.Tensor, # (B, L, V) + lam_del: torch.Tensor, # (B, L) + lam_sub: torch.Tensor, # (B, L) + logits_sub: torch.Tensor, # (B, L, V) + z_t: torch.Tensor, # (B, N) aligned, with eps_id + z_1: torch.Tensor, # (B, N) aligned target, with eps_id + x_t: torch.Tensor, # (B, L) + valid_mask: torch.Tensor, # (B, L) bool, False==padding, True==valid + precomputed_weight: torch.Tensor, # (B,) or () = kappa_dot/(1-kappa) + eps_id: int, + bos_id: int, + eos_id: int, + ) -> torch.Tensor: + """ + Implements Eq. 23 style loss for Edit Flows. + + We: + 1) penalize total outgoing rate on x_t + 2) for every column in (z_t, z_1) that still differs, we map it to an x_t edit + and add - w * log u_required + + BOS/EOS handling: + - z_t and z_1 are aligned and already contain BOS/EOS + - we do NOT allow edits that delete/replace BOS/EOS + - we also skip columns whose target token is BOS/EOS (they should already match) + """ + device = lam_ins.device + B, L = x_t.shape + _, N = z_t.shape + + # 1. ----- RATE TERM ----- + # valid_mask: True = real token, False = pad + valid_f = valid_mask.to(lam_ins.dtype) # (B, L) 1.0 for valid, 0.0 for pad + + # total outgoing rate at each position + total_rate_pos = lam_ins + lam_del + lam_sub # (B, L) + total_rate_pos = total_rate_pos * valid_f # zero out pads + loss_rate = total_rate_pos.sum(dim=1) # (B,) + + # 2. ----- EDIT TERM ----- + # precompute log-softmax for tokens (better numerics) + logp_ins = F.log_softmax(logits_ins, dim=-1) # (B, L, V) + logp_sub = F.log_softmax(logits_sub, dim=-1) # (B, L, V) + + # Make weight shape nice + if precomputed_weight.dim() == 0: + precomputed_weight = precomputed_weight.view(1).expand(B).to(device) + else: + precomputed_weight = precomputed_weight.to(device) + + # use float for accumulation + loss_edit = torch.zeros(B, dtype=torch.float32, device=device) + + for b in range(B): + # how many *valid* tokens in x_t[b] + valid_len = int(valid_mask[b].sum().item()) + + prefix_non_eps = 0 # = number of non-ε seen in z_t[b, :i] + + for i in range(N): + zt = int(z_t[b, i].item()) + z1 = int(z_1[b, i].item()) + + # map this column to an x_t position + if zt != eps_id: + # this column corresponds to x_t[b, prefix_non_eps] + x_pos = prefix_non_eps + is_token = True + prefix_non_eps += 1 + else: + # this is a gap column, sits BETWEEN tokens + x_pos = prefix_non_eps + is_token = False + + # if already matched -> no term + if zt == z1: + continue + + # --- BOS/EOS guards on the target side --- + # if the target token is BOS or EOS, we shouldn't try to force an edit here + if z1 == bos_id or z1 == eos_id: + # skip this column; aligned BOS/EOS should already match + continue + + # figure out which edit we need + # CASE 1: deletion: z_t has token, z_1 has ε + if is_token and (z1 == eps_id): + # delete token at x_pos + if x_pos >= valid_len: + # out of range (shouldn't happen if alignment & mask match) + raise NotImplementedError + + x_token = int(x_t[b, x_pos].item()) + # do NOT delete BOS/EOS + if x_token == bos_id or x_token == eos_id: + continue + + lam = lam_del[b, x_pos].clamp_min(1e-12) + log_u_req = torch.log(lam) + + # CASE 2: substitution: token -> different token + elif is_token and (z1 != eps_id) and (zt != z1): + if x_pos >= valid_len: + raise NotImplementedError + + x_token = int(x_t[b, x_pos].item()) + # do NOT substitute BOS/EOS + if x_token == bos_id or x_token == eos_id: + continue + + lam = lam_sub[b, x_pos].clamp_min(1e-12) + logp_tok = logp_sub[b, x_pos, z1] + log_u_req = torch.log(lam) + logp_tok + + # CASE 3: insertion: ε -> token + elif (not is_token) and (z1 != eps_id): + # insertion in the gap after token (x_pos - 1) + # if x_pos == 0, we insert after "BOS"/at start -> map to position 0 + ins_pos = x_pos - 1 + if ins_pos < 0: + ins_pos = 0 + + # clamp to last valid position if needed + if valid_len == 0: + # degenerate, but avoid -1 + ins_pos = 0 + elif ins_pos >= valid_len: + ins_pos = valid_len - 1 + + # also don't insert "after" EOS if ins_pos currently points to EOS + x_token = int(x_t[b, ins_pos].item()) + if x_token == eos_id: + # simplest policy: skip this insertion supervision + continue + + lam = lam_ins[b, ins_pos].clamp_min(1e-12) + logp_tok = logp_ins[b, ins_pos, z1] + log_u_req = torch.log(lam) + logp_tok + + else: + # unknown pattern (shouldn't happen with proper alignment) + raise NotImplementedError + + w = precomputed_weight[b] + loss_edit[b] += - w * log_u_req + + # 3. ----- COMBINE ----- + loss = loss_rate + loss_edit # (B,) + loss = loss.mean() + return loss + + +class EditFlowsLossReParam(_Loss): + """ + Edit Flows loss (Eq. 23) but parameterized by: + - lam_total: (B, L) >= 0 + - pi_type: (B, L, 3) prob over {ins, del, sub}, sums to 1 + + Token distributions are still given by logits_ins/logits_sub. + + Required-edit log rate: + del: log u = log lam_total[pos] + log pi_del[pos] + sub: log u = log lam_total[pos] + log pi_sub[pos] + log Q_sub(token) + ins: log u = log lam_total[pos] + log pi_ins[pos] + log Q_ins(token) + """ + + def __init__(self, reduction: str = "mean", gamma_rate: float = 1.0, eps: float = 1e-12) -> None: + super().__init__(None, None, reduction) + self.gamma_rate = float(gamma_rate) + self.eps = float(eps) + + def forward( + self, + lam_total: torch.Tensor, # (B, L) + pi_type: torch.Tensor, # (B, L, 3) probabilities over {ins, del, sub} + logits_ins: torch.Tensor, # (B, L, V) + logits_sub: torch.Tensor, # (B, L, V) + z_t: torch.Tensor, # (B, N) aligned, with eps_id + z_1: torch.Tensor, # (B, N) aligned target, with eps_id + x_t: torch.Tensor, # (B, L) + valid_mask: torch.Tensor, # (B, L) bool, False==padding, True==valid (same as your old loss) + precomputed_weight: torch.Tensor, # (B,) or () = kappa_dot/(1-kappa) + eps_id: int, + bos_id: int, + eos_id: int, + ) -> torch.Tensor: + device = lam_total.device + B, L = x_t.shape + _, N = z_t.shape + + # --- RATE TERM --- + valid_f = valid_mask.to(lam_total.dtype) # True=valid -> 1.0 + loss_rate = (lam_total * valid_f).sum(dim=1) # (B,) + loss_rate = self.gamma_rate * loss_rate + + # --- EDIT TERM --- + logp_ins = F.log_softmax(logits_ins, dim=-1) # (B, L, V) + logp_sub = F.log_softmax(logits_sub, dim=-1) # (B, L, V) + + # normalize pi_type defensively (in case caller passes slightly off-simplex) + pi_sum = pi_type.sum(dim=-1, keepdim=True).clamp_min(self.eps) + pi_type = pi_type / pi_sum + + log_lam = torch.log(lam_total.clamp_min(self.eps)) # (B, L) + log_pi = torch.log(pi_type.clamp_min(self.eps)) # (B, L, 3) + + # weight shape + if precomputed_weight.dim() == 0: + precomputed_weight = precomputed_weight.view(1).expand(B).to(device) + else: + precomputed_weight = precomputed_weight.to(device) + + loss_edit = torch.zeros(B, dtype=torch.float32, device=device) + + for b in range(B): + valid_len = int(valid_mask[b].sum().item()) + prefix_non_eps = 0 + + for i in range(N): + zt = int(z_t[b, i].item()) + z1 = int(z_1[b, i].item()) + + # map this alignment column to an x_t position + if zt != eps_id: + x_pos = prefix_non_eps + is_token = True + prefix_non_eps += 1 + else: + x_pos = prefix_non_eps + is_token = False + + # already matched -> no supervision + if zt == z1: + continue + + # don't force edits where target is BOS/EOS + if z1 == bos_id or z1 == eos_id: + continue + + # CASE 1: deletion (token -> eps) + if is_token and (z1 == eps_id): + if x_pos >= valid_len: + raise NotImplementedError + + x_token = int(x_t[b, x_pos].item()) + if x_token == bos_id or x_token == eos_id: + continue + + log_u_req = log_lam[b, x_pos] + log_pi[b, x_pos, 1] # del index = 1 + + # CASE 2: substitution (token -> different token) + elif is_token and (z1 != eps_id) and (zt != z1): + if x_pos >= valid_len: + raise NotImplementedError + + x_token = int(x_t[b, x_pos].item()) + if x_token == bos_id or x_token == eos_id: + continue + + log_u_req = ( + log_lam[b, x_pos] + + log_pi[b, x_pos, 2] # sub index = 2 + + logp_sub[b, x_pos, z1] + ) + + # CASE 3: insertion (eps -> token), map gap to "after token (x_pos-1)" + elif (not is_token) and (z1 != eps_id): + ins_pos = x_pos - 1 + if ins_pos < 0: + ins_pos = 0 + + if valid_len == 0: + ins_pos = 0 + elif ins_pos >= valid_len: + ins_pos = valid_len - 1 + + x_token = int(x_t[b, ins_pos].item()) + if x_token == eos_id: + continue + + log_u_req = ( + log_lam[b, ins_pos] + + log_pi[b, ins_pos, 0] # ins index = 0 + + logp_ins[b, ins_pos, z1] + ) + + else: + raise NotImplementedError + + w = precomputed_weight[b] + loss_edit[b] += -w * log_u_req + + loss = loss_rate + loss_edit + return loss.mean() \ No newline at end of file diff --git a/flow_matching/loss/generalized_loss.py b/flow_matching/loss/generalized_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..cc1507e031d102825d9eb06e9bb7a0f248d16952 --- /dev/null +++ b/flow_matching/loss/generalized_loss.py @@ -0,0 +1,80 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch import Tensor +from torch.nn.modules.loss import _Loss + +from flow_matching.path import MixtureDiscreteProbPath + + +class MixturePathGeneralizedKL(_Loss): + r"""A generalized KL loss for discrete flow matching. + A class that measures the generalized KL of a discrete flow model :math:`p_{1|t}` w.r.t. a probability path given by ``path``. Note: this class is assuming that the model is trained on the same path. + + For a model trained on a space :math:`\mathcal{S} = \mathcal{T}^d`, :math:`\mathcal{T} = [K] = \set{1,2,\ldots,K}`, the loss is given by + + .. math:: + \ell_i(x_1, x_t, t) = -\frac{\dot{\kappa}_t}{1-\kappa_t} \biggr[ p_{1|t}(x_t^i|x_t) -\delta_{x^i_1}(x_t^i) + (1-\delta_{x^i_1}(x_t^i))\left(\log p_{1|t}(x_1^i|x_t)\right)\biggr], + + where :math:`\kappa_t` is the scheduler associated with ``path``. + + Args: + path (MixtureDiscreteProbPath): Probability path (x-prediction training). + reduction (str, optional): Specify the reduction to apply to the output ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction is applied to the output, ``'mean'``: the output is reduced by mean over sequence elements, ``'sum'``: the output is reduced by sum over sequence elements. Defaults to 'mean'. + """ + + def __init__(self, path: MixtureDiscreteProbPath, reduction: str = "mean") -> None: + super().__init__(None, None, reduction) + self.path = path + + def forward(self, logits: Tensor, x_1: Tensor, x_t: Tensor, t: Tensor) -> Tensor: + r"""Evaluates the generalized KL loss. + + Args: + logits (Tensor): posterior model output (i.e., softmax(``logits``) :math:`=p_{1|t}(x|x_t)`), shape (batch, d, K). + x_1 (Tensor): target data point :math:`x_1 \sim q`, shape (batch, d). + x_t (Tensor): conditional sample at :math:`x_t \sim p_t(\cdot|x_1)`, shape (batch, d). + t (Tensor): times in :math:`[0,1]`, shape (batch). + + Raises: + ValueError: reduction value must be one of ``'none'`` | ``'mean'`` | ``'sum'``. + + Returns: + Tensor: Generalized KL loss. + """ + x_1_shape = x_1.shape + + # extract x_1 value of log(p_{1|t}(x|x_t)). + log_p_1t = torch.log_softmax(logits, dim=-1) + log_p_1t_x1 = torch.gather(log_p_1t, dim=-1, index=x_1.unsqueeze(-1)) + log_p_1t_x1 = log_p_1t_x1.view(*x_1_shape) + + # extract x_t value of p_{1|t}(x|x_t). + p_1t = torch.exp(log_p_1t) + p_1t_xt = torch.gather(p_1t, dim=-1, index=x_t.unsqueeze(-1)) + p_1t_xt = p_1t_xt.view(*x_1_shape) + + scheduler_output = self.path.scheduler(t) + + jump_coefficient = ( + scheduler_output.d_alpha_t / (1 - scheduler_output.alpha_t) + )[(...,) + (None,) * (x_1.dim() - 1)] + jump_coefficient = jump_coefficient.repeat(1, *x_1_shape[1:]) + delta_x1_xt = (x_t == x_1).to(log_p_1t.dtype) + + loss = -jump_coefficient * ( + p_1t_xt - delta_x1_xt + (1 - delta_x1_xt) * log_p_1t_x1 + ) + + if self.reduction == "mean": + return torch.mean(loss) + elif self.reduction == "sum": + return torch.sum(loss) + elif self.reduction == "none": + return loss + else: + raise ValueError(f"{self.reduction} is not a valid value for reduction") diff --git a/flow_matching/path/.DS_Store b/flow_matching/path/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..b2a7aa41a9aa92d105c1c88facc5fe2914357fee Binary files /dev/null and b/flow_matching/path/.DS_Store differ diff --git a/flow_matching/path/__init__.py b/flow_matching/path/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..293c78e1eb084c0abbfb7352b7e9d1e67a7caf47 --- /dev/null +++ b/flow_matching/path/__init__.py @@ -0,0 +1,24 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from .affine import AffineProbPath, CondOTProbPath +from .geodesic import GeodesicProbPath +from .mixture import MixtureDiscreteProbPath +from .path import ProbPath +from .path_sample import DiscretePathSample, PathSample +from .editflows_adapter import EditFlowsPathAdapter + + +__all__ = [ + "ProbPath", + "AffineProbPath", + "CondOTProbPath", + "MixtureDiscreteProbPath", + "GeodesicProbPath", + "PathSample", + "DiscretePathSample", + "EditFlowsPathAdapter", +] diff --git a/flow_matching/path/affine.py b/flow_matching/path/affine.py new file mode 100644 index 0000000000000000000000000000000000000000..81cb7ed31f2434d03424ea9a5571a36bfc9f2681 --- /dev/null +++ b/flow_matching/path/affine.py @@ -0,0 +1,260 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from torch import Tensor + +from flow_matching.path.path import ProbPath +from flow_matching.path.path_sample import PathSample +from flow_matching.path.scheduler.scheduler import CondOTScheduler, Scheduler +from flow_matching.utils import expand_tensor_like + + +class AffineProbPath(ProbPath): + r"""The ``AffineProbPath`` class represents a specific type of probability path where the transformation between distributions is affine. + An affine transformation can be represented as: + + .. math:: + + X_t = \alpha_t X_1 + \sigma_t X_0, + + where :math:`X_t` is the transformed data point at time `t`. :math:`X_0` and :math:`X_1` are the source and target data points, respectively. :math:`\alpha_t` and :math:`\sigma_t` are the parameters of the affine transformation at time `t`. + + The scheduler is responsible for providing the time-dependent parameters :math:`\alpha_t` and :math:`\sigma_t`, as well as their derivatives, which define the affine transformation at any given time `t`. + + Using ``AffineProbPath`` in the flow matching framework: + + .. code-block:: python + + # Instantiates a probability path + my_path = AffineProbPath(...) + mse_loss = torch.nn.MSELoss() + + for x_1 in dataset: + # Sets x_0 to random noise + x_0 = torch.randn() + + # Sets t to a random value in [0,1] + t = torch.rand() + + # Samples the conditional path X_t ~ p_t(X_t|X_0,X_1) + path_sample = my_path.sample(x_0=x_0, x_1=x_1, t=t) + + # Computes the MSE loss w.r.t. the velocity + loss = mse_loss(path_sample.dx_t, my_model(x_t, t)) + loss.backward() + + Args: + scheduler (Scheduler): An instance of a scheduler that provides the parameters :math:`\alpha_t`, :math:`\sigma_t`, and their derivatives over time. + + """ + + def __init__(self, scheduler: Scheduler): + self.scheduler = scheduler + + def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor) -> PathSample: + r"""Sample from the affine probability path: + + | given :math:`(X_0,X_1) \sim \pi(X_0,X_1)` and a scheduler :math:`(\alpha_t,\sigma_t)`. + | return :math:`X_0, X_1, X_t = \alpha_t X_1 + \sigma_t X_0`, and the conditional velocity at :math:`X_t, \dot{X}_t = \dot{\alpha}_t X_1 + \dot{\sigma}_t X_0`. + + Args: + x_0 (Tensor): source data point, shape (batch_size, ...). + x_1 (Tensor): target data point, shape (batch_size, ...). + t (Tensor): times in [0,1], shape (batch_size). + + Returns: + PathSample: a conditional sample at :math:`X_t \sim p_t`. + """ + self.assert_sample_shape(x_0=x_0, x_1=x_1, t=t) + + scheduler_output = self.scheduler(t) + + alpha_t = expand_tensor_like( + input_tensor=scheduler_output.alpha_t, expand_to=x_1 + ) + sigma_t = expand_tensor_like( + input_tensor=scheduler_output.sigma_t, expand_to=x_1 + ) + d_alpha_t = expand_tensor_like( + input_tensor=scheduler_output.d_alpha_t, expand_to=x_1 + ) + d_sigma_t = expand_tensor_like( + input_tensor=scheduler_output.d_sigma_t, expand_to=x_1 + ) + + # construct xt ~ p_t(x|x1). + x_t = sigma_t * x_0 + alpha_t * x_1 + dx_t = d_sigma_t * x_0 + d_alpha_t * x_1 + + return PathSample(x_t=x_t, dx_t=dx_t, x_1=x_1, x_0=x_0, t=t) + + def target_to_velocity(self, x_1: Tensor, x_t: Tensor, t: Tensor) -> Tensor: + r"""Convert from x_1 representation to velocity. + + | given :math:`X_1`. + | return :math:`\dot{X}_t`. + + Args: + x_1 (Tensor): target data point. + x_t (Tensor): path sample at time t. + t (Tensor): time in [0,1]. + + Returns: + Tensor: velocity. + """ + scheduler_output = self.scheduler(t) + + alpha_t = scheduler_output.alpha_t + d_alpha_t = scheduler_output.d_alpha_t + sigma_t = scheduler_output.sigma_t + d_sigma_t = scheduler_output.d_sigma_t + + a_t = d_sigma_t / sigma_t + b_t = (d_alpha_t * sigma_t - d_sigma_t * alpha_t) / sigma_t + + return a_t * x_t + b_t * x_1 + + def epsilon_to_velocity(self, epsilon: Tensor, x_t: Tensor, t: Tensor) -> Tensor: + r"""Convert from epsilon representation to velocity. + + | given :math:`\epsilon`. + | return :math:`\dot{X}_t`. + + Args: + epsilon (Tensor): noise in the path sample. + x_t (Tensor): path sample at time t. + t (Tensor): time in [0,1]. + + Returns: + Tensor: velocity. + """ + scheduler_output = self.scheduler(t) + + alpha_t = scheduler_output.alpha_t + d_alpha_t = scheduler_output.d_alpha_t + sigma_t = scheduler_output.sigma_t + d_sigma_t = scheduler_output.d_sigma_t + + a_t = d_alpha_t / alpha_t + b_t = (d_sigma_t * alpha_t - d_alpha_t * sigma_t) / alpha_t + + return a_t * x_t + b_t * epsilon + + def velocity_to_target(self, velocity: Tensor, x_t: Tensor, t: Tensor) -> Tensor: + r"""Convert from velocity to x_1 representation. + + | given :math:`\dot{X}_t`. + | return :math:`X_1`. + + Args: + velocity (Tensor): velocity at the path sample. + x_t (Tensor): path sample at time t. + t (Tensor): time in [0,1]. + + Returns: + Tensor: target data point. + """ + scheduler_output = self.scheduler(t) + + alpha_t = scheduler_output.alpha_t + d_alpha_t = scheduler_output.d_alpha_t + sigma_t = scheduler_output.sigma_t + d_sigma_t = scheduler_output.d_sigma_t + + a_t = -d_sigma_t / (d_alpha_t * sigma_t - d_sigma_t * alpha_t) + b_t = sigma_t / (d_alpha_t * sigma_t - d_sigma_t * alpha_t) + + return a_t * x_t + b_t * velocity + + def epsilon_to_target(self, epsilon: Tensor, x_t: Tensor, t: Tensor) -> Tensor: + r"""Convert from epsilon representation to x_1 representation. + + | given :math:`\epsilon`. + | return :math:`X_1`. + + Args: + epsilon (Tensor): noise in the path sample. + x_t (Tensor): path sample at time t. + t (Tensor): time in [0,1]. + + Returns: + Tensor: target data point. + """ + scheduler_output = self.scheduler(t) + + alpha_t = scheduler_output.alpha_t + sigma_t = scheduler_output.sigma_t + + a_t = 1 / alpha_t + b_t = -sigma_t / alpha_t + + return a_t * x_t + b_t * epsilon + + def velocity_to_epsilon(self, velocity: Tensor, x_t: Tensor, t: Tensor) -> Tensor: + r"""Convert from velocity to noise representation. + + | given :math:`\dot{X}_t`. + | return :math:`\epsilon`. + + Args: + velocity (Tensor): velocity at the path sample. + x_t (Tensor): path sample at time t. + t (Tensor): time in [0,1]. + + Returns: + Tensor: noise in the path sample. + """ + scheduler_output = self.scheduler(t) + + alpha_t = scheduler_output.alpha_t + d_alpha_t = scheduler_output.d_alpha_t + sigma_t = scheduler_output.sigma_t + d_sigma_t = scheduler_output.d_sigma_t + + a_t = -d_alpha_t / (d_sigma_t * alpha_t - d_alpha_t * sigma_t) + b_t = alpha_t / (d_sigma_t * alpha_t - d_alpha_t * sigma_t) + + return a_t * x_t + b_t * velocity + + def target_to_epsilon(self, x_1: Tensor, x_t: Tensor, t: Tensor) -> Tensor: + r"""Convert from x_1 representation to velocity. + + | given :math:`X_1`. + | return :math:`\epsilon`. + + Args: + x_1 (Tensor): target data point. + x_t (Tensor): path sample at time t. + t (Tensor): time in [0,1]. + + Returns: + Tensor: noise in the path sample. + """ + scheduler_output = self.scheduler(t) + + alpha_t = scheduler_output.alpha_t + sigma_t = scheduler_output.sigma_t + + a_t = 1 / sigma_t + b_t = -alpha_t / sigma_t + + return a_t * x_t + b_t * x_1 + + +class CondOTProbPath(AffineProbPath): + r"""The ``CondOTProbPath`` class represents a conditional optimal transport probability path. + + This class is a specialized version of the ``AffineProbPath`` that uses a conditional optimal transport scheduler to determine the parameters of the affine transformation. + + The parameters :math:`\alpha_t` and :math:`\sigma_t` for the conditional optimal transport path are defined as: + + .. math:: + + \alpha_t = t \quad \text{and} \quad \sigma_t = 1 - t. + """ + + def __init__(self): + self.scheduler = CondOTScheduler() diff --git a/flow_matching/path/editflows_adapter.py b/flow_matching/path/editflows_adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..338b04bd2c7db08f464e7dc1cfcdef2d35fa468b --- /dev/null +++ b/flow_matching/path/editflows_adapter.py @@ -0,0 +1,85 @@ +from dataclasses import dataclass +from typing import List, Sequence, Optional, Tuple + +import torch +from torch import Tensor +import pdb + +from flow_matching.path import MixtureDiscreteProbPath + + +@dataclass +class EditFlowsPathSampleRagged: + # aligned inputs/outputs (ragged) + z0_list: List[Tensor] # list[(T_i,)] + z1_list: List[Tensor] # list[(T_i,)] + zt_list: List[Tensor] # list[(T_i,)] + # projected current sequences (ragged X_t) + x_t_list: List[Tensor] # list[(n_i,)] + # targets/masks for EF loss (ragged) + need_delete_list: List[Tensor] # list[(n_i,)] bool + need_substitute_list: List[Tensor] # list[(n_i,)] bool + sub_target_list: List[Tensor] # list[(n_i,)] long + ins_slot_idx_list: List[Tensor] # list[(K_i,)] long indices in [0..n_i] + ins_target_list: List[Tensor] # list[(K_i,)] long + # time + t: Tensor # (B,) + + +class EditFlowsPathAdapter: + """ + Ragged adapter for Edit Flows: + - Accepts aligned Z sequences as tensors (B, N) (may include eps_id). + - Applies MixtureDiscreteProbPath's scheduler (σ_t) to sample z_t by + flipping z0→z1 with prob (1-σ_t) per aligned token. + - You can later project z_t -> x_t by removing ε. + """ + def __init__(self, mixture_path, eps_id: int): + self.path = mixture_path + self.eps_id = int(eps_id) + + # expose scheduler so training can precompute weights κ̇/(1-κ) + def scheduler(self, t: Tensor): + return self.path.scheduler(t) + + @torch.no_grad() + def sample( + self, + z0: Tensor, # (B, N) + z1: Tensor, # (B, N) + t: Tensor, # () or (B,) + ) -> Tensor: + """ + Sample z_t from (z0, z1) using the discrete mixture path: + - σ_t = scheduler(t) + - with prob σ_t keep z0 + - with prob (1 - σ_t) take z1 + We do this per-column. + BOS/EOS are already aligned -> z0 == z1 there -> safe. + """ + device = z0.device + B, N = z0.shape + + # get sigma_t from the path + sigma = self.path.scheduler(t).sigma_t # could be scalar or (B,) + if sigma.dim() == 0: + sigma = sigma.expand(B) # (B,) + else: + # ensure shape is (B,) + sigma = sigma.view(B) + # reshape to broadcast over columns + sigma = sigma.view(B, 1) # (B, 1) + + # uniform noise per position + u = torch.rand(B, N, device=device) + + # we only need to flip where z0 != z1 + diff_mask = (z0 != z1) # (B, N) + + # flip where u > sigma (i.e. prob 1 - sigma) AND tokens differ + flip = (u > sigma) & diff_mask # (B, N) + + # choose from z1 when flip, else z0 + z_t = torch.where(flip, z1, z0) # (B, N) + + return z_t \ No newline at end of file diff --git a/flow_matching/path/geodesic.py b/flow_matching/path/geodesic.py new file mode 100644 index 0000000000000000000000000000000000000000..d04bf6770a98804f9bbcc7447c352be6ddec65dc --- /dev/null +++ b/flow_matching/path/geodesic.py @@ -0,0 +1,100 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from torch import Tensor +from torch.func import jvp, vmap + +from flow_matching.path.path import ProbPath + +from flow_matching.path.path_sample import PathSample +from flow_matching.path.scheduler import ConvexScheduler +from flow_matching.utils import expand_tensor_like + +from flow_matching.utils.manifolds import geodesic, Manifold + + +class GeodesicProbPath(ProbPath): + r"""The ``GeodesicProbPath`` class represents a specific type of probability path where the transformation between distributions is defined through the geodesic path. + Mathematically, a geodesic path can be represented as: + + .. math:: + + X_t = \psi_t(X_0 | X_1) = \exp_{X_1}(\kappa_t \log_{X_1}(X_0)), + + where :math:`X_t` is the transformed data point at time `t`, :math:`X_0` and :math:`X_1` are the source and target data points, respectively, and :math:`\kappa_t` is a scheduler. + + The scheduler is responsible for providing the time-dependent :math:`\kappa_t` and must be differentiable. + + Using ``GeodesicProbPath`` in the flow matching framework: + + .. code-block:: python + # Instantiates a manifold + manifold = FlatTorus() + + # Instantiates a scheduler + scheduler = CondOTScheduler() + + # Instantiates a probability path + my_path = GeodesicProbPath(scheduler, manifold) + mse_loss = torch.nn.MSELoss() + + for x_1 in dataset: + # Sets x_0 to random noise + x_0 = torch.randn() + + # Sets t to a random value in [0,1] + t = torch.rand() + + # Samples the conditional path :math:`X_t \sim p_t(X_t|X_0,X_1)` + path_sample = my_path.sample(x_0=x_0, x_1=x_1, t=t) + + # Computes the MSE loss w.r.t. the velocity + loss = mse_loss(path_sample.dx_t, my_model(x_t, t)) + loss.backward() + + Args: + scheduler (ConvexScheduler): The scheduler that provides :math:`\kappa_t`. + manifold (Manifold): The manifold on which the probability path is defined. + + """ + + def __init__(self, scheduler: ConvexScheduler, manifold: Manifold): + self.scheduler = scheduler + self.manifold = manifold + + def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor) -> PathSample: + r"""Sample from the Riemannian probability path with geodesic interpolation: + + | given :math:`(X_0,X_1) \sim \pi(X_0,X_1)` and a scheduler :math:`\kappa_t`. + | return :math:`X_0, X_1, X_t = \exp_{X_1}(\kappa_t \log_{X_1}(X_0))`, and the conditional velocity at :math:`X_t, \dot{X}_t`. + + Args: + x_0 (Tensor): source data point, shape (batch_size, ...). + x_1 (Tensor): target data point, shape (batch_size, ...). + t (Tensor): times in [0,1], shape (batch_size). + + Returns: + PathSample: A conditional sample at :math:`X_t \sim p_t`. + """ + self.assert_sample_shape(x_0=x_0, x_1=x_1, t=t) + t = expand_tensor_like(input_tensor=t, expand_to=x_1[..., 0:1]).clone() + + def cond_u(x_0, x_1, t): + path = geodesic(self.manifold, x_0, x_1) + x_t, dx_t = jvp( + lambda t: path(self.scheduler(t).alpha_t), + (t,), + (torch.ones_like(t).to(t),), + ) + return x_t, dx_t + + x_t, dx_t = vmap(cond_u)(x_0, x_1, t) + x_t = x_t.reshape_as(x_1) + dx_t = dx_t.reshape_as(x_1) + + return PathSample(x_t=x_t, dx_t=dx_t, x_1=x_1, x_0=x_0, t=t) diff --git a/flow_matching/path/mixture.py b/flow_matching/path/mixture.py new file mode 100644 index 0000000000000000000000000000000000000000..28b4043ebe89a35408ada3a28cf33291ae277cd1 --- /dev/null +++ b/flow_matching/path/mixture.py @@ -0,0 +1,117 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn.functional as F + +from torch import Tensor + +from flow_matching.path.path import ProbPath + +from flow_matching.path.path_sample import DiscretePathSample +from flow_matching.path.scheduler import ConvexScheduler +from flow_matching.utils import expand_tensor_like, unsqueeze_to_match + + +class MixtureDiscreteProbPath(ProbPath): + r"""The ``MixtureDiscreteProbPath`` class defines a factorized discrete probability path. + + This path remains constant at the source data point :math:`X_0` until a random time, determined by the scheduler, when it flips to the target data point :math:`X_1`. + The scheduler determines the flip probability using the parameter :math:`\sigma_t`, which is a function of time `t`. Specifically, :math:`\sigma_t` represents the probability of remaining at :math:`X_0`, while :math:`1 - \sigma_t` is the probability of flipping to :math:`X_1`: + + .. math:: + + P(X_t = X_0) = \sigma_t \quad \text{and} \quad P(X_t = X_1) = 1 - \sigma_t, + + where :math:`\sigma_t` is provided by the scheduler. + + Example: + + .. code-block:: python + + >>> x_0 = torch.zeros((1, 3, 3)) + >>> x_1 = torch.ones((1, 3, 3)) + + >>> path = MixtureDiscreteProbPath(PolynomialConvexScheduler(n=1.0)) + >>> result = path.sample(x_0, x_1, t=torch.tensor([0.1])).x_t + >>> result + tensor([[[0.0, 0.0, 0.0], + [0.0, 0.0, 1.0], + [0.0, 0.0, 0.0]]]) + + >>> result = path.sample(x_0, x_1, t=torch.tensor([0.5])).x_t + >>> result + tensor([[[1.0, 0.0, 1.0], + [0.0, 1.0, 0.0], + [0.0, 1.0, 0.0]]]) + + >>> result = path.sample(x_0, x_1, t=torch.tensor([1.0])).x_t + >>> result + tensor([[[1.0, 1.0, 1.0], + [1.0, 1.0, 1.0], + [1.0, 1.0, 1.0]]]) + + Args: + scheduler (ConvexScheduler): The scheduler that provides :math:`\sigma_t`. + """ + + def __init__(self, scheduler: ConvexScheduler): + assert isinstance( + scheduler, ConvexScheduler + ), "Scheduler for ConvexProbPath must be a ConvexScheduler." + + self.scheduler = scheduler + + def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor) -> DiscretePathSample: + r"""Sample from the affine probability path: + | given :math:`(X_0,X_1) \sim \pi(X_0,X_1)` and a scheduler :math:`(\alpha_t,\sigma_t)`. + | return :math:`X_0, X_1, t`, and :math:`X_t \sim p_t`. + Args: + x_0 (Tensor): source data point, shape (batch_size, ...). + x_1 (Tensor): target data point, shape (batch_size, ...). + t (Tensor): times in [0,1], shape (batch_size). + + Returns: + DiscretePathSample: a conditional sample at :math:`X_t ~ p_t`. + """ + self.assert_sample_shape(x_0=x_0, x_1=x_1, t=t) + + sigma_t = self.scheduler(t).sigma_t + + sigma_t = expand_tensor_like(input_tensor=sigma_t, expand_to=x_1) + + source_indices = torch.rand(size=x_1.shape, device=x_1.device) < sigma_t + x_t = torch.where(condition=source_indices, input=x_0, other=x_1) + + return DiscretePathSample(x_t=x_t, x_1=x_1, x_0=x_0, t=t) + + def posterior_to_velocity( + self, posterior_logits: Tensor, x_t: Tensor, t: Tensor + ) -> Tensor: + r"""Convert the factorized posterior to velocity. + + | given :math:`p(X_1|X_t)`. In the factorized case: :math:`\prod_i p(X_1^i | X_t)`. + | return :math:`u_t`. + + Args: + posterior_logits (Tensor): logits of the x_1 posterior conditional on x_t, shape (..., vocab size). + x_t (Tensor): path sample at time t, shape (...). + t (Tensor): time in [0,1]. + + Returns: + Tensor: velocity. + """ + posterior = torch.softmax(posterior_logits, dim=-1) + vocabulary_size = posterior.shape[-1] + x_t = F.one_hot(x_t, num_classes=vocabulary_size) + t = unsqueeze_to_match(source=t, target=x_t) + + scheduler_output = self.scheduler(t) + + kappa_t = scheduler_output.alpha_t + d_kappa_t = scheduler_output.d_alpha_t + + return (d_kappa_t / (1 - kappa_t)) * (posterior - x_t) diff --git a/flow_matching/path/path.py b/flow_matching/path/path.py new file mode 100644 index 0000000000000000000000000000000000000000..94cc39e7ce3e0e1dadd5553e8110c8e43bd35ffe --- /dev/null +++ b/flow_matching/path/path.py @@ -0,0 +1,64 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from abc import ABC, abstractmethod + +from torch import Tensor + +from flow_matching.path.path_sample import PathSample + + +class ProbPath(ABC): + r"""Abstract class, representing a probability path. + + A probability path transforms the distribution :math:`p(X_0)` into :math:`p(X_1)` over :math:`t=0\rightarrow 1`. + + The ``ProbPath`` class is designed to support model training in the flow matching framework. It supports two key functionalities: (1) sampling the conditional probability path and (2) conversion between various training objectives. + Here is a high-level example + + .. code-block:: python + + # Instantiate a probability path + my_path = ProbPath(...) + + for x_0, x_1 in dataset: + # Sets t to a random value in [0,1] + t = torch.rand() + + # Samples the conditional path X_t ~ p_t(X_t|X_0,X_1) + path_sample = my_path.sample(x_0=x_0, x_1=x_1, t=t) + + # Optimizes the model. The loss function varies, depending on model and path. + loss(path_sample, my_model(x_t, t)).backward() + + """ + + @abstractmethod + def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor) -> PathSample: + r"""Sample from an abstract probability path: + + | given :math:`(X_0,X_1) \sim \pi(X_0,X_1)`. + | returns :math:`X_0, X_1, X_t \sim p_t(X_t)`, and a conditional target :math:`Y`, all objects are under ``PathSample``. + + Args: + x_0 (Tensor): source data point, shape (batch_size, ...). + x_1 (Tensor): target data point, shape (batch_size, ...). + t (Tensor): times in [0,1], shape (batch_size). + + Returns: + PathSample: a conditional sample. + """ + + def assert_sample_shape(self, x_0, x_1, t): + # x_0, x_1: list[Tensor] of variable length (ragged) + # t: Tensor of shape (B,) + assert isinstance(x_0, (list, tuple)), f"x_0 must be a list or tuple of ragged Tensors. Got {type(x_0)}" + assert isinstance(x_1, (list, tuple)), f"x_1 must be a list or tuple of ragged Tensors. Got {type(x_1)}" + assert isinstance(t, Tensor), f"t must be a torch.Tensor. Got {type(t)}" + batch_size = len(x_0) + assert len(x_1) == batch_size, f"x_1 must match x_0 batch size ({batch_size}). Got {len(x_1)}" + assert t.ndim == 1, f"The time vector t must have shape [batch_size]. Got {t.shape}." + assert t.shape[0] == batch_size, f"Time t dimension must match the batch size [{batch_size}]. Got {t.shape}" diff --git a/flow_matching/path/path_sample.py b/flow_matching/path/path_sample.py new file mode 100644 index 0000000000000000000000000000000000000000..3f18bb233790d9af401a86e6f0b3c3774d363163 --- /dev/null +++ b/flow_matching/path/path_sample.py @@ -0,0 +1,53 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from dataclasses import dataclass, field + +from torch import Tensor + + +@dataclass +class PathSample: + r"""Represents a sample of a conditional-flow generated probability path. + + Attributes: + x_1 (Tensor): the target sample :math:`X_1`. + x_0 (Tensor): the source sample :math:`X_0`. + t (Tensor): the time sample :math:`t`. + x_t (Tensor): samples :math:`X_t \sim p_t(X_t)`, shape (batch_size, ...). + dx_t (Tensor): conditional target :math:`\frac{\partial X}{\partial t}`, shape: (batch_size, ...). + + """ + + x_1: Tensor = field(metadata={"help": "target samples X_1 (batch_size, ...)."}) + x_0: Tensor = field(metadata={"help": "source samples X_0 (batch_size, ...)."}) + t: Tensor = field(metadata={"help": "time samples t (batch_size, ...)."}) + x_t: Tensor = field( + metadata={"help": "samples x_t ~ p_t(X_t), shape (batch_size, ...)."} + ) + dx_t: Tensor = field( + metadata={"help": "conditional target dX_t, shape: (batch_size, ...)."} + ) + + +@dataclass +class DiscretePathSample: + r""" + Represents a sample of a conditional-flow generated discrete probability path. + + Attributes: + x_1 (Tensor): the target sample :math:`X_1`. + x_0 (Tensor): the source sample :math:`X_0`. + t (Tensor): the time sample :math:`t`. + x_t (Tensor): the sample along the path :math:`X_t \sim p_t`. + """ + + x_1: Tensor = field(metadata={"help": "target samples X_1 (batch_size, ...)."}) + x_0: Tensor = field(metadata={"help": "source samples X_0 (batch_size, ...)."}) + t: Tensor = field(metadata={"help": "time samples t (batch_size, ...)."}) + x_t: Tensor = field( + metadata={"help": "samples X_t ~ p_t(X_t), shape (batch_size, ...)."} + ) diff --git a/flow_matching/path/scheduler/.DS_Store b/flow_matching/path/scheduler/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..5008ddfcf53c02e82d7eee2e57c38e5672ef89f6 Binary files /dev/null and b/flow_matching/path/scheduler/.DS_Store differ diff --git a/flow_matching/path/scheduler/__init__.py b/flow_matching/path/scheduler/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f3b1a43f08b3c85a9c69ebe87b4fdeaf9eff29b1 --- /dev/null +++ b/flow_matching/path/scheduler/__init__.py @@ -0,0 +1,29 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from .schedule_transform import ScheduleTransformedModel +from .scheduler import ( + CondOTScheduler, + ConvexScheduler, + CosineScheduler, + LinearVPScheduler, + PolynomialConvexScheduler, + Scheduler, + SchedulerOutput, + VPScheduler, +) + +__all__ = [ + "CondOTScheduler", + "CosineScheduler", + "ConvexScheduler", + "PolynomialConvexScheduler", + "ScheduleTransformedModel", + "Scheduler", + "VPScheduler", + "LinearVPScheduler", + "SchedulerOutput", +] diff --git a/flow_matching/path/scheduler/schedule_transform.py b/flow_matching/path/scheduler/schedule_transform.py new file mode 100644 index 0000000000000000000000000000000000000000..a366f19c9eafd961d5e40db28f553b409ff82520 --- /dev/null +++ b/flow_matching/path/scheduler/schedule_transform.py @@ -0,0 +1,148 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from torch import Tensor + +from flow_matching.path.scheduler.scheduler import Scheduler +from flow_matching.utils import ModelWrapper + + +class ScheduleTransformedModel(ModelWrapper): + """ + Change of scheduler for a velocity model. + + This class wraps a given velocity model and transforms its scheduling + to a new scheduler function. It modifies the time + dynamics of the model according to the new scheduler while maintaining + the original model's behavior. + + Example: + + .. code-block:: python + + import torch + from flow_matching.path.scheduler import CondOTScheduler, CosineScheduler, ScheduleTransformedModel + from flow_matching.solver import ODESolver + + # Initialize the model and schedulers + model = ... + + original_scheduler = CondOTScheduler() + new_scheduler = CosineScheduler() + + # Create the transformed model + transformed_model = ScheduleTransformedModel( + velocity_model=model, + original_scheduler=original_scheduler, + new_scheduler=new_scheduler + ) + + # Set up the solver + solver = ODESolver(velocity_model=transformed_model) + + x_0 = torch.randn([10, 2]) # Example initial condition + + x_1 = solver.sample( + time_steps=torch.tensor([0.0, 1.0]), + x_init=x_0, + step_size=1/1000 + )[1] + + Args: + velocity_model (ModelWrapper): The original velocity model to be transformed. + original_scheduler (Scheduler): The scheduler used by the original model. Must implement the snr_inverse function. + new_scheduler (Scheduler): The new scheduler to be applied to the model. + """ + + def __init__( + self, + velocity_model: ModelWrapper, + original_scheduler: Scheduler, + new_scheduler: Scheduler, + ): + super().__init__(model=velocity_model) + self.original_scheduler = original_scheduler + self.new_scheduler = new_scheduler + + assert hasattr(self.original_scheduler, "snr_inverse") and callable( + getattr(self.original_scheduler, "snr_inverse") + ), "The original scheduler must have a callable 'snr_inverse' method." + + def forward(self, x: Tensor, t: Tensor, **extras) -> Tensor: + r""" + Compute the transformed marginal velocity field for a new scheduler. + This method implements a post-training velocity scheduler change for + affine conditional flows. It transforms a generating marginal velocity + field :math:`u_t(x)` based on an original scheduler to a new marginal velocity + field :math:`\bar{u}_r(x)` based on a different scheduler, while maintaining + the same data coupling. + The transformation is based on the scale-time (ST) transformation + between the two conditional flows, defined as: + + .. math:: + + \bar{X}_r = s_r X_{t_r}, + + where :math:`X_t` and :math:`\bar{X}_r` are defined by their respective schedulers. + The ST transformation is computed as: + + .. math:: + + t_r = \rho^{-1}(\bar{\rho}(r)) \quad \text{and} \quad s_r = \frac{\bar{\sigma}_r}{\sigma_{t_r}}. + + Here, :math:`\rho(t)` is the signal-to-noise ratio (SNR) defined as: + + .. math:: + + \rho(t) = \frac{\alpha_t}{\sigma_t}. + + :math:`\bar{\rho}(r)` is similarly defined for the new scheduler. + The marginal velocity for the new scheduler is then given by: + + .. math:: + + \bar{u}_r(x) = \left(\frac{\dot{s}_r}{s_r}\right) x + s_r \dot{t}_r u_{t_r}\left(\frac{x}{s_r}\right). + + Args: + x (Tensor): :math:`x_t`, the input tensor. + t (Tensor): The time tensor (denoted as :math:`r` above). + **extras: Additional arguments for the model. + Returns: + Tensor: The transformed velocity. + """ + r = t + + r_scheduler_output = self.new_scheduler(t=r) + + alpha_r = r_scheduler_output.alpha_t + sigma_r = r_scheduler_output.sigma_t + d_alpha_r = r_scheduler_output.d_alpha_t + d_sigma_r = r_scheduler_output.d_sigma_t + + t = self.original_scheduler.snr_inverse(alpha_r / sigma_r) + + t_scheduler_output = self.original_scheduler(t=t) + + alpha_t = t_scheduler_output.alpha_t + sigma_t = t_scheduler_output.sigma_t + d_alpha_t = t_scheduler_output.d_alpha_t + d_sigma_t = t_scheduler_output.d_sigma_t + + s_r = sigma_r / sigma_t + + dt_r = ( + sigma_t + * sigma_t + * (sigma_r * d_alpha_r - alpha_r * d_sigma_r) + / (sigma_r * sigma_r * (sigma_t * d_alpha_t - alpha_t * d_sigma_t)) + ) + + ds_r = (sigma_t * d_sigma_r - sigma_r * d_sigma_t * dt_r) / (sigma_t * sigma_t) + + u_t = self.model(x=x / s_r, t=t, **extras) + u_r = ds_r * x / s_r + dt_r * s_r * u_t + + return u_r diff --git a/flow_matching/path/scheduler/scheduler.py b/flow_matching/path/scheduler/scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..2055e95ce359c7aa6a7c2441362dc0756f3d900b --- /dev/null +++ b/flow_matching/path/scheduler/scheduler.py @@ -0,0 +1,199 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field + +from typing import Union + +import torch + +from torch import Tensor + + +@dataclass +class SchedulerOutput: + r"""Represents a sample of a conditional-flow generated probability path. + + Attributes: + alpha_t (Tensor): :math:`\alpha_t`, shape (...). + sigma_t (Tensor): :math:`\sigma_t`, shape (...). + d_alpha_t (Tensor): :math:`\frac{\partial}{\partial t}\alpha_t`, shape (...). + d_sigma_t (Tensor): :math:`\frac{\partial}{\partial t}\sigma_t`, shape (...). + + """ + + alpha_t: Tensor = field(metadata={"help": "alpha_t"}) + sigma_t: Tensor = field(metadata={"help": "sigma_t"}) + d_alpha_t: Tensor = field(metadata={"help": "Derivative of alpha_t."}) + d_sigma_t: Tensor = field(metadata={"help": "Derivative of sigma_t."}) + + +class Scheduler(ABC): + """Base Scheduler class.""" + + @abstractmethod + def __call__(self, t: Tensor) -> SchedulerOutput: + r""" + Args: + t (Tensor): times in [0,1], shape (...). + + Returns: + SchedulerOutput: :math:`\alpha_t,\sigma_t,\frac{\partial}{\partial t}\alpha_t,\frac{\partial}{\partial t}\sigma_t` + """ + ... + + @abstractmethod + def snr_inverse(self, snr: Tensor) -> Tensor: + r""" + Computes :math:`t` from the signal-to-noise ratio :math:`\frac{\alpha_t}{\sigma_t}`. + + Args: + snr (Tensor): The signal-to-noise, shape (...) + + Returns: + Tensor: t, shape (...) + """ + ... + + +class ConvexScheduler(Scheduler): + @abstractmethod + def __call__(self, t: Tensor) -> SchedulerOutput: + r"""Scheduler for convex paths. + + Args: + t (Tensor): times in [0,1], shape (...). + + Returns: + SchedulerOutput: :math:`\alpha_t,\sigma_t,\frac{\partial}{\partial t}\alpha_t,\frac{\partial}{\partial t}\sigma_t` + """ + ... + + @abstractmethod + def kappa_inverse(self, kappa: Tensor) -> Tensor: + r""" + Computes :math:`t` from :math:`\kappa_t`. + + Args: + kappa (Tensor): :math:`\kappa`, shape (...) + + Returns: + Tensor: t, shape (...) + """ + ... + + def snr_inverse(self, snr: Tensor) -> Tensor: + r""" + Computes :math:`t` from the signal-to-noise ratio :math:`\frac{\alpha_t}{\sigma_t}`. + + Args: + snr (Tensor): The signal-to-noise, shape (...) + + Returns: + Tensor: t, shape (...) + """ + kappa_t = snr / (1.0 + snr) + + return self.kappa_inverse(kappa=kappa_t) + + +class CondOTScheduler(ConvexScheduler): + """CondOT Scheduler.""" + + def __call__(self, t: Tensor) -> SchedulerOutput: + return SchedulerOutput( + alpha_t=t, + sigma_t=1 - t, + d_alpha_t=torch.ones_like(t), + d_sigma_t=-torch.ones_like(t), + ) + + def kappa_inverse(self, kappa: Tensor) -> Tensor: + return kappa + + +class PolynomialConvexScheduler(ConvexScheduler): + """Polynomial Scheduler.""" + + def __init__(self, n: Union[float, int]) -> None: + assert isinstance( + n, (float, int) + ), f"`n` must be a float or int. Got {type(n)=}." + assert n > 0, f"`n` must be positive. Got {n=}." + + self.n = n + + def __call__(self, t: Tensor) -> SchedulerOutput: + return SchedulerOutput( + alpha_t=t**self.n, + sigma_t=1 - t**self.n, + d_alpha_t=self.n * (t ** (self.n - 1)), + d_sigma_t=-self.n * (t ** (self.n - 1)), + ) + + def kappa_inverse(self, kappa: Tensor) -> Tensor: + return torch.pow(kappa, 1.0 / self.n) + + +class VPScheduler(Scheduler): + """Variance Preserving Scheduler.""" + + def __init__(self, beta_min: float = 0.1, beta_max: float = 20.0) -> None: + self.beta_min = beta_min + self.beta_max = beta_max + super().__init__() + + def __call__(self, t: Tensor) -> SchedulerOutput: + b = self.beta_min + B = self.beta_max + T = 0.5 * (1 - t) ** 2 * (B - b) + (1 - t) * b + dT = -(1 - t) * (B - b) - b + + return SchedulerOutput( + alpha_t=torch.exp(-0.5 * T), + sigma_t=torch.sqrt(1 - torch.exp(-T)), + d_alpha_t=-0.5 * dT * torch.exp(-0.5 * T), + d_sigma_t=0.5 * dT * torch.exp(-T) / torch.sqrt(1 - torch.exp(-T)), + ) + + def snr_inverse(self, snr: Tensor) -> Tensor: + T = -torch.log(snr**2 / (snr**2 + 1)) + b = self.beta_min + B = self.beta_max + t = 1 - ((-b + torch.sqrt(b**2 + 2 * (B - b) * T)) / (B - b)) + return t + + +class LinearVPScheduler(Scheduler): + """Linear Variance Preserving Scheduler.""" + + def __call__(self, t: Tensor) -> SchedulerOutput: + return SchedulerOutput( + alpha_t=t, + sigma_t=(1 - t**2) ** 0.5, + d_alpha_t=torch.ones_like(t), + d_sigma_t=-t / (1 - t**2) ** 0.5, + ) + + def snr_inverse(self, snr: Tensor) -> Tensor: + return torch.sqrt(snr**2 / (1 + snr**2)) + + +class CosineScheduler(Scheduler): + """Cosine Scheduler.""" + + def __call__(self, t: Tensor) -> SchedulerOutput: + pi = torch.pi + return SchedulerOutput( + alpha_t=torch.sin(pi / 2 * t), + sigma_t=torch.cos(pi / 2 * t), + d_alpha_t=pi / 2 * torch.cos(pi / 2 * t), + d_sigma_t=-pi / 2 * torch.sin(pi / 2 * t), + ) + + def snr_inverse(self, snr: Tensor) -> Tensor: + return 2.0 * torch.atan(snr) / torch.pi diff --git a/flow_matching/solver/.DS_Store b/flow_matching/solver/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..5008ddfcf53c02e82d7eee2e57c38e5672ef89f6 Binary files /dev/null and b/flow_matching/solver/.DS_Store differ diff --git a/flow_matching/solver/__init__.py b/flow_matching/solver/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0b33155d40c64c622a39271344ca61be546ff9d1 --- /dev/null +++ b/flow_matching/solver/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from .discrete_solver import MixtureDiscreteEulerSolver +from .ode_solver import ODESolver +from .riemannian_ode_solver import RiemannianODESolver +from .solver import Solver +from .editflows_solver import EditFlowsEulerSolver + +__all__ = [ + "ODESolver", + "Solver", + "ModelWrapper", + "MixtureDiscreteEulerSolver", + "RiemannianODESolver", + "EditFlowsEulerSolver", +] diff --git a/flow_matching/solver/discrete_solver.py b/flow_matching/solver/discrete_solver.py new file mode 100644 index 0000000000000000000000000000000000000000..8ca84cdb1210fcff09b9a603273b27881c53d308 --- /dev/null +++ b/flow_matching/solver/discrete_solver.py @@ -0,0 +1,260 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext +from math import ceil +from typing import Callable, Optional, Union + +import torch +from torch import Tensor + +from torch.nn import functional as F + +from flow_matching.path import MixtureDiscreteProbPath + +from flow_matching.solver.solver import Solver +from flow_matching.utils import categorical, ModelWrapper +from .utils import get_nearest_times + +try: + from tqdm import tqdm + + TQDM_AVAILABLE = True +except ImportError: + TQDM_AVAILABLE = False + + +class MixtureDiscreteEulerSolver(Solver): + r"""Solver that simulates the CTMC process :math:`(X_t)_{t_{\text{init}}\leq t\leq t_{\text{final}}}` defined by :math:`p_t` the marginal probability path of ``path``. + Given :math:`X_t \sim p_t`, the algorithm of solver step from :math:`t` to :math:`t+h` for the i-th coordinate is: + + .. math:: + + \begin{align*} + & X_1^i \sim p_{1|t}^i(\cdot|X_t)\\ + & \lambda^i \gets \sum_{x^i\ne X_t^i} u_t^i(x^i, X_t^i|X_1^i)\\ + & Z^i_{\text{change}} \sim U[0,1]\\ + & X_{t+h}^i \sim \begin{cases} + \frac{u_t^i(\cdot, X_t^i|X_1^i)}{\lambda^i}(1-\delta_{X_t^i}(\cdot)) \text{ if $Z^i_{\text{change}}\le 1-e^{-h\lambda^i}$}\\ + \delta_{X_t^i}(\cdot) \text{ else } + \end{cases} + \end{align*} + + Where :math:`p_{1|t}(\cdot|X_t)` is the output of ``model``, and the conditional probability velocity is of the mixture probability path is: + + .. math:: + + u_t^i(x^i, y^i|x_1^i) = \hat{u}_t^i(x^i, y^i|x_1^i) + c_{\text{div\_free}}\left[\hat{u}_t^i(x^i, y^i|x_1^i) - \check{u}_t^i(x^i, y^i|x_1^i) \right], + + where + + .. math:: + \hat{u}_t^i(x^i, y^i|x_1^i) = \frac{\dot{\kappa}_t}{1-\kappa_t} \left[ \delta_{x_1^i}(x^i) - \delta_{y^i}(x^i) \right], + + and + + .. math:: + + \check{u}_t^i(x^i, y^i|x_1^i) = \frac{\dot{\kappa}_t}{\kappa_t}\left[ \delta_{y^i}(x^i) - p(x^i) \right]. + + The source distribution :math:`p(x^i)` is given by ``p``. + + Args: + model (ModelWrapper): trained with x-prediction, outputting posterior probabilities (in the range :math:`[0,1]`), output must be [..., vocabulary_size]. + path (MixtureDiscreteProbPath): Probability path used for x-prediction training. + vocabulary_size (int): size of the discrete vocabulary. + source_distribution_p (Optional[Tensor], optional): Source distribution, must be of shape [vocabulary_size]. Required only when divergence-free term for the probability velocity is non-zero. Defaults to None. + """ + + def __init__( + self, + model: ModelWrapper, + path: MixtureDiscreteProbPath, + vocabulary_size: int, + source_distribution_p: Optional[Tensor] = None, + ): + super().__init__() + self.model = model + self.path = path + self.vocabulary_size = vocabulary_size + + if source_distribution_p is not None: + assert source_distribution_p.shape == torch.Size( + [vocabulary_size] + ), f"Source distribution p dimension must match the vocabulary size {vocabulary_size}. Got {source_distribution_p.shape}." + + self.source_distribution_p = source_distribution_p + + @torch.no_grad() + def sample( + self, + x_init: Tensor, + step_size: Optional[float], + div_free: Union[float, Callable[[float], float]] = 0.0, + dtype_categorical: torch.dtype = torch.float32, + time_grid: Tensor = torch.tensor([0.0, 1.0]), + return_intermediates: bool = False, + verbose: bool = False, + **model_extras, + ) -> Tensor: + """ + Sample a sequence of discrete values from the given model. + + .. code-block:: python + + import torch + from flow_matching.utils import ModelWrapper + from flow_matching.solver import MixtureDiscreteEulerSolver + + class DummyModel(ModelWrapper): + def __init__(self): + super().__init__(None) + def forward(self, x: torch.Tensor, t: torch.Tensor, **extras) -> torch.Tensor: + return ... + + model = DummyModel() + solver = MixtureDiscreteEulerSolver(model=model) + + x_init = torch.LongTensor([122, 725]) + step_size = 0.001 + time_grid = torch.tensor([0.0, 1.0]) + + result = solver.sample(x_init=x_init, step_size=step_size, time_grid=time_grid) + + Args: + x_init (Tensor): The initial state. + step_size (Optional[float]): If float then time discretization is uniform with the given step size. If None then time discretization is set to be time_grid. + div_free (Union[float, Callable[[float], float]]): The coefficient of the divergence-free term in the probability velocity. Can be either a float or a time dependent function. Defaults to 0.0. + dtype_categorical (torch.dtype): Precision to use for categorical sampler. Defaults to torch.float32. + time_grid (Tensor): The CTMC process is solved in the interval [time_grid[0], time_grid[-1]] and if step_size is None then time discretization is set by the time grid. Defaults to torch.tensor([0.0,1.0]). + return_intermediates (bool): If True then return intermediate time steps according to time_grid. Defaults to False. + verbose (bool): Whether to print progress bars. Defaults to False. + **model_extras: Additional input for the model. + + Returns: + Tensor: The sampled sequence of discrete values. + + Raises: + ImportError: To run in verbose mode, tqdm must be installed. + """ + if not div_free == 0.0: + assert ( + self.source_distribution_p is not None + ), "Source distribution p must be specified in order to add a divergence-free term to the probability velocity." + + # Initialize the current state `x_t` with the initial state `X_0`. + time_grid = time_grid.to(device=x_init.device) + + if step_size is None: + # If step_size is None then set the t discretization to time_grid. + t_discretization = time_grid + n_steps = len(time_grid) - 1 + else: + # If step_size is float then t discretization is uniform with step size set by step_size. + t_init = time_grid[0].item() + t_final = time_grid[-1].item() + assert ( + t_final - t_init + ) > step_size, f"Time interval [time_grid[0], time_grid[-1]] must be larger than step_size. Got a time interval [{t_init}, {t_final}] and step_size {step_size}." + + n_steps = ceil((t_final - t_init) / step_size) + t_discretization = torch.tensor( + [t_init + step_size * i for i in range(n_steps)] + [t_final], + device=x_init.device, + ) + + if return_intermediates: + # get order of intermediate steps: + order = torch.argsort(time_grid) + # Compute intermediate steps to return via nearest points in t_discretization to time_grid. + time_grid = get_nearest_times( + time_grid=time_grid, t_discretization=t_discretization + ) + + x_t = x_init.clone() + steps_counter = 0 + res = [] + + if return_intermediates: + res = [x_init.clone()] + + if verbose: + if not TQDM_AVAILABLE: + raise ImportError( + "tqdm is required for verbose mode. Please install it." + ) + ctx = tqdm(total=t_final, desc=f"NFE: {steps_counter}") + else: + ctx = nullcontext() + + with ctx: + for i in range(n_steps): + t = t_discretization[i : i + 1] + h = t_discretization[i + 1 : i + 2] - t_discretization[i : i + 1] + + # Sample x_1 ~ p_1|t( \cdot |x_t) + p_1t = self.model(x=x_t, t=t.repeat(x_t.shape[0]), **model_extras) + x_1 = categorical(p_1t.to(dtype=dtype_categorical)) + + # Checks if final step + if i == n_steps - 1: + x_t = x_1 + else: + # Compute u_t(x|x_t,x_1) + scheduler_output = self.path.scheduler(t=t) + + k_t = scheduler_output.alpha_t + d_k_t = scheduler_output.d_alpha_t + + delta_1 = F.one_hot(x_1, num_classes=self.vocabulary_size).to( + k_t.dtype + ) + u = d_k_t / (1 - k_t) * delta_1 + + # Add divergence-free part + div_free_t = div_free(t) if callable(div_free) else div_free + + if div_free_t > 0: + p_0 = self.source_distribution_p[(None,) * x_t.dim()] + u = u + div_free_t * d_k_t / (k_t * (1 - k_t)) * ( + (1 - k_t) * p_0 + k_t * delta_1 + ) + + # Set u_t(x_t|x_t,x_1) = 0 + delta_t = F.one_hot(x_t, num_classes=self.vocabulary_size) + u = torch.where( + delta_t.to(dtype=torch.bool), torch.zeros_like(u), u + ) + + # Sample x_t ~ u_t( \cdot |x_t,x_1) + intensity = u.sum(dim=-1) # Assuming u_t(xt|xt,x1) := 0 + mask_jump = torch.rand( + size=x_t.shape, device=x_t.device + ) < 1 - torch.exp(-h * intensity) + + if mask_jump.sum() > 0: + x_t[mask_jump] = categorical( + u[mask_jump].to(dtype=dtype_categorical) + ) + + steps_counter += 1 + t = t + h + + if return_intermediates and (t in time_grid): + res.append(x_t.clone()) + + if verbose: + ctx.n = t.item() + ctx.refresh() + ctx.set_description(f"NFE: {steps_counter}") + + if return_intermediates: + if step_size is None: + return torch.stack(res, dim=0) + else: + return torch.stack(res, dim=0)[order] + else: + return x_t diff --git a/flow_matching/solver/editflows_solver.py b/flow_matching/solver/editflows_solver.py new file mode 100644 index 0000000000000000000000000000000000000000..e6189cead9770f045ccfadeed651ab085ae15fd4 --- /dev/null +++ b/flow_matching/solver/editflows_solver.py @@ -0,0 +1,168 @@ +# Copyright (c) Meta Platforms, Inc. +# All rights reserved. + +from typing import Callable, List, Sequence, Tuple, Optional + +import torch +import torch.nn.functional as F +from torch import Tensor + + +class EditFlowsEulerSolver: + """ + Ragged EditFlows CTMC solver with Euler thinning (≤1 jump per step). + + At each step h: + - query model for (λ_ins[slots], Q_ins[slots], λ_del[tokens], λ_sub[tokens], Q_sub[tokens]) + - total intensity Λ = sum(λ_ins) + sum(λ_del) + sum(λ_sub) + - with prob 1 - exp(-h Λ), take ONE jump: + * sample which event from concatenated rates + * if ins/sub: sample token from the Q at that position/slot + * apply edit to the ragged sequence (insert/delete/replace) + - advance t ← t + h + + This is faithful to the CTMC in the small-h regime used by Euler schemes. + """ + + def __init__( + self, + model, # callable: (x_list, t_vec) -> ragged heads + scheduler: Callable[[Tensor], object], # path.scheduler(t) providing alpha_t, d_alpha_t, sigma_t if needed + vocab_size: int, + dtype_categorical: torch.dtype = torch.float64, + generator: Optional[torch.Generator] = None, + ) -> None: + self.model = model + self.scheduler = scheduler + self.vocab_size = vocab_size + self.dtype_cat = dtype_categorical + self.gen = generator + + @torch.no_grad() + def _sample_categorical_logits(self, logits: Tensor) -> Tensor: + # logits: (V,) -> sample id + probs = F.softmax(logits, dim=-1, dtype=self.dtype_cat) + return torch.multinomial(probs, num_samples=1, replacement=True, generator=self.gen).squeeze(0) + + @torch.no_grad() + def _maybe_one_jump( + self, + x: Tensor, # (n,) + lam_ins: Tensor, # (n+1,) + logits_ins: Tensor, # (n+1, V) + lam_del: Tensor, # (n,) + lam_sub: Tensor, # (n,) + logits_sub: Tensor, # (n, V) + h: float, + device: torch.device, + ) -> Tensor: + """ + Apply at most one edit event to x with step size h. + Returns possibly modified x (1-D LongTensor). + """ + n = x.numel() + # Total intensity + sum_ins = lam_ins.sum() if lam_ins.numel() else torch.tensor(0.0, device=device) + sum_del = lam_del.sum() if lam_del.numel() else torch.tensor(0.0, device=device) + sum_sub = lam_sub.sum() if lam_sub.numel() else torch.tensor(0.0, device=device) + total = (sum_ins + sum_del + sum_sub).clamp(min=0.0) + + if total.item() <= 0: + return x # no possible jumps + + # Bernoulli: does a jump occur? + p_jump = 1.0 - torch.exp(torch.tensor(-h, device=device) * total) + if torch.rand((), device=device, generator=self.gen) >= p_jump: + return x # no jump this step + + # Build event vector: [ins(0..n), del(0..n-1), sub(0..n-1)] + parts = [] + if lam_ins.numel(): parts.append(lam_ins) + if lam_del.numel(): parts.append(lam_del) + if lam_sub.numel(): parts.append(lam_sub) + lam_cat = torch.cat(parts, dim=0) # (n+1 + n + n,) + probs = (lam_cat / total).to(self.dtype_cat) + + # Sample which event + e_idx = torch.multinomial(probs, num_samples=1, generator=self.gen).item() + + # Decode which branch + offset = 0 + if e_idx < lam_ins.numel(): + # INSERT at slot s + s = e_idx + # sample token from slot distribution + y = self._sample_categorical_logits(logits_ins[s]) # id ∈ [0..V-1] + # perform insertion at slot s: x' = [x[:s], y, x[s:]] + if s == 0: + x_new = torch.cat([y.view(1), x], dim=0) + elif s == n: + x_new = torch.cat([x, y.view(1)], dim=0) + else: + x_new = torch.cat([x[:s], y.view(1), x[s:]], dim=0) + return x_new + + e_idx -= lam_ins.numel() + if e_idx < lam_del.numel(): + # DELETE token j + j = e_idx + if n == 0: + return x + if j == 0: + return x[1:] + elif j == n - 1: + return x[:-1] + else: + return torch.cat([x[:j], x[j + 1 :]], dim=0) + + e_idx -= lam_del.numel() + # SUBSTITUTE at token j + j = e_idx + if n == 0: + return x + y = self._sample_categorical_logits(logits_sub[j]) + x_new = x.clone() + x_new[j] = y + return x_new + + @torch.no_grad() + def sample( + self, + x_list: Sequence[Tensor], # ragged: list[(n_i,)] + n_steps: int, + t0: float = 0.0, + t1: float = 1.0, + verbose: bool = False, + ) -> List[Tensor]: + """ + Evolve all sequences independently with shared time grid t_k. + """ + assert n_steps >= 1 + B = len(x_list) + device = x_list[0].device if B > 0 else torch.device("cpu") + + x = [xi.clone() for xi in x_list] + t = torch.full((B,), float(t0), device=device) + h = (t1 - t0) / float(n_steps) + + for k in range(n_steps): + # Query model once at time t_k (vector of size B) + lam_ins, logits_ins, lam_del, lam_sub, logits_sub = self.model(x, t) + + # Apply at most one jump per sample + for i in range(B): + x[i] = self._maybe_one_jump( + x=x[i], + lam_ins=lam_ins[i].squeeze(0) if lam_ins[i].dim() == 2 else lam_ins[i], + logits_ins=logits_ins[i].squeeze(0) if logits_ins[i].dim() == 3 else logits_ins[i], + lam_del=lam_del[i].squeeze(0) if lam_del[i].dim() == 2 else lam_del[i], + lam_sub=lam_sub[i].squeeze(0) if lam_sub[i].dim() == 2 else lam_sub[i], + logits_sub=logits_sub[i].squeeze(0) if logits_sub[i].dim() == 3 else logits_sub[i], + h=h, + device=device, + ) + + # advance time + t = t + h + + return x diff --git a/flow_matching/solver/ode_solver.py b/flow_matching/solver/ode_solver.py new file mode 100644 index 0000000000000000000000000000000000000000..2ee8a17b23f726bdb04b5d8bc1e5b89ca8b65d52 --- /dev/null +++ b/flow_matching/solver/ode_solver.py @@ -0,0 +1,197 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Callable, Optional, Sequence, Tuple, Union + +import torch +from torch import Tensor +from torchdiffeq import odeint + +from flow_matching.solver.solver import Solver +from flow_matching.utils import gradient, ModelWrapper + + +class ODESolver(Solver): + """A class to solve ordinary differential equations (ODEs) using a specified velocity model. + + This class utilizes a velocity field model to solve ODEs over a given time grid using numerical ode solvers. + + Args: + velocity_model (Union[ModelWrapper, Callable]): a velocity field model receiving :math:`(x,t)` and returning :math:`u_t(x)` + """ + + def __init__(self, velocity_model: Union[ModelWrapper, Callable]): + super().__init__() + self.velocity_model = velocity_model + + def sample( + self, + x_init: Tensor, + step_size: Optional[float], + method: str = "euler", + atol: float = 1e-5, + rtol: float = 1e-5, + time_grid: Tensor = torch.tensor([0.0, 1.0]), + return_intermediates: bool = False, + enable_grad: bool = False, + **model_extras, + ) -> Union[Tensor, Sequence[Tensor]]: + r"""Solve the ODE with the velocity field. + + Example: + + .. code-block:: python + + import torch + from flow_matching.utils import ModelWrapper + from flow_matching.solver import ODESolver + + class DummyModel(ModelWrapper): + def __init__(self): + super().__init__(None) + + def forward(self, x: torch.Tensor, t: torch.Tensor, **extras) -> torch.Tensor: + return torch.ones_like(x) * 3.0 * t**2 + + velocity_model = DummyModel() + solver = ODESolver(velocity_model=velocity_model) + x_init = torch.tensor([0.0, 0.0]) + step_size = 0.001 + time_grid = torch.tensor([0.0, 1.0]) + + result = solver.sample(x_init=x_init, step_size=step_size, time_grid=time_grid) + + Args: + x_init (Tensor): initial conditions (e.g., source samples :math:`X_0 \sim p`). Shape: [batch_size, ...]. + step_size (Optional[float]): The step size. Must be None for adaptive step solvers. + method (str): A method supported by torchdiffeq. Defaults to "euler". Other commonly used solvers are "dopri5", "midpoint" and "heun3". For a complete list, see torchdiffeq. + atol (float): Absolute tolerance, used for adaptive step solvers. + rtol (float): Relative tolerance, used for adaptive step solvers. + time_grid (Tensor): The process is solved in the interval [min(time_grid, max(time_grid)] and if step_size is None then time discretization is set by the time grid. May specify a descending time_grid to solve in the reverse direction. Defaults to torch.tensor([0.0, 1.0]). + return_intermediates (bool, optional): If True then return intermediate time steps according to time_grid. Defaults to False. + enable_grad (bool, optional): Whether to compute gradients during sampling. Defaults to False. + **model_extras: Additional input for the model. + + Returns: + Union[Tensor, Sequence[Tensor]]: The last timestep when return_intermediates=False, otherwise all values specified in time_grid. + """ + + time_grid = time_grid.to(x_init.device) + + def ode_func(t, x): + return self.velocity_model(x=x, t=t, **model_extras) + + ode_opts = {"step_size": step_size} if step_size is not None else {} + + with torch.set_grad_enabled(enable_grad): + # Approximate ODE solution with numerical ODE solver + sol = odeint( + ode_func, + x_init, + time_grid, + method=method, + options=ode_opts, + atol=atol, + rtol=rtol, + ) + + if return_intermediates: + return sol + else: + return sol[-1] + + def compute_likelihood( + self, + x_1: Tensor, + log_p0: Callable[[Tensor], Tensor], + step_size: Optional[float], + method: str = "euler", + atol: float = 1e-5, + rtol: float = 1e-5, + time_grid: Tensor = torch.tensor([1.0, 0.0]), + return_intermediates: bool = False, + exact_divergence: bool = False, + enable_grad: bool = False, + **model_extras, + ) -> Union[Tuple[Tensor, Tensor], Tuple[Sequence[Tensor], Tensor]]: + r"""Solve for log likelihood given a target sample at :math:`t=0`. + + Works similarly to sample, but solves the ODE in reverse to compute the log-likelihood. The velocity model must be differentiable with respect to x. + The function assumes log_p0 is the log probability of the source distribution at :math:`t=0`. + + Args: + x_1 (Tensor): target sample (e.g., samples :math:`X_1 \sim p_1`). + log_p0 (Callable[[Tensor], Tensor]): Log probability function of the source distribution. + step_size (Optional[float]): The step size. Must be None for adaptive step solvers. + method (str): A method supported by torchdiffeq. Defaults to "euler". Other commonly used solvers are "dopri5", "midpoint" and "heun3". For a complete list, see torchdiffeq. + atol (float): Absolute tolerance, used for adaptive step solvers. + rtol (float): Relative tolerance, used for adaptive step solvers. + time_grid (Tensor): If step_size is None then time discretization is set by the time grid. Must start at 1.0 and end at 0.0, otherwise the likelihood computation is not valid. Defaults to torch.tensor([1.0, 0.0]). + return_intermediates (bool, optional): If True then return intermediate time steps according to time_grid. Otherwise only return the final sample. Defaults to False. + exact_divergence (bool): Whether to compute the exact divergence or use the Hutchinson estimator. + enable_grad (bool, optional): Whether to compute gradients during sampling. Defaults to False. + **model_extras: Additional input for the model. + + Returns: + Union[Tuple[Tensor, Tensor], Tuple[Sequence[Tensor], Tensor]]: Samples at time_grid and log likelihood values of given x_1. + """ + assert ( + time_grid[0] == 1.0 and time_grid[-1] == 0.0 + ), f"Time grid must start at 1.0 and end at 0.0. Got {time_grid}" + + # Fix the random projection for the Hutchinson divergence estimator + if not exact_divergence: + z = (torch.randn_like(x_1).to(x_1.device) < 0) * 2.0 - 1.0 + + def ode_func(x, t): + return self.velocity_model(x=x, t=t, **model_extras) + + def dynamics_func(t, states): + xt = states[0] + with torch.set_grad_enabled(True): + xt.requires_grad_() + ut = ode_func(xt, t) + + if exact_divergence: + # Compute exact divergence + div = 0 + for i in range(ut.flatten(1).shape[1]): + div += gradient(ut[:, i], xt, create_graph=True)[:, i].detach() + else: + # Compute Hutchinson divergence estimator E[z^T D_x(ut) z] + ut_dot_z = torch.einsum( + "ij,ij->i", ut.flatten(start_dim=1), z.flatten(start_dim=1) + ) + grad_ut_dot_z = gradient(ut_dot_z, xt) + div = torch.einsum( + "ij,ij->i", + grad_ut_dot_z.flatten(start_dim=1), + z.flatten(start_dim=1), + ) + + return ut.detach(), div.detach() + + y_init = (x_1, torch.zeros(x_1.shape[0], device=x_1.device)) + ode_opts = {"step_size": step_size} if step_size is not None else {} + + with torch.set_grad_enabled(enable_grad): + sol, log_det = odeint( + dynamics_func, + y_init, + time_grid, + method=method, + options=ode_opts, + atol=atol, + rtol=rtol, + ) + + x_source = sol[-1] + source_log_p = log_p0(x_source) + + if return_intermediates: + return sol, source_log_p + log_det[-1] + else: + return sol[-1], source_log_p + log_det[-1] diff --git a/flow_matching/solver/riemannian_ode_solver.py b/flow_matching/solver/riemannian_ode_solver.py new file mode 100644 index 0000000000000000000000000000000000000000..a6ff0fd84b0eb27e25e4b28fabb56c63b3ca5602 --- /dev/null +++ b/flow_matching/solver/riemannian_ode_solver.py @@ -0,0 +1,261 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +import math +from typing import Callable + +import torch +from torch import Tensor + +from flow_matching.solver.solver import Solver +from flow_matching.utils import ModelWrapper +from flow_matching.utils.manifolds import geodesic, Manifold + +try: + from tqdm import tqdm + + TQDM_AVAILABLE = True +except ImportError: + TQDM_AVAILABLE = False + + +class RiemannianODESolver(Solver): + r"""Riemannian ODE solver + Initialize the ``RiemannianODESolver``. + + Args: + manifold (Manifold): the manifold to solve on. + velocity_model (ModelWrapper): a velocity field model receiving :math:`(x,t)` + and returning :math:`u_t(x)` which is assumed to lie on the tangent plane at `x`. + """ + + def __init__(self, manifold: Manifold, velocity_model: ModelWrapper): + super().__init__() + self.manifold = manifold + self.velocity_model = velocity_model + + def sample( + self, + x_init: Tensor, + step_size: float, + projx: bool = True, + proju: bool = True, + method: str = "euler", + time_grid: Tensor = torch.tensor([0.0, 1.0]), + return_intermediates: bool = False, + verbose: bool = False, + enable_grad: bool = False, + **model_extras, + ) -> Tensor: + r"""Solve the ODE with the `velocity_field` on the manifold. + + Args: + x_init (Tensor): initial conditions (e.g., source samples :math:`X_0 \sim p`). + step_size (float): The step size. + projx (bool): Whether to project the point onto the manifold at each step. Defaults to True. + proju (bool): Whether to project the vector field onto the tangent plane at each step. Defaults to True. + method (str): One of ["euler", "midpoint", "rk4"]. Defaults to "euler". + time_grid (Tensor, optional): The process is solved in the interval [min(time_grid, max(time_grid)] and if step_size is None then time discretization is set by the time grid. Defaults to torch.tensor([0.0,1.0]). + return_intermediates (bool, optional): If True then return intermediate time steps according to time_grid. Defaults to False. + verbose (bool, optional): Whether to print progress bars. Defaults to False. + enable_grad (bool, optional): Whether to compute gradients during sampling. Defaults to False. + **model_extras: Additional input for the model. + + Returns: + Tensor: The sampled sequence. Defaults to returning samples at :math:`t=1`. + + Raises: + ImportError: To run in verbose mode, tqdm must be installed. + """ + step_fns = { + "euler": _euler_step, + "midpoint": _midpoint_step, + "rk4": _rk4_step, + } + assert method in step_fns.keys(), f"Unknown method {method}" + step_fn = step_fns[method] + + def velocity_func(x, t): + return self.velocity_model(x=x, t=t, **model_extras) + + # --- Factor this out. + time_grid = torch.sort(time_grid.to(device=x_init.device)).values + + if step_size is None: + # If step_size is None then set the t discretization to time_grid. + t_discretization = time_grid + n_steps = len(time_grid) - 1 + else: + # If step_size is float then t discretization is uniform with step size set by step_size. + t_init = time_grid[0].item() + t_final = time_grid[-1].item() + assert ( + t_final - t_init + ) > step_size, f"Time interval [min(time_grid), max(time_grid)] must be larger than step_size. Got a time interval [{t_init}, {t_final}] and step_size {step_size}." + + n_steps = math.ceil((t_final - t_init) / step_size) + t_discretization = torch.tensor( + [step_size * i for i in range(n_steps)] + [t_final], + device=x_init.device, + ) + # --- + t0s = t_discretization[:-1] + + if verbose: + if not TQDM_AVAILABLE: + raise ImportError( + "tqdm is required for verbose mode. Please install it." + ) + t0s = tqdm(t0s) + + if return_intermediates: + xts = [] + i_ret = 0 + + with torch.set_grad_enabled(enable_grad): + xt = x_init + for t0, t1 in zip(t0s, t_discretization[1:]): + dt = t1 - t0 + xt_next = step_fn( + velocity_func, + xt, + t0, + dt, + manifold=self.manifold, + projx=projx, + proju=proju, + ) + if return_intermediates: + while ( + i_ret < len(time_grid) + and t0 <= time_grid[i_ret] + and time_grid[i_ret] <= t1 + ): + xts.append( + interp(self.manifold, xt, xt_next, t0, t1, time_grid[i_ret]) + ) + i_ret += 1 + xt = xt_next + + if return_intermediates: + return torch.stack(xts, dim=0) + else: + return xt + + +def interp(manifold, xt, xt_next, t, t_next, t_ret): + return geodesic(manifold, xt, xt_next)( + (t_ret - t) / (t_next - t).reshape(1) + ).reshape_as(xt) + + +def _euler_step( + velocity_model: Callable, + xt: Tensor, + t0: Tensor, + dt: Tensor, + manifold: Manifold, + projx: bool = True, + proju: bool = True, +) -> Tensor: + r"""Perform an Euler step on a manifold. + + Args: + velocity_model (Callable): the velocity model + xt (Tensor): tensor containing the state at time t0 + t0 (Tensor): the time at which this step is taken + dt (Tensor): the step size + manifold (Manifold): a manifold object + projx (bool, optional): whether to project the state onto the manifold. Defaults to True. + proju (bool, optional): whether to project the velocity onto the tangent plane. Defaults to True. + + Returns: + Tensor: tensor containing the state after the step + """ + velocity_fn = lambda x, t: ( + manifold.proju(x, velocity_model(x, t)) if proju else velocity_model(x, t) + ) + projx_fn = lambda x: manifold.projx(x) if projx else x + + vt = velocity_fn(xt, t0) + + xt = xt + dt * vt + + return projx_fn(xt) + + +def _midpoint_step( + velocity_model: Callable, + xt: Tensor, + t0: Tensor, + dt: Tensor, + manifold: Manifold, + projx: bool = True, + proju: bool = True, +) -> Tensor: + r"""Perform a midpoint step on a manifold. + + Args: + velocity_model (Callable): the velocity model + xt (Tensor): tensor containing the state at time t0 + t0 (Tensor): the time at which this step is taken + dt (Tensor): the step size + manifold (Manifold): a manifold object + projx (bool, optional): whether to project the state onto the manifold. Defaults to True. + proju (bool, optional): whether to project the velocity onto the tangent plane. Defaults to True. + + Returns: + Tensor: tensor containing the state after the step + """ + velocity_fn = lambda x, t: ( + manifold.proju(x, velocity_model(x, t)) if proju else velocity_model(x, t) + ) + projx_fn = lambda x: manifold.projx(x) if projx else x + + half_dt = 0.5 * dt + vt = velocity_fn(xt, t0) + x_mid = xt + half_dt * vt + x_mid = projx_fn(x_mid) + + xt = xt + dt * velocity_fn(x_mid, t0 + half_dt) + + return projx_fn(xt) + + +def _rk4_step( + velocity_model: Callable, + xt: Tensor, + t0: Tensor, + dt: Tensor, + manifold: Manifold, + projx: bool = True, + proju: bool = True, +) -> Tensor: + r"""Perform an RK4 step on a manifold. + + Args: + velocity_model (Callable): the velocity model + xt (Tensor): tensor containing the state at time t0 + t0 (Tensor): the time at which this step is taken + dt (Tensor): the step size + manifold (Manifold): a manifold object + projx (bool, optional): whether to project the state onto the manifold. Defaults to True. + proju (bool, optional): whether to project the velocity onto the tangent plane. Defaults to True. + + Returns: + Tensor: tensor containing the state after the step + """ + velocity_fn = lambda x, t: ( + manifold.proju(x, velocity_model(x, t)) if proju else velocity_model(x, t) + ) + projx_fn = lambda x: manifold.projx(x) if projx else x + + k1 = velocity_fn(xt, t0) + k2 = velocity_fn(projx_fn(xt + dt * k1 / 3), t0 + dt / 3) + k3 = velocity_fn(projx_fn(xt + dt * (k2 - k1 / 3)), t0 + dt * 2 / 3) + k4 = velocity_fn(projx_fn(xt + dt * (k1 - k2 + k3)), t0 + dt) + + return projx_fn(xt + (k1 + 3 * (k2 + k3) + k4) * dt * 0.125) diff --git a/flow_matching/solver/solver.py b/flow_matching/solver/solver.py new file mode 100644 index 0000000000000000000000000000000000000000..4819e1c0ff148c411648c6e248380f471efd0a6c --- /dev/null +++ b/flow_matching/solver/solver.py @@ -0,0 +1,17 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from abc import ABC, abstractmethod + +from torch import nn, Tensor + + +class Solver(ABC, nn.Module): + """Abstract base class for solvers.""" + + @abstractmethod + def sample(self, x_0: Tensor = None) -> Tensor: + ... diff --git a/flow_matching/solver/utils.py b/flow_matching/solver/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f3a34ee4b0ca4c3ef64dfa1b55cb13da1187b965 --- /dev/null +++ b/flow_matching/solver/utils.py @@ -0,0 +1,19 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch import Tensor + + +def get_nearest_times(time_grid: Tensor, t_discretization: Tensor) -> Tensor: + distances = torch.cdist( + time_grid.unsqueeze(1), + t_discretization.unsqueeze(1), + compute_mode="donot_use_mm_for_euclid_dist", + ) + nearest_indices = distances.argmin(dim=1) + + return t_discretization[nearest_indices] diff --git a/flow_matching/utils/.DS_Store b/flow_matching/utils/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..1765e53704c05f68c7324afe4446fd17bfc92089 Binary files /dev/null and b/flow_matching/utils/.DS_Store differ diff --git a/flow_matching/utils/__init__.py b/flow_matching/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0085c44c5e044d5a7dd0c2bde9755e9d3c18c0b1 --- /dev/null +++ b/flow_matching/utils/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from .categorical_sampler import categorical +from .model_wrapper import ModelWrapper +from .utils import expand_tensor_like, gradient, unsqueeze_to_match + +__all__ = [ + "unsqueeze_to_match", + "expand_tensor_like", + "gradient", + "categorical", + "ModelWrapper", +] diff --git a/flow_matching/utils/categorical_sampler.py b/flow_matching/utils/categorical_sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..70937af5af1da4c3c2b18a85745197e240724a96 --- /dev/null +++ b/flow_matching/utils/categorical_sampler.py @@ -0,0 +1,23 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch import Tensor + + +def categorical(probs: Tensor) -> Tensor: + r"""Categorical sampler according to weights in the last dimension of ``probs`` using :func:`torch.multinomial`. + + Args: + probs (Tensor): probabilities. + + Returns: + Tensor: Samples. + """ + + return torch.multinomial(probs.flatten(0, -2), 1, replacement=True).view( + *probs.shape[:-1] + ) diff --git a/flow_matching/utils/manifolds/.DS_Store b/flow_matching/utils/manifolds/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..5008ddfcf53c02e82d7eee2e57c38e5672ef89f6 Binary files /dev/null and b/flow_matching/utils/manifolds/.DS_Store differ diff --git a/flow_matching/utils/manifolds/__init__.py b/flow_matching/utils/manifolds/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1148872f29c5e42e2578eab31751abaf6def51da --- /dev/null +++ b/flow_matching/utils/manifolds/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from .manifold import Euclidean, Manifold +from .sphere import Sphere +from .torus import FlatTorus +from .utils import geodesic + +__all__ = [ + "Euclidean", + "Manifold", + "Sphere", + "FlatTorus", + "geodesic", +] diff --git a/flow_matching/utils/manifolds/manifold.py b/flow_matching/utils/manifolds/manifold.py new file mode 100644 index 0000000000000000000000000000000000000000..52a6a1bc6ea7d83e3aeefd335edb19a842aeca21 --- /dev/null +++ b/flow_matching/utils/manifolds/manifold.py @@ -0,0 +1,93 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +import abc + +import torch.nn as nn +from torch import Tensor + + +class Manifold(nn.Module, metaclass=abc.ABCMeta): + """A manifold class that contains projection operations and logarithm and exponential maps.""" + + @abc.abstractmethod + def expmap(self, x: Tensor, u: Tensor) -> Tensor: + r"""Computes exponential map :math:`\exp_x(u)`. + + Args: + x (Tensor): point on the manifold + u (Tensor): tangent vector at point :math:`x` + + Raises: + NotImplementedError: if not implemented + + Returns: + Tensor: transported point + """ + raise NotImplementedError + + @abc.abstractmethod + def logmap(self, x: Tensor, y: Tensor) -> Tensor: + r"""Computes logarithmic map :math:`\log_x(y)`. + + Args: + x (Tensor): point on the manifold + y (Tensor): point on the manifold + + Raises: + NotImplementedError: if not implemented + + Returns: + Tensor: tangent vector at point :math:`x` + """ + raise NotImplementedError + + @abc.abstractmethod + def projx(self, x: Tensor) -> Tensor: + """Project point :math:`x` on the manifold. + + Args: + x (Tensor): point to be projected + + Raises: + NotImplementedError: if not implemented + + Returns: + Tensor: projected point on the manifold + """ + raise NotImplementedError + + @abc.abstractmethod + def proju(self, x: Tensor, u: Tensor) -> Tensor: + """Project vector :math:`u` on a tangent space for :math:`x`. + + Args: + x (Tensor): point on the manifold + u (Tensor): vector to be projected + + Raises: + NotImplementedError: if not implemented + + Returns: + Tensor: projected tangent vector + """ + raise NotImplementedError + + +class Euclidean(Manifold): + """The Euclidean manifold.""" + + def expmap(self, x: Tensor, u: Tensor) -> Tensor: + return x + u + + def logmap(self, x: Tensor, y: Tensor) -> Tensor: + return y - x + + def projx(self, x: Tensor) -> Tensor: + return x + + def proju(self, x: Tensor, u: Tensor) -> Tensor: + return u diff --git a/flow_matching/utils/manifolds/sphere.py b/flow_matching/utils/manifolds/sphere.py new file mode 100644 index 0000000000000000000000000000000000000000..76bf748f54b055b10cd4c4a085bf8204df6f47a0 --- /dev/null +++ b/flow_matching/utils/manifolds/sphere.py @@ -0,0 +1,45 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch import Tensor + +from flow_matching.utils.manifolds import Manifold + + +class Sphere(Manifold): + """Represents a hyperpshere in :math:`R^D`. Isometric to the product of 1-D spheres.""" + + EPS = {torch.float32: 1e-4, torch.float64: 1e-7} + + def expmap(self, x: Tensor, u: Tensor) -> Tensor: + norm_u = u.norm(dim=-1, keepdim=True) + exp = x * torch.cos(norm_u) + u * torch.sin(norm_u) / norm_u + retr = self.projx(x + u) + cond = norm_u > self.EPS[norm_u.dtype] + + return torch.where(cond, exp, retr) + + def logmap(self, x: Tensor, y: Tensor) -> Tensor: + u = self.proju(x, y - x) + dist = self.dist(x, y, keepdim=True) + cond = dist.gt(self.EPS[x.dtype]) + result = torch.where( + cond, + u * dist / u.norm(dim=-1, keepdim=True).clamp_min(self.EPS[x.dtype]), + u, + ) + return result + + def projx(self, x: Tensor) -> Tensor: + return x / x.norm(dim=-1, keepdim=True) + + def proju(self, x: Tensor, u: Tensor) -> Tensor: + return u - (x * u).sum(dim=-1, keepdim=True) * x + + def dist(self, x: Tensor, y: Tensor, *, keepdim=False) -> Tensor: + inner = (x * y).sum(-1, keepdim=keepdim) + return torch.acos(inner) diff --git a/flow_matching/utils/manifolds/torus.py b/flow_matching/utils/manifolds/torus.py new file mode 100644 index 0000000000000000000000000000000000000000..3587ed7567e3f55bb28edb176ab042882048d858 --- /dev/null +++ b/flow_matching/utils/manifolds/torus.py @@ -0,0 +1,28 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +import math + +import torch +from torch import Tensor + +from flow_matching.utils.manifolds import Manifold + + +class FlatTorus(Manifold): + r"""Represents a flat torus on the :math:`[0, 2\pi]^D` subspace. Isometric to the product of 1-D spheres.""" + + def expmap(self, x: Tensor, u: Tensor) -> Tensor: + return (x + u) % (2 * math.pi) + + def logmap(self, x: Tensor, y: Tensor) -> Tensor: + return torch.atan2(torch.sin(y - x), torch.cos(y - x)) + + def projx(self, x: Tensor) -> Tensor: + return x % (2 * math.pi) + + def proju(self, x: Tensor, u: Tensor) -> Tensor: + return u diff --git a/flow_matching/utils/manifolds/utils.py b/flow_matching/utils/manifolds/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b83d2faf8ac4f182ca8621f25cafaefca32f5091 --- /dev/null +++ b/flow_matching/utils/manifolds/utils.py @@ -0,0 +1,45 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Callable + +import torch +from torch import Tensor + +from flow_matching.utils.manifolds import Manifold + + +def geodesic( + manifold: Manifold, start_point: Tensor, end_point: Tensor +) -> Callable[[Tensor], Tensor]: + """Generate parameterized function for geodesic curve. + + Args: + manifold (Manifold): the manifold to compute geodesic on. + start_point (Tensor): point on the manifold at :math:`t=0`. + end_point (Tensor): point on the manifold at :math:`t=1`. + + Returns: + Callable[[Tensor], Tensor]: a function that takes in :math:`t` and outputs the geodesic at time :math:`t`. + """ + + shooting_tangent_vec = manifold.logmap(start_point, end_point) + + def path(t: Tensor) -> Tensor: + """Generate parameterized function for geodesic curve. + + Args: + t (Tensor): Times at which to compute points of the geodesics. + + Returns: + Tensor: geodesic path evaluated at time t. + """ + tangent_vecs = torch.einsum("i,...k->...ik", t, shooting_tangent_vec) + points_at_time_t = manifold.expmap(start_point.unsqueeze(-2), tangent_vecs) + + return points_at_time_t + + return path diff --git a/flow_matching/utils/model_wrapper.py b/flow_matching/utils/model_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..ac7d932c851b3d2194bd63ce597ef924dc819a7f --- /dev/null +++ b/flow_matching/utils/model_wrapper.py @@ -0,0 +1,43 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from abc import ABC + +from torch import nn, Tensor + + +class ModelWrapper(ABC, nn.Module): + """ + This class is used to wrap around another model, adding custom forward pass logic. + """ + + def __init__(self, model: nn.Module): + super().__init__() + self.model = model + + def forward(self, x: Tensor, t: Tensor, **extras) -> Tensor: + r""" + This method defines how inputs should be passed through the wrapped model. + Here, we're assuming that the wrapped model takes both :math:`x` and :math:`t` as input, + along with any additional keyword arguments. + + Optional things to do here: + - check that t is in the dimensions that the model is expecting. + - add a custom forward pass logic. + - call the wrapped model. + + | given x, t + | returns the model output for input x at time t, with extra information `extra`. + + Args: + x (Tensor): input data to the model (batch_size, ...). + t (Tensor): time (batch_size). + **extras: additional information forwarded to the model, e.g., text condition. + + Returns: + Tensor: model output. + """ + return self.model(x=x, t=t, **extras) diff --git a/flow_matching/utils/utils.py b/flow_matching/utils/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..beb31ff27343c40955ec52513c44f75a6ab5d665 --- /dev/null +++ b/flow_matching/utils/utils.py @@ -0,0 +1,90 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Optional + +import torch +from torch import Tensor + + +def unsqueeze_to_match(source: Tensor, target: Tensor, how: str = "suffix") -> Tensor: + """ + Unsqueeze the source tensor to match the dimensionality of the target tensor. + + Args: + source (Tensor): The source tensor to be unsqueezed. + target (Tensor): The target tensor to match the dimensionality of. + how (str, optional): Whether to unsqueeze the source tensor at the beginning + ("prefix") or end ("suffix"). Defaults to "suffix". + + Returns: + Tensor: The unsqueezed source tensor. + """ + assert ( + how == "prefix" or how == "suffix" + ), f"{how} is not supported, only 'prefix' and 'suffix' are supported." + + dim_diff = target.dim() - source.dim() + + for _ in range(dim_diff): + if how == "prefix": + source = source.unsqueeze(0) + elif how == "suffix": + source = source.unsqueeze(-1) + + return source + + +def expand_tensor_like(input_tensor: Tensor, expand_to: Tensor) -> Tensor: + """`input_tensor` is a 1d vector of length equal to the batch size of `expand_to`, + expand `input_tensor` to have the same shape as `expand_to` along all remaining dimensions. + + Args: + input_tensor (Tensor): (batch_size,). + expand_to (Tensor): (batch_size, ...). + + Returns: + Tensor: (batch_size, ...). + """ + assert input_tensor.ndim == 1, "Input tensor must be a 1d vector." + assert ( + input_tensor.shape[0] == expand_to.shape[0] + ), f"The first (batch_size) dimension must match. Got shape {input_tensor.shape} and {expand_to.shape}." + + dim_diff = expand_to.ndim - input_tensor.ndim + + t_expanded = input_tensor.clone() + t_expanded = t_expanded.reshape(-1, *([1] * dim_diff)) + + return t_expanded.expand_as(expand_to) + + +def gradient( + output: Tensor, + x: Tensor, + grad_outputs: Optional[Tensor] = None, + create_graph: bool = False, +) -> Tensor: + """ + Compute the gradient of the inner product of output and grad_outputs w.r.t :math:`x`. + + Args: + output (Tensor): [N, D] Output of the function. + x (Tensor): [N, d_1, d_2, ... ] input + grad_outputs (Optional[Tensor]): [N, D] Gradient of outputs, if `None`, + then will use a tensor of ones + create_graph (bool): If True, graph of the derivative will be constructed, allowing + to compute higher order derivative products. Defaults to False. + Returns: + Tensor: [N, d_1, d_2, ... ]. the gradient w.r.t x. + """ + + if grad_outputs is None: + grad_outputs = torch.ones_like(output).detach() + grad = torch.autograd.grad( + output, x, grad_outputs=grad_outputs, create_graph=create_graph + )[0] + return grad diff --git a/logic/.DS_Store b/logic/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..5008ddfcf53c02e82d7eee2e57c38e5672ef89f6 Binary files /dev/null and b/logic/.DS_Store differ diff --git a/logic/__init__.py b/logic/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..36d719522180cd6c67a6b7eb05ee12e1cff6f386 --- /dev/null +++ b/logic/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. diff --git a/logic/evaluate.py b/logic/evaluate.py new file mode 100644 index 0000000000000000000000000000000000000000..3e8d977c92b7fb65552289dbf8a3bb887ebf9aec --- /dev/null +++ b/logic/evaluate.py @@ -0,0 +1,183 @@ +# Copyright (c) Meta Platforms, Inc. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. +# Part of this implementation is adapted from https://github.com/louaaron/Score-Entropy-Discrete-Diffusion +# which is released under MIT license + +import math +from collections import Counter +from typing import Iterable, List, Optional, Sequence, Tuple + +import torch +import torch.nn.functional as F +from torch import Tensor +from torch.utils.data import DataLoader +from tqdm import tqdm + + +# --------------------------- +# Helpers for ragged handling +# --------------------------- + +def _batch_iter(seq: Sequence, batch_size: int) -> Iterable[Sequence]: + for i in range(0, len(seq), batch_size): + yield seq[i : i + batch_size] + + +def _pad_ragged_to_dense( + seqs: Sequence[Tensor], + pad_id: int, +) -> Tuple[Tensor, Tensor]: + """ + Pad a list of 1-D LongTensors -> (B, Lmax) LongTensor and an attention mask (B, Lmax) bool. + pad_id is written into padded positions; attention_mask == 1 for real tokens. + """ + if len(seqs) == 0: + return torch.empty(0, 0, dtype=torch.long), torch.empty(0, 0, dtype=torch.bool) + + device = seqs[0].device + lens = [int(s.numel()) for s in seqs] + Lmax = max(lens) + B = len(seqs) + + out = torch.full((B, Lmax), pad_id, dtype=torch.long, device=device) + attn = torch.zeros((B, Lmax), dtype=torch.bool, device=device) + + for i, s in enumerate(seqs): + L = s.numel() + if L > 0: + out[i, :L] = s + attn[i, :L] = True + + return out, attn + + +# --------------------------- +# Perplexity (ragged-friendly) +# --------------------------- + +@torch.no_grad() +def compute_perplexity( + samples: List[Tensor], + *, + lm_model, # a causal LM compatible with your tokenizer/vocab (e.g., HF AutoModelForCausalLM) + pad_id: int, + batch_size: int = 16, +) -> Tensor: + """ + Compute perplexity on ragged samples using a provided *compatible* causal LM. + + Args: + samples: list of 1-D LongTensors (ragged, token ids in the *same* vocab the LM expects) + lm_model: a causal LM (e.g., transformers.AutoModelForCausalLM) on the right device + pad_id: integer id used for padding during evaluation + batch_size: eval batch size + + Returns: + Scalar tensor: mean perplexity across sequences. + """ + if len(samples) == 0: + return torch.tensor(float("nan")) + + device = samples[0].device + lm_model.eval() + + ppl_vals: List[Tensor] = [] + for chunk in _batch_iter(samples, batch_size): + # pad ragged batch + x, attn = _pad_ragged_to_dense(chunk, pad_id=pad_id) # (B, L), (B, L) + + # shift for next-token prediction + x_in = x[:, :-1] + x_tgt = x[:, 1:] + attn_in = attn[:, :-1] # (B, L-1) + + # forward LM + out = lm_model(input_ids=x_in, attention_mask=attn_in, use_cache=False) + logits = out.logits # (B, L-1, V) + + # compute token-wise NLL with ignore on pads + # set ignored targets to -100 per HF convention + labels = x_tgt.clone() + labels[~attn[:, 1:]] = -100 + + # cross-entropy averaged over *valid* tokens; we want per-sequence mean → then exp + ce = F.cross_entropy( + logits.transpose(-1, -2), # (B, V, L-1) + labels, + reduction="none", + ignore_index=-100, + ) # (B, L-1) + + # per-sequence mean over valid positions + valid = (labels != -100).float() + token_counts = valid.sum(dim=1).clamp_min(1.0) # avoid div-by-zero + ce_seq = (ce * valid).sum(dim=1) / token_counts # (B,) + ppl_seq = ce_seq.exp() # (B,) + + ppl_vals.append(ppl_seq) + + ppl_all = torch.cat(ppl_vals, dim=0) # (N,) + return ppl_all.mean() + + +# --------------------------- +# Entropy (ragged-friendly) +# --------------------------- + +def _sample_entropy_1d(sample: Tensor) -> float: + """ + Entropy (base-2) of a 1-D LongTensor by empirical token histogram. + """ + if sample.numel() == 0: + return 0.0 + histogram = Counter(sample.tolist()) + total = float(sum(histogram.values())) + ent = 0.0 + for c in histogram.values(): + p = c / total + ent -= p * math.log2(max(p, 1e-12)) + return ent + + +@torch.no_grad() +def compute_entropy(samples: List[Tensor]) -> Tensor: + """ + Mean per-sequence token entropy (base-2) for ragged samples. + """ + if len(samples) == 0: + return torch.tensor(float("nan")) + device = samples[0].device + vals = [_sample_entropy_1d(s) for s in samples] + return torch.tensor(sum(vals) / len(vals), device=device) + + +# --------------------------- +# Likelihood / ELBO (EF note) +# --------------------------- + +@torch.no_grad() +def estimate_likelihood( + model: torch.nn.Module, + dataloader: DataLoader, + source_distribution, + path, + n_discretization: int, + device: torch.device, + batch_size: int = 32, + epsilon: float = 1e-3, +) -> Tensor: + """ + Placeholder: ELBO-style likelihood estimation used for DFM does not + directly apply to Edit Flows (different objective / dynamics). + If you need a quantitative likelihood-like diagnostic for EF, + consider CTMC pathwise estimators or reverse-time simulators instead. + + We raise NotImplementedError to avoid silently reporting a mismatched metric. + """ + raise NotImplementedError( + "ELBO/likelihood estimation used for discrete flow matching is not applicable to Edit Flows. " + "Use task metrics or CTMC-based diagnostics instead." + ) diff --git a/logic/flow.py b/logic/flow.py new file mode 100644 index 0000000000000000000000000000000000000000..97c77858c05e52cde4c8d21f7a8e341821cad32e --- /dev/null +++ b/logic/flow.py @@ -0,0 +1,439 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +from abc import ABC +from typing import List, Optional, Tuple, Union + +import torch +from flow_matching.loss import MixturePathGeneralizedKL, EditFlowsLoss, EditFlowsLossReParam +from flow_matching.path import MixtureDiscreteProbPath # for the scheduler only +from flow_matching.path.scheduler import PolynomialConvexScheduler +from flow_matching.path.editflows_adapter import EditFlowsPathAdapter # <-- NEW import +from torch import Tensor +from torch.nn.modules.loss import _Loss + + +class SourceDistribution(ABC): + def __init__(self) -> None: + ... + + def sample(self, tensor_size: Tuple[int, ...], device: torch.device) -> Tensor: + ... + + def sample_like(self, tensor_like: Tensor) -> Tensor: + ... + + +class MaskedSourceDistribution(SourceDistribution): + def __init__(self, mask_token: int) -> None: + self.mask_token = mask_token + + @property + def masked(self) -> bool: + return True + + def sample(self, tensor_size: Tuple[int, ...], device: torch.device) -> Tensor: + return torch.zeros(tensor_size, device=device).fill_(self.mask_token).long() + + def sample_like(self, tensor_like: Tensor) -> Tensor: + return torch.zeros_like(tensor_like).fill_(self.mask_token).long() + + +import torch +from typing import List, Tuple, Optional, Union +Tensor = torch.Tensor + +class UniformSourceDistribution: + def __init__(self, vocab_size, special_token_ids = None): + self.vocab_size = vocab_size + self.special_token_ids = set(special_token_ids) if special_token_ids is not None else set() + # Compute allowed tokens by removing all special tokens from vocab + self._allowed_tokens = [i for i in range(vocab_size) if i not in self.special_token_ids] + if len(self._allowed_tokens) == 0: + raise ValueError(f"All tokens are special tokens: {special_token_ids}") + + @property + def masked(self) -> bool: + return False + + @torch.no_grad() + def _sample_from_allowed(self, shape, device, allowed_tokens = None, generator = None): + """Sample uniformly from allowed tokens with given shape.""" + if allowed_tokens is None: + allowed_tokens = self._allowed_tokens + if len(allowed_tokens) == 0: + raise ValueError("No allowed tokens provided") + + allowed_tensor = torch.tensor(allowed_tokens, dtype=torch.long, device=device) + num_allowed = len(allowed_tokens) + + if shape.numel() == 0: + # Return an empty tensor (length-0 sequence) + return torch.empty(shape, dtype=torch.long, device=device) + + # Sample indices in [0, num_allowed) + indices = torch.randint( + low=0, high=num_allowed, size=shape, device=device, generator=generator + ) + return allowed_tensor[indices] + + @torch.no_grad() + def sample(self, tensor_size, device, allowed_tokens = None, generator = None): + return self._sample_from_allowed(tensor_size, device, allowed_tokens, generator) + + @torch.no_grad() + def sample_like(self, tensor_like, allowed_tokens = None, generator = None): + """ + Keep original semantics: sample tokens with the SAME SHAPE(S) as tensor_like. + """ + if isinstance(tensor_like, (list, tuple)): + return [ + self._sample_from_allowed(seq.shape, seq.device, allowed_tokens, generator) + for seq in tensor_like + ] + return self._sample_from_allowed(tensor_like.shape, tensor_like.device, allowed_tokens, generator) + + @torch.no_grad() + def sample_like( + self, + tensor_like, + allowed_tokens: Optional[List[int]] = None, + min_len_factor: float = 0.0, + max_len_factor: float = 2.0, + generator: Optional[torch.Generator] = None, + ): + """ + For each reference x1, sample x0 with LENGTH L ~ Uniform{ floor(min_len_factor*N) .. floor(max_len_factor*N) }, + defaulting to [0, 2*N]. Returns an empty tensor if L == 0. + + Supports Tensor (1D) or List[Tensor] (ragged). If you pass a 2D tensor, we + interpret N as the last dimension and return a 1D sequence for that tensor. + """ + def _one(x1): + assert x1.dim() >= 1, "x1 must be at least 1D" + device = x1.device + dtype = torch.long + + N = int(x1.shape[-1]) # use last-dim length as reference + lo = int(max(0, int(min_len_factor * N))) + hi = int(max(0, int(max_len_factor * N))) + # randint is [low, high), so make high inclusive with +1 + L = int(torch.randint(low=lo, high=hi + 1, size=(1,), device=device, generator=generator).item()) + if L == 0: + return torch.empty((0,), dtype=dtype, device=device) + return self._sample_from_allowed(torch.Size([L]), device, allowed_tokens, generator) + + if isinstance(tensor_like, (list, tuple)): + return [_one(seq) for seq in tensor_like] + else: + return _one(tensor_like) + + def sample_x0_from_x1(self, x1, pad_id, allowed_tokens, scale_size = 1.0, bos_id = 0, eos_id = 2): + """ + For each sequence in x1, sample an x0 whose *core* length (excluding BOS/EOS) + is in [0, scale_size * len_valid(x1)], where len_valid(x1) counts only tokens that are + NOT {BOS, EOS, PAD}. + + Rules: + - x0 always starts with BOS and ends with EOS + - x0 core tokens are sampled uniformly from vocab excluding {BOS, EOS, PAD} + - We batch and pad x0 to a common length (B, L0) with pad_id + - "length of x0 does not account for BOS and EOS" = the sampled core length + + Returns: + x0: (B, L0) Long, padded with pad_id + """ + device = x1.device + B, L1 = x1.shape + + # compute valid length of x1 per sequence + # valid = not pad, not BOS, not EOS + valid_mask_x1 = (x1 != pad_id) & (x1 != bos_id) & (x1 != eos_id) + valid_len = valid_mask_x1.sum(dim=1) # (B,) + + # we will store all sequences here before padding + x0_seqs = [] + + for b in range(B): + max_core_len = int((1 + scale_size) * valid_len[b].item()) # may be 0 + min_core_len = int((1 - scale_size) * valid_len[b].item()) + + # sample core length in [0, max_core_len] + core_len = int(torch.randint(low=min_core_len, high=max_core_len + 1, size=(1,), device=device).item()) + + # 4) sample core tokens + if core_len > 0: + idx = torch.randint(0, allowed_tokens.size(0), (core_len,), device=device) + core_tokens = allowed_tokens[idx] # (core_len,) + else: + core_tokens = torch.empty(0, dtype=torch.long, device=device) + + # 5) build full x0: [BOS] + core + [EOS] + seq = torch.cat([ + torch.tensor([bos_id], device=device, dtype=torch.long), + core_tokens, + torch.tensor([eos_id], device=device, dtype=torch.long), + ], dim=0) # (1 + core_len + 1,) + + x0_seqs.append(seq) + + x0 = torch.nn.utils.rnn.pad_sequence(x0_seqs, batch_first=True, padding_value=pad_id) + + return x0 + + +# class UniformSourceDistribution(SourceDistribution): +# def __init__(self, vocab_size: int, special_token_ids: Optional[List[int]] = None) -> None: +# self.vocab_size = vocab_size +# self.special_token_ids = set(special_token_ids) if special_token_ids is not None else set() +# # Compute allowed tokens by removing all special tokens from vocab +# self._allowed_tokens = [i for i in range(vocab_size) if i not in self.special_token_ids] + +# if len(self._allowed_tokens) == 0: +# raise ValueError(f"All tokens are special tokens: {special_token_ids}") + +# @property +# def masked(self) -> bool: +# return False + +# def _sample_from_allowed( +# self, shape: torch.Size, device: torch.device, allowed_tokens: Optional[List[int]] = None +# ) -> Tensor: +# """Sample uniformly from allowed tokens.""" +# if allowed_tokens is None: +# # Use default allowed tokens (vocab minus special tokens) +# allowed_tokens = self._allowed_tokens + +# if len(allowed_tokens) == 0: +# raise ValueError(f"No allowed tokens provided") + +# # print(f"allowed_tokens: {allowed_tokens}") + +# # Sample indices into allowed_tokens, then map to actual token IDs +# allowed_tensor = torch.tensor(allowed_tokens, dtype=torch.long, device=device) +# num_allowed = len(allowed_tokens) + +# # Sample indices in [0, num_allowed) +# indices = torch.randint(size=shape, high=num_allowed, device=device) +# return allowed_tensor[indices] + +# def sample( +# self, +# tensor_size: Tuple[int, ...], +# device: torch.device, +# allowed_tokens: Optional[List[int]] = None +# ) -> Tensor: +# return self._sample_from_allowed(tensor_size, device, allowed_tokens) + +# def sample_like( +# self, +# tensor_like: Union[Tensor, List[Tensor]], +# allowed_tokens: Optional[List[int]] = None +# ) -> Union[Tensor, List[Tensor]]: +# """ +# Sample uniform tokens matching the shape of tensor_like. +# Supports both Tensor and List[Tensor] for ragged inputs. + +# Args: +# tensor_like: Either a Tensor or List[Tensor] (for ragged inputs) +# allowed_tokens: Optional list of allowed token IDs. If None, uses vocab minus special tokens. + +# Returns: +# Tensor or List[Tensor] matching the input shape(s) +# """ +# # Handle ragged input (list of tensors) +# if isinstance(tensor_like, (list, tuple)): +# return [ +# self._sample_from_allowed(seq.shape, seq.device, allowed_tokens) +# for seq in tensor_like +# ] + +# # Handle regular tensor input +# return self._sample_from_allowed(tensor_like.shape, tensor_like.device, allowed_tokens) + + +class EmpiricalSourceDistribution(SourceDistribution): + def __init__(self, vocab_size: int, probs: torch.Tensor, length: int): + self.vocab_size = vocab_size + self.registered = probs / probs.sum() + self.length = length # e.g., 100 tokens + + @property + def masked(self) -> bool: + return False + + def sample(self, tensor_size: Tuple[int, ...], device: torch.device) -> Tensor: + B = tensor_size[0] + idx = torch.multinomial(self.registered.to(device), num_samples=self.length, replacement=True) + return idx.view(1, self.length).repeat(B, 1) + + def sample_like(self, tensor_like: Tensor) -> Tensor: + B = tensor_like.shape[0] + idx = torch.multinomial(self.registered.to(tensor_like.device), num_samples=self.length, replacement=True) + return idx.view(1, self.length).repeat(B, 1) + + +# NOTE: return type changed to the adapter (ragged). We only rely on .scheduler(t) and .sample(...). +def get_path(scheduler_type: str, exponent: Optional[float] = None, eps_id: int = -1) -> EditFlowsPathAdapter: + if scheduler_type == "polynomial": + # paper uses cubic => exponent=3 + scheduler = PolynomialConvexScheduler(n=exponent) + else: + raise ValueError(f"{scheduler_type} is not supported") + + # MixtureDiscreteProbPath carries the scheduler; the adapter will sample ragged z_t itself + mixture = MixtureDiscreteProbPath(scheduler=scheduler) + return EditFlowsPathAdapter(mixture_path=mixture, eps_id=eps_id) + + +def get_source_distribution( + source_distribution: str, + p_emp: Optional[Tensor] = None, + length: Optional[int] = None, + vocab_size: Optional[int] = None, + special_token_ids: Optional[List[int]] = None, +) -> SourceDistribution: + if p_emp is not None: + assert vocab_size is not None and length is not None, "Empirical source requires vocab_size and length" + return EmpiricalSourceDistribution(vocab_size=vocab_size, probs=p_emp, length=length) + + if source_distribution == "mask": + assert vocab_size is not None, "Masked source requires vocab_size" + return MaskedSourceDistribution(mask_token=vocab_size) + elif source_distribution == "uniform": + assert vocab_size is not None, "Uniform source requires vocab_size" + return UniformSourceDistribution(vocab_size=vocab_size, special_token_ids=special_token_ids) + else: + raise ValueError(f"{source_distribution} is not supported") + + +def get_loss_function(loss_function: str, path: Optional[Union[MixtureDiscreteProbPath, EditFlowsPathAdapter]] = None) -> _Loss: + if loss_function == "cross_entropy": + return torch.nn.CrossEntropyLoss() + elif loss_function == "generalized_kl": + assert path is not None + # Generalized KL still expects a (dense) path; fine for DFM experiments + return MixturePathGeneralizedKL(path=path) + elif loss_function == "editflows": + # Ragged EF loss does NOT need the path; training.step precomputes the weight + return EditFlowsLoss(reduction="mean") + elif loss_function == "editflows_reparam": + return EditFlowsLossReParam(reduction="mean") + else: + raise ValueError(f"{loss_function} is not supported") + + +# # Copyright (c) Meta Platforms, Inc. and affiliates. +# # All rights reserved. +# # +# # This source code is licensed under the CC-by-NC license found in the +# # LICENSE file in the root directory of this source tree. + +# from abc import ABC +# from typing import Optional, Tuple + +# import torch +# from flow_matching.loss import MixturePathGeneralizedKL, EditFlowsLoss +# from flow_matching.path import MixtureDiscreteProbPath, ProbPath, EditFlowsPathAdapter +# from flow_matching.path.scheduler import PolynomialConvexScheduler +# from torch import Tensor +# from torch.nn.modules.loss import _Loss + + +# class SourceDistribution(ABC): +# def __init__( +# self, +# ) -> None: +# ... + +# def sample(self, tensor_size: Tuple[int, ...], device: torch.device) -> Tensor: +# ... + +# def sample_like(self, tensor_like: Tensor) -> Tensor: +# ... + + +# class MaskedSourceDistribution(SourceDistribution): +# def __init__(self, mask_token: int) -> None: +# self.mask_token = mask_token + +# @property +# def masked(self) -> bool: +# return True + +# def sample(self, tensor_size: Tuple[int, ...], device: torch.device) -> Tensor: +# return torch.zeros(tensor_size, device=device).fill_(self.mask_token).long() + +# def sample_like(self, tensor_like: Tensor) -> Tensor: +# return torch.zeros_like(tensor_like).fill_(self.mask_token).long() + + +# class UniformSourceDistribution(SourceDistribution): +# def __init__(self, vocab_size: int) -> None: +# self.vocab_size = vocab_size + +# @property +# def masked(self) -> bool: +# return False + +# def sample(self, tensor_size: Tuple[int, ...], device: torch.device) -> Tensor: +# return torch.randint(size=tensor_size, high=self.vocab_size, device=device) + +# def sample_like(self, tensor_like: Tensor) -> Tensor: +# return torch.randint_like(tensor_like, high=self.vocab_size) + +# class EmpiricalSourceDistribution(SourceDistribution): +# def __init__(self, vocab_size: int, probs: torch.Tensor, length: int): +# self.vocab_size = vocab_size +# self.registered = probs / probs.sum() +# self.length = length # e.g., 100 tokens as in the paper’s variant + +# @property +# def masked(self) -> bool: return False + +# def sample(self, tensor_size: Tuple[int, ...], device: torch.device) -> Tensor: +# B = tensor_size[0] +# idx = torch.multinomial(self.registered.to(device), num_samples=self.length, replacement=True) +# return idx.view(1, self.length).repeat(B, 1) + +# def sample_like(self, tensor_like: Tensor) -> Tensor: +# B = tensor_like.shape[0] +# idx = torch.multinomial(self.registered.to(tensor_like.device), num_samples=self.length, replacement=True) +# return idx.view(1, self.length).repeat(B, 1) + + +# def get_path(scheduler_type: str, exponent: Optional[float] = None) -> ProbPath: +# if scheduler_type == "polynomial": +# scheduler = PolynomialConvexScheduler(n=exponent) +# else: +# raise ValueError(f"{scheduler_type} is not supported") + +# return EditFlowsPathAdapter(path=MixtureDiscreteProbPath(scheduler=scheduler)) # still need to decide (1) whether to implement the z_0, z_1 creation here, and (2) how to pass the eps_id to the adapter + + +# def get_source_distribution( +# source_distribution: str, p_emp: Optional[Tensor] = None, length: Optional[int] = None, vocab_size: int = None +# ) -> SourceDistribution: +# if p_emp is not None: +# return EmpiricalSourceDistribution(vocab_size=vocab_size, probs=p_emp, length=length) +# if source_distribution == "mask": +# return MaskedSourceDistribution(mask_token=vocab_size) +# elif source_distribution == "uniform": +# return UniformSourceDistribution(vocab_size=vocab_size) +# else: +# raise ValueError(f"{source_distribution} is not supported") + + +# def get_loss_function(loss_function: str, path: Optional[ProbPath] = None) -> _Loss: +# if loss_function == "cross_entropy": +# return torch.nn.CrossEntropyLoss() +# elif loss_function == "generalized_kl": +# assert path is not None + +# return MixturePathGeneralizedKL(path=path) +# elif loss_function == "editflows": +# assert path is not None + +# return EditFlowsLoss(path=path) +# else: +# raise ValueError(f"{loss_function} is not supported") diff --git a/logic/generate.py b/logic/generate.py new file mode 100644 index 0000000000000000000000000000000000000000..e6b0548f707fbf436f8038d52f1d09f8e0ed50ae --- /dev/null +++ b/logic/generate.py @@ -0,0 +1,103 @@ +# Copyright (c) Meta Platforms, Inc. +# All rights reserved. + +from pathlib import Path +from typing import List, Optional, Sequence, Tuple + +import torch +from torch import nn, Tensor + +from transformers.tokenization_utils import PreTrainedTokenizer + +from flow_matching.path import ProbPath # adapter exposing .scheduler(t) +from flow_matching.solver import EditFlowsEulerSolver +from .flow import SourceDistribution + + +class WrappedEFModel: + """ + Thin wrapper so the solver has a stable interface: + forward(x_list, t_scalar or t_vec) -> lam_ins, logits_ins, lam_del, lam_sub, logits_sub + """ + def __init__(self, model: nn.Module) -> None: + self.model = model + + @torch.no_grad() + def __call__(self, x_list: Sequence[Tensor], t: Tensor): + # Model already returns ragged lists (lam_ins, logits_ins, lam_del, lam_sub, logits_sub) + return self.model(x_t=list(x_list), time=t) + + +def _rows_to_ragged(x_dense: Tensor) -> List[Tensor]: + # Start from a (B, L) dense init -> ragged list of (L,) + B, L = x_dense.shape + return [x_dense[b].clone() for b in range(B)] + + +def _ragged_to_text( + seqs: Sequence[Tensor], + tokenizer: PreTrainedTokenizer, +) -> List[str]: + # Decode each sequence independently (no padding required) + out = [] + for s in seqs: + ids = s.detach().tolist() + out.append(tokenizer.decode(ids, skip_special_tokens=False)) + return out + + +def generate_samples( + model: nn.Module, + step: int, + vocab_size: int, + tokenizer: PreTrainedTokenizer, + rank: int, + device: torch.device, + path: ProbPath, # adapter or object exposing .scheduler(t) + source_distribution: SourceDistribution, + sample_batch_size: int, + sequence_length: int, + sampling_steps: int, + time_epsilon: float = 0.0, # not used in Option A; kept for API parity + sample_dir: Optional[Path] = None, + dtype_categorical: torch.dtype = torch.float64, +) -> List[Tensor]: + """ + EditFlows ragged generation with Euler thinning (≤1 jump per step). + Returns list[LongTensor] of final sequences. + """ + wrapped = WrappedEFModel(model=model) + + # Initial sequences (dense) -> ragged + x_init_dense = source_distribution.sample( + tensor_size=(sample_batch_size, sequence_length), device=device + ).long() + x_list = _rows_to_ragged(x_init_dense) + + # New ragged solver + solver = EditFlowsEulerSolver( + model=wrapped, + scheduler=path.scheduler, + vocab_size=vocab_size, + dtype_categorical=dtype_categorical, + ) + + # Simulate from t=0 → t=1 + x_final = solver.sample( + x_list=x_list, + n_steps=sampling_steps, + t0=0.0, + t1=1.0, + verbose=False, + ) + + # Decode to text (optional, for logging/inspection) + sentences = _ragged_to_text(x_final, tokenizer) + if sample_dir is not None: + file_name = sample_dir / f"iter_{step}" / f"sample_{rank}.txt" + file_name.parents[0].mkdir(exist_ok=True, parents=True) + with open(file_name, "w") as f: + for s in sentences: + f.write(s + "\n" + "=" * 20 + " New sample " + "=" * 20 + "\n") + + return x_final diff --git a/logic/state.py b/logic/state.py new file mode 100644 index 0000000000000000000000000000000000000000..ebafbb8064f4da6094e31c295429fd4c63c0a903 --- /dev/null +++ b/logic/state.py @@ -0,0 +1,89 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +import logging +from pathlib import Path + +import torch + +from torch import nn +from torch.optim import Optimizer + + +class TrainState: + def __init__( + self, + model: nn.Module, + optimizer: Optimizer, + step: int, + # data_state: DataState, + ): + self._model = model + self._optimizer = optimizer + self._step = step + # self._data_state = data_state + + @property + def step(self) -> int: + return self._step + + @step.setter + def step(self, value: int) -> None: + self._step = value + + @property + def optimizer(self) -> Optimizer: + return self._optimizer + + @property + def model(self) -> nn.Module: + return self._model + + # @property + # def data_state(self) -> DataState: + # return self._data_state + + def compile_model(self) -> None: + self._model = torch.compile(self._model) + + def restore_checkpoint( + self, ckpt_dir: Path, device: torch.device, rank: int + ) -> None: + if ckpt_dir.exists(): + loaded_state = torch.load(ckpt_dir, map_location=device, weights_only=True) + + self.optimizer.load_state_dict(loaded_state["optimizer"]) + self.model.module.load_state_dict(loaded_state["model"]) + self.step = loaded_state["step"] + # self._data_state.test.load_state_dict(loaded_state["test_sampler"]) + # self._data_state.train.sampler.load_state_dict( + # loaded_state["train_sampler"] + # ) + else: + ckpt_dir.parent.mkdir(exist_ok=True, parents=True) + + if rank == 0: + logging.warning( + f"No checkpoint found at {ckpt_dir}. Returned the same state as input" + ) + + def save_checkpoint(self, ckpt_dir: str, rank: int) -> None: + saved_state = { + "optimizer": self.optimizer.state_dict(), + "model": self.model.module.state_dict(), + "step": self.step, + # "train_sampler": self._data_state.train.sampler.state_dict(), + # "test_sampler": self._data_state.test.sampler.state_dict(), + } + + if rank == 0: + torch.save(saved_state, ckpt_dir) + + def eval(self) -> None: + self.train(training=False) + + def train(self, training: bool = True) -> None: + self._model.train(mode=training) diff --git a/logic/training.py b/logic/training.py new file mode 100644 index 0000000000000000000000000000000000000000..f2dd7fceb9304af177265d4cbcfb7df1491ec44f --- /dev/null +++ b/logic/training.py @@ -0,0 +1,134 @@ +import math +from contextlib import nullcontext +from typing import Optional + +import torch +from flow_matching.loss import MixturePathGeneralizedKL, EditFlowsLoss +from flow_matching.path import ProbPath +from omegaconf.dictconfig import DictConfig +from torch import nn, Tensor +from torch.cuda.amp import GradScaler +from torch.nn.utils.rnn import pad_sequence + +from torch.utils.data import DataLoader +from utils.logging import TrainLogger + +from .flow import SourceDistribution +from .state import TrainState +from ..model.utils import build_z0_z1_with_alignment, remove_eps + +from dataclasses import dataclass +from typing import List, Tuple, Optional +import torch +from torch import Tensor + +import pdb + + +def step( + state: TrainState, + loss_fn: nn.Module, # EditFlowsLoss + path: ProbPath, # EditFlowsPathAdapter (exposes .scheduler) + scaler: GradScaler, + iterator: DataLoader, + device: torch.device, + source_distribution: SourceDistribution, + logger: TrainLogger, + training: bool, + optim_params: Optional[DictConfig], + pad_id: int, + bos_id: int, + eos_id: int, +) -> Tensor: + assert (training and (optim_params is not None)) or (not training) + state.train() if training else state.eval() + + batch = next(iterator) + # x_1 = pad_sequence(batch['input_ids'], batch_first=True, padding_value=pad_id).to(device) + x_1 = torch.tensor(batch["input_ids"]).to(device) + B = x_1.shape[0] + + # === Source & time === + with torch.no_grad(): + eps_id = getattr(path, "eps_id", -1) + allowed_tokens = torch.tensor([tok for tok in source_distribution._allowed_tokens if tok != eps_id]).to(device) + + x_0 = source_distribution.sample_x0_from_x1(x_1, pad_id=pad_id, allowed_tokens=allowed_tokens, scale_size=2, bos_id = bos_id, eos_id = eos_id) + t = torch.rand(B, device=device) + + sched = path.scheduler(t) + precomputed_weight = sched.d_alpha_t / sched.sigma_t # (B,) + + z_0, z_1 = build_z0_z1_with_alignment(x_0, x_1, eps_id, pad_id, bos_id, eos_id, p_optimal=0.6) + + z_t = path.sample(z_0, z_1, t=t) + x_t, mask = remove_eps(z_t, eps_id, pad_id) + + ctx = torch.amp.autocast('cuda', dtype=torch.float16) if training else torch.no_grad() + with ctx: + # pdb.set_trace() + lam_ins, logits_ins, lam_del, lam_sub, logits_sub = state.model(x_t=x_t, mask=mask,t=t) + + loss = loss_fn(lam_ins, logits_ins, lam_del, lam_sub, logits_sub, + z_t, z_1, x_t, mask, precomputed_weight, eps_id, bos_id, eos_id) + + if training: + optimization_step( + state=state, + loss=loss, + scaler=scaler, + optim_params=optim_params, + logger=logger, + ) + + return loss.detach() + + + +def _get_lr(lr: float, step: int, warmup: int, n_iters: int, eta_min_ratio: float): + if step < warmup: + # Linear warmup + return lr * (step / warmup) + else: + # Cosine annealing + total_steps = n_iters + eta_min = eta_min_ratio * lr + cosine_decay = 0.5 * ( + 1 + math.cos(math.pi * (step - warmup) / (total_steps - warmup)) + ) + return eta_min + (lr - eta_min) * cosine_decay + +def optimization_step( + state: TrainState, + scaler: GradScaler, + loss: Tensor, + optim_params: DictConfig, + logger: TrainLogger, +) -> None: + scaler.scale(loss).backward() + scaler.unscale_(state.optimizer) + + lr = _get_lr( + lr=optim_params.lr, + step=state.step, + warmup=optim_params.warmup, + n_iters=optim_params.n_iters, + eta_min_ratio=optim_params.eta_min_ratio, + ) + + # Update learning rate in optimizer + for g in state.optimizer.param_groups: + g["lr"] = lr + + if state.step % optim_params.log_lr_every == 0: + logger.log_lr(value=lr, step=state.step) + + if optim_params.grad_clip >= 0: + torch.nn.utils.clip_grad_norm_( + state.model.parameters(), max_norm=optim_params.grad_clip + ) + + scaler.step(state.optimizer) + scaler.update() + + state.optimizer.zero_grad() diff --git a/model/.DS_Store b/model/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..5008ddfcf53c02e82d7eee2e57c38e5672ef89f6 Binary files /dev/null and b/model/.DS_Store differ diff --git a/model/__init__.py b/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..544f5e9521e0e88f67e62464e35461b6ad8e29a5 --- /dev/null +++ b/model/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +# from .transformer import Transformer +# from .rotary import Rotary + +# __all__ = [ +# "Transformer", +# "Rotary", +# ] diff --git a/model/base_models.py b/model/base_models.py new file mode 100644 index 0000000000000000000000000000000000000000..242a72fea2da48c87faf7a45b2a1031373219cbd --- /dev/null +++ b/model/base_models.py @@ -0,0 +1,520 @@ +from typing import Optional, Tuple +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import EsmModel +import pytorch_lightning as pl + +from .utils import build_z0_z1_with_alignment, remove_eps +import pdb + +# ---------- Utilities ---------- + +def exists(x): return x is not None + +def default(val, d): + return val if exists(val) else d + +# ---------- Timestep embedding (sinusoidal -> MLP) ---------- + +class TimeEmbedding(nn.Module): + """ + Sinusoidal time embedding followed by a small MLP. + Accepts t of shape (B,) or scalar; outputs (B, d_model) and broadcasts over L. + """ + def __init__(self, d_model: int, hidden: Optional[int] = None, max_period: int = 10000): + super().__init__() + self.d_model = d_model + self.max_period = max_period + hidden = default(hidden, d_model * 4) + # Use even dim for sin/cos pairing + pe_dim = d_model if d_model % 2 == 0 else d_model - 1 + self.pe_dim = pe_dim + self.mlp = nn.Sequential( + nn.Linear(pe_dim, hidden), + nn.SiLU(), + nn.Linear(hidden, d_model), + ) + + def forward(self, t: torch.Tensor, batch_size: Optional[int] = None) -> torch.Tensor: + """ + t: (B,) or () in [0,1] + returns: (B, d_model) + """ + if t.dim() == 0: + # scalar -> expand to batch + if batch_size is None: + raise ValueError("When t is scalar, provide batch_size.") + t = t.expand(batch_size) + B = t.shape[0] + device = t.device + half = self.pe_dim // 2 + # frequencies + freqs = torch.exp( + torch.arange(half, device=device, dtype=t.dtype) * (-math.log(self.max_period) / (half - 1 + 1e-8)) + ) + angles = t[:, None] * freqs[None, :] * math.pi # (B, half) + pe = torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1) # (B, pe_dim) + if self.pe_dim < self.d_model: # rare case when d_model is odd + pe = F.pad(pe, (0, 1), value=0.0) + return self.mlp(pe) # (B, d_model) + +# ---------- RoPE (rotary position embedding) ---------- + +def apply_rotary(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """ + q, k: (B, h, L, d_head) + cos, sin: (L, d_head) broadcastable to (B, h, L, d_head) + """ + # split last dim into pairs + d = q.shape[-1] + if d % 2 != 0: + # pad to even + q = F.pad(q, (0, 1), value=0.0) + k = F.pad(k, (0, 1), value=0.0) + d += 1 + + q1, q2 = q[..., :d//2], q[..., d//2:] + k1, k2 = k[..., :d//2], k[..., d//2:] + + # broadcast cos/sin + while cos.dim() < q1.dim(): + cos = cos.unsqueeze(0) + sin = sin.unsqueeze(0) + + rq = torch.cat([q1 * cos - q2 * sin, q1 * sin + q2 * cos], dim=-1) + rk = torch.cat([k1 * cos - k2 * sin, k1 * sin + k2 * cos], dim=-1) + return rq[..., :q.shape[-1]], rk[..., :k.shape[-1]] + +class RotaryPositionalEmbedding(nn.Module): + """ + Precomputes cos/sin for RoPE given max_len and head_dim. + """ + def __init__(self, head_dim: int, max_len: int = 8192, base: int = 10000): + super().__init__() + if head_dim % 2 != 0: + # allow odd by padding inside apply_rotary, but prefer even + pass + inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.max_len = max_len + self.head_dim = head_dim + self._cached_len = 0 + self.register_buffer("cos_cached", torch.empty(0), persistent=False) + self.register_buffer("sin_cached", torch.empty(0), persistent=False) + + def _update_cache(self, seq_len: int, device, dtype): + if seq_len <= self._cached_len and self.cos_cached.device == device and self.cos_cached.dtype == dtype: + return + t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) + freqs = torch.einsum('n,d->nd', t, self.inv_freq) # (L, head_dim/2) + # emb = torch.cat((freqs, freqs), dim=-1) # (L, head_dim) + self.cos_cached = freqs.cos().to(dtype=dtype) + self.sin_cached = freqs.sin().to(dtype=dtype) + self._cached_len = seq_len + + def forward(self, L: int, device, dtype): + self._update_cache(L, device, dtype) + return self.cos_cached[:L], self.sin_cached[:L] + +# ---------- Rotary MHA + Transformer block ---------- + +class RotaryMHA(nn.Module): + def __init__(self, d_model: int, n_heads: int, rope: RotaryPositionalEmbedding, attn_dropout: float = 0.0, proj_dropout: float = 0.0): + super().__init__() + assert d_model % n_heads == 0, "d_model must be divisible by n_heads" + self.qkv = nn.Linear(d_model, 3 * d_model, bias=False) + self.proj = nn.Linear(d_model, d_model, bias=False) + self.n_heads = n_heads + self.d_head = d_model // n_heads + self.attn_dropout = nn.Dropout(attn_dropout) + self.proj_dropout = nn.Dropout(proj_dropout) + self.rope = rope + + def forward(self, x: torch.Tensor, key_padding_mask: Optional[torch.Tensor] = None): + """ + x: (B, L, d_model) + key_padding_mask: (B, L) bool, True=pad (masked) + """ + B, L, D = x.shape + qkv = self.qkv(x) # (B, L, 3D) + q, k, v = qkv.chunk(3, dim=-1) + q = q.view(B, L, self.n_heads, self.d_head).transpose(1, 2) # (B, h, L, d) + k = k.view(B, L, self.n_heads, self.d_head).transpose(1, 2) + v = v.view(B, L, self.n_heads, self.d_head).transpose(1, 2) + + cos, sin = self.rope(L, x.device, x.dtype) + q, k = apply_rotary(q, k, cos, sin) + + attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_head) # (B, h, L, L) + + if exists(key_padding_mask): + # mask: True means pad -> set attention to -inf + mask = key_padding_mask[:, None, None, :].to(dtype=torch.bool) # (B,1,1,L) + attn_scores = attn_scores.masked_fill(mask, float('-inf')) + + attn = torch.softmax(attn_scores, dim=-1) + attn = self.attn_dropout(attn) + out = torch.matmul(attn, v) # (B, h, L, d) + out = out.transpose(1, 2).contiguous().view(B, L, D) + return self.proj_dropout(self.proj(out)) + +class TransformerBlock(nn.Module): + def __init__(self, d_model: int, n_heads: int, mlp_ratio: float = 4.0, attn_dropout: float = 0.0, proj_dropout: float = 0.0, rope: Optional[RotaryPositionalEmbedding] = None): + super().__init__() + self.norm1 = nn.LayerNorm(d_model) + self.attn = RotaryMHA(d_model, n_heads, rope=rope, attn_dropout=attn_dropout, proj_dropout=proj_dropout) + self.norm2 = nn.LayerNorm(d_model) + hidden = int(d_model * mlp_ratio) + self.mlp = nn.Sequential( + nn.Linear(d_model, hidden), + nn.SiLU(), + nn.Linear(hidden, d_model), + ) + + def forward(self, x: torch.Tensor, key_padding_mask: Optional[torch.Tensor] = None): + x = x + self.attn(self.norm1(x), key_padding_mask=key_padding_mask) + x = x + self.mlp(self.norm2(x)) + return x + +class ProteinEditFlowModel(nn.Module): + """ + Inputs: + x_t: (B, L) Long + mask: (B, L) bool, True=pad (i.e., should be ignored) + t: (B,) or scalar in [0,1] + + Outputs: + lam_ins: (B, L) >= 0 + logits_ins: (B, L, V) + lam_del: (B, L) >= 0 + lam_sub: (B, L) >= 0 + logits_sub: (B, L, V) + """ + def __init__(self, vocab_size, pad_id, config): + super().__init__() + self.d_model = getattr(config, "d_model", 768) + self.n_layers = getattr(config, "n_layers", 12) + self.n_heads = getattr(config, "n_heads", 12) + self.mlp_ratio = getattr(config, "mlp_ratio", 4) + self.max_len = getattr(config, "max_len", 2048) + self.dropout = getattr(config, "dropout", 0.1) + self.attn_dropout = getattr(config, "attn_dropout", 0) + self.proj_dropout = getattr(config, "proj_dropout", 0) + self.vocab_size = vocab_size + self.pad_id = pad_id + + # --- Embedding --- + self.esm_emb = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D") + for param in self.esm_emb.parameters(): + param.requires_grad = False + self.time_emb = TimeEmbedding(d_model=self.d_model) + + self.tok_embed_to_hidden = nn.Linear(1280, self.d_model) + + # --- RoPE shared by attention blocks --- + rope = RotaryPositionalEmbedding(head_dim=self.d_model // self.n_heads, max_len=self.max_len) + + # --- Encoder --- + self.blocks = nn.ModuleList([ + TransformerBlock( + d_model=self.d_model, + n_heads=self.n_heads, + mlp_ratio=self.mlp_ratio, + attn_dropout=self.attn_dropout, + proj_dropout=self.proj_dropout, + rope=rope + ) + for _ in range(self.n_layers) + ]) + self.final_norm = nn.LayerNorm(self.d_model) + + # --- Heads --- + # We use small MLP heads for rates; logits are linear. + self.lam_ins_head = nn.Sequential(nn.Linear(self.d_model, self.d_model//2), nn.SiLU(), nn.Linear(self.d_model//2, 1)) + self.lam_del_head = nn.Sequential(nn.Linear(self.d_model, self.d_model//2), nn.SiLU(), nn.Linear(self.d_model//2, 1)) + self.lam_sub_head = nn.Sequential(nn.Linear(self.d_model, self.d_model//2), nn.SiLU(), nn.Linear(self.d_model//2, 1)) + self.logits_ins_head = nn.Linear(self.d_model, vocab_size, bias=False) + self.logits_sub_head = nn.Linear(self.d_model, vocab_size, bias=False) + + # nonnegativity via softplus (safer than exp) + self.softplus = nn.Softplus(beta=1.0) + + def forward( + self, + x_t: torch.LongTensor, + mask: torch.BoolTensor, + t: torch.Tensor, + ): + """ + x_t: (B, L) long tokens + mask: (B, L) bool, True = PAD (ignored) + t: (B,) or scalar in [0,1] + esmbed: optional (B, L, esm2_embed_dim) if use_real_esm2=True + """ + B, L = x_t.shape + # pdb.set_trace() + + # --- Embedding --- + h = self.esm_emb(x_t, mask).last_hidden_state + h = self.tok_embed_to_hidden(h) + + # --- Add time embedding (broadcast across length) --- + t_emb = self.time_emb(t, batch_size=B) # (B, d_model) + h = h + t_emb.unsqueeze(1) # (B, L, d_model) + + # --- Encoder blocks with key padding mask --- + for blk in self.blocks: + h = blk(h, key_padding_mask=(~mask)) + + h = self.final_norm(h) # (B, L, d_model) + + # --- Heads --- + lam_ins = self.softplus(self.lam_ins_head(h)).squeeze(-1) # (B, L) + lam_del = self.softplus(self.lam_del_head(h)).squeeze(-1) # (B, L) + lam_sub = self.softplus(self.lam_sub_head(h)).squeeze(-1) # (B, L) + logits_ins = self.logits_ins_head(h) # (B, L, V) + logits_sub = self.logits_sub_head(h) # (B, L, V) + + # --- Zero-out padded positions so they contribute nothing downstream --- + if exists(mask): + # For lambdas: force to 0 on pads + pad_mask_f = mask.to(h.dtype) # True=valid -> 1.0 + lam_ins = lam_ins * pad_mask_f + lam_del = lam_del * pad_mask_f + lam_sub = lam_sub * pad_mask_f + + # kill logits on pads + neg_val = torch.tensor(-1e4, device=h.device, dtype=h.dtype) + logits_ins = logits_ins.masked_fill((~mask).unsqueeze(-1), neg_val) + logits_sub = logits_sub.masked_fill((~mask).unsqueeze(-1), neg_val) + + return lam_ins, logits_ins, lam_del, lam_sub, logits_sub + +class SMILESEditFlowModel(nn.Module): + """ + Inputs: + x_t: (B, L) Long + mask: (B, L) bool, True=pad (i.e., should be ignored) + t: (B,) or scalar in [0,1] + + Outputs: + lam_ins: (B, L) >= 0 + logits_ins: (B, L, V) + lam_del: (B, L) >= 0 + lam_sub: (B, L) >= 0 + logits_sub: (B, L, V) + """ + def __init__(self, vocab_size, pad_id, config): + super().__init__() + self.d_model = getattr(config, "d_model", 768) + self.n_layers = getattr(config, "n_layers", 12) + self.n_heads = getattr(config, "n_heads", 12) + self.mlp_ratio = getattr(config, "mlp_ratio", 4) + self.max_len = getattr(config, "max_len", 2048) + self.dropout = getattr(config, "dropout", 0.1) + self.attn_dropout = getattr(config, "attn_dropout", 0) + self.proj_dropout = getattr(config, "proj_dropout", 0) + self.vocab_size = vocab_size + self.pad_id = pad_id + + # --- Embedding --- + self.seq_emb = nn.Embedding(self.vocab_size, self.d_model, padding_idx=self.pad_id) + self.time_emb = TimeEmbedding(d_model=self.d_model) + + # --- RoPE shared by attention blocks --- + rope = RotaryPositionalEmbedding(head_dim=self.d_model // self.n_heads, max_len=self.max_len) + + # --- Encoder --- + self.blocks = nn.ModuleList([ + TransformerBlock( + d_model=self.d_model, + n_heads=self.n_heads, + mlp_ratio=self.mlp_ratio, + attn_dropout=self.attn_dropout, + proj_dropout=self.proj_dropout, + rope=rope + ) + for _ in range(self.n_layers) + ]) + self.final_norm = nn.LayerNorm(self.d_model) + + # --- Heads --- + # We use small MLP heads for rates; logits are linear. + self.lam_ins_head = nn.Sequential(nn.Linear(self.d_model, self.d_model//2), nn.SiLU(), nn.Linear(self.d_model//2, 1)) + self.lam_del_head = nn.Sequential(nn.Linear(self.d_model, self.d_model//2), nn.SiLU(), nn.Linear(self.d_model//2, 1)) + self.lam_sub_head = nn.Sequential(nn.Linear(self.d_model, self.d_model//2), nn.SiLU(), nn.Linear(self.d_model//2, 1)) + self.logits_ins_head = nn.Linear(self.d_model, vocab_size, bias=False) + self.logits_sub_head = nn.Linear(self.d_model, vocab_size, bias=False) + + # nonnegativity via softplus (safer than exp) + self.softplus = nn.Softplus(beta=1.0) + + def forward( + self, + x_t: torch.LongTensor, + mask: torch.BoolTensor, + t: torch.Tensor, + ): + """ + x_t: (B, L) long tokens + mask: (B, L) bool, True = PAD (ignored) + t: (B,) or scalar in [0,1] + esmbed: optional (B, L, esm2_embed_dim) if use_real_esm2=True + """ + B, L = x_t.shape + + # --- Embedding --- + h = self.seq_emb(x_t) + + # --- Add time embedding (broadcast across length) --- + t_emb = self.time_emb(t, batch_size=B) # (B, d_model) + h = h + t_emb.unsqueeze(1) # (B, L, d_model) + + # --- Encoder blocks with key padding mask --- + for blk in self.blocks: + h = blk(h, key_padding_mask=(~mask)) + + h = self.final_norm(h) # (B, L, d_model) + + # --- Heads --- + lam_ins = self.softplus(self.lam_ins_head(h)).squeeze(-1) # (B, L) + lam_del = self.softplus(self.lam_del_head(h)).squeeze(-1) # (B, L) + lam_sub = self.softplus(self.lam_sub_head(h)).squeeze(-1) # (B, L) + logits_ins = self.logits_ins_head(h) # (B, L, V) + logits_sub = self.logits_sub_head(h) # (B, L, V) + + # --- Zero-out padded positions so they contribute nothing downstream --- + if exists(mask): + # For lambdas: force to 0 on pads + pad_mask_f = mask.to(h.dtype) # True=valid -> 1.0 + lam_ins = lam_ins * pad_mask_f + lam_del = lam_del * pad_mask_f + lam_sub = lam_sub * pad_mask_f + + # kill logits on pads + neg_val = torch.tensor(-1e4, device=h.device, dtype=h.dtype) + logits_ins = logits_ins.masked_fill((~mask).unsqueeze(-1), neg_val) + logits_sub = logits_sub.masked_fill((~mask).unsqueeze(-1), neg_val) + + return lam_ins, logits_ins, lam_del, lam_sub, logits_sub + +class EditFlow(pl.LightningModule): + def __init__(self, + model, + loss_fn, + path, + source_distribution, + pad_id, + bos_id, + eos_id, + config, + ): + super().__init__() + + self.cfg = config + + self.source_distribution = source_distribution + self.path = path + self.model = model + self.loss_fn = loss_fn + + self.bos_id = bos_id + self.eos_id = eos_id + self.pad_id = pad_id + self.eps_id = getattr(self.path, "eps_id", -1) + + self._total_steps = None + + def configure_optimizers(self): + opt = torch.optim.AdamW( + self.parameters(), + lr=float(self.cfg.optim.lr), + betas=(self.cfg.optim.beta1, self.cfg.optim.beta2), + eps=float(self.cfg.optim.eps), + weight_decay=self.cfg.optim.weight_decay, + fused=self.cfg.optim.fused, + ) + + warmup_ratio = getattr(self.cfg.optim, "warmup_ratio", 0.1) + min_scale = 0.1 + + def lr_lambda(global_step: int): + # until on_train_start runs we just return 1.0 + if self._total_steps is None or self._total_steps == 0: + return 1.0 + + total_steps = self._total_steps + warmup_steps = max(1, int(warmup_ratio * total_steps)) + + if global_step < warmup_steps: + # linear warmup: 0.1 -> 1.0 + alpha = (global_step + 1) / warmup_steps + return 0.1 + 0.9 * alpha + else: + # cosine from 1.0 down to min_scale + progress = (global_step - warmup_steps) / max(1, total_steps - warmup_steps) + cosine = 0.5 * (1 + math.cos(math.pi * progress)) # 1 -> 0 + return min_scale + (1.0 - min_scale) * cosine + + sch = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda=lr_lambda) + + return { + "optimizer": opt, + "lr_scheduler": { + "scheduler": sch, + "interval": "step", # <- per-step + "frequency": 1, + }, + } + + def preparation(self, x_1): + B = x_1.shape[0] + + with torch.no_grad(): + allowed_tokens = torch.tensor([tok for tok in self.source_distribution._allowed_tokens if tok != self.eps_id]).to(self.device) + + x_0 = self.source_distribution.sample_x0_from_x1(x_1, pad_id=self.pad_id, allowed_tokens=allowed_tokens, scale_size=self.cfg.model.scale_size, bos_id = self.bos_id, eos_id = self.eos_id) + t = torch.rand(B, device=self.device) + + sched = self.path.scheduler(t) + weight = sched.d_alpha_t / sched.sigma_t # (B,) + + z_0, z_1 = build_z0_z1_with_alignment(x_0, x_1, self.eps_id, self.pad_id, self.bos_id, self.eos_id, p_optimal=self.cfg.model.p_optimal) + + z_t = self.path.sample(z_0, z_1, t=t) + x_t, mask = remove_eps(z_t, self.eps_id, self.pad_id) + + lam_ins, logits_ins, lam_del, lam_sub, logits_sub = self.model(x_t=x_t, mask=mask,t=t) + + return lam_ins, logits_ins, lam_del, lam_sub, logits_sub, z_t, z_1, x_t, mask, weight + + + def training_step(self, batch, batch_idx): + x_1 = torch.tensor(batch["input_ids"]).to(self.device) + B = x_1.shape[0] + + lam_ins, logits_ins, lam_del, lam_sub, logits_sub, z_t, z_1, x_t, mask, weight = self.preparation(x_1) + + loss = self.loss_fn(lam_ins, logits_ins, lam_del, lam_sub, logits_sub, + z_t, z_1, x_t, mask, weight, self.eps_id, self.bos_id, self.eos_id) + self.log("train_loss", loss, prog_bar=True, on_step=True, on_epoch=True, batch_size=B, sync_dist=True) + + return loss + + def validation_step(self, batch, batch_idx): + x_1 = torch.tensor(batch["input_ids"]).to(self.device) + B = x_1.shape[0] + + lam_ins, logits_ins, lam_del, lam_sub, logits_sub, z_t, z_1, x_t, mask, weight = self.preparation(x_1) + + loss = self.loss_fn(lam_ins, logits_ins, lam_del, lam_sub, logits_sub, + z_t, z_1, x_t, mask, weight, self.eps_id, self.bos_id, self.eos_id) + self.log("val_loss", loss, prog_bar=True, on_step=True, on_epoch=True, batch_size=B, sync_dist=True) + + return loss + + def on_train_start(self): + # how many optimizer steps we will take in this fit + self._total_steps = self.trainer.estimated_stepping_batches \ No newline at end of file diff --git a/model/reparam_models.py b/model/reparam_models.py new file mode 100644 index 0000000000000000000000000000000000000000..e5ef157ce2193031d5956f3bb65734821945e5e4 --- /dev/null +++ b/model/reparam_models.py @@ -0,0 +1,540 @@ +from typing import Optional, Tuple +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import EsmModel +import pytorch_lightning as pl + +from .utils import build_z0_z1_with_alignment, remove_eps +import pdb + +# ---------- Utilities ---------- + +def exists(x): return x is not None + +def default(val, d): + return val if exists(val) else d + +# ---------- Timestep embedding (sinusoidal -> MLP) ---------- + +class TimeEmbedding(nn.Module): + """ + Sinusoidal time embedding followed by a small MLP. + Accepts t of shape (B,) or scalar; outputs (B, d_model) and broadcasts over L. + """ + def __init__(self, d_model: int, hidden: Optional[int] = None, max_period: int = 10000): + super().__init__() + self.d_model = d_model + self.max_period = max_period + hidden = default(hidden, d_model * 4) + # Use even dim for sin/cos pairing + pe_dim = d_model if d_model % 2 == 0 else d_model - 1 + self.pe_dim = pe_dim + self.mlp = nn.Sequential( + nn.Linear(pe_dim, hidden), + nn.SiLU(), + nn.Linear(hidden, d_model), + ) + + def forward(self, t: torch.Tensor, batch_size: Optional[int] = None) -> torch.Tensor: + """ + t: (B,) or () in [0,1] + returns: (B, d_model) + """ + if t.dim() == 0: + # scalar -> expand to batch + if batch_size is None: + raise ValueError("When t is scalar, provide batch_size.") + t = t.expand(batch_size) + B = t.shape[0] + device = t.device + half = self.pe_dim // 2 + # frequencies + freqs = torch.exp( + torch.arange(half, device=device, dtype=t.dtype) * (-math.log(self.max_period) / (half - 1 + 1e-8)) + ) + angles = t[:, None] * freqs[None, :] * math.pi # (B, half) + pe = torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1) # (B, pe_dim) + if self.pe_dim < self.d_model: # rare case when d_model is odd + pe = F.pad(pe, (0, 1), value=0.0) + return self.mlp(pe) # (B, d_model) + +# ---------- RoPE (rotary position embedding) ---------- + +def apply_rotary(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """ + q, k: (B, h, L, d_head) + cos, sin: (L, d_head) broadcastable to (B, h, L, d_head) + """ + # split last dim into pairs + d = q.shape[-1] + if d % 2 != 0: + # pad to even + q = F.pad(q, (0, 1), value=0.0) + k = F.pad(k, (0, 1), value=0.0) + d += 1 + + q1, q2 = q[..., :d//2], q[..., d//2:] + k1, k2 = k[..., :d//2], k[..., d//2:] + + # broadcast cos/sin + while cos.dim() < q1.dim(): + cos = cos.unsqueeze(0) + sin = sin.unsqueeze(0) + + rq = torch.cat([q1 * cos - q2 * sin, q1 * sin + q2 * cos], dim=-1) + rk = torch.cat([k1 * cos - k2 * sin, k1 * sin + k2 * cos], dim=-1) + return rq[..., :q.shape[-1]], rk[..., :k.shape[-1]] + +class RotaryPositionalEmbedding(nn.Module): + """ + Precomputes cos/sin for RoPE given max_len and head_dim. + """ + def __init__(self, head_dim: int, max_len: int = 8192, base: int = 10000): + super().__init__() + if head_dim % 2 != 0: + # allow odd by padding inside apply_rotary, but prefer even + pass + inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.max_len = max_len + self.head_dim = head_dim + self._cached_len = 0 + self.register_buffer("cos_cached", torch.empty(0), persistent=False) + self.register_buffer("sin_cached", torch.empty(0), persistent=False) + + def _update_cache(self, seq_len: int, device, dtype): + if seq_len <= self._cached_len and self.cos_cached.device == device and self.cos_cached.dtype == dtype: + return + t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) + freqs = torch.einsum('n,d->nd', t, self.inv_freq) # (L, head_dim/2) + # emb = torch.cat((freqs, freqs), dim=-1) # (L, head_dim) + self.cos_cached = freqs.cos().to(dtype=dtype) + self.sin_cached = freqs.sin().to(dtype=dtype) + self._cached_len = seq_len + + def forward(self, L: int, device, dtype): + self._update_cache(L, device, dtype) + return self.cos_cached[:L], self.sin_cached[:L] + +# ---------- Rotary MHA + Transformer block ---------- + +class RotaryMHA(nn.Module): + def __init__(self, d_model: int, n_heads: int, rope: RotaryPositionalEmbedding, attn_dropout: float = 0.0, proj_dropout: float = 0.0): + super().__init__() + assert d_model % n_heads == 0, "d_model must be divisible by n_heads" + self.qkv = nn.Linear(d_model, 3 * d_model, bias=False) + self.proj = nn.Linear(d_model, d_model, bias=False) + self.n_heads = n_heads + self.d_head = d_model // n_heads + self.attn_dropout = nn.Dropout(attn_dropout) + self.proj_dropout = nn.Dropout(proj_dropout) + self.rope = rope + + def forward(self, x: torch.Tensor, key_padding_mask: Optional[torch.Tensor] = None): + """ + x: (B, L, d_model) + key_padding_mask: (B, L) bool, True=pad (masked) + """ + B, L, D = x.shape + qkv = self.qkv(x) # (B, L, 3D) + q, k, v = qkv.chunk(3, dim=-1) + q = q.view(B, L, self.n_heads, self.d_head).transpose(1, 2) # (B, h, L, d) + k = k.view(B, L, self.n_heads, self.d_head).transpose(1, 2) + v = v.view(B, L, self.n_heads, self.d_head).transpose(1, 2) + + cos, sin = self.rope(L, x.device, x.dtype) + q, k = apply_rotary(q, k, cos, sin) + + attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_head) # (B, h, L, L) + + if exists(key_padding_mask): + # mask: True means pad -> set attention to -inf + mask = key_padding_mask[:, None, None, :].to(dtype=torch.bool) # (B,1,1,L) + attn_scores = attn_scores.masked_fill(mask, float('-inf')) + + attn = torch.softmax(attn_scores, dim=-1) + attn = self.attn_dropout(attn) + out = torch.matmul(attn, v) # (B, h, L, d) + out = out.transpose(1, 2).contiguous().view(B, L, D) + return self.proj_dropout(self.proj(out)) + +class TransformerBlock(nn.Module): + def __init__(self, d_model: int, n_heads: int, mlp_ratio: float = 4.0, attn_dropout: float = 0.0, proj_dropout: float = 0.0, rope: Optional[RotaryPositionalEmbedding] = None): + super().__init__() + self.norm1 = nn.LayerNorm(d_model) + self.attn = RotaryMHA(d_model, n_heads, rope=rope, attn_dropout=attn_dropout, proj_dropout=proj_dropout) + self.norm2 = nn.LayerNorm(d_model) + hidden = int(d_model * mlp_ratio) + self.mlp = nn.Sequential( + nn.Linear(d_model, hidden), + nn.SiLU(), + nn.Linear(hidden, d_model), + ) + + def forward(self, x: torch.Tensor, key_padding_mask: Optional[torch.Tensor] = None): + x = x + self.attn(self.norm1(x), key_padding_mask=key_padding_mask) + x = x + self.mlp(self.norm2(x)) + return x + + +class ProteinEditFlowModel(nn.Module): + """ + Inputs: + x_t: (B, L) Long + mask: (B, L) bool, True=pad (i.e., should be ignored) + t: (B,) or scalar in [0,1] + + Outputs: + lam_ins: (B, L) >= 0 + logits_ins: (B, L, V) + lam_del: (B, L) >= 0 + lam_sub: (B, L) >= 0 + logits_sub: (B, L, V) + """ + def __init__(self, vocab_size, pad_id, config): + super().__init__() + self.d_model = getattr(config, "d_model", 768) + self.n_layers = getattr(config, "n_layers", 12) + self.n_heads = getattr(config, "n_heads", 12) + self.mlp_ratio = getattr(config, "mlp_ratio", 4) + self.max_len = getattr(config, "max_len", 2048) + self.dropout = getattr(config, "dropout", 0.1) + self.attn_dropout = getattr(config, "attn_dropout", 0) + self.proj_dropout = getattr(config, "proj_dropout", 0) + self.vocab_size = vocab_size + self.pad_id = pad_id + + # --- Embedding --- + esm_model_name = getattr(config, "esm_model_name", "facebook/esm2_t33_650M_UR50D") + freeze_esm = getattr(config, "freeze_esm", True) + self.esm_emb = EsmModel.from_pretrained(esm_model_name) + if freeze_esm: + for param in self.esm_emb.parameters(): + param.requires_grad = False + self.time_emb = TimeEmbedding(d_model=self.d_model) + + self.tok_embed_to_hidden = nn.Linear(self.esm_emb.config.hidden_size, self.d_model) + + # --- RoPE shared by attention blocks --- + rope = RotaryPositionalEmbedding(head_dim=self.d_model // self.n_heads, max_len=self.max_len) + + # --- Encoder --- + self.blocks = nn.ModuleList([ + TransformerBlock( + d_model=self.d_model, + n_heads=self.n_heads, + mlp_ratio=self.mlp_ratio, + attn_dropout=self.attn_dropout, + proj_dropout=self.proj_dropout, + rope=rope + ) + for _ in range(self.n_layers) + ]) + self.final_norm = nn.LayerNorm(self.d_model) + + # --- Heads --- + # We use small MLP heads for rates; logits are linear. + self.lam_total_head = nn.Sequential(nn.Linear(self.d_model, self.d_model//2), nn.SiLU(), nn.Linear(self.d_model//2, 1)) + self.logits_type_head = nn.Sequential(nn.Linear(self.d_model, self.d_model//2), nn.SiLU(), nn.Linear(self.d_model//2, 3)) + self.logits_ins_head = nn.Linear(self.d_model, vocab_size, bias=False) + self.logits_sub_head = nn.Linear(self.d_model, vocab_size, bias=False) + + # nonnegativity via softplus (safer than exp) + self.softplus = nn.Softplus(beta=1.0) + + def forward( + self, + x_t: torch.LongTensor, + mask: torch.BoolTensor, + t: torch.Tensor, + ): + """ + x_t: (B, L) long tokens + mask: (B, L) bool, True=valid (i.e., should be kept) + t: (B,) or scalar in [0,1] + """ + B, L = x_t.shape + + # --- Embedding --- + h = self.esm_emb(x_t, mask).last_hidden_state # (B, L, esm_hidden) + h = self.tok_embed_to_hidden(h) # (B, L, d_model) + + # --- Add time embedding (broadcast across length) --- + t_emb = self.time_emb(t, batch_size=B) # (B, d_model) + h = h + t_emb.unsqueeze(1) # (B, L, d_model) + + # --- Encoder blocks with key padding mask --- + for blk in self.blocks: + h = blk(h, key_padding_mask=(~mask)) + + h = self.final_norm(h) # (B, L, d_model) + + # --- Heads --- + lam_total = self.softplus(self.lam_total_head(h)).squeeze(-1) # (B, L) >= 0 + logits_type = self.logits_type_head(h) # (B, L, 3) + logits_ins = self.logits_ins_head(h) # (B, L, V) + logits_sub = self.logits_sub_head(h) # (B, L, V) + + # --- Compute type probs and per-type rates --- + pi_type = torch.softmax(logits_type, dim=-1) # (B, L, 3) + lam_ins = lam_total * pi_type[..., 0] # (B, L) + lam_del = lam_total * pi_type[..., 1] # (B, L) + lam_sub = lam_total * pi_type[..., 2] # (B, L) + + # --- Zero-out padded positions so they contribute nothing downstream --- + if exists(mask): + pad_mask_f = mask.to(h.dtype) # True=valid -> 1.0 + lam_total = lam_total * pad_mask_f + lam_ins = lam_ins * pad_mask_f + lam_del = lam_del * pad_mask_f + lam_sub = lam_sub * pad_mask_f + + neg_val = torch.tensor(-1e4, device=h.device, dtype=h.dtype) + logits_type = logits_type.masked_fill((~mask).unsqueeze(-1), neg_val) + logits_ins = logits_ins.masked_fill((~mask).unsqueeze(-1), neg_val) + logits_sub = logits_sub.masked_fill((~mask).unsqueeze(-1), neg_val) + + return lam_ins, logits_ins, lam_del, lam_sub, logits_sub, lam_total, pi_type + + + +class SMILESEditFlowModel(nn.Module): + """ + Inputs: + x_t: (B, L) Long + mask: (B, L) bool, True=pad (i.e., should be ignored) + t: (B,) or scalar in [0,1] + + Outputs: + lam_ins: (B, L) >= 0 + logits_ins: (B, L, V) + lam_del: (B, L) >= 0 + lam_sub: (B, L) >= 0 + logits_sub: (B, L, V) + """ + def __init__(self, vocab_size, pad_id, config): + super().__init__() + self.d_model = getattr(config, "d_model", 768) + self.n_layers = getattr(config, "n_layers", 12) + self.n_heads = getattr(config, "n_heads", 12) + self.mlp_ratio = getattr(config, "mlp_ratio", 4) + self.max_len = getattr(config, "max_len", 2048) + self.dropout = getattr(config, "dropout", 0.1) + self.attn_dropout = getattr(config, "attn_dropout", 0) + self.proj_dropout = getattr(config, "proj_dropout", 0) + self.vocab_size = vocab_size + self.pad_id = pad_id + + # --- Embedding --- + self.seq_emb = nn.Embedding(self.vocab_size, self.d_model, padding_idx=self.pad_id) + self.time_emb = TimeEmbedding(d_model=self.d_model) + + # --- RoPE shared by attention blocks --- + rope = RotaryPositionalEmbedding(head_dim=self.d_model // self.n_heads, max_len=self.max_len) + + # --- Encoder --- + self.blocks = nn.ModuleList([ + TransformerBlock( + d_model=self.d_model, + n_heads=self.n_heads, + mlp_ratio=self.mlp_ratio, + attn_dropout=self.attn_dropout, + proj_dropout=self.proj_dropout, + rope=rope + ) + for _ in range(self.n_layers) + ]) + self.final_norm = nn.LayerNorm(self.d_model) + + # --- Heads (REPARAM) --- + # total rate λ_total and type distribution π_type over {ins, del, sub} + self.lam_total_head = nn.Sequential( + nn.Linear(self.d_model, self.d_model // 2), + nn.SiLU(), + nn.Linear(self.d_model // 2, 1), + ) + self.type_logits_head = nn.Sequential( + nn.Linear(self.d_model, self.d_model // 2), + nn.SiLU(), + nn.Linear(self.d_model // 2, 3), # 3 edit types + ) + + # token logits (unchanged) + self.logits_ins_head = nn.Linear(self.d_model, vocab_size, bias=False) + self.logits_sub_head = nn.Linear(self.d_model, vocab_size, bias=False) + + self.softplus = nn.Softplus(beta=1.0) + + def forward( + self, + x_t: torch.LongTensor, + mask: torch.BoolTensor, + t: torch.Tensor, + ): + """ + x_t: (B, L) long tokens + mask: (B, L) bool, True = PAD (ignored) + t: (B,) or scalar in [0,1] + esmbed: optional (B, L, esm2_embed_dim) if use_real_esm2=True + """ + B, L = x_t.shape + + # --- Embedding --- + h = self.seq_emb(x_t) + + # --- Add time embedding (broadcast across length) --- + t_emb = self.time_emb(t, batch_size=B) # (B, d_model) + h = h + t_emb.unsqueeze(1) # (B, L, d_model) + + # --- Encoder blocks with key padding mask --- + for blk in self.blocks: + h = blk(h, key_padding_mask=(~mask)) + + h = self.final_norm(h) # (B, L, d_model) + + # --- Heads (REPARAM) --- + lam_total = self.softplus(self.lam_total_head(h)).squeeze(-1) # (B, L), >=0 + type_logits = self.type_logits_head(h) # (B, L, 3) + pi_type = torch.softmax(type_logits, dim=-1) # (B, L, 3), sums to 1 + + lam_ins = lam_total * pi_type[..., 0] # (B, L) + lam_del = lam_total * pi_type[..., 1] + lam_sub = lam_total * pi_type[..., 2] + + logits_ins = self.logits_ins_head(h) # (B, L, V) + logits_sub = self.logits_sub_head(h) # (B, L, V) + + # --- Masking (keep your existing behavior) --- + if exists(mask): + pad_mask_f = mask.to(h.dtype) + lam_total = lam_total * pad_mask_f + lam_ins = lam_ins * pad_mask_f + lam_del = lam_del * pad_mask_f + lam_sub = lam_sub * pad_mask_f + + neg_val = torch.tensor(-1e4, device=h.device, dtype=h.dtype) + logits_ins = logits_ins.masked_fill((~mask).unsqueeze(-1), neg_val) + logits_sub = logits_sub.masked_fill((~mask).unsqueeze(-1), neg_val) + + return lam_ins, logits_ins, lam_del, lam_sub, logits_sub, lam_total, pi_type + +class EditFlow(pl.LightningModule): + def __init__(self, + model, + loss_fn, + path, + source_distribution, + pad_id, + bos_id, + eos_id, + config, + ): + super().__init__() + + self.cfg = config + + self.source_distribution = source_distribution + self.path = path + self.model = model + self.loss_fn = loss_fn + + self.bos_id = bos_id + self.eos_id = eos_id + self.pad_id = pad_id + self.eps_id = getattr(self.path, "eps_id", -1) + + self._total_steps = None + + def configure_optimizers(self): + opt = torch.optim.AdamW( + self.parameters(), + lr=float(self.cfg.optim.lr), + betas=(self.cfg.optim.beta1, self.cfg.optim.beta2), + eps=float(self.cfg.optim.eps), + weight_decay=self.cfg.optim.weight_decay, + fused=self.cfg.optim.fused, + ) + + warmup_ratio = getattr(self.cfg.optim, "warmup_ratio", 0.1) + min_scale = 0.1 + + def lr_lambda(global_step: int): + # until on_train_start runs we just return 1.0 + if self._total_steps is None or self._total_steps == 0: + return 1.0 + + total_steps = self._total_steps + warmup_steps = max(1, int(warmup_ratio * total_steps)) + + if global_step < warmup_steps: + # linear warmup: 0.1 -> 1.0 + alpha = (global_step + 1) / warmup_steps + return 0.1 + 0.9 * alpha + else: + # cosine from 1.0 down to min_scale + progress = (global_step - warmup_steps) / max(1, total_steps - warmup_steps) + cosine = 0.5 * (1 + math.cos(math.pi * progress)) # 1 -> 0 + return min_scale + (1.0 - min_scale) * cosine + + sch = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda=lr_lambda) + + return { + "optimizer": opt, + "lr_scheduler": { + "scheduler": sch, + "interval": "step", # <- per-step + "frequency": 1, + }, + } + + def preparation(self, x_1): + B = x_1.shape[0] + + with torch.no_grad(): + allowed_tokens = torch.tensor([tok for tok in self.source_distribution._allowed_tokens if tok != self.eps_id]).to(self.device) + + x_0 = self.source_distribution.sample_x0_from_x1(x_1, pad_id=self.pad_id, allowed_tokens=allowed_tokens, scale_size=self.cfg.model.scale_size, bos_id = self.bos_id, eos_id = self.eos_id) + t = torch.rand(B, device=self.device) + + sched = self.path.scheduler(t) + weight = sched.d_alpha_t / sched.sigma_t # (B,) + + z_0, z_1 = build_z0_z1_with_alignment(x_0, x_1, self.eps_id, self.pad_id, self.bos_id, self.eos_id, p_optimal=self.cfg.model.p_optimal) + + z_t = self.path.sample(z_0, z_1, t=t) + x_t, mask = remove_eps(z_t, self.eps_id, self.pad_id) + + lam_ins, logits_ins, lam_del, lam_sub, logits_sub, lam_total, pi_type = self.model(x_t=x_t, mask=mask,t=t) + + return lam_total, pi_type, logits_ins, logits_sub, z_t, z_1, x_t, mask, weight + + + def training_step(self, batch, batch_idx): + x_1 = torch.tensor(batch["input_ids"]).to(self.device) + B = x_1.shape[0] + + lam_total, pi_type, logits_ins, logits_sub, z_t, z_1, x_t, mask, weight = self.preparation(x_1) + + loss = self.loss_fn(lam_total, pi_type, logits_ins, logits_sub, + z_t, z_1, x_t, mask, weight, self.eps_id, self.bos_id, self.eos_id) + self.log("train_loss", loss, prog_bar=True, on_step=True, on_epoch=True, batch_size=B, sync_dist=True) + + return loss + + def validation_step(self, batch, batch_idx): + x_1 = torch.tensor(batch["input_ids"]).to(self.device) + B = x_1.shape[0] + + lam_total, pi_type, logits_ins, logits_sub, z_t, z_1, x_t, mask, weight = self.preparation(x_1) + + loss = self.loss_fn(lam_total, pi_type, logits_ins, logits_sub, + z_t, z_1, x_t, mask, weight, self.eps_id, self.bos_id, self.eos_id) + self.log("val_loss", loss, prog_bar=True, on_step=True, on_epoch=True, batch_size=B, sync_dist=True) + + return loss + + def on_train_start(self): + # how many optimizer steps we will take in this fit + self._total_steps = self.trainer.estimated_stepping_batches \ No newline at end of file diff --git a/model/rotary.py b/model/rotary.py new file mode 100644 index 0000000000000000000000000000000000000000..3913cc3473e4ec73985dd52d0bf0eb6dd05842cb --- /dev/null +++ b/model/rotary.py @@ -0,0 +1,115 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. +# Part of this implementation is adapted from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/layers/rotary.py#L20 +# which is released under BSD-3 license +# Part of this implementation is adapted from https://github.com/louaaron/Score-Entropy-Discrete-Diffusion +# which is released under MIT license + +from typing import Tuple, Optional + +import torch +from einops import repeat +from torch import Tensor + + +class Rotary(torch.nn.Module): + """ + From: https://github.com/louaaron/Score-Entropy-Discrete-Diffusion + """ + + def __init__(self, dim: int, base: int = 10_000): + super().__init__() + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) + self.register_buffer("inv_freq", inv_freq) + self.seq_len_cached = None + self.cos_cached = None + self.sin_cached = None + + def positions_like(self, lengths: torch.Tensor) -> torch.Tensor: + # returns (T,) positions 0..len_i-1 concatenated across sequences + return torch.cat([torch.arange(int(L), device=lengths.device) for L in lengths], dim=0) + + def forward(self, x: Tensor, seq_dim: int = 1) -> Tuple[Tensor, Tensor]: + seq_len = x.shape[seq_dim] + if seq_len != self.seq_len_cached: + self.seq_len_cached = seq_len + t = torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq) + freqs = torch.einsum("i,j->ij", t, self.inv_freq.clone()) + emb = torch.cat((freqs, freqs), dim=-1).to(x.device) + + # dims are: batch, seq_len, qkv, head, dim + self.cos_cached = emb.cos()[None, :, None, None, :].repeat(1, 1, 3, 1, 1) + self.sin_cached = emb.sin()[None, :, None, None, :].repeat(1, 1, 3, 1, 1) + + # This makes the transformation on v an identity. + self.cos_cached[:, :, 2, :, :].fill_(1.0) + self.sin_cached[:, :, 2, :, :].fill_(0.0) + + return self.cos_cached, self.sin_cached + + +def rotate_half(x: Tensor) -> Tensor: + x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] + + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_emb_torch(x, cos, sin, interleaved: bool = False, rotary_dim: Optional[int] = None): + # cos/sin cached as in your Rotary.forward + cos = cos[0, :, 0, 0, : cos.shape[-1] // 2].to(dtype=x.dtype, device=x.device) # (S, d/2) + sin = sin[0, :, 0, 0, : sin.shape[-1] // 2].to(dtype=x.dtype, device=x.device) # (S, d/2) + + # Decide how many head dims get RoPE + if rotary_dim is None: + rotary_dim = cos.shape[-1] * 2 # default: full cached rotary span + rotary_dim = min(rotary_dim, x.shape[-1]) + assert rotary_dim % 2 == 0, "rotary_dim must be even" + + # Broadcast to (..., S, 1, rotary_dim) + cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") + sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") + + x_rot = x[..., :rotary_dim] + x_tail = x[..., rotary_dim:] + + x_rot = x_rot * cos + rotate_half(x_rot) * sin + return torch.cat([x_rot, x_tail], dim=-1) + +def apply_rotary_emb_ragged( + q: torch.Tensor, k: torch.Tensor, + cos: torch.Tensor, sin: torch.Tensor, + positions: torch.Tensor, head_dim: Optional[int] = None, + interleaved: bool = False, inplace: bool = True +): + # q,k: (T, H, Dh), positions: (T,) + cos_half = cos[0, :, 0, 0, : cos.shape[-1] // 2].to(dtype=q.dtype, device=q.device) # (S, d/2) + sin_half = sin[0, :, 0, 0, : sin.shape[-1] // 2].to(dtype=q.dtype, device=q.device) + assert positions.max().item() < cos_half.size(0), "position index exceeds cached length" + + cos_sel = cos_half.index_select(0, positions) # (T, d/2) + sin_sel = sin_half.index_select(0, positions) # (T, d/2) + + cos_sel = repeat(cos_sel, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") # (T,1,d) + sin_sel = repeat(sin_sel, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") + + ro_dim = head_dim or cos_sel.shape[-1] + ro_dim = min(ro_dim, q.shape[-1]) + assert ro_dim % 2 == 0, "rotary_dim/head_dim must be even" + + if not inplace: + q = q.clone() + k = k.clone() + + def rot(x): + x1, x2 = x[..., : ro_dim // 2], x[..., ro_dim // 2: ro_dim] + return torch.cat([-x2, x1], dim=-1) + + q_head = q[..., :ro_dim] + k_head = k[..., :ro_dim] + q[..., :ro_dim] = q_head * cos_sel + rot(q_head) * sin_sel + k[..., :ro_dim] = k_head * cos_sel + rot(k_head) * sin_sel + return q, k + diff --git a/model/transformer.py b/model/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..c0ff17989157ef6679a997f265da7506bd77e01d --- /dev/null +++ b/model/transformer.py @@ -0,0 +1,471 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. +# Part of this implementation is adapted from https://github.com/facebookresearch/DiT +# which is released under NonCommercial-4.0 license +# Part of this implementation is adapted from https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py +# which is released under MIT license +# Part of this implementation is adapted from https://github.com/louaaron/Score-Entropy-Discrete-Diffusion +# which is released under MIT license + +import math +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from torch.nn.attention.flex_attention import flex_attention + +from einops import rearrange +from omegaconf import OmegaConf +from omegaconf.dictconfig import DictConfig + +from torch import nn, Tensor +from transformers import EsmModel + +from . import rotary + +import pdb + +def lengths_to_offsets(lengths: torch.Tensor) -> torch.Tensor: + # lengths: (B,) long -> offsets: (B+1,) long + return F.pad(lengths.cumsum(0), (1, 0)) + +def build_seq_ids(lengths: torch.Tensor, device=None) -> torch.Tensor: + # lengths: (B,) -> seq_ids: (T,) where T=sum(lengths) + device = device or lengths.device + return torch.repeat_interleave(torch.arange(lengths.numel(), device=device), lengths) + +def make_score_mod_for_intra_sequence_only(lengths: torch.Tensor): + seq_ids = build_seq_ids(lengths, device=lengths.device) # (T,) + + def score_mod(scores, b, h, q_idx, k_idx): + # scores: (..., q_block, k_block) + same = (seq_ids[q_idx] == seq_ids[k_idx]) # bool, broadcastable to scores + # Set cross-sequence scores to a very negative value *without* in-place ops + neg_large = torch.finfo(scores.dtype).min # dtype-safe "−large" (no -inf issues) + return torch.where(same, scores, torch.full_like(scores, neg_large)) + + return score_mod + +# def make_score_mod_for_intra_sequence_only(lengths: torch.Tensor): +# """ +# Returns a score_mod callback for FlexAttention that sets logits to -inf +# when query and key belong to different sequences (ragged, no padding). +# """ +# seq_ids = build_seq_ids(lengths, device=lengths.device) # (T,) + +# def score_mod(scores, b, h, q_idx, k_idx): +# # scores: (..., q_block, k_block); q_idx/k_idx: flat token indices in [0..T-1] +# original_type = scores.dtype +# scores = scores.float() +# same = (seq_ids[q_idx] == seq_ids[k_idx]).to(scores.dtype) +# scores += (same - 1) * 1e9 # subtract a large number for cross-seq pairs +# return scores.to(original_type) + +# return score_mod + + + +def bias_dropout_add_scale( + x: Tensor, scale: Tensor, residual: Optional[Tensor], prob: float, training: bool +) -> Tensor: + return residual + scale * F.dropout(x, p=prob, training=training) + + +def modulate(x: Tensor, shift: Tensor, scale: Tensor) -> Tensor: + return x * (1 + scale) + shift + + +class LayerNorm(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.weight = nn.Parameter(torch.ones([dim])) + self.dim = dim + + def forward(self, x: Tensor) -> Tensor: + with torch.amp.autocast("cuda", enabled=False): + y = F.layer_norm(x.float(), [self.dim]) + if y.dim() == 3: # (B,S,H) + scale = self.weight[None, None, :] + elif y.dim() == 2: # (T,H) + scale = self.weight[None, :] + else: + raise ValueError(f"LayerNorm expects 2D/3D, got {y.shape}") + return y * scale + + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(time: Tensor, dim: int, max_period: int = 10000) -> Tensor: + """ + Create sinusoidal timestep embeddings. + :param t: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an (N, D) Tensor of positional embeddings. + """ + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(start=0, end=half, dtype=torch.float32) + / half + ).to(device=time.device) + args = time[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat( + [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 + ) + return embedding + + def forward(self, time: Tensor) -> Tensor: + t_freq = self.timestep_embedding(time=time, dim=self.frequency_embedding_size) + t_emb = self.mlp(t_freq) + return t_emb + + +class DDiTBlock(nn.Module): + def __init__( + self, + dim: int, + n_heads: int, + cond_dim: int, + mlp_ratio: int = 4, + dropout: float = 0.1, + ): + super().__init__() + assert dim % n_heads == 0, "dim must be devisable by n_heads" + + self.n_heads = n_heads + self.dim = dim + self.dropout = dropout + + self.head_dim = self.dim // self.n_heads + + self.norm1 = LayerNorm(dim=dim) + + self.qw = nn.Linear(dim, dim, bias=False) + self.kw = nn.Linear(dim, dim, bias=False) + self.vw = nn.Linear(dim, dim, bias=False) + + self.attn_out = nn.Linear(dim, dim, bias=False) + self.dropout1 = nn.Dropout(dropout) + + self.norm2 = LayerNorm(dim=dim) + self.mlp = nn.Sequential( + nn.Linear(dim, mlp_ratio * dim, bias=True), + nn.GELU(approximate="tanh"), + nn.Linear(mlp_ratio * dim, dim, bias=True), + ) + + self.adaLN_modulation = nn.Linear(cond_dim, 6 * dim, bias=True) + self.adaLN_modulation.weight.data.zero_() + self.adaLN_modulation.bias.data.zero_() + + def forward( + self, + x: Tensor, # (T, H) flat tokens; T = sum(lengths) + rotary_caches: tuple, # (cos, sin) + c: Tensor # (B, cond_dim) + ) -> Tensor: + + ( + shift_msa, + scale_msa, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + ) = self.adaLN_modulation(c)[:, None].chunk(6, dim=2) + + seq_ids = build_seq_ids(lengths, device=x.device) # (T,) + + shift_msa_b = shift_msa.squeeze(1) + scale_msa_b = scale_msa.squeeze(1) + gate_msa_b = gate_msa.squeeze(1) + shift_mlp_b = shift_mlp.squeeze(1) + scale_mlp_b = scale_mlp.squeeze(1) + gate_mlp_b = gate_mlp.squeeze(1) + + # Token-wise (T,H) + shift_msa_t = shift_msa_b[seq_ids] + scale_msa_t = scale_msa_b[seq_ids] + gate_msa_t = gate_msa_b[seq_ids] + shift_mlp_t = shift_mlp_b[seq_ids] + scale_mlp_t = scale_mlp_b[seq_ids] + gate_mlp_t = gate_mlp_b[seq_ids] + + x_skip = x + x = modulate(self.norm1(x), shift=shift_msa_t, scale=scale_msa_t) # (T, H) + + q = self.qw(x); k = self.kw(x); v = self.vw(x) # (T, H) + T = x.shape[0] + q = q.view(T, self.n_heads, self.head_dim) + k = k.view(T, self.n_heads, self.head_dim) + v = v.view(T, self.n_heads, self.head_dim) + + with torch.amp.autocast("cuda", enabled=False): + ################## + # TODO: + ################## + q = q.to(dtype); k = k.to(dtype) + + # fold heads into batch + q = q.transpose(0, 1).contiguous() # (Hh, T, Dh) + k = k.transpose(0, 1).contiguous() + v = v.transpose(0, 1).contiguous() + + q = q.unsqueeze(0) # -> (1, Hh, T, Dh) + k = k.unsqueeze(0) # -> (1, Hh, T, Dh) + v = v.unsqueeze(0) # -> (1, Hh, T, Dh) + + score_mod = make_score_mod_for_intra_sequence_only(lengths) + attn_out = flex_attention(q, k, v, score_mod=score_mod) # (Hh, T, Dh) + attn_out = attn_out.squeeze(0) # -> (Hh, T, Dh) + + x = attn_out.transpose(0, 1).contiguous().view(T, self.dim) # (T,H) + x = bias_dropout_add_scale(self.attn_out(x), gate_msa_t, x_skip, self.dropout, self.training) + + x = bias_dropout_add_scale( + self.mlp(modulate(self.norm2(x), shift=shift_mlp_t, scale=scale_mlp_t)), + gate_mlp_t, x, self.dropout, self.training + ) + return x + + + +class DDitFinalLayer(nn.Module): + def __init__(self, hidden_size: int, out_channels: int, cond_dim: int): + super().__init__() + self.norm_final = LayerNorm(hidden_size) + self.linear = nn.Linear(hidden_size, out_channels) + self.linear.weight.data.zero_() + self.linear.bias.data.zero_() + + self.adaLN_modulation = nn.Linear(cond_dim, 2 * hidden_size, bias=True) + self.adaLN_modulation.weight.data.zero_() + self.adaLN_modulation.bias.data.zero_() + + def forward(self, x: Tensor, c: Tensor) -> Tensor: + shift, scale = self.adaLN_modulation(c)[:, None].chunk(2, dim=2) + x = modulate(x=self.norm_final(x), shift=shift, scale=scale) + x = self.linear(x) + + return x + +# ------------------------------------------------------------- + +# NEW: small helper to keep λ ≥ 0 +class Positive(nn.Module): + def __init__(self): + super().__init__() + self.bias = nn.Parameter(torch.tensor(0.0)) + def forward(self, x: Tensor) -> Tensor: + return F.softplus(x + self.bias) + +# NEW: project token hidden states to "slot" hidden states (n+1) +class SlotProjector(nn.Module): + """ + Builds n+1 slot states from n token states using learnable boundaries. + slot i uses left h_{i-1} and right h_i (with learned BOS/EOS). + """ + def __init__(self, hidden_size: int): + super().__init__() + self.bos = nn.Parameter(torch.zeros(hidden_size)) + self.eos = nn.Parameter(torch.zeros(hidden_size)) + self.proj = nn.Linear(2 * hidden_size, hidden_size, bias=True) + self.act = nn.GELU() + + def forward(self, h_tok: Tensor) -> Tensor: + # h_tok: (B, n, H) + B, n, H = h_tok.shape + bos = self.bos.expand(B, 1, H) # (B, 1, H) + eos = self.eos.expand(B, 1, H) # (B, 1, H) + + # left/right neighbors for the n+1 between-token slots + left = torch.cat([bos, h_tok], dim=1) # (B, n+1, H) + right = torch.cat([h_tok, eos], dim=1) # (B, n+1, H) + + slots = torch.cat([left, right], dim=-1) # (B, n+1, 2H) + return self.act(self.proj(slots)) # (B, n+1, H) + +# NEW: multi-head output layer for edit flows +class EditFlowsHead(nn.Module): + """ + Produces: + - λ_ins: (B, n+1), Q_ins: (B, n+1, V) + - λ_del: (B, n), λ_sub: (B, n), Q_sub: (B, n, V) + """ + def __init__(self, hidden_size: int, vocab_size: int, cond_dim: int): + super().__init__() + self.norm = LayerNorm(hidden_size) + self.cond = nn.Linear(cond_dim, 2 * hidden_size, bias=True) + self.cond.weight.data.zero_() + self.cond.bias.data.zero_() + + # token-position heads (n) + self.lambda_del = nn.Linear(hidden_size, 1, bias=True) + self.lambda_sub = nn.Linear(hidden_size, 1, bias=True) + self.q_sub = nn.Linear(hidden_size, vocab_size, bias=True) + + # slot-position heads (n+1) + self.slot_proj = SlotProjector(hidden_size) + self.lambda_ins = nn.Linear(hidden_size, 1, bias=True) + self.q_ins = nn.Linear(hidden_size, vocab_size, bias=True) + + self.to_positive = Positive() + + # init: keep outputs small at start + for m in [self.lambda_del, self.lambda_sub, self.lambda_ins, self.q_sub, self.q_ins]: + nn.init.zeros_(m.weight); nn.init.zeros_(m.bias) + + def forward(self, h_tok: Tensor, c: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + # FiLM-style modulation like your DDitFinalLayer + shift, scale = self.cond(c)[:, None].chunk(2, dim=-1) + h_tok = modulate(self.norm(h_tok), shift=shift, scale=scale) # (B, n, H) + + # token positions (n) + lam_del = self.to_positive(self.lambda_del(h_tok)).squeeze(-1) # (B, n) + lam_sub = self.to_positive(self.lambda_sub(h_tok)).squeeze(-1) # (B, n) + q_sub = self.q_sub(h_tok) # (B, n, V) -- removed softmax here, may need to adjust tensor shape + + # slot positions (n+1) + h_slot = self.slot_proj(h_tok) # (B, n+1, H) + lam_ins = self.to_positive(self.lambda_ins(h_slot)).squeeze(-1) # (B, n+1) + q_ins = self.q_ins(h_slot) # (B, n+1, V) -- removed softmax here, may need to adjust tensor shape + + return lam_ins, q_ins, lam_del, lam_sub, q_sub + + +class Transformer(nn.Module): + def __init__(self, vocab_size: int, masked: bool, config: DictConfig): + super().__init__() + if isinstance(config, dict): + config = OmegaConf.create(config) + self.config = config + self.vocab_size = vocab_size + self.pad_id = getattr(config, "pad_id", 0) + + add_token = 1 if masked else 0 # keep if you need a mask token elsewhere + + # ESM-2 embedding approach (similar to gpm_model.py) + esm_model_name = getattr(config, "esm_model_name", "facebook/esm2_t33_650M_UR50D") + freeze_esm = getattr(config, "freeze_esm", True) + + if esm_model_name is not None: + self.tok_embedder = EsmModel.from_pretrained(esm_model_name) + tok_embed_dim = self.tok_embedder.config.hidden_size + + if freeze_esm: + for param in self.tok_embedder.parameters(): + param.requires_grad = False + self.tok_embedder.eval() + + # Project from ESM hidden size to model hidden size + self.tok_embed_to_hidden = nn.Linear(tok_embed_dim, config.hidden_size) + self.vocab_embed = None # Not needed when using ESM + else: + # Fallback to embedding layer if ESM is not used + self.tok_embedder = None + self.tok_embed_to_hidden = None # Not needed when using standard embedding + self.vocab_embed = nn.Embedding(self.vocab_size + add_token, config.hidden_size) + + self.time_embedding = TimestepEmbedder(hidden_size=config.cond_dim) + self.rotary_emb = rotary.Rotary(dim=config.hidden_size // config.n_heads) + + self.blocks = nn.ModuleList( + [ + DDiTBlock( + dim=config.hidden_size, + n_heads=config.n_heads, + cond_dim=config.cond_dim, + dropout=config.dropout, + ) + for _ in range(config.n_blocks) + ] + ) + + # CHANGED: use EditFlowsHead instead of DDitFinalLayer + self.output_layer = EditFlowsHead( + hidden_size=config.hidden_size, + vocab_size=vocab_size + add_token, + cond_dim=config.cond_dim, + ) + + def _embed_ragged(self, x_t, mask): + # Handle device - use ESM device or vocab_embed device + if self.tok_embedder is not None: + device = next(self.tok_embedder.parameters()).device + else: + device = self.vocab_embed.weight.device + + if self.tok_embedder is not None: + x_t_emb = self.tok_embedder(x_t, mask).last_hidden_state + x = self.tok_embed_to_hidden(x_t_emb) + else: + # Fallback to standard embedding + x = self.vocab_embed(x_t) # (T, H) + + # positions from RoPE helper: + B, S = x_t.shape + lengths = torch.full((B,), S, device=device, dtype=torch.long) + positions = self.rotary_emb.positions_like(lengths) # (T,) + return x, lengths, positions + + + def forward(self, x_t, time: torch.Tensor): + """ + x_t: List[LongTensor] of variable lengths OR a (B,S) tensor + time: (B,) float in [0,1] + returns lists (ragged): lam_ins, q_ins, lam_del, lam_sub, q_sub + """ + # 0) pad x_t + x_t = torch.nn.utils.rnn.pad_sequence(x_t, batch_first=True, padding_value=self.pad_id) + mask = (x_t != self.pad_id).long() + + # 1) embed ragged + x_emb = self._embed_ragged(x_t, mask) + + # 2) time conditioning + B = x_t.shape[0] + t_emb = F.silu(self.time_embedding(time)) # (B, cond_dim) + + # 3) rotary caches (you can keep your existing cache logic) + ############### + # TODO: + ############### + + # 4) transformer blocks + with torch.amp.autocast("cuda", dtype=torch.bfloat16): + for blk in self.blocks: + x = blk( + x=x_emb, + mask=mask, + rotary_caches=rotary_caches, + c=t_emb, + ) + + ############# + # TODO: + ############# + lam_ins, q_ins, lam_del, lam_sub, q_sub = zip(*outs) + + # 6) return ragged lists (keep padding outside the model if needed) + return list(lam_ins), list(q_ins), list(lam_del), list(lam_sub), list(q_sub) diff --git a/model/utils.py b/model/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2e4c8777c527262274cca7da9bf3f907e22f862f --- /dev/null +++ b/model/utils.py @@ -0,0 +1,533 @@ +from dataclasses import dataclass +from typing import List, Tuple, Optional, Sequence + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.utils.rnn import pad_sequence +from rdkit import Chem + +Tensor = torch.Tensor + + +def _optimal_align_core(core0: torch.Tensor, core1: torch.Tensor, eps_id: int): + """ + Edit-distance alignment on the *core* (no BOS/EOS). + Returns two python lists of ints of the same length, using eps_id for gaps. + """ + L0 = core0.size(0) + L1 = core1.size(0) + + dp = torch.zeros((L0 + 1, L1 + 1), dtype=torch.long, device=core0.device) + for i in range(1, L0 + 1): + dp[i, 0] = i + for j in range(1, L1 + 1): + dp[0, j] = j + + for i in range(1, L0 + 1): + for j in range(1, L1 + 1): + cost_sub = 0 if core0[i-1].item() == core1[j-1].item() else 1 + dp[i, j] = min( + dp[i-1, j] + 1, # delete core0[i-1] + dp[i, j-1] + 1, # insert core1[j-1] + dp[i-1, j-1] + cost_sub # match/sub + ) + + z0_core = [] + z1_core = [] + i, j = L0, L1 + while i > 0 or j > 0: + if i > 0 and j > 0: + cost_sub = 0 if core0[i-1].item() == core1[j-1].item() else 1 + if dp[i, j].item() == dp[i-1, j-1].item() + cost_sub: + z0_core.append(int(core0[i-1].item())) + z1_core.append(int(core1[j-1].item())) + i -= 1 + j -= 1 + continue + if i > 0 and dp[i, j].item() == dp[i-1, j].item() + 1: + z0_core.append(int(core0[i-1].item())) + z1_core.append(eps_id) + i -= 1 + continue + if j > 0 and dp[i, j].item() == dp[i, j-1].item() + 1: + z0_core.append(eps_id) + z1_core.append(int(core1[j-1].item())) + j -= 1 + continue + + z0_core.reverse() + z1_core.reverse() + return z0_core, z1_core + + +def _suboptimal_align_core(core0: torch.Tensor, core1: torch.Tensor, eps_id: int): + """ + Left-align cores; pad the shorter core with eps_id. + """ + L0 = core0.size(0) + L1 = core1.size(0) + N = max(L0, L1) + z0_core, z1_core = [], [] + for k in range(N): + tok0 = int(core0[k].item()) if k < L0 else eps_id + tok1 = int(core1[k].item()) if k < L1 else eps_id + z0_core.append(tok0) + z1_core.append(tok1) + return z0_core, z1_core + + +def build_z0_z1_with_alignment( + x0: torch.Tensor, # (B, L0), padded with pad_id, contains BOS/EOS + x1: torch.Tensor, # (B, L1), padded with pad_id, contains BOS/EOS + eps_id: int, + pad_id: int, + bos_id: int, + eos_id: int, + p_optimal: float = 0.6, +): + """ + Align x0 and x1 such that: + - BOS aligns with BOS + - EOS aligns with EOS + - between BOS and EOS we align with eps_id + - after EOS we pad with pad_id + + Returns: + z0: (B, N_max) + z1: (B, N_max) + """ + device = x0.device + B = x0.size(0) + + z0_list = [] + z1_list = [] + max_len = 0 + + rand = torch.rand(B, device=device) + + for b in range(B): + # strip pads + seq0 = x0[b][x0[b] != pad_id] # e.g. [BOS, ..., EOS] + seq1 = x1[b][x1[b] != pad_id] + + # find BOS/EOS positions (assume 1 each, in order) + # usually BOS is at index 0, but let's be safe + bos_pos0 = (seq0 == bos_id).nonzero(as_tuple=False)[0, 0].item() + bos_pos1 = (seq1 == bos_id).nonzero(as_tuple=False)[0, 0].item() + eos_pos0 = (seq0 == eos_id).nonzero(as_tuple=False)[0, 0].item() + eos_pos1 = (seq1 == eos_id).nonzero(as_tuple=False)[0, 0].item() + + # cores: everything between BOS and EOS + core0 = seq0[bos_pos0 + 1 : eos_pos0] # may be empty + core1 = seq1[bos_pos1 + 1 : eos_pos1] + + # pick alignment strategy for the core + if rand[b].item() < p_optimal: + core0_aligned, core1_aligned = _optimal_align_core(core0, core1, eps_id) + else: + core0_aligned, core1_aligned = _suboptimal_align_core(core0, core1, eps_id) + + # rebuild full aligned sequences: [BOS] + core_aligned + [EOS] + aligned0 = [bos_id] + core0_aligned + [eos_id] + aligned1 = [bos_id] + core1_aligned + [eos_id] + + cur_len = len(aligned0) + assert cur_len == len(aligned1) + if cur_len > max_len: + max_len = cur_len + + z0_list.append(aligned0) + z1_list.append(aligned1) + + # pad with pad_id AFTER eos + z0 = torch.full((B, max_len), pad_id, dtype=torch.long, device=device) + z1 = torch.full((B, max_len), pad_id, dtype=torch.long, device=device) + + for b in range(B): + cur = len(z0_list[b]) + z0[b, :cur] = torch.tensor(z0_list[b], device=device, dtype=torch.long) + z1[b, :cur] = torch.tensor(z1_list[b], device=device, dtype=torch.long) + + return z0, z1 + +def remove_eps( + z_t: torch.Tensor, # (B, N) + eps_id: int, + pad_id: int, + return_mask: bool = True, +): + device = z_t.device + B, N = z_t.shape + + x_t = [] + for b in range(B): + seq = z_t[b] + core = seq[seq != eps_id] # remove eps + x_t.append(core) + + x_t = pad_sequence(x_t, batch_first=True, padding_value=pad_id) + mask = (x_t != pad_id).bool() + + if return_mask: + return x_t, mask + return x_t + +@torch.no_grad() +def generate_from_x0( + model, + x0: torch.Tensor, # (B, L) long, has BOS/EOS, padded with pad_id + *, + pad_id: int, + bos_id: int, + eos_id: int, + allowed_tokens: torch.Tensor = None, # 1D tensor of vocab ids we can generate + num_steps: int = 32, + max_len_cap: int = None, + op_temperature: float = 1.0, # temperature for choosing insert vs delete vs sub + token_temperature: float = 1.0, # temperature for choosing the token to insert/sub + device: torch.device = None, +): + """ + Discrete edit sampler for Edit Flows with temperature on: + - operation choice (insert/delete/sub) + - token choice (for insert/sub) + + At each step we apply at most ONE edit per sequence. + """ + if device is None: + device = x0.device + x = x0.clone().to(device) + B = x.size(0) + + def sample_token_from_logits(logits_row: torch.Tensor) -> int: + """ + logits_row: (V,) + Apply temperature + allowed_tokens filtering, then sample. + """ + logit = logits_row + if allowed_tokens is not None: + mask = torch.zeros_like(logit, dtype=torch.bool) + mask[allowed_tokens] = True + logit = logit.masked_fill(~mask, -1e4) + + if token_temperature is not None and token_temperature > 0.0: + logit = logit / token_temperature + + probs = F.softmax(logit, dim=-1) + # multinomial expects probs >= 0 and sum=1 + idx = torch.multinomial(probs, num_samples=1) + return int(idx.item()) + + for step in range(num_steps): + # t in [0,1] + t = torch.full((B,), float(step) / float(max(1, num_steps - 1)), device=device) + + # build mask: True = valid, False = pad + mask = (x != pad_id) + + # forward through model + lam_ins, logits_ins, lam_del, lam_sub, logits_sub, lam_total, pi_type = model(x_t=x, mask=mask, t=t) + + # collect new sequences + new_seqs = [] + max_len_this_round = 0 + + for b in range(B): + seq = x[b] + valid = (seq != pad_id) + tokens = seq[valid].tolist() # python list + + if len(tokens) == 0: + new_seq = torch.tensor([], device=device, dtype=torch.long) + new_seqs.append(new_seq) + continue + + # find EOS pos + try: + eos_pos = tokens.index(eos_id) + except ValueError: + eos_pos = len(tokens) - 1 + + length_b = valid.sum().item() + lam_ins_b = lam_ins[b] + lam_del_b = lam_del[b] + lam_sub_b = lam_sub[b] + logits_ins_b = logits_ins[b] + logits_sub_b = logits_sub[b] + + # --- collect best candidate per op --- + + # insertion: pick position with highest lambda, but skip after EOS + best_ins_pos = None + best_ins_val = 0.0 + for i in range(length_b): + if tokens[i] == eos_id: + continue + val = lam_ins_b[i].item() + if val > best_ins_val: + best_ins_val = val + best_ins_pos = i + + # deletion: pick position with highest lambda, skip BOS/EOS + best_del_pos = None + best_del_val = 0.0 + for i in range(length_b): + if tokens[i] == bos_id or tokens[i] == eos_id: + continue + val = lam_del_b[i].item() + if val > best_del_val: + best_del_val = val + best_del_pos = i + + # substitution: pick position with highest lambda, skip BOS/EOS + best_sub_pos = None + best_sub_val = 0.0 + for i in range(length_b): + if tokens[i] == bos_id or tokens[i] == eos_id: + continue + val = lam_sub_b[i].item() + if val > best_sub_val: + best_sub_val = val + best_sub_pos = i + + # --- choose which operation to apply --- + # we form a 3-vector of op "scores" = the lambdas + op_scores = torch.tensor( + [best_ins_val, best_del_val, best_sub_val], + device=device, + dtype=torch.float32, + ) + + # if all zero-ish, just keep sequence + if torch.all(op_scores <= 1e-6): + new_seq = torch.tensor(tokens, device=device, dtype=torch.long) + new_seqs.append(new_seq) + max_len_this_round = max(max_len_this_round, new_seq.size(0)) + continue + + # temperature over ops + if op_temperature is not None and op_temperature > 0.0: + op_logits = op_scores / op_temperature + op_probs = F.softmax(op_logits, dim=0) + op_idx = int(torch.multinomial(op_probs, 1).item()) + else: + op_idx = int(torch.argmax(op_scores).item()) + + # 0 -> insert, 1 -> delete, 2 -> sub + if op_idx == 0: + # insertion + pos = best_ins_pos + if pos is not None: + ins_tok = sample_token_from_logits(logits_ins_b[pos]) + tokens = tokens[:pos + 1] + [ins_tok] + tokens[pos + 1:] + + elif op_idx == 1: + # deletion + pos = best_del_pos + if pos is not None: + tokens = tokens[:pos] + tokens[pos + 1:] + + else: + # substitution + pos = best_sub_pos + if pos is not None: + sub_tok = sample_token_from_logits(logits_sub_b[pos]) + tokens = tokens[:pos] + [sub_tok] + tokens[pos + 1:] + + # make sure we still end with EOS + if len(tokens) == 0 or tokens[-1] != eos_id: + tokens.append(eos_id) + + # enforce max_len_cap + if max_len_cap is not None and len(tokens) > max_len_cap: + tokens = tokens[:max_len_cap] + if tokens[-1] != eos_id: + tokens[-1] = eos_id + + new_seq = torch.tensor(tokens, device=device, dtype=torch.long) + new_seqs.append(new_seq) + max_len_this_round = max(max_len_this_round, new_seq.size(0)) + + # pad batch back to tensor + x = x.new_full((B, max_len_this_round), pad_id) + for b, seq_b in enumerate(new_seqs): + x[b, :seq_b.size(0)] = seq_b + + return x + +def generate_from_x0_multi_edit( + model, + x0: torch.Tensor, # (B, L) long, has BOS/EOS, padded with pad_id + *, + pad_id: int, + bos_id: int, + eos_id: int, + allowed_tokens: torch.Tensor = None, # 1D tensor of vocab ids we can generate + num_steps: int = 32, + max_len_cap: int = None, + op_temperature: float = 1.0, # temperature for choosing insert vs delete vs sub + token_temperature: float = 1.0, # temperature for choosing the token to insert/sub + device: torch.device = None, +): + """ + Multi-edit discrete edit sampler for Edit Flows. + + At each step: + - For each position i, independently "fire" an edit with probability + p_i = 1 - exp(-delta * lambda_i), + where lambda_i = lam_ins[i] + lam_del[i] + lam_sub[i] (after masking illegal ops). + - If fired, sample ONE op type at that position (ins/del/sub) proportional to rates, + with optional op_temperature. + - For ins/sub, sample token from logits with optional token_temperature and allowed_tokens. + - Apply edits in a single left-to-right pass (avoids index-shift headaches). + """ + if device is None: + device = x0.device + x = x0.clone().to(device) + B = x.size(0) + + # User-requested: delta = 1 / num_steps + delta = 1.0 / float(max(1, num_steps)) + + def sample_token_from_logits(logits_row: torch.Tensor) -> int: + """ + logits_row: (V,) + Apply temperature + allowed_tokens filtering, then sample. + """ + logit = logits_row + if allowed_tokens is not None: + mask = torch.zeros_like(logit, dtype=torch.bool) + mask[allowed_tokens] = True + logit = logit.masked_fill(~mask, -1e4) + + if token_temperature is not None and token_temperature > 0.0: + logit = logit / token_temperature + + probs = F.softmax(logit, dim=-1) + idx = torch.multinomial(probs, num_samples=1) + return int(idx.item()) + + for step in range(num_steps): + # t in [0,1] + t = torch.full((B,), float(step) / float(max(1, num_steps - 1)), device=device) + + # mask: True = valid (non-pad) + mask = (x != pad_id) + + # forward + lam_ins, logits_ins, lam_del, lam_sub, logits_sub, lam_total, pi_type = model(x_t=x, mask=mask, t=t) + + new_seqs = [] + max_len_this_round = 0 + + for b in range(B): + seq = x[b] + valid = (seq != pad_id) + tokens = seq[valid].tolist() + + if len(tokens) == 0: + new_seq = torch.tensor([], device=device, dtype=torch.long) + new_seqs.append(new_seq) + continue + + # Ensure there's an EOS somewhere (fallback: append later) + if eos_id not in tokens: + tokens = tokens + [eos_id] + + Lb = len(tokens) + + lam_ins_b = lam_ins[b][:Lb].clone() + lam_del_b = lam_del[b][:Lb].clone() + lam_sub_b = lam_sub[b][:Lb].clone() + logits_ins_b = logits_ins[b][:Lb] + logits_sub_b = logits_sub[b][:Lb] + + # --- operation legality masks at current positions --- + tok_tensor = torch.tensor(tokens, device=device, dtype=torch.long) + is_bos = (tok_tensor == bos_id) + is_eos = (tok_tensor == eos_id) + + # insertion not allowed at EOS + lam_ins_b = lam_ins_b.masked_fill(is_eos, 0.0) + # deletion/substitution not allowed at BOS/EOS + lam_del_b = lam_del_b.masked_fill(is_bos | is_eos, 0.0) + lam_sub_b = lam_sub_b.masked_fill(is_bos | is_eos, 0.0) + + lam_pos_total = lam_ins_b + lam_del_b + lam_sub_b + + # fire prob per position + # p_i = 1 - exp(-delta * lambda_i) + p_fire = 1.0 - torch.exp(-delta * lam_pos_total.clamp(min=0.0)) + + # sample fired positions + fired = (torch.rand(Lb, device=device) < p_fire) & (lam_pos_total > 1e-12) + + # sample op type (0=ins, 1=del, 2=sub) for ALL positions (we'll use only where fired) + rates = torch.stack([lam_ins_b, lam_del_b, lam_sub_b], dim=-1) # (Lb, 3) + + # temperature over ops: probs ∝ rate^(1/temp) == softmax(log(rate)/temp) + if op_temperature is not None and op_temperature > 0.0: + op_logits = torch.log(rates + 1e-20) / op_temperature + op_probs = F.softmax(op_logits, dim=-1) + else: + # greedy: pick max-rate op; represent as one-hot probs for multinomial compatibility + op_idx_greedy = torch.argmax(rates, dim=-1) # (Lb,) + op_probs = F.one_hot(op_idx_greedy, num_classes=3).float() + + # multinomial per row + # torch.multinomial accepts (n, m) -> (n, num_samples) + op_idx = torch.multinomial(op_probs, num_samples=1).squeeze(-1) # (Lb,) + + # pre-sample tokens for fired ins/sub positions (loop only over fired positions) + ins_tok_map = {} + sub_tok_map = {} + fired_idx = fired.nonzero(as_tuple=True)[0].tolist() + for i in fired_idx: + oi = int(op_idx[i].item()) + if oi == 0: + # insertion + ins_tok_map[i] = sample_token_from_logits(logits_ins_b[i]) + elif oi == 2: + # substitution + sub_tok_map[i] = sample_token_from_logits(logits_sub_b[i]) + + # apply edits in one pass (left-to-right) + out = [] + for i in range(Lb): + tok = tokens[i] + if fired[i]: + oi = int(op_idx[i].item()) + if oi == 1: + # deletion (already masked for BOS/EOS) + continue + elif oi == 2: + # substitution + tok = sub_tok_map.get(i, tok) + + out.append(tok) + + # insertion happens AFTER this token (and never after EOS, due to masking) + if fired[i] and int(op_idx[i].item()) == 0: + out.append(ins_tok_map.get(i)) + + # ensure EOS at end + if len(out) == 0 or out[-1] != eos_id: + out.append(eos_id) + + # enforce max_len_cap + if max_len_cap is not None and len(out) > max_len_cap: + out = out[:max_len_cap] + if out[-1] != eos_id: + out[-1] = eos_id + + new_seq = torch.tensor(out, device=device, dtype=torch.long) + new_seqs.append(new_seq) + max_len_this_round = max(max_len_this_round, new_seq.size(0)) + + # pad batch + x = x.new_full((B, max_len_this_round), pad_id) + for b, seq_b in enumerate(new_seqs): + x[b, :seq_b.size(0)] = seq_b + + return x \ No newline at end of file diff --git a/paths.py b/paths.py new file mode 100644 index 0000000000000000000000000000000000000000..86e73155b94d3fab1269f41367cff18d90d55243 --- /dev/null +++ b/paths.py @@ -0,0 +1,64 @@ +"""Repo-root path helpers. All scripts should import this instead of hardcoding hosts.""" +from __future__ import annotations + +import os +import shutil +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parent + + +def add_repo_to_sys_path() -> Path: + root = str(REPO_ROOT) + if root not in sys.path: + sys.path.insert(0, root) + return REPO_ROOT + + +def resolve_path(path: str | Path) -> Path: + p = Path(path).expanduser() + if p.is_absolute(): + return p + return (REPO_ROOT / p).resolve() + + +def peptiverse_root() -> Path: + env = os.environ.get("PEPTIVERSE_ROOT") + if env: + root = Path(env).expanduser().resolve() + else: + root = (REPO_ROOT.parent / "PeptiVerse").resolve() + if not root.is_dir(): + raise FileNotFoundError( + "PeptiVerse was not found. Clone it first and either place it next " + f"to this repo ({REPO_ROOT.parent / 'PeptiVerse'}) or set PEPTIVERSE_ROOT.\n" + " git clone https://huggingface.co/ChatterjeeLab/PeptiVerse" + ) + return root + + +def mafft_path() -> str: + env = os.environ.get("MAFFT_PATH") + if env: + if Path(env).is_file() and os.access(env, os.X_OK): + return env + raise FileNotFoundError(f"MAFFT_PATH is set but not an executable: {env}") + found = shutil.which("mafft") + if found: + return found + raise FileNotFoundError( + "MAFFT was not found on PATH. Install it and make sure `mafft` is " + "executable, or set MAFFT_PATH to the binary. GFP FPredX oracles require MAFFT." + ) + + +def smiles_vocab_files() -> tuple[Path, Path]: + return ( + REPO_ROOT / "smiles_tokenizer" / "new_vocab.txt", + REPO_ROOT / "smiles_tokenizer" / "new_splits.txt", + ) + + +def selfies_vocab_path() -> Path: + return REPO_ROOT / "smiles_tokenizer" / "selfies_vocab.json" diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..885e17bcbb442bc984f64126f7a3ca0106a1fd04 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,29 @@ +# Core Edit Flow / pCoMole +torch>=2.1.0 +pytorch-lightning>=2.0.0 +lightning>=2.0.0 +transformers==4.46.0 +datasets>=2.14.0 +easydict +pyyaml +numpy>=1.24.0 +pandas>=2.0.0 +tqdm +einops +omegaconf +joblib>=1.3.0 +scikit-learn>=1.3.0 + +# Peptidomimetic chemistry +rdkit>=2023.3.1 +selfies +SmilesPE +networkx +torch-geometric + +# Oracles used by pCoMole +xgboost>=2.0.0 +chemprop>=2.0.0 + +# Optional logging +# wandb diff --git a/smiles_tokenizer/.DS_Store b/smiles_tokenizer/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..5008ddfcf53c02e82d7eee2e57c38e5672ef89f6 Binary files /dev/null and b/smiles_tokenizer/.DS_Store differ diff --git a/smiles_tokenizer/my_tokenizers.py b/smiles_tokenizer/my_tokenizers.py new file mode 100644 index 0000000000000000000000000000000000000000..14787878748642bc0a44586b74c042a82569b095 --- /dev/null +++ b/smiles_tokenizer/my_tokenizers.py @@ -0,0 +1,468 @@ +import collections +import logging +import os +import re +import codecs +import unicodedata +from typing import List, Optional +from transformers import PreTrainedTokenizer +from SmilesPE.tokenizer import SPE_Tokenizer +import torch +import torch.nn.utils.rnn as rnn_utils +import pdb + +def load_vocab(vocab_file): + """Loads a vocabulary file into a dictionary.""" + vocab = collections.OrderedDict() + with open(vocab_file, "r", encoding="utf-8") as reader: + tokens = reader.readlines() + for index, token in enumerate(tokens): + token = token.rstrip("\n") + vocab[token] = index + return vocab + +class Atomwise_Tokenizer(object): + """Run atom-level SMILES tokenization""" + + def __init__(self): + """ Constructs a atom-level Tokenizer. + """ + # self.regex_pattern = r"(\[[^\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\|\/|:|~|@|\?|>>?|\*|\$|\%[0-9]{2}|[0-9])" + self.regex_pattern = r"(\([^\(\)]{0,4}\)|\[[^\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\|\/\/?|:|~|@|\?|>>?|\*|\$|\%[0-9]{2}|[0-9])" + + self.regex = re.compile(self.regex_pattern) + + def tokenize(self, text): + """ Basic Tokenization of a SMILES. + """ + tokens = [token for token in self.regex.findall(text)] + return tokens + +class SMILES_SPE_Tokenizer(PreTrainedTokenizer): + r""" + Constructs a SMILES tokenizer. Based on SMILES Pair Encoding (https://github.com/XinhaoLi74/SmilesPE). + This tokenizer inherits from :class:`~transformers.PreTrainedTokenizer` which contains most of the methods. Users + should refer to the superclass for more information regarding methods. + Args: + vocab_file (:obj:`string`): + File containing the vocabulary. + spe_file (:obj:`string`): + File containing the trained SMILES Pair Encoding vocabulary. + unk_token (:obj:`string`, `optional`, defaults to "[UNK]"): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + sep_token (:obj:`string`, `optional`, defaults to "[SEP]"): + The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences + for sequence classification or for a text and a question for question answering. + It is also used as the last token of a sequence built with special tokens. + pad_token (:obj:`string`, `optional`, defaults to "[PAD]"): + The token used for padding, for example when batching sequences of different lengths. + cls_token (:obj:`string`, `optional`, defaults to "[CLS]"): + The classifier token which is used when doing sequence classification (classification of the whole + sequence instead of per-token classification). It is the first token of the sequence when built with + special tokens. + mask_token (:obj:`string`, `optional`, defaults to "[MASK]"): + The token used for masking values. This is the token used when training this model with masked language + modeling. This is the token which the model will try to predict. + """ + + def __init__(self, vocab_file, spe_file, + unk_token="[UNK]", + sep_token="[SEP]", + pad_token="[PAD]", + cls_token="[CLS]", + mask_token="[MASK]", + **kwargs): + if not os.path.isfile(vocab_file): + raise ValueError("Can't find a vocabulary file at path '{}'.".format(vocab_file)) + if not os.path.isfile(spe_file): + raise ValueError("Can't find a SPE vocabulary file at path '{}'.".format(spe_file)) + + self.vocab = load_vocab(vocab_file) + self.spe_vocab = open(spe_file, 'r', encoding='utf-8') + self.ids_to_tokens = collections.OrderedDict([(ids, tok) for tok, ids in self.vocab.items()]) + self.spe_tokenizer = SPE_Tokenizer(self.spe_vocab) + + super().__init__( + unk_token=unk_token, + sep_token=sep_token, + pad_token=pad_token, + cls_token=cls_token, + mask_token=mask_token, + **kwargs) + + @property + def vocab_size(self): + return len(self.vocab) + + def get_vocab(self): + return dict(self.vocab, **self.added_tokens_encoder) + + def _tokenize(self, text): + return self.spe_tokenizer.tokenize(text).split(' ') + + def _convert_token_to_id(self, token): + """ Converts a token (str) in an id using the vocab. """ + return self.vocab.get(token, self.vocab.get(self.unk_token)) + + # changed encode and decode functions + def encode(self, token_array): + token_ids = [] + token_ids.append(2) + for token in token_array: + id = self._convert_token_to_id(token) + token_ids.append(id) + token_ids.append(3) + token_ids = torch.tensor([token_ids]) + attn_mask = torch.ones_like(token_ids) + return {'input_ids': token_ids, 'attention_mask': attn_mask} + + def batch_encode(self, sequences): + input_ids_list = [] + attention_masks_list = [] + for sequence in sequences: + res = self.encode(sequence) + input_ids_list.append(res['input_ids'][0]) + attention_masks_list.append(res['attention_mask'][0]) + + input_ids_tensor = rnn_utils.pad_sequence( + input_ids_list, + batch_first=True, + padding_value=0 + ) + + attention_mask_tensor = rnn_utils.pad_sequence( + attention_masks_list, + batch_first=True, + padding_value=0 # Pad attention mask with 0s + ) + + return { + 'input_ids': input_ids_tensor, + 'attention_mask': attention_mask_tensor + } + + def decode(self, token_ids, skip_special_tokens=True): + token_ids = token_ids.squeeze(0).cpu().tolist() + token_array = [] + for idx in token_ids: + if idx == 3: # Stop decoding when token ID 3 is encountered + break + if skip_special_tokens and idx in self.all_special_ids: + continue + token = self._convert_id_to_token(idx) + token_array.append(token) + sequence = "".join(token_array) + return sequence + + def batch_decode(self, batch_token_ids, skip_special_tokens=True): + sequences = [] + for token_ids in batch_token_ids: + sequences.append(self.decode(token_ids)) + return sequences + + def get_token_split(self, token_ids): + if isinstance(token_ids, torch.Tensor): + token_ids = token_ids.cpu().tolist() + + token_array = [] + for seq_ids in token_ids: + seq_array = [] + for id in seq_ids: + token = self._convert_id_to_token(id) + seq_array.append(token) + token_array.append(seq_array) + + return token_array + + def _convert_id_to_token(self, index): + """Converts an index (integer) in a token (str) using the vocab.""" + return self.ids_to_tokens.get(index, self.unk_token) + + def convert_tokens_to_string(self, tokens): + """ Converts a sequence of tokens (string) in a single string. """ + out_string = " ".join(tokens).replace(" ##", "").strip() + return out_string + + def build_inputs_with_special_tokens( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None + ) -> List[int]: + """ + Build model inputs from a sequence or a pair of sequence for sequence classification tasks + by concatenating and adding special tokens. + A BERT sequence has the following format: + - single sequence: ``[CLS] X [SEP]`` + - pair of sequences: ``[CLS] A [SEP] B [SEP]`` + Args: + token_ids_0 (:obj:`List[int]`): + List of IDs to which the special tokens will be added + token_ids_1 (:obj:`List[int]`, `optional`, defaults to :obj:`None`): + Optional second list of IDs for sequence pairs. + Returns: + :obj:`List[int]`: list of `input IDs <../glossary.html#input-ids>`__ with the appropriate special tokens. + """ + if token_ids_1 is None: + return [self.cls_token_id] + token_ids_0 + [self.sep_token_id] + cls = [self.cls_token_id] + sep = [self.sep_token_id] + return cls + token_ids_0 + sep + token_ids_1 + sep + + def get_special_tokens_mask( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False + ) -> List[int]: + """ + Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding + special tokens using the tokenizer ``prepare_for_model`` method. + Args: + token_ids_0 (:obj:`List[int]`): + List of ids. + token_ids_1 (:obj:`List[int]`, `optional`, defaults to :obj:`None`): + Optional second list of IDs for sequence pairs. + already_has_special_tokens (:obj:`bool`, `optional`, defaults to :obj:`False`): + Set to True if the token list is already formatted with special tokens for the model + Returns: + :obj:`List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token. + """ + + if already_has_special_tokens: + if token_ids_1 is not None: + raise ValueError( + "You should not supply a second sequence if the provided sequence of " + "ids is already formated with special tokens for the model." + ) + return list(map(lambda x: 1 if x in [self.sep_token_id, self.cls_token_id] else 0, token_ids_0)) + + if token_ids_1 is not None: + return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1] + return [1] + ([0] * len(token_ids_0)) + [1] + + def create_token_type_ids_from_sequences( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None + ) -> List[int]: + """ + Creates a mask from the two sequences passed to be used in a sequence-pair classification task. + A BERT sequence pair mask has the following format: + :: + 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 + | first sequence | second sequence | + if token_ids_1 is None, only returns the first portion of the mask (0's). + Args: + token_ids_0 (:obj:`List[int]`): + List of ids. + token_ids_1 (:obj:`List[int]`, `optional`, defaults to :obj:`None`): + Optional second list of IDs for sequence pairs. + Returns: + :obj:`List[int]`: List of `token type IDs <../glossary.html#token-type-ids>`_ according to the given + sequence(s). + """ + sep = [self.sep_token_id] + cls = [self.cls_token_id] + if token_ids_1 is None: + return len(cls + token_ids_0 + sep) * [0] + return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1] + + def save_vocabulary(self, vocab_path): + """ + Save the sentencepiece vocabulary (copy original file) and special tokens file to a directory. + Args: + vocab_path (:obj:`str`): + The directory in which to save the vocabulary. + Returns: + :obj:`Tuple(str)`: Paths to the files saved. + """ + index = 0 + if os.path.isdir(vocab_path): + vocab_file = os.path.join(vocab_path, VOCAB_FILES_NAMES["vocab_file"]) + else: + vocab_file = vocab_path + with open(vocab_file, "w", encoding="utf-8") as writer: + for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]): + if index != token_index: + logger.warning( + "Saving vocabulary to {}: vocabulary indices are not consecutive." + " Please check that the vocabulary is not corrupted!".format(vocab_file) + ) + index = token_index + writer.write(token + "\n") + index += 1 + return (vocab_file,) + +class SMILES_Atomwise_Tokenizer(PreTrainedTokenizer): + r""" + Constructs a SMILES tokenizer. Based on SMILES Pair Encoding (https://github.com/XinhaoLi74/SmilesPE). + This tokenizer inherits from :class:`~transformers.PreTrainedTokenizer` which contains most of the methods. Users + should refer to the superclass for more information regarding methods. + Args: + vocab_file (:obj:`string`): + File containing the vocabulary. + unk_token (:obj:`string`, `optional`, defaults to "[UNK]"): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + sep_token (:obj:`string`, `optional`, defaults to "[SEP]"): + The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences + for sequence classification or for a text and a question for question answering. + It is also used as the last token of a sequence built with special tokens. + pad_token (:obj:`string`, `optional`, defaults to "[PAD]"): + The token used for padding, for example when batching sequences of different lengths. + cls_token (:obj:`string`, `optional`, defaults to "[CLS]"): + The classifier token which is used when doing sequence classification (classification of the whole + sequence instead of per-token classification). It is the first token of the sequence when built with + special tokens. + mask_token (:obj:`string`, `optional`, defaults to "[MASK]"): + The token used for masking values. This is the token used when training this model with masked language + modeling. This is the token which the model will try to predict. + """ + + def __init__( + self, + vocab_file, + unk_token="[UNK]", + sep_token="[SEP]", + pad_token="[PAD]", + cls_token="[CLS]", + mask_token="[MASK]", + **kwargs + ): + super().__init__( + unk_token=unk_token, + sep_token=sep_token, + pad_token=pad_token, + cls_token=cls_token, + mask_token=mask_token, + **kwargs, + ) + + if not os.path.isfile(vocab_file): + raise ValueError( + "Can't find a vocabulary file at path '{}'.".format(vocab_file) + ) + self.vocab = load_vocab(vocab_file) + self.ids_to_tokens = collections.OrderedDict([(ids, tok) for tok, ids in self.vocab.items()]) + self.tokenizer = Atomwise_Tokenizer() + + @property + def vocab_size(self): + return len(self.vocab) + + def get_vocab(self): + return dict(self.vocab, **self.added_tokens_encoder) + + + def _tokenize(self, text): + return self.tokenizer.tokenize(text) + + def _convert_token_to_id(self, token): + """ Converts a token (str) in an id using the vocab. """ + return self.vocab.get(token, self.vocab.get(self.unk_token)) + + def _convert_id_to_token(self, index): + """Converts an index (integer) in a token (str) using the vocab.""" + return self.ids_to_tokens.get(index, self.unk_token) + + def convert_tokens_to_string(self, tokens): + """ Converts a sequence of tokens (string) in a single string. """ + out_string = " ".join(tokens).replace(" ##", "").strip() + return out_string + + def build_inputs_with_special_tokens( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None + ) -> List[int]: + """ + Build model inputs from a sequence or a pair of sequence for sequence classification tasks + by concatenating and adding special tokens. + A BERT sequence has the following format: + - single sequence: ``[CLS] X [SEP]`` + - pair of sequences: ``[CLS] A [SEP] B [SEP]`` + Args: + token_ids_0 (:obj:`List[int]`): + List of IDs to which the special tokens will be added + token_ids_1 (:obj:`List[int]`, `optional`, defaults to :obj:`None`): + Optional second list of IDs for sequence pairs. + Returns: + :obj:`List[int]`: list of `input IDs <../glossary.html#input-ids>`__ with the appropriate special tokens. + """ + if token_ids_1 is None: + return [self.cls_token_id] + token_ids_0 + [self.sep_token_id] + cls = [self.cls_token_id] + sep = [self.sep_token_id] + return cls + token_ids_0 + sep + token_ids_1 + sep + + def get_special_tokens_mask( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False + ) -> List[int]: + """ + Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding + special tokens using the tokenizer ``prepare_for_model`` method. + Args: + token_ids_0 (:obj:`List[int]`): + List of ids. + token_ids_1 (:obj:`List[int]`, `optional`, defaults to :obj:`None`): + Optional second list of IDs for sequence pairs. + already_has_special_tokens (:obj:`bool`, `optional`, defaults to :obj:`False`): + Set to True if the token list is already formatted with special tokens for the model + Returns: + :obj:`List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token. + """ + + if already_has_special_tokens: + if token_ids_1 is not None: + raise ValueError( + "You should not supply a second sequence if the provided sequence of " + "ids is already formated with special tokens for the model." + ) + return list(map(lambda x: 1 if x in [self.sep_token_id, self.cls_token_id] else 0, token_ids_0)) + + if token_ids_1 is not None: + return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1] + return [1] + ([0] * len(token_ids_0)) + [1] + + def create_token_type_ids_from_sequences( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None + ) -> List[int]: + """ + Creates a mask from the two sequences passed to be used in a sequence-pair classification task. + A BERT sequence pair mask has the following format: + :: + 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 + | first sequence | second sequence | + if token_ids_1 is None, only returns the first portion of the mask (0's). + Args: + token_ids_0 (:obj:`List[int]`): + List of ids. + token_ids_1 (:obj:`List[int]`, `optional`, defaults to :obj:`None`): + Optional second list of IDs for sequence pairs. + Returns: + :obj:`List[int]`: List of `token type IDs <../glossary.html#token-type-ids>`_ according to the given + sequence(s). + """ + sep = [self.sep_token_id] + cls = [self.cls_token_id] + if token_ids_1 is None: + return len(cls + token_ids_0 + sep) * [0] + return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1] + + def save_vocabulary(self, vocab_path): + """ + Save the sentencepiece vocabulary (copy original file) and special tokens file to a directory. + Args: + vocab_path (:obj:`str`): + The directory in which to save the vocabulary. + Returns: + :obj:`Tuple(str)`: Paths to the files saved. + """ + index = 0 + if os.path.isdir(vocab_path): + vocab_file = os.path.join(vocab_path, VOCAB_FILES_NAMES["vocab_file"]) + else: + vocab_file = vocab_path + with open(vocab_file, "w", encoding="utf-8") as writer: + for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]): + if index != token_index: + logger.warning( + "Saving vocabulary to {}: vocabulary indices are not consecutive." + " Please check that the vocabulary is not corrupted!".format(vocab_file) + ) + index = token_index + writer.write(token + "\n") + index += 1 + return (vocab_file,) diff --git a/smiles_tokenizer/new_splits.txt b/smiles_tokenizer/new_splits.txt new file mode 100644 index 0000000000000000000000000000000000000000..e48784ed46b8515935fe27feaeff177b63e5da5b --- /dev/null +++ b/smiles_tokenizer/new_splits.txt @@ -0,0 +1,159 @@ +c 1 +c 2 +c 3 +c 4 +c 5 +c 6 +c 7 +c 8 +c 9 +( c1 +( c2 +c1 ) +c2 ) +n 1 +n 2 +n 3 +n 4 +n 5 +n 6 +n 7 +n 8 +n 9 +( n1 +( n2 +n1 ) +n2 ) +O 1 +O 2 +O 3 +O 4 +O 5 +O 6 +O 7 +O 8 +O 9 +( O1 +( O2 +O2 ) +O2 ) += O += C += c += N += n +=C C +=C N +=C c +=c c +=N C +=N c +=n C +=n c +# N +# C +#N C +#C C +#C N +#N N +( C +C ) +( O +O ) +( N +N ) +Br c +( =O +(=O ) +C (=O) +C =O +C =N +C #N +C #C +C C +CC C +CC N +CC O +CC S +CC c +CC n +C N +CN C +CN c +C O +CO C +CO N +CO c +C S +CS C +CS S +CS c +C c +Cl c +C n +F c +N C +NC C +NC c +N N +N O +N c +N n +O C +OC C +OC O +OC c +O N +O O +O c +S C +SC C +SC c +S S +S c +c c +cc c +cc n +cc o +cc s +cc cc +c n +cn c +cn n +c o +co c +c s +cs c +cs n +n c +nc c +nc n +nc o +nc s +n n +nn c +nn n +n o +no c +no n +n s +ns c +ns n +o c +oc c +o n +s c +sc c +sc n +s n +N P +P N +C P +P C +N S +S N +C S +S C +S P +P S +C I \ No newline at end of file diff --git a/smiles_tokenizer/new_vocab.txt b/smiles_tokenizer/new_vocab.txt new file mode 100644 index 0000000000000000000000000000000000000000..ce613b0ccabbffa3cecaef1d8e944eb404c30345 --- /dev/null +++ b/smiles_tokenizer/new_vocab.txt @@ -0,0 +1,586 @@ +[PAD] +[UNK] +[CLS] +[SEP] +[MASK] +# +% +( +) ++ +- +/ +0 +1 +2 +3 +4 +5 +6 +7 +8 +9 += +@ +A +B +Br +Brc +C +CC +CCC +CCN +CCO +CCS +CCc +CCn +CN +CNC +CNc +CO +COC +CON +COc +CS +CSC +CSS +CSc +Cc +Cl +Clc +Cn +F +Fc +H +I +K +L +M +N +NC +NCC +NCc +NN +NO +Nc +Nn +O +OC +OCC +OCO +OCc +ON +OO +Oc +P +R +S +SC +SCC +SCc +SS +Sc +T +X +Z +[ +\\ +(/ +] +a +b +c +cc +ccc +ccn +cco +ccs +cn +cnc +cnn +co +coc +cs +csc +csn +e +g +i +l +n +nc +ncc +ncn +nco +ncs +nn +nnc +nnn +no +noc +non +ns +nsc +nsn +o +oc +occ +on +p +r +s +sc +scc +scn +sn +t +c1 +c2 +c3 +c4 +c5 +c6 +c7 +c8 +c9 +n1 +n2 +n3 +n4 +n5 +n6 +n7 +n8 +n9 +O1 +O2 +O3 +O4 +O5 +O6 +O7 +O8 +O9 +(c1 +(c2 +c1) +c2) +(n1 +(n2 +n1) +n2) +(O1 +(O2 +O2) +=O +=C +=c +=N +=n +=CC +=CN +=Cc +=cc +=NC +=Nc +=nC +=nc +#C +#CC +#CN +#N +#NC +#NN +(C +C) +(O +O) +(N +N) +NP +PN +CP +PC +NS +SN +SP +PS +C(=O) +(/Br) +(/C#N) +(/C) +(/C=N) +(/C=O) +(/CBr) +(/CC) +(/CCC) +(/CCF) +(/CCN) +(/CCO) +(/CCl) +(/CI) +(/CN) +(/CO) +(/CS) +(/Cl) +(/F) +(/I) +(/N) +(/NC) +(/NCC) +(/NO) +(/O) +(/OC) +(/OCC) +(/S) +(/SC) +(=C) +(=C/C) +(=C/F) +(=C/I) +(=C/N) +(=C/O) +(=CBr) +(=CC) +(=CCF) +(=CCN) +(=CCO) +(=CCl) +(=CF) +(=CI) +(=CN) +(=CO) +(=C\\C) +(=C\\F) +(=C\\I) +(=C\\N) +(=C\\O) +(=N) +(=N/C) +(=N/N) +(=N/O) +(=NBr) +(=NC) +(=NCC) +(=NCl) +(=NN) +(=NO) +(=NOC) +(=N\\C) +(=N\\N) +(=N\\O) +(=O) +(=S) +(B) +(Br) +(C#C) +(C#CC) +(C#CI) +(C#CO) +(C#N) +(C#SN) +(C) +(C=C) +(C=CF) +(C=CI) +(C=N) +(C=NN) +(C=NO) +(C=O) +(C=S) +(CBr) +(CC#C) +(CC#N) +(CC) +(CC=C) +(CC=O) +(CCBr) +(CCC) +(CCCC) +(CCCF) +(CCCI) +(CCCN) +(CCCO) +(CCCS) +(CCCl) +(CCF) +(CCI) +(CCN) +(CCNC) +(CCNN) +(CCNO) +(CCO) +(CCOC) +(CCON) +(CCS) +(CCSC) +(CCl) +(CF) +(CI) +(CN) +(CN=O) +(CNC) +(CNCC) +(CNCO) +(CNN) +(CNNC) +(CNO) +(CNOC) +(CO) +(COC) +(COCC) +(COCI) +(COCN) +(COCO) +(COF) +(CON) +(COO) +(CS) +(CSC) +(CSCC) +(CSCF) +(CSO) +(Cl) +(F) +(I) +(N) +(N=N) +(N=NO) +(N=O) +(N=S) +(NBr) +(NC#N) +(NC) +(NC=N) +(NC=O) +(NC=S) +(NCBr) +(NCC) +(NCCC) +(NCCF) +(NCCN) +(NCCO) +(NCCS) +(NCCl) +(NCNC) +(NCO) +(NCS) +(NCl) +(NN) +(NN=O) +(NNC) +(NO) +(NOC) +(O) +(OC#N) +(OC) +(OC=C) +(OC=O) +(OC=S) +(OCBr) +(OCC) +(OCCC) +(OCCF) +(OCCI) +(OCCN) +(OCCO) +(OCCS) +(OCCl) +(OCF) +(OCI) +(OCO) +(OCOC) +(OCON) +(OCSC) +(OCl) +(OI) +(ON) +(OO) +(OOC) +(OOCC) +(OOSN) +(OSC) +(P) +(S) +(SC#N) +(SC) +(SCC) +(SCCC) +(SCCF) +(SCCN) +(SCCO) +(SCCS) +(SCCl) +(SCF) +(SCN) +(SCOC) +(SCSC) +(SCl) +(SI) +(SN) +(SN=O) +(SO) +(SOC) +(SOOO) +(SS) +(SSC) +(SSCC) +([At]) +([O-]) +([O]) +([S-]) +(\\Br) +(\\C#N) +(\\C) +(\\C=N) +(\\C=O) +(\\CBr) +(\\CC) +(\\CCC) +(\\CCO) +(\\CCl) +(\\CF) +(\\CN) +(\\CNC) +(\\CO) +(\\COC) +(\\Cl) +(\\F) +(\\I) +(\\N) +(\\NC) +(\\NCC) +(\\NN) +(\\NO) +(\\NOC) +(\\O) +(\\OC) +(\\OCC) +(\\ON) +(\\S) +(\\SC) +(\\SCC) +[Ag+] +[Ag-4] +[Ag] +[Al-3] +[Al] +[As+] +[AsH3] +[AsH] +[As] +[At] +[B-] +[B@-] +[B@@-] +[BH-] +[BH2-] +[BH3-] +[B] +[Ba] +[Br+2] +[BrH] +[Br] +[C+] +[C-] +[C@@H] +[C@@] +[C@H] +[C@] +[CH-] +[CH2] +[CH3] +[CH] +[C] +[CaH2] +[Ca] +[Cl+2] +[Cl+3] +[Cl+] +[Cs] +[FH] +[F] +[H] +[He] +[I+2] +[I+3] +[I+] +[IH] +[I] +[K] +[Kr] +[Li+] +[LiH] +[MgH2] +[Mg] +[N+] +[N-] +[N@+] +[N@@+] +[N@@] +[N@] +[NH+] +[NH-] +[NH2+] +[NH3] +[NH] +[N] +[Na] +[O+] +[O-] +[OH+] +[OH2] +[OH] +[O] +[P+] +[P@+] +[P@@+] +[P@@] +[P@] +[PH2] +[PH] +[P] +[Ra] +[Rb] +[S+] +[S-] +[S@+] +[S@@+] +[S@@] +[S@] +[SH+] +[SH2] +[SH] +[S] +[Se+] +[Se-2] +[SeH2] +[SeH] +[Se] +[Si@] +[SiH2] +[SiH] +[Si] +[SrH2] +[TeH] +[Te] +[Xe] +[Zn+2] +[Zn-2] +[Zn] +[b-] +[c+] +[c-] +[cH-] +[cH] +[c] +[n+] +[n-] +[nH] +[n] +[o+] +[s+] +[se+] +[se] +[te+] +[te] \ No newline at end of file diff --git a/smiles_tokenizer/selfies_tokenizer_config.json b/smiles_tokenizer/selfies_tokenizer_config.json new file mode 100644 index 0000000000000000000000000000000000000000..467ad8ef1476bd0862480ef0671e6e1a5cb7f3da --- /dev/null +++ b/smiles_tokenizer/selfies_tokenizer_config.json @@ -0,0 +1,10 @@ +{ + "type": "selfies-wordlevel", + "version": "2.1.1", + "special_tokens": { + "pad": "", + "bos": "", + "eos": "" + }, + "semantic_constraints": null +} \ No newline at end of file diff --git a/smiles_tokenizer/selfies_tokenizers.py b/smiles_tokenizer/selfies_tokenizers.py new file mode 100644 index 0000000000000000000000000000000000000000..732a6ce32744f27f838138dc760d9ca3107d7330 --- /dev/null +++ b/smiles_tokenizer/selfies_tokenizers.py @@ -0,0 +1,186 @@ +import os +import math +import pandas as pd +from datasets import Dataset, DatasetDict +from typing import List, Dict, Any, Iterable, Set +from tqdm import tqdm +import numpy as np +import selfies as sf +import json +from rdkit import Chem +from rdkit import RDLogger +import pdb + +PAD_TOKEN = "" +BOS_TOKEN = "" +EOS_TOKEN = "" + +class SelfiesTokenizer: + """ + Minimal SELFIES tokenizer: + - build_vocab(selfies_list) -> token->id mapping + - encode(smiles_or_selfies, already_selfies=False) -> {'input_ids': List[int]} + - decode(ids) -> selfies string (no BOS/EOS returned) + """ + def __init__(self, vocab: Dict[str,int] = None): + self.vocab = vocab or {} + # filled after build_vocab if not provided + self.inv_vocab = {i:t for t,i in (self.vocab or {}).items()} + self.vocab_size = len(self.vocab) + + @staticmethod + def smiles_to_selfies(s: str) -> str: + return sf.encoder(s) + + @staticmethod + def selfies_tokens(s: str) -> List[str]: + return list(sf.split_selfies(s)) + + # def vocab_size(self) -> int: + # return len(self.vocab) + + def build_vocab(self, selfies_corpus: Iterable[str], add_semantic_default: bool = False): + tokens: Set[str] = set() + for s in selfies_corpus: + tokens.update(self.selfies_tokens(s)) + + if add_semantic_default: + tokens.update(sf.get_semantic_robust_alphabet()) + + # special tokens first with fixed ids + vocab_list = [PAD_TOKEN, BOS_TOKEN, EOS_TOKEN] + sorted(tokens) + self.vocab = {tok: i for i, tok in enumerate(vocab_list)} + self.inv_vocab = {i: tok for tok, i in self.vocab.items()} + self.vocab_size = len(self.vocab) + + def encode(self, text: str, already_selfies: bool = True, add_bos_eos: bool = True) -> Dict[str, List[int]]: + if not already_selfies: + # optionally validate SMILES with RDKit to avoid weird cases + if text and Chem.MolFromSmiles(text) is None: + return {"input_ids": []} # skip invalid + text = sf.encoder(text) + + toks = self.selfies_tokens(text) + ids = [] + if add_bos_eos: + ids.append(self.vocab[BOS_TOKEN]) + for t in toks: + if t not in self.vocab: + # unseen token -> extend vocab or skip; here we skip sample + return {"input_ids": []} + ids.append(self.vocab[t]) + if add_bos_eos: + ids.append(self.vocab[EOS_TOKEN]) + + attention_mask = [1] * len(ids) + return {"input_ids": ids, "attention_mask": attention_mask} + + def save(self, folder, constraints = None, version = None): + os.makedirs(folder, exist_ok=True) + meta = { + "type": "selfies-wordlevel", + "version": version or getattr(sf, "__version__", None), + "special_tokens": { + "pad": PAD_TOKEN, "bos": BOS_TOKEN, "eos": EOS_TOKEN + }, + "semantic_constraints": constraints or None, + } + with open(os.path.join(folder, "vocab.json"), "w") as f: + json.dump(self.vocab, f, indent=2, sort_keys=True) + with open(os.path.join(folder, "tokenizer_config.json"), "w") as f: + json.dump(meta, f, indent=2) + + @classmethod + def load(cls, file) -> "SelfiesTokenizer": + with open(file) as f: + vocab = json.load(f) + tok = cls(vocab=vocab) + return tok + + def _ids_to_selfies_tokens( + self, + ids: list[int], + stop_at_eos: bool = True, + ) -> list[str]: + """Map ids -> SELFIES tokens, filtering PAD/BOS/EOS and (optionally) stopping at EOS.""" + toks = [] + for i in ids: + tok = self.inv_vocab.get(i, None) + if tok is None: + continue + if tok == '': + continue + if tok == '': + continue + if tok == '': + if stop_at_eos: + break + else: + continue + toks.append(tok) + return toks + + def decode( + self, + ids: list[int], + return_smiles: bool = True, + stop_at_eos: bool = True, + sanitize_smiles: bool = False, + ) -> str: + """ + Decode a sequence of token ids. + + Args: + ids: list of token ids. + return_smiles: if True, return canonical SMILES; else return SELFIES string. + stop_at_eos: stop decoding when is encountered. + sanitize_smiles: if True (and RDKit available), canonicalize decoded SMILES. + + Returns: + str: SMILES (default) or SELFIES string. + """ + # 1) ids -> SELFIES tokens -> SELFIES string + toks = self._ids_to_selfies_tokens(ids, stop_at_eos=stop_at_eos) + selfies_str = "".join(toks) + # print(selfies_str) + # pdb.set_trace() + + if not return_smiles: + return selfies_str + + # 2) SELFIES -> SMILES + try: + smiles = sf.decoder(selfies_str) + except Exception: + return "" # decoding failed + + if not sanitize_smiles: + return smiles + + # 3) Optional RDKit sanitization / canonicalization + try: + mol = Chem.MolFromSmiles(smiles) + if mol is None: + return "" + # Canonical SMILES (Kekule off, isomeric on to keep stereo) + can = Chem.MolToSmiles(mol, canonical=True, isomericSmiles=True) + return can + except Exception: + return smiles # fall back to raw decoded SMILES + + def batch_decode( + self, + batch_ids: list[list[int]], + return_smiles: bool = True, + stop_at_eos: bool = True, + sanitize_smiles: bool = True, + ) -> list[str]: + return [ + self.decode( + ids, + return_smiles=return_smiles, + stop_at_eos=stop_at_eos, + sanitize_smiles=sanitize_smiles, + ) + for ids in batch_ids + ] \ No newline at end of file diff --git a/smiles_tokenizer/selfies_vocab.json b/smiles_tokenizer/selfies_vocab.json new file mode 100644 index 0000000000000000000000000000000000000000..a52c28786347a35bb445e64b2c2afca65a9b529e --- /dev/null +++ b/smiles_tokenizer/selfies_vocab.json @@ -0,0 +1,46 @@ +{ + "": 1, + "": 2, + "": 0, + "[#Branch1]": 3, + "[#Branch2]": 4, + "[#C]": 5, + "[#N]": 6, + "[/-Ring1]": 7, + "[/C]": 8, + "[/N]": 9, + "[=Branch1]": 10, + "[=Branch2]": 11, + "[=C]": 12, + "[=N]": 13, + "[=O]": 14, + "[=Ring1]": 15, + "[=Ring2]": 16, + "[=S]": 17, + "[Br]": 18, + "[Branch1]": 19, + "[Branch2]": 20, + "[C@@H1]": 21, + "[C@@]": 22, + "[C@H1]": 23, + "[C@]": 24, + "[CH1-1]": 25, + "[C]": 26, + "[Cl]": 27, + "[F]": 28, + "[I]": 29, + "[N+1]": 30, + "[N-1]": 31, + "[NH1]": 32, + "[N]": 33, + "[O-1]": 34, + "[O]": 35, + "[P]": 36, + "[Ring1]": 37, + "[Ring2]": 38, + "[S]": 39, + "[\\-Ring1]": 40, + "[\\-Ring2]": 41, + "[\\C]": 42, + "[\\N]": 43 +} \ No newline at end of file diff --git a/train.py b/train.py new file mode 100644 index 0000000000000000000000000000000000000000..368521272bef904bb054e7e98b8e16f268914343 --- /dev/null +++ b/train.py @@ -0,0 +1,146 @@ +import argparse +import os + +import pytorch_lightning as pl +import torch +import yaml +from datasets import load_from_disk +from easydict import EasyDict as edict +from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint +from torch.utils.data import DataLoader +from transformers import EsmTokenizer + +from logic import flow +from model.reparam_models import EditFlow, ProteinEditFlowModel, SMILESEditFlowModel +from paths import add_repo_to_sys_path, resolve_path, selfies_vocab_path, smiles_vocab_files +from smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer +from smiles_tokenizer.selfies_tokenizers import SelfiesTokenizer + +add_repo_to_sys_path() + + +def load_config(config_path: str) -> edict: + with open(resolve_path(config_path), "r") as f: + return edict(yaml.safe_load(f)) + + +def build_editflow(cfg, device=None): + if cfg.task == "protein": + tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D") + vocab_size = 24 + source_distribution = flow.get_source_distribution( + source_distribution=cfg.flow.source_distribution, + vocab_size=vocab_size, + special_token_ids=[0, 1, 2, 3], + ) + pad_id, bos_id, eos_id = 1, 0, 2 + model = ProteinEditFlowModel(vocab_size=vocab_size, pad_id=pad_id, config=cfg.model) + elif cfg.task == "smiles": + vocab_path, splits_path = smiles_vocab_files() + tokenizer = SMILES_SPE_Tokenizer(str(vocab_path), str(splits_path)) + vocab_size = 586 + source_distribution = flow.get_source_distribution( + source_distribution=cfg.flow.source_distribution, + vocab_size=vocab_size, + special_token_ids=[0, 1, 2, 3, 4], + ) + pad_id, bos_id, eos_id = 0, 2, 3 + model = SMILESEditFlowModel(vocab_size=vocab_size, pad_id=pad_id, config=cfg.model) + elif cfg.task == "selfies": + tokenizer = SelfiesTokenizer.load(str(selfies_vocab_path())) + vocab_size = 44 + source_distribution = flow.get_source_distribution( + source_distribution=cfg.flow.source_distribution, + vocab_size=vocab_size, + special_token_ids=[0, 1, 2], + ) + pad_id, bos_id, eos_id = 0, 1, 2 + model = SMILESEditFlowModel(vocab_size=vocab_size, pad_id=pad_id, config=cfg.model) + else: + raise NotImplementedError(f"Unsupported task: {cfg.task}") + + if device is not None: + model = model.to(device) + + eps_id = getattr(cfg.flow, "eps_id", -1) + path = flow.get_path(scheduler_type=cfg.flow.scheduler_type, exponent=cfg.flow.exponent, eps_id=eps_id) + loss_fn = flow.get_loss_function(loss_function=cfg.flow.loss_function, path=path) + editflow = EditFlow(model, loss_fn, path, source_distribution, pad_id, bos_id, eos_id, cfg) + return editflow, tokenizer, source_distribution, pad_id, bos_id, eos_id, eps_id + + +def build_dataloaders(cfg): + train_path = resolve_path(cfg.data.train_path) + val_path = resolve_path(cfg.data.val_path) + if not train_path.exists() or not val_path.exists(): + raise FileNotFoundError( + "Training data was not found.\n" + f" train: {train_path}\n" + f" val: {val_path}\n" + "Update data.train_path / data.val_path in the config. " + "The shipped SELFIES peptidomimetic dataset lives at data/selfies/28k_mimetics." + ) + num_workers = int(getattr(getattr(cfg, "data", {}), "num_workers", 4) or 4) + train_dataloader = DataLoader(load_from_disk(str(train_path)), batch_size=None, shuffle=True, num_workers=num_workers) + val_dataloader = DataLoader(load_from_disk(str(val_path)), batch_size=None, shuffle=False, num_workers=num_workers) + return train_dataloader, val_dataloader + + +def main(): + parser = argparse.ArgumentParser(description="Train an Edit Flow model") + parser.add_argument("--config", type=str, required=True, help="Path to YAML config") + parser.add_argument("--wandb", action="store_true", help="Log to Weights & Biases") + args = parser.parse_args() + + cfg = load_config(args.config) + run_name = ( + f"reparam_{cfg.task}_lr{cfg.optim.lr}_epoch{cfg.optim.n_epochs}" + f"_scale{cfg.model.scale_size}_optimal{cfg.model.p_optimal}" + ) + workdir = resolve_path(getattr(cfg, "work_dir", "outputs")) / run_name + os.makedirs(workdir, exist_ok=True) + + pl.seed_everything(cfg.training.seed, workers=True) + editflow, _, _, _, _, _, _ = build_editflow(cfg) + train_dataloader, val_dataloader = build_dataloaders(cfg) + + ckpt = ModelCheckpoint( + dirpath=os.path.join(workdir, "checkpoint"), + monitor="val_loss", + mode="min", + save_top_k=3, + save_last=True, + filename="epoch{epoch:04d}-val{val_loss:.2f}", + auto_insert_metric_name=False, + ) + callbacks = [ckpt, LearningRateMonitor(logging_interval="step")] + + logger = False + if args.wandb: + from pytorch_lightning.loggers import WandbLogger + + logger = WandbLogger( + project=getattr(getattr(cfg, "logging", {}), "project", "pCoMole"), + name=run_name, + entity=getattr(getattr(cfg, "logging", {}), "entity", None), + ) + + trainer = pl.Trainer( + default_root_dir=str(workdir), + accelerator="gpu" if torch.cuda.is_available() else "cpu", + devices=cfg.compute.ngpus, + strategy="ddp" if cfg.compute.ngpus > 1 else "auto", + precision="bf16-mixed", + max_epochs=cfg.optim.n_epochs, + log_every_n_steps=10, + callbacks=callbacks, + enable_checkpointing=True, + gradient_clip_val=1.0, + deterministic=False, + logger=logger, + ) + trainer.fit(editflow, train_dataloader, val_dataloader) + + +if __name__ == "__main__": + main() diff --git a/utils/.DS_Store b/utils/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..5008ddfcf53c02e82d7eee2e57c38e5672ef89f6 Binary files /dev/null and b/utils/.DS_Store differ diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..36d719522180cd6c67a6b7eb05ee12e1cff6f386 --- /dev/null +++ b/utils/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. diff --git a/utils/checkpointing.py b/utils/checkpointing.py new file mode 100644 index 0000000000000000000000000000000000000000..d50f6a02b20ae5cd4cff726784691c3b0a3bb2ac --- /dev/null +++ b/utils/checkpointing.py @@ -0,0 +1,76 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. +# Part of this implementation is adapted from https://github.com/louaaron/Score-Entropy-Discrete-Diffusion +# which is released under MIT license + +from dataclasses import dataclass, field +from pathlib import Path + +import torch +from logic.flow import SourceDistribution +from model import Transformer +from omegaconf import OmegaConf +from torch import nn +from torch.nn.parallel import DistributedDataParallel as DDP + + +def load_cfg_from_path(work_dir: str) -> OmegaConf: + work_dir = Path(work_dir) + + root_dir = work_dir if work_dir.is_dir() else work_dir.parents[1] + + cfg_path = root_dir / ".hydra/config.yaml" + + return OmegaConf.load(cfg_path) + + +def load_model_from_path( + work_dir: str, + source_distribution: SourceDistribution, + device: torch.device, + vocab_size: int, + cfg: OmegaConf, +) -> nn.Module: + work_dir = Path(work_dir) + + if work_dir.is_dir(): + root_dir = work_dir + ckpt_dir = work_dir / "checkpoints" / "checkpoint.pth" + else: + root_dir = work_dir.parents[1] + ckpt_dir = work_dir + + model = Transformer( + config=cfg, vocab_size=vocab_size, masked=source_distribution.masked + ).to(device) + model = DDP(model, device_ids=[device]) + + ckpt_dir = root_dir / "checkpoints" / "checkpoint.pth" + loaded_state = torch.load(ckpt_dir, map_location=device, weights_only=True) + + model.module.load_state_dict(loaded_state["model"]) + + return model + + +@dataclass +class WorkDirectory: + root: Path = field(metadata={"help": "Root work directory"}) + checkpoint: Path = field(metadata={"help": "Checkpoint directory"}) + samples: Path = field(metadata={"help": "Samples directory"}) + + +def get_work_dirs(work_dir: str, rank: int) -> WorkDirectory: + work_dir = Path(work_dir) + + sample_dir = work_dir / "samples" + checkpoint_dir = work_dir / "checkpoints" / "checkpoint.pth" + + if rank == 0: + sample_dir.mkdir(exist_ok=True) + checkpoint_dir.parents[0].mkdir(exist_ok=True) + + return WorkDirectory(root=work_dir, checkpoint=checkpoint_dir, samples=sample_dir) diff --git a/utils/logging.py b/utils/logging.py new file mode 100644 index 0000000000000000000000000000000000000000..7b73067ffbce83551b5dacee334cd0ae57ebbcbe --- /dev/null +++ b/utils/logging.py @@ -0,0 +1,122 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the CC-by-NC license found in the +# LICENSE file in the root directory of this source tree. + +import logging +import os +from logging import Logger +from pathlib import Path +from typing import Optional + +import torch +import wandb +from omegaconf import OmegaConf + + +def get_logger(log_path: str, rank: int): + if rank != 0: + return logging.getLogger("dummy") + + logger = logging.getLogger() + default_level = logging.INFO + + if logger.hasHandlers(): + logger.handlers.clear() + + logger.setLevel(default_level) + + formatter = logging.Formatter( + "%(levelname)s | %(asctime)s | %(message)s", "%Y-%m-%d %H:%M:%S" + ) + + info_file_handler = logging.FileHandler(log_path, mode="a") + info_file_handler.setLevel(default_level) + info_file_handler.setFormatter(formatter) + logger.addHandler(info_file_handler) + + console_handler = logging.StreamHandler() + console_handler.setLevel(default_level) + console_handler.setFormatter(formatter) + logger.addHandler(console_handler) + + return logger + + +class TrainLogger: + def __init__(self, log_dir: Path, rank: int, cfg: bool = False): + self.log_dir = log_dir + self.cfg = cfg + + self._init_text_logger(rank=rank) + + self.enable_wandb = self.cfg.logging.enable_wandb and (rank == 0) + + if self.enable_wandb: + self._init_wandb() + + def _init_text_logger(self, rank: int): + log_path = self.log_dir / self.cfg.logging.log_file_name + self._logger = get_logger(log_path=log_path, rank=rank) + + def _init_wandb( + self, + ): + wandb_run_id_path = self.log_dir / "wandb_run.id" + + try: + wandb_run_id = wandb_run_id_path.read_text() + except FileNotFoundError: + wandb_run_id = wandb.util.generate_id() + wandb_run_id_path.write_text(wandb_run_id) + + self.wandb_logger = wandb.init( + id=wandb_run_id, + project=self.cfg.logging.project, + group=self.cfg.logging.group, + dir=self.log_dir, + entity=self.cfg.logging.entity, + resume="allow", + config=OmegaConf.to_container(self.cfg, resolve=True), + ) + + def log_metric(self, value: float, name: str, stage: bool, step: int) -> None: + self._logger.info(f"[{step}] {stage} {name}: {value:.3f}") + + if self.enable_wandb: + self.wandb_logger.log(data={f"{stage}/{name}": value}, step=step) + + def log_lr(self, value: float, step: int) -> None: + if self.enable_wandb: + self.wandb_logger.log(data={"Optimization/LR": value}, step=step) + + def info(self, msg: str, step: Optional[int] = None) -> None: + step_str = f"[{step}] " if step else "" + self._logger.info(f"{step_str}{msg}") + + def warning(self, msg: str) -> None: + self._logger.warning(msg) + + def finish(self) -> None: + for handler in self._logger.handlers: + if isinstance(handler, logging.FileHandler): + handler.close() + + if self.enable_wandb: + wandb.finish() + + @staticmethod + def log_devices(device: torch.device, logger: Logger) -> None: + if device.type == "cuda": + logger.info("Found {} CUDA devices.".format(torch.cuda.device_count())) + for i in range(torch.cuda.device_count()): + props = torch.cuda.get_device_properties(i) + logger.info( + "{} \t Memory: {:.2f}GB".format( + props.name, props.total_memory / (1024**3) + ) + ) + else: + logger.warning("WARNING: Using device {}".format(device)) + logger.info(f"Found {os.cpu_count()} total number of CPUs.")