Add Cas9 task

Adds the Cas9 vertical alongside gfp and peptidomimetics: the pCoMole
editor, its three objectives (Cas9-likeness, deletion count, PAM
matching), constraints, data-prep scripts, run scripts, and config.

- cas9/model/ and cas9/logic/ are vendored because the Cas9 Edit Flow
keeps the reparameterized models inside base_models.py, whereas the
repo root splits them into reparam_models.py. The two are not
interchangeable, so Cas9 imports its own copies.
- objectives.py is trimmed to the three Cas9 classes; the ADMET oracles
already ship in peptidomimetics/objectives.py.
- Oracle and checkpoint locations resolve through paths.py
(PROTEIN2PAM_ROOT, CAS_PREDICTOR_ROOT, CAS9_CLASSIFIER_CKPT,
CAS9_CLASSIFIER_CONFIG, CAS9_HMM_DIR) instead of absolute paths.
- Ten Cas9 ortholog seeds ship as cas9/data/orthologs.fasta.

Checkpoints and oracle packages are not included in this commit;
cas9/ckpt/, cas9/classifier_ckpt/ and cas9/hmm/ are placeholders.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015JoDLXiG2Vc8rwe9shw78R

README.md

@@ -26,8 +26,26 @@ pCoMole/
     ckpt/SMILES_BindEvaluator.ckpt
     ckpt/admetica/                # LD50, solubility, Caco-2, half-life
     ckpt/deepdtagen/
+  cas9/                           # Cas9 generation + pCoMole
+    pcomole.py                    # official Cas9 editor
+    generate_cas9_csvs.py
+    objectives.py                 # Cas9Classification, DeletionCount, PAMMatching
+    constraints.py
+    domain_detector.py            # HNH / RuvC completeness (HMMER)
+    pam_detector_hmm.py           # PAM-interacting domain masking
+    model/ logic/                 # Cas9 Edit Flow variant (see note below)
+    data/orthologs.fasta          # 10 Cas9 ortholog seeds
+    ckpt/last.ckpt
+    classifier_ckpt/              # Cas9-likeness classifier
+    hmm/                          # HMM databases
 ```
 
+> **Note.** `cas9/model/` and `cas9/logic/` are a *fork* of the repo-root
+> packages, not duplicates: the Cas9 stack keeps the reparameterized Edit Flow
+> models inside `base_models.py`, whereas the root splits them into
+> `reparam_models.py`. The two are not interchangeable, so the Cas9 task
+> imports its own copies (`from cas9.model...`).
+
 ## 1. Environment
 
 A CUDA GPU is strongly recommended.
@@ -76,6 +94,21 @@ export MAFFT_PATH=/path/to/mafft
 
 Confirm with `mafft --version` before running GFP editing.
 
+### HMMER and protein2pam (required for Cas9 pCoMole)
+
+The Cas9 domain-completeness and PAM-masking oracles shell out to `hmmscan`:
+
+```bash
+conda install -c bioconda hmmer
+```
+
+The PAM objective needs the `protein2pam` model class (weights download from the
+Hub as `Profluent-Bio/protein2pam-cas9_full`):
+
+```bash
+pip install protein2pam==0.2.0
+```
+
 ## 2. Train, evaluate, and sample Edit Flows
 
 Run every command from the **repository root**.
@@ -172,6 +205,40 @@ bash peptidomimetics/scripts/pcomole.sh
 
 `--target` and `--motifs` are required for affinity and motif oracles. `--specificity` adds the seventh objective.
 
+### Cas9
+
+Official editor: `cas9/pcomole.py`. Three objectives — Cas9-likeness, deletion
+count, and PAM matching — under Cas9-score, PAM-matching and length constraints.
+
+Cas9 code is a package, so run it from the repository root with `-m`:
+
+```bash
+ORTHOLOG=cjcas9 bash cas9/scripts/pcomole.sh
+```
+
+Equivalent command:
+
+```bash
+python -m cas9.pcomole \
+  --pcomol_config configs/config_cas9.yaml \
+  --ckpt cas9/ckpt/last.ckpt \
+  --input "$(awk -v n='>cjcas9' '$1==n{f=1;next}/^>/{f=0}f{printf "%s",$0}' cas9/data/orthologs.fasta)" \
+  --cas9_objective --deletion_count_obj --pam_distr_objective \
+  --objective_weights 1 1 1 \
+  --target_pam matching \
+  --num_steps 5 --num_rollouts 5 --num_candidates 10 --num_sequences 50 \
+  --max_deletion_percentage 0.2 --deletion_rate_scale 2500 \
+  --pam_min_confidence 0.7 --cas9_score_threshold 0.5 \
+  --max_target_length 930 --min_target_length 890 \
+  --use_ce_loss --allow_equal_length --zero_lam_ins \
+  --output_file outputs/cas9_cjcas9_pcomole.csv
+```
+
+Ten ortholog seeds ship in `cas9/data/orthologs.fasta` (SpCas9, SaCas9, CjCas9,
+St1/St3Cas9, GeoCas9, NmCas9, iSpyMac, SpCas9++, SpRYc). Set `ORTHOLOG` to any
+of them. `--target_pam matching` optimizes toward the seed's own PAM; pass an
+explicit 10-nt `ACGTN` string to retarget.
+
 ## 4. Included checkpoints
 
 | Path | Role |
@@ -183,6 +250,9 @@ bash peptidomimetics/scripts/pcomole.sh
 | `peptidomimetics/ckpt/SMILES_BindEvaluator.ckpt` | motif / specificity |
 | `peptidomimetics/ckpt/admetica/*.ckpt` | Admetica blending oracles |
 | `peptidomimetics/ckpt/deepdtagen/` | DeepDTAGen affinity |
+| `cas9/ckpt/last.ckpt` | Cas9 Edit Flow |
+| `cas9/classifier_ckpt/` | Cas9-likeness classifier (objective + hard constraint) |
+| `cas9/hmm/` | HNH / RuvC and PAM-interacting HMM databases |
 
 Together these are about 10 GB. HuggingFace uploads should use Git LFS for the `.ckpt` / `.pt` / `.pth` files.
 
@@ -193,13 +263,18 @@ Together these are about 10 GB. HuggingFace uploads should use Git LFS for the `
 | `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 |
+| `PROTEIN2PAM_ROOT` | protein2pam package root (PAM oracle model class) |
+| `CAS_PREDICTOR_ROOT` | cas_predictor package root (Cas9 classifier code) |
+| `CAS9_CLASSIFIER_CKPT` | override for the Cas9 classifier checkpoint |
+| `CAS9_CLASSIFIER_CONFIG` | optional Cas9 classifier config |
+| `CAS9_HMM_DIR` | directory holding the Cas9 HMM databases |
 
 ## 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.
+- Cas9 objectives fail at init if `hmmscan`, `protein2pam`, or `cas_predictor` are missing.
 - Training writes to `outputs/` by default.
 
 ## Citation

cas9/ckpt/README.md

@@ -0,0 +1,3 @@
+Tier B artifact slot — not yet populated.
+
+See the repository README, section 4, for what belongs here.

cas9/classifier_ckpt/README.md

@@ -0,0 +1,3 @@
+Tier B artifact slot — not yet populated.
+
+See the repository README, section 4, for what belongs here.

cas9/constraints.py

@@ -0,0 +1,705 @@
+# constraints.py
+
+import os
+import hashlib
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Sequence, Tuple
+
+import torch
+from rdkit import Chem
+from rdkit.Chem import Descriptors, rdMolDescriptors
+# Make selfies optional (only needed for SMILES constraints, not protein constraints)
+try:
+    import selfies as sf
+except ImportError:
+    sf = None
+
+# from .is_peptidomimetic import is_peptidomimetic_not_natural
+import pdb
+
+# ---- Optional import: Cas9 domain detector (HNH + RuvC) ----
+# Expects domain_detector.py to live in the same package as this constraints.py
+try:
+    from .domain_detector import DomainDetector, DomainDetectorConfig, HitCriteria
+except Exception as e:
+    import traceback
+    print(f"Warning: Failed to import DomainDetector: {e}")
+    traceback.print_exc()
+    DomainDetector = None
+    DomainDetectorConfig = None
+    HitCriteria = None
+
+
+def detokenize_output(x, cfg, tokenizer, bos_id, eos_id, pad_id):
+    """
+    Convert a single generated sequence (1, L) back to string.
+    """
+    seq = x[0].tolist()
+    # strip padding
+    seq = [tok for tok in seq if tok != pad_id]
+    # strip BOS/EOS
+    if len(seq) > 0 and seq[0] == bos_id:
+        seq = seq[1:]
+    if len(seq) > 0 and seq[-1] == eos_id:
+        seq = seq[:-1]
+
+    if cfg.task == 'protein':
+        # esm tokenizer has batch_decode
+        return tokenizer.batch_decode([seq], skip_special_tokens=True)[0]
+    elif cfg.task in ('smiles', 'selfies'):
+        return tokenizer.decode(seq)
+    else:
+        return " ".join(map(str, seq))
+
+
+class Peptidomimetic:
+    def __init__(self):
+        pass
+
+    def __call__(self, smiles_tokens, smiles_seq):
+        flag, audit = is_peptidomimetic_not_natural(smiles_seq)
+        return 1 if flag else 0
+
+
+# SMARTS
+_AMIDE       = Chem.MolFromSmarts("[CX3](=[OX1])[NX3]")               # amide C(=O)-N (includes carbamates/ureas as centers)
+_SULFONAMIDE = Chem.MolFromSmarts("[SX4](=[OX1])(=[OX1])[NX3]")       # sulfonamide S(=O)2-N  (optional isostere)
+
+
+def _centers_and_N(mol, include_sulfonamide=True):
+    centers, n2c = set(), {}
+    for c_idx, _, n_idx in mol.GetSubstructMatches(_AMIDE):
+        centers.add(c_idx)
+        n2c.setdefault(n_idx, set()).add(c_idx)
+    if include_sulfonamide:
+        for s_idx, _, _, n_idx in mol.GetSubstructMatches(_SULFONAMIDE):
+            centers.add(s_idx)
+            n2c.setdefault(n_idx, set()).add(s_idx)
+    return centers, n2c
+
+
+def _is_sp3_carbon(atom):
+    return atom.GetAtomicNum() == 6 and atom.GetHybridization() == Chem.HybridizationType.SP3
+
+
+def _backbone_edges(mol, centers, n2c):
+    """
+    Edge between centers i--j exists if there is a single-bond path:
+      center_i — N — C(sp3) — center_j    (length 3)
+      or center_i — N — (O|S|N) — C(sp3) — center_j  (length 4)
+    """
+    edges = set()
+    for n_idx, left_centers in n2c.items():
+        n = mol.GetAtomWithIdx(n_idx)
+        for b1 in n.GetBonds():
+            if b1.GetBondType() != Chem.BondType.SINGLE:
+                continue
+            a1 = b1.GetOtherAtom(n)
+
+            # case A: N—C(sp3)
+            if _is_sp3_carbon(a1):
+                for b2 in a1.GetBonds():
+                    if b2.GetBondType() != Chem.BondType.SINGLE:
+                        continue
+                    other = b2.GetOtherAtom(a1)
+                    j = other.GetIdx()
+                    if j in centers:
+                        for i in left_centers:
+                            if i != j:
+                                edges.add(tuple(sorted((i, j))))
+
+            # case B: N—(O|S|N)—C(sp3)
+            if a1.GetAtomicNum() in (8, 16, 7):  # O, S, N
+                for b2 in a1.GetBonds():
+                    if b2.GetBondType() != Chem.BondType.SINGLE:
+                        continue
+                    a2 = b2.GetOtherAtom(a1)
+                    if _is_sp3_carbon(a2):
+                        for b3 in a2.GetBonds():
+                            if b3.GetBondType() != Chem.BondType.SINGLE:
+                                continue
+                            other = b3.GetOtherAtom(a2)
+                            j = other.GetIdx()
+                            if j in centers:
+                                for i in left_centers:
+                                    if i != j:
+                                        edges.add(tuple(sorted((i, j))))
+    return edges
+
+
+def _largest_component(centers, edges):
+    if not centers:
+        return set()
+    adj = {i: set() for i in centers}
+    for i, j in edges:
+        if i in adj and j in adj:
+            adj[i].add(j)
+            adj[j].add(i)
+    seen, best = set(), set()
+    for s in adj:
+        if s in seen:
+            continue
+        stack, comp = [s], set()
+        while stack:
+            u = stack.pop()
+            if u in comp:
+                continue
+            comp.add(u)
+            seen.add(u)
+            stack.extend(v for v in adj[u] if v not in comp)
+        if len(comp) > len(best):
+            best = comp
+    return best
+
+
+def backbone_metrics(smiles: str, include_sulfonamide=True):
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        raise ValueError(f"Invalid SMILES: {smiles}")
+    centers, n2c = _centers_and_N(mol, include_sulfonamide)
+    edges = _backbone_edges(mol, centers, n2c)
+    main_nodes = _largest_component(centers, edges)
+    main_edges = {e for e in edges if e[0] in main_nodes and e[1] in main_nodes}
+    return {
+        "backbone_centers_main": len(main_nodes),
+        "backbone_links_main": len(main_edges),
+        "MW": Descriptors.MolWt(mol),
+        "HAC": mol.GetNumAtoms(),
+        "TPSA": rdMolDescriptors.CalcTPSA(mol),
+        "rotatable_bonds": rdMolDescriptors.CalcNumRotatableBonds(mol),
+    }
+
+
+def compare_shorter_smaller(orig_smiles: str, cand_smiles: str, include_sulfonamide=True,
+                            mw_margin: float = 0.0, hac_margin: int = 0):
+    o = backbone_metrics(orig_smiles, include_sulfonamide)
+    c = backbone_metrics(cand_smiles, include_sulfonamide)
+    shorter_backbone = (c["backbone_centers_main"] < o["backbone_centers_main"]) or \
+                       (c["backbone_links_main"]   < o["backbone_links_main"])
+    smaller_bulk = (c["MW"] <= o["MW"] - mw_margin) and (c["HAC"] <= o["HAC"] - hac_margin)
+    return {"original": o, "candidate": c, "decisions": {
+        "shorter_backbone": shorter_backbone,
+        "smaller_bulk": smaller_bulk,
+        "overall_shorter_or_smaller": (shorter_backbone or smaller_bulk),
+        "require_both_shorter_and_smaller": (shorter_backbone and smaller_bulk),
+    }}
+
+
+class Length:
+    def __init__(self, orig_smiles, cfg, smiles_tokenizer):
+        self.cfg = cfg
+        self.smiles_tokenizer = smiles_tokenizer
+
+        self.orig_smiles = sf.decoder(sf.encoder(orig_smiles))
+        self.L0 = smiles_tokenizer(orig_smiles, return_tensors='pt')['input_ids'].shape[1]
+        print("Initial Length: ", self.L0)
+
+    def __call__(self, smiles_tokens, smiles_seq):
+        res = compare_shorter_smaller(self.orig_smiles, smiles_seq)
+        return int(res["decisions"]["overall_shorter_or_smaller"] and smiles_tokens.shape[1] >= self.L0 // 2)
+        # return int(smiles_tokens.shape[1] <= self.L0 and smiles_tokens.shape[1] >= self.L0 // 2)
+
+
+class ProteinLength:
+    """
+    Length constraint for protein sequences.
+    
+    Identical logic to the SMILES Length constraint, but using protein sequence length.
+    By default requires candidate to be strictly shorter than original AND at least half
+    the original length. If allow_equal_length is True, candidates equal to the original
+    length also pass the upper bound (still must be at least L0 // 2).
+    """
+    def __init__(self, orig_protein_seq, cfg, tokenizer, allow_equal_length: bool = False):
+        self.cfg = cfg
+        self.tokenizer = tokenizer
+        self.allow_equal_length = bool(allow_equal_length)
+        
+        self.orig_protein_seq = orig_protein_seq.strip().replace(" ", "").replace("\n", "")
+        # Store length in amino acids (equivalent to tokenized length for proteins)
+        self.L0 = len(self.orig_protein_seq)
+        print("Initial Length: ", self.L0)
+        if self.allow_equal_length:
+            print("ProteinLength: allowing equal length (candidates may have length == L0)")
+
+    def __call__(self, protein_tokens, protein_seq):
+        """
+        Check protein sequence length against original.
+        
+        - Candidate must be at least half the original length (L0 // 2).
+        - Upper bound: strictly shorter than L0 by default; if allow_equal_length, L_cand <= L0.
+        
+        Args:
+            protein_tokens: Unused (kept for interface compatibility, could use for tokenized length)
+            protein_seq: Protein sequence string (single sequence)
+        
+        Returns:
+            int: 1 if constraints pass, 0 otherwise
+        """
+        seq_clean = (protein_seq or "").strip().replace(" ", "").replace("\n", "")
+        L_cand = len(seq_clean)
+        
+        if self.allow_equal_length:
+            is_within_upper = L_cand <= self.L0
+        else:
+            is_within_upper = L_cand < self.L0
+        is_at_least_half = L_cand >= self.L0 // 2
+        
+        return int(is_within_upper and is_at_least_half)
+
+
+class TargetLength:
+    """
+    Target length constraint for protein sequences.
+    
+    Requires candidate sequence length to exactly match the target length.
+    """
+    def __init__(self, target_length: int):
+        """
+        Args:
+            target_length: The exact target length in amino acids
+        """
+        if target_length <= 0:
+            raise ValueError(f"target_length must be positive, got {target_length}")
+        self.target_length = int(target_length)
+        print(f"Target Length: {self.target_length}")
+
+    def __call__(self, protein_tokens, protein_seq):
+        """
+        Check if protein sequence length exactly matches the target length.
+        
+        Args:
+            protein_tokens: Unused (kept for interface compatibility)
+            protein_seq: Protein sequence string (single sequence)
+        
+        Returns:
+            int: 1 if length matches target, 0 otherwise
+        """
+        seq_clean = (protein_seq or "").strip().replace(" ", "").replace("\n", "")
+        L_cand = len(seq_clean)
+        
+        return int(L_cand == self.target_length)
+    
+    def predict_batch(self, seqs: Sequence[str]) -> List[bool]:
+        """
+        Batch prediction helper for evaluation / faster integration.
+        
+        Args:
+            seqs: List of protein sequence strings
+        
+        Returns:
+            List of bool: True if length matches target, False otherwise
+        """
+        clean = [(s or "").strip().replace(" ", "").replace("\n", "") for s in seqs]
+        return [len(s) == self.target_length for s in clean]
+
+
+class MinTargetLength:
+    """
+    Minimum target length constraint for protein sequences.
+
+    Requires candidate sequence length to be greater than or equal to the minimum target length.
+    """
+    def __init__(self, min_target_length: int):
+        """
+        Args:
+            min_target_length: The minimum allowed length in amino acids (inclusive)
+        """
+        if min_target_length <= 0:
+            raise ValueError(f"min_target_length must be positive, got {min_target_length}")
+        self.min_target_length = int(min_target_length)
+        print(f"Min Target Length: {self.min_target_length}")
+
+    def __call__(self, protein_tokens, protein_seq):
+        """
+        Check if protein sequence length is greater than or equal to the minimum target length.
+
+        Args:
+            protein_tokens: Unused (kept for interface compatibility)
+            protein_seq: Protein sequence string (single sequence)
+
+        Returns:
+            int: 1 if length >= min_target_length, 0 otherwise
+        """
+        seq_clean = (protein_seq or "").strip().replace(" ", "").replace("\n", "")
+        L_cand = len(seq_clean)
+
+        return int(L_cand >= self.min_target_length)
+
+    def predict_batch(self, seqs: Sequence[str]) -> List[bool]:
+        """
+        Batch prediction helper for evaluation / faster integration.
+
+        Args:
+            seqs: List of protein sequence strings
+
+        Returns:
+            List of bool: True if length >= min_target_length, False otherwise
+        """
+        clean = [(s or "").strip().replace(" ", "").replace("\n", "") for s in seqs]
+        return [len(s) >= self.min_target_length for s in clean]
+
+
+class MaxTargetLength:
+    """
+    Maximum target length constraint for protein sequences.
+    
+    Requires candidate sequence length to be less than or equal to the maximum target length.
+    """
+    def __init__(self, max_target_length: int):
+        """
+        Args:
+            max_target_length: The maximum allowed length in amino acids (inclusive)
+        """
+        if max_target_length <= 0:
+            raise ValueError(f"max_target_length must be positive, got {max_target_length}")
+        self.max_target_length = int(max_target_length)
+        print(f"Max Target Length: {self.max_target_length}")
+
+    def __call__(self, protein_tokens, protein_seq):
+        """
+        Check if protein sequence length is less than or equal to the maximum target length.
+        
+        Args:
+            protein_tokens: Unused (kept for interface compatibility)
+            protein_seq: Protein sequence string (single sequence)
+        
+        Returns:
+            int: 1 if length <= max_target_length, 0 otherwise
+        """
+        seq_clean = (protein_seq or "").strip().replace(" ", "").replace("\n", "")
+        L_cand = len(seq_clean)
+        
+        return int(L_cand <= self.max_target_length)
+    
+    def predict_batch(self, seqs: Sequence[str]) -> List[bool]:
+        """
+        Batch prediction helper for evaluation / faster integration.
+        
+        Args:
+            seqs: List of protein sequence strings
+        
+        Returns:
+            List of bool: True if length <= max_target_length, False otherwise
+        """
+        clean = [(s or "").strip().replace(" ", "").replace("\n", "") for s in seqs]
+        return [len(s) <= self.max_target_length for s in clean]
+
+
+class PAMMatchingConstraint:
+    """
+    Terminal constraint for protein sequences: predicted PAM must match target PAM.
+    
+    Requires that the predicted PAM of the generated sequence exactly matches the target PAM.
+    For positions where target PAM is 'N', any nucleotide is acceptable.
+    """
+    def __init__(self, pam_matching_obj):
+        """
+        Args:
+            pam_matching_obj: PAMMatching objective object that has predict_pam method and target_pam attribute
+        """
+        if not hasattr(pam_matching_obj, 'predict_pam'):
+            raise ValueError("pam_matching_obj must have a predict_pam method")
+        if not hasattr(pam_matching_obj, 'target_pam'):
+            raise ValueError("pam_matching_obj must have a target_pam attribute")
+        self.pam_matching_obj = pam_matching_obj
+        self.target_pam = pam_matching_obj.target_pam.upper()
+        print(f"PAM Matching Constraint: sequences must have predicted PAM matching target: {self.target_pam}")
+
+    def _pam_matches(self, predicted_pam: str, target_pam: str) -> bool:
+        """
+        Check if predicted PAM matches target PAM.
+        For positions where target PAM is 'N', any nucleotide is acceptable.
+        
+        Args:
+            predicted_pam: Predicted PAM string (10 nucleotides)
+            target_pam: Target PAM string (10 nucleotides, may contain 'N')
+        
+        Returns:
+            bool: True if predicted PAM matches target PAM, False otherwise
+        """
+        # if len(predicted_pam) != len(target_pam):
+        #     return False
+        
+        # for pred_nuc, target_nuc in zip(predicted_pam, target_pam):
+        #     if target_nuc == 'N':
+        #         # N in target means any nucleotide is acceptable
+        #         continue
+        #     if pred_nuc != target_nuc:
+        #         return False
+        
+        # return True
+        return predicted_pam.upper() == target_pam.upper()
+
+    def __call__(self, protein_tokens, protein_seq):
+        """
+        Check if protein sequence's predicted PAM matches the target PAM.
+        
+        Args:
+            protein_tokens: Unused (kept for interface compatibility)
+            protein_seq: Protein sequence string (single sequence)
+        
+        Returns:
+            int: 1 if predicted PAM matches target PAM, 0 otherwise
+        """
+        try:
+            # Check if the model is available
+            if not hasattr(self.pam_matching_obj, 'model') or self.pam_matching_obj.model is None:
+                print(f"[WARNING] PAMMatchingConstraint: model is not initialized. Cannot predict PAM.")
+                return 0
+            
+            if not hasattr(self.pam_matching_obj, 'tokenizer') or self.pam_matching_obj.tokenizer is None:
+                print(f"[WARNING] PAMMatchingConstraint: tokenizer is not initialized. Cannot predict PAM.")
+                return 0
+            
+            predicted_pams = self.pam_matching_obj.predict_pam([protein_seq])
+            if not predicted_pams:
+                print(f"[WARNING] PAMMatchingConstraint: predict_pam returned empty list for sequence.")
+                return 0
+            
+            predicted_pam = predicted_pams[0]
+            matches = self._pam_matches(predicted_pam, self.target_pam)
+            if not matches:
+                # Debug output: show what was predicted vs expected
+                print(f"[DEBUG] PAMMatchingConstraint: predicted '{predicted_pam}' != target '{self.target_pam}'")
+            return int(matches)
+        except Exception as e:
+            import traceback
+            print(f"[ERROR] PAMMatchingConstraint failed: {e}")
+            traceback.print_exc()
+            return 0
+    
+    def predict_batch(self, seqs: Sequence[str]) -> List[bool]:
+        """
+        Batch prediction helper for evaluation / faster integration.
+        
+        Args:
+            seqs: List of protein sequence strings
+        
+        Returns:
+            List of bool: True if predicted PAM matches target PAM, False otherwise
+        """
+        try:
+            # Check if the model is available
+            if not hasattr(self.pam_matching_obj, 'model') or self.pam_matching_obj.model is None:
+                print(f"[WARNING] PAMMatchingConstraint.predict_batch: model is not initialized. Cannot predict PAM.")
+                return [False] * len(seqs)
+            
+            if not hasattr(self.pam_matching_obj, 'tokenizer') or self.pam_matching_obj.tokenizer is None:
+                print(f"[WARNING] PAMMatchingConstraint.predict_batch: tokenizer is not initialized. Cannot predict PAM.")
+                return [False] * len(seqs)
+            
+            predicted_pams = self.pam_matching_obj.predict_pam(seqs)
+            results = []
+            for i, predicted_pam in enumerate(predicted_pams):
+                if not predicted_pam:
+                    print(f"[WARNING] PAMMatchingConstraint.predict_batch: empty prediction for sequence {i}")
+                    results.append(False)
+                else:
+                    results.append(self._pam_matches(predicted_pam, self.target_pam))
+            return results
+        except Exception as e:
+            import traceback
+            print(f"[ERROR] PAMMatchingConstraint.predict_batch failed: {e}")
+            traceback.print_exc()
+            return [False] * len(seqs)
+
+
+class Cas9ScoreThreshold:
+    """
+    Terminal constraint for protein sequences: Cas9 classifier score must be above a threshold.
+    
+    Requires that the Cas9 classifier score (probability) is strictly greater than the threshold.
+    """
+    def __init__(self, cas9_classifier_obj, threshold: float):
+        """
+        Args:
+            cas9_classifier_obj: Cas9Classification objective object that has get_scores method
+            threshold: Minimum Cas9 score threshold (0-1). Sequences must have score > threshold.
+        """
+        if not hasattr(cas9_classifier_obj, 'get_scores'):
+            raise ValueError("cas9_classifier_obj must have a get_scores method")
+        if not (0.0 <= threshold <= 1.0):
+            raise ValueError(f"threshold must be in [0, 1], got {threshold}")
+        self.cas9_classifier_obj = cas9_classifier_obj
+        self.threshold = float(threshold)
+        print(f"Cas9 Score Threshold Constraint: sequences must have Cas9 score > {self.threshold}")
+
+    def __call__(self, protein_tokens, protein_seq):
+        """
+        Check if protein sequence's Cas9 score is above the threshold.
+        
+        Args:
+            protein_tokens: Unused (kept for interface compatibility)
+            protein_seq: Protein sequence string (single sequence)
+        
+        Returns:
+            int: 1 if Cas9 score > threshold, 0 otherwise
+        """
+        scores = self.cas9_classifier_obj.get_scores([protein_seq])
+        if not scores:
+            return 0
+        
+        score = scores[0]
+        return int(score > self.threshold)
+    
+    def predict_batch(self, seqs: Sequence[str]) -> List[bool]:
+        """
+        Batch prediction helper for evaluation / faster integration.
+        
+        Args:
+            seqs: List of protein sequence strings
+        
+        Returns:
+            List of bool: True if Cas9 score > threshold, False otherwise
+        """
+        scores = self.cas9_classifier_obj.get_scores(seqs)
+        return [score > self.threshold for score in scores]
+
+
+class PAMMatchingProbabilityThreshold:
+    """
+    Terminal constraint for protein sequences: PAM matching probability must be above a threshold.
+
+    Uses the PAMMatching objective model to compute probability for its configured target PAM and
+    requires probability > threshold.
+    """
+    def __init__(self, pam_matching_obj, threshold: float):
+        """
+        Args:
+            pam_matching_obj: PAMMatching-like object with target_pam and get_score_for_pam methods
+            threshold: Minimum PAM matching probability threshold (0-1). Sequences must have score > threshold.
+        """
+        if not hasattr(pam_matching_obj, "get_score_for_pam"):
+            raise ValueError("pam_matching_obj must have a get_score_for_pam method")
+        if not hasattr(pam_matching_obj, "target_pam"):
+            raise ValueError("pam_matching_obj must have a target_pam attribute")
+        if not (0.0 <= threshold <= 1.0):
+            raise ValueError(f"threshold must be in [0, 1], got {threshold}")
+
+        self.pam_matching_obj = pam_matching_obj
+        self.threshold = float(threshold)
+        self.target_pam = pam_matching_obj.target_pam
+        print(
+            "PAM Matching Probability Threshold Constraint: sequences must have "
+            f"PAM matching probability > {self.threshold} for target PAM {self.target_pam}"
+        )
+
+    def __call__(self, protein_tokens, protein_seq):
+        """
+        Check if protein sequence's PAM matching probability is above the threshold.
+
+        Args:
+            protein_tokens: Unused (kept for interface compatibility)
+            protein_seq: Protein sequence string (single sequence)
+
+        Returns:
+            int: 1 if PAM matching probability > threshold, 0 otherwise
+        """
+        scores = self.pam_matching_obj.get_score_for_pam(
+            [protein_seq], self.target_pam, use_temperature_scaling=False
+        )
+        if not scores:
+            return 0
+        score = scores[0]
+        return int(score > self.threshold)
+
+    def predict_batch(self, seqs: Sequence[str]) -> List[bool]:
+        """
+        Batch prediction helper for evaluation / faster integration.
+
+        Args:
+            seqs: List of protein sequence strings
+
+        Returns:
+            List of bool: True if PAM matching probability > threshold, False otherwise
+        """
+        scores = self.pam_matching_obj.get_score_for_pam(
+            seqs, self.target_pam, use_temperature_scaling=False
+        )
+        return [score > self.threshold for score in scores]
+
+
+# -------------------------------------------------------------------------
+# Cas9 nuclease domain completeness detector (HNH + RuvC), recall-first
+# -------------------------------------------------------------------------
+
+def _sha1(s: str) -> str:
+    return hashlib.sha1(s.encode("utf-8")).hexdigest()
+
+
+class Cas9DomainCompleteness:
+    """
+    Hard constraint for protein sequences:
+      pass iff sequence contains BOTH HNH and RuvC domains (domain presence),
+      using the Option-B detector (Stage1 fast scan + Stage2 --max rescue).
+
+    Intended use:
+      - As a *domain presence* gate (high recall)
+      - Combine with an MLP classifier to enforce "is Cas9" / other plausibility checks.
+    """
+    def __init__(
+        self,
+        hmm_path: str,
+        cpu: int = 16,
+        # Stage1 acceptance
+        ev1: float = 1e-2,
+        minlen1: int = 35,
+        # Stage2 acceptance (recall-first)
+        ev2: float = 10.0,
+        minlen2: int = 20,
+        # Optional: relax Stage1 HMMER filters (usually leave False)
+        stage1_relax_filters: bool = False,
+        # Cache results by sequence content to avoid repeated hmmscan calls
+        enable_cache: bool = True,
+    ):
+        if DomainDetector is None:
+            raise ImportError(
+                "Could not import DomainDetector. Ensure domain_detector.py is in the same package as constraints.py "
+                "and that its imports succeed."
+            )
+        self.hmm_path = hmm_path
+        self.enable_cache = enable_cache
+        self._cache: Dict[str, int] = {}
+
+        cfg = DomainDetectorConfig(
+            hmm_path=hmm_path,
+            cpu=cpu,
+            stage1=HitCriteria(ievalue_max=ev1, min_dom_len=minlen1),
+            stage2=HitCriteria(ievalue_max=ev2, min_dom_len=minlen2),
+        )
+
+        if stage1_relax_filters:
+            cfg = DomainDetectorConfig(
+                hmm_path=hmm_path,
+                cpu=cpu,
+                stage1=HitCriteria(ievalue_max=ev1, min_dom_len=minlen1),
+                stage2=HitCriteria(ievalue_max=ev2, min_dom_len=minlen2),
+                stage1_F1=1.0, stage1_F2=1.0, stage1_F3=1.0,
+                stage1_report_E=1e6, stage1_report_domE=1e6,
+            )
+
+        self.detector = DomainDetector(cfg)
+
+    def __call__(self, protein_tokens, protein_seq: str) -> int:
+        # protein_tokens is unused; kept to match existing constraint signature
+        seq = (protein_seq or "").strip().replace(" ", "").replace("\n", "")
+        if not seq:
+            return 0
+
+        if self.enable_cache:
+            key = _sha1(seq)
+            if key in self._cache:
+                return self._cache[key]
+
+        passed = int(self.detector.predict_one(seq, id_="query"))
+
+        if self.enable_cache:
+            self._cache[_sha1(seq)] = passed
+        return passed
+
+    def predict_batch(self, seqs: Sequence[str], ids: Optional[Sequence[str]] = None) -> List[bool]:
+        """
+        Batch prediction helper for evaluation / faster integration.
+        """
+        clean = [(s or "").strip().replace(" ", "").replace("\n", "") for s in seqs]
+        return self.detector.predict_batch(clean, ids=ids)

cas9/data/__init__.py

@@ -0,0 +1,9 @@
+# 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 .data import DataState
+
+__all__ = ["DataState"]

cas9/data/cas9_pre_batching.py

@@ -0,0 +1,160 @@
+import os
+import random
+import pandas as pd
+from datasets import Dataset, DatasetDict
+
+# import selfies as sf
+# from rdkit import Chem
+# from rdkit import RDLogger
+# RDLogger.DisableLog('rdApp.*')
+
+import sys
+
+from transformers import EsmTokenizer
+
+# ---------------------------------------------------------------------
+# CONFIG
+# ---------------------------------------------------------------------
+# Set PROJECT_ROOT or edit paths below
+PROJECT_ROOT = os.environ.get('PROJECT_ROOT', 'path/to/project')
+FASTA_PATH = os.path.join(PROJECT_ROOT, 'path/to/protein2pam_db_v1.0/cas9_full.fasta')
+OUTPUT_DIR = os.path.join(PROJECT_ROOT, 'data/cas9/cas9_dataset_esm2_tokenized_under_1125')
+
+MAX_TOKENS_PER_BATCH = 8192
+TRAIN_RATIO, VAL_RATIO, TEST_RATIO = 0.8, 0.1, 0.1
+
+# ------------------------------------------------------------
+# 1) Load sequences from FASTA file
+# ------------------------------------------------------------
+from Bio import SeqIO
+
+sequences = []
+print("Loading sequences from FASTA and validating…")
+for record in SeqIO.parse(FASTA_PATH, "fasta"):
+    s = str(record.seq)
+    # skip empty sequences
+    if not s:
+        continue
+    sequences.append(s)
+
+# remove all duplicates 
+sequences = list(set(sequences))
+
+## filter out any sequence that contains non-natural amino acids
+print(f"len(sequences) before filtering: {len(sequences)}")
+sequences = [
+    seq for seq in sequences
+    if all(aa in "ACDEFGHIKLMNPQRSTVWY" for aa in seq) and len(seq) <= 1125
+]
+print(f"len(sequences) after filtering: {len(sequences)}")
+print(f"Average length of sequences: {sum(len(seq) for seq in sequences) / len(sequences)}")
+
+# ---- Initialize ESM tokenizer ----
+tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
+print(f"Vocabulary Size: {tokenizer.vocab_size}")
+
+# shuffle before splitting
+random.shuffle(sequences)
+n = len(sequences)
+n_train = int(n * TRAIN_RATIO)
+n_val = int(n * VAL_RATIO)
+
+train_sequences = sequences[:n_train]
+val_sequences   = sequences[n_train:n_train + n_val]
+test_sequences  = sequences[n_train + n_val:]
+
+# ------------------------------------------------------------
+# 2) helper: tokenize sequences using the ESM tokenizer
+# ------------------------------------------------------------
+def tokenize_sequences(sequences):
+    enc = {"input_ids": [], "attention_mask": []}
+    for s in sequences:
+        res = tokenizer(s, add_special_tokens=True)
+        enc["input_ids"].append(res["input_ids"])
+        enc["attention_mask"].append(res["attention_mask"])
+    return enc
+
+
+# ------------------------------------------------------------
+# 3) build batched dataset (token-based batching: groups sequences by length, packs up to max_tokens)
+# ------------------------------------------------------------
+from tqdm import tqdm
+
+def build_batched_dataset(sequences, max_tokens=8192):
+    # tokenize all first
+    toks = tokenize_sequences(sequences)
+    input_ids_list = toks["input_ids"]
+    attn_mask_list = toks["attention_mask"]
+
+    # add lengths
+    items = []
+    for ids, mask in tqdm(zip(input_ids_list, attn_mask_list), total=len(input_ids_list), desc="Collecting items with length"):
+        items.append({
+                "input_ids": ids,
+                "attention_mask": mask,
+                "length": len(ids),
+        })
+
+    # sort by length
+    items.sort(key=lambda x: x["length"])
+
+    batched_input_ids = []
+    batched_attention_masks = []
+    batched_lengths = []
+    batched_batch_sizes = []
+
+    i = 0
+    n = len(items)
+    while i < n:
+        L = items[i]["length"]
+        # how many of length L can we pack?
+        max_bs = max_tokens // L
+        if max_bs < 1:
+            max_bs = 1
+
+        # collect up to max_bs items with same length
+        cur_ids = []
+        cur_masks = []
+        taken = 0
+        j = i
+        while j < n and items[j]["length"] == L and taken < max_bs:
+            cur_ids.append(items[j]["input_ids"])         # length L
+            cur_masks.append(items[j]["attention_mask"])  # length L
+            j += 1
+            taken += 1
+
+        # now this is one batch: shape (B, L)
+        batched_input_ids.append(cur_ids)
+        batched_attention_masks.append(cur_masks)
+        batched_lengths.append(L)
+        batched_batch_sizes.append(taken)
+
+        i = j
+
+    ds = Dataset.from_dict({
+        "input_ids": batched_input_ids,           # list of (B, L)
+            "attention_mask": batched_attention_masks,
+        "seq_length": batched_lengths,
+            "batch_size": batched_batch_sizes,
+    })
+    return ds
+
+
+# ------------------------------------------------------------
+# 4) build batched datasets for each split
+# ------------------------------------------------------------
+train_ds = build_batched_dataset(train_sequences, max_tokens=MAX_TOKENS_PER_BATCH)
+val_ds   = build_batched_dataset(val_sequences,   max_tokens=MAX_TOKENS_PER_BATCH)
+test_ds  = build_batched_dataset(test_sequences,  max_tokens=MAX_TOKENS_PER_BATCH)
+
+dsdict = DatasetDict(
+    {
+        "train": train_ds,
+        "validation": val_ds,
+        "test": test_ds,
+    }
+)
+
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+dsdict.save_to_disk(OUTPUT_DIR)
+print(f"saved to {OUTPUT_DIR}")

cas9/data/cas9_pre_batching_csv.py

@@ -0,0 +1,138 @@
+import os
+import random
+import pandas as pd
+from datasets import Dataset, DatasetDict
+from transformers import EsmTokenizer
+
+import selfies as sf
+from rdkit import Chem
+from rdkit import RDLogger
+RDLogger.DisableLog('rdApp.*')
+
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
+from smiles_tokenizer.selfies_tokenizers import SelfiesTokenizer
+
+# ---------------------------------------------------------------------
+# CONFIG
+# ---------------------------------------------------------------------
+PROJECT_ROOT = os.environ.get('PROJECT_ROOT', 'path/to/project')
+CSV_PATH = os.path.join(PROJECT_ROOT, "data/cas9/cas9.csv")
+OUTPUT_DIR = os.path.join(PROJECT_ROOT, "data/cas9")
+MAX_TOKENS_PER_BATCH = 8192
+TRAIN_RATIO, VAL_RATIO, TEST_RATIO = 0.8, 0.1, 0.1
+
+
+# ------------------------------------------------------------
+# 1) Load CSV and split 8:1:1
+# ------------------------------------------------------------
+df = pd.read_csv(CSV_PATH)
+sequences = df["Sequence"].tolist()
+
+# ---- Build vocab ----
+tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
+
+random.shuffle(sequences)
+n = len(sequences)
+n_train = int(n * TRAIN_RATIO)
+n_val = int(n * VAL_RATIO)
+
+train_df = sequences[:n_train]
+val_df   = sequences[n_train:n_train + n_val]
+test_df  = sequences[n_train + n_val:]
+
+
+# ------------------------------------------------------------
+# 2) Tokenizer
+# ------------------------------------------------------------
+
+def tokenize_selfies(sequences):
+    enc = {"input_ids": [], "attention_mask": []}
+    for s in sequences:
+        res = tokenizer(s)
+        enc["input_ids"].append(res["input_ids"])
+        enc["attention_mask"].append(res["attention_mask"])
+    # enc["input_ids"] is a list of list[int]
+    return enc
+
+
+# ------------------------------------------------------------
+# helper: build batches (one row = one batch)
+# ------------------------------------------------------------
+def build_batched_dataset(sequences, max_tokens=9182):
+    # tokenize all first
+    toks = tokenize_selfies(sequences)
+    input_ids_list = toks["input_ids"]
+    attn_mask_list = toks["attention_mask"]
+
+    # add lengths
+    items = []
+    for ids, mask in zip(input_ids_list, attn_mask_list):
+        items.append({
+            "input_ids": ids,
+            "attention_mask": mask,
+            "length": len(ids),
+        })
+
+    # sort by length
+    items.sort(key=lambda x: x["length"])
+
+    batched_input_ids = []
+    batched_attention_masks = []
+    batched_lengths = []
+    batched_batch_sizes = []
+
+    i = 0
+    n = len(items)
+    while i < n:
+        L = items[i]["length"]
+        # how many of length L can we pack?
+        max_bs = max_tokens // L
+        if max_bs < 1:
+            max_bs = 1
+
+        # collect up to max_bs items with same length
+        cur_ids = []
+        cur_masks = []
+        taken = 0
+        j = i
+        while j < n and items[j]["length"] == L and taken < max_bs:
+            cur_ids.append(items[j]["input_ids"])         # length L
+            cur_masks.append(items[j]["attention_mask"])  # length L
+            j += 1
+            taken += 1
+
+        # now this is one batch: shape (B, L)
+        batched_input_ids.append(cur_ids)
+        batched_attention_masks.append(cur_masks)
+        batched_lengths.append(L)
+        batched_batch_sizes.append(taken)
+
+        i = j
+
+    ds = Dataset.from_dict({
+        "input_ids": batched_input_ids,           # list of (B, L)
+        "attention_mask": batched_attention_masks,
+        "seq_length": batched_lengths,
+        "batch_size": batched_batch_sizes,
+    })
+    return ds
+
+
+# ------------------------------------------------------------
+# 3) build batched datasets for each split
+# ------------------------------------------------------------
+train_ds = build_batched_dataset(train_df, MAX_TOKENS_PER_BATCH)
+val_ds   = build_batched_dataset(val_df,   MAX_TOKENS_PER_BATCH)
+test_ds  = build_batched_dataset(test_df,  MAX_TOKENS_PER_BATCH)
+
+dsdict = DatasetDict({
+    "train": train_ds,
+    "validation": val_ds,
+    "test": test_ds,
+})
+
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+dsdict.save_to_disk(OUTPUT_DIR)
+print(f"saved to {OUTPUT_DIR}")

cas9/data/data.py

@@ -0,0 +1,309 @@
+# 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
+
+from dataclasses import dataclass, field
+from typing import Dict, Iterable, List, Tuple, Optional
+
+import os
+
+from datasets import Dataset as HFDataset, DatasetDict, load_dataset
+from omegaconf import OmegaConf
+
+import torch
+from torch.utils.data import DataLoader, Dataset
+from transformers import AutoTokenizer
+try:
+    from transformers import EsmTokenizer
+    _HAVE_ESM = True
+except Exception:
+    _HAVE_ESM = False
+
+from cas9.data.tokenizer import wt_detokenizer
+from cas9.data.utils import cycle_loader, StatefulDistributedSampler
+
+
+# -------------------------
+# Tokenizer helper
+# -------------------------
+
+def _get_tokenizer(name: Optional[str]):
+    """
+    Returns a HF tokenizer by name. Supports ESM when available.
+    """
+    tok_name = name or "gpt2"
+    if _HAVE_ESM and tok_name.startswith("facebook/esm"):
+        return EsmTokenizer.from_pretrained(tok_name)
+    return AutoTokenizer.from_pretrained(tok_name)
+
+
+# -------------------------
+# HuggingFace datasets (ragged)
+# -------------------------
+
+def _get_hf_dataset_ragged(
+    name: str,
+    mode: str,
+    cache_dir: Optional[str],
+    num_proc: int,
+    tokenizer_name: Optional[str],
+) -> HFDataset:
+    """
+    Loads an HF dataset split and tokenizes into variable-length sequences.
+    No chunking to fixed block_size; each example is one tokenized text (+ EOS if available).
+    Returns a HuggingFace Dataset with a single column "input_ids" (list[int]) per row.
+    """
+    detokenizer = None
+    if name == "wikitext103":
+        ds = load_dataset("wikitext", name="wikitext-103-raw-v1", cache_dir=cache_dir)[mode]
+        detokenizer = wt_detokenizer
+    elif name == "fineweb-edu":
+        ds = load_dataset("HuggingFaceFW/fineweb-edu", name="CC-MAIN-2024-10", cache_dir=cache_dir)[mode]
+    else:
+        ds = load_dataset(name, cache_dir=cache_dir)[mode]
+
+    tok = _get_tokenizer(tokenizer_name)
+    eos_id = None
+    if hasattr(tok, "eos_token") and tok.eos_token is not None:
+        eos_id = tok.convert_tokens_to_ids(tok.eos_token)
+        if isinstance(eos_id, list):
+            eos_id = eos_id[0]
+
+    def _apply_detok(batch: Dict):
+        texts = batch["text"]
+        if detokenizer is not None:
+            texts = [detokenizer(t) for t in texts]
+        return {"text": texts}
+
+    if "text" in ds.column_names:
+        ds = ds.map(_apply_detok, batched=True, num_proc=num_proc, load_from_cache_file=True)
+
+    def _tokenize(batch: Dict):
+        # Use fast tokenizer; do not return attention masks
+        ids = tok(batch["text"], add_special_tokens=True, return_attention_mask=False)["input_ids"]
+        # if eos_id is not None:
+        #     for arr in ids:
+        #         arr.append(eos_id)
+        return {"input_ids": ids}
+
+    # Tokenize
+    tokenized = ds.map(
+        _tokenize,
+        batched=True,
+        num_proc=num_proc,
+        remove_columns=[c for c in ds.column_names if c != "text"],
+        load_from_cache_file=True,
+    )
+
+    # Keep only input_ids
+    if "text" in tokenized.column_names:
+        tokenized = tokenized.remove_columns("text")
+
+    # For fineweb there can be extra meta columns; keep only input_ids
+    extra_cols = [c for c in tokenized.column_names if c != "input_ids"]
+    if extra_cols:
+        tokenized = tokenized.remove_columns(extra_cols)
+
+    tokenized = tokenized.with_format("python")  # return plain python objects; collate will tensorize
+    return tokenized
+
+
+# -------------------------
+# Local FASTA dataset (ragged)
+# -------------------------
+
+class FASTADataset(Dataset):
+    """
+    Loads a local FASTA file and tokenizes each sequence to variable-length ids.
+    One HF-like item per sequence: {"input_ids": List[int]}.
+    """
+
+    def __init__(self, fasta_path: str, tokenizer_name: Optional[str]):
+        assert os.path.isfile(fasta_path), f"FASTA file not found: {fasta_path}"
+        self.fasta_path = fasta_path
+        self.tok = _get_tokenizer(tokenizer_name)
+        self.eos_id = None
+        if hasattr(self.tok, "eos_token") and self.tok.eos_token is not None:
+            eid = self.tok.convert_tokens_to_ids(self.tok.eos_token)
+            self.eos_id = eid[0] if isinstance(eid, list) else eid
+
+        # Parse FASTA (simple parser; no Biopython dependency)
+        self._seqs: List[str] = []
+        with open(fasta_path, "r") as f:
+            cur = []
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                if line.startswith(">"):
+                    if cur:
+                        self._seqs.append("".join(cur))
+                        cur = []
+                else:
+                    cur.append(line)
+            if cur:
+                self._seqs.append("".join(cur))
+
+    def __len__(self) -> int:
+        return len(self._seqs)
+
+    def __getitem__(self, idx: int) -> Dict[str, List[int]]:
+        seq = self._seqs[idx]
+        ids = self.tok(seq, add_special_tokens=True, return_attention_mask=False)["input_ids"]
+        if isinstance(ids, list) and len(ids) > 0 and isinstance(ids[0], list):
+            # Some tokenizers may return nested lists for per-char tokenization; flatten
+            ids = [t for sub in ids for t in (sub if isinstance(sub, list) else [sub])]
+        # if self.eos_id is not None:
+        #     ids.append(self.eos_id)
+        return {"input_ids": ids}
+
+
+# -------------------------
+# Ragged collate
+# -------------------------
+
+def ragged_collate(batch: List[Dict[str, List[int]]]) -> Dict[str, List[torch.Tensor]]:
+    """
+    Collate a list of {"input_ids": List[int]} into {"input_ids": List[LongTensor]}.
+    No padding; tensorization only.
+    """
+    out: List[torch.Tensor] = []
+    for item in batch:
+        ids = item["input_ids"]
+        out.append(torch.tensor(ids, dtype=torch.long))
+    return {"input_ids": out}
+
+
+# -------------------------
+# Public dataclasses
+# -------------------------
+
+@dataclass
+class DatasetWrap:
+    dataset: Dataset  # HF Dataset (python format) OR torch Dataset
+    sampler: StatefulDistributedSampler
+
+
+@dataclass
+class DataState:
+    train: DatasetWrap
+    test: DatasetWrap
+
+
+# -------------------------
+# Builders
+# -------------------------
+
+def _get_dataset(
+    name: str,
+    mode: str,
+    cache_dir: Optional[str],
+    num_proc: int,
+    batch_size: int,
+    ngpus: int,
+    tokenizer_name: Optional[str],
+    fasta_path: Optional[str] = None,
+) -> DatasetWrap:
+    """
+    Build either an HF dataset (ragged) or a local FASTA dataset.
+    - name == "fasta" -> requires fasta_path
+    - otherwise -> HF dataset by name
+    """
+    assert batch_size % ngpus == 0, f"{mode} batch size must be divisible by number of gpus."
+
+    if name.lower() == "fasta":
+        assert fasta_path is not None, "For name='fasta', you must set data.fasta_path in the config."
+        dataset = FASTADataset(fasta_path=fasta_path, tokenizer_name=tokenizer_name)
+    else:
+        dataset = _get_hf_dataset_ragged(
+            name=name,
+            mode=mode,
+            cache_dir=cache_dir,
+            num_proc=num_proc,
+            tokenizer_name=tokenizer_name,
+        )
+
+    sampler = StatefulDistributedSampler(dataset=dataset)
+    return DatasetWrap(dataset=dataset, sampler=sampler)
+
+
+def get_data_state(config: OmegaConf) -> DataState:
+    """
+    Expects (typical):
+      config.data.train / config.data.valid : dataset name (e.g., "wikitext103", "fineweb-edu", "fasta", or any HF hub dataset id)
+      config.data.cache_dir
+      config.data.num_workers
+      config.data.tokenizer (string, e.g., "gpt2" or "facebook/esm2_t33_650M_UR50S")
+      config.data.fasta_path (optional, used when name == "fasta")
+      config.training.batch_size
+      config.eval.batch_size
+      config.compute.ngpus
+    """
+    tok_name = getattr(config.data, "tokenizer", None)
+
+    train = _get_dataset(
+        name=config.data.train,
+        mode="train",
+        cache_dir=config.data.cache_dir,
+        num_proc=config.data.num_workers,
+        batch_size=config.training.batch_size,
+        ngpus=config.compute.ngpus,
+        tokenizer_name=tok_name,
+        fasta_path=getattr(config.data, "fasta_path", None),
+    )
+
+    valid = _get_dataset(
+        name=config.data.valid,
+        mode="validation",
+        cache_dir=config.data.cache_dir,
+        num_proc=config.data.num_workers,
+        batch_size=config.eval.batch_size,
+        ngpus=config.compute.ngpus,
+        tokenizer_name=tok_name,
+        fasta_path=getattr(config.data, "fasta_path", None),
+    )
+
+    return DataState(train=train, test=valid)
+
+
+def get_data_loaders(
+    config: OmegaConf,
+    data_state: DataState,
+) -> Tuple[Iterable, Iterable]:
+    """
+    Returns infinite iterators over train/valid using ragged_collate.
+    Each batch item is a dict: {"input_ids": List[LongTensor]}.
+    """
+    per_gpu_train_bs = config.training.batch_size // config.compute.ngpus
+    per_gpu_valid_bs = config.eval.batch_size // config.compute.ngpus
+
+    train_loader = cycle_loader(
+        DataLoader(
+            data_state.train.dataset,
+            batch_size=per_gpu_train_bs,
+            sampler=data_state.train.sampler,
+            num_workers=config.data.num_workers,
+            pin_memory=True,
+            shuffle=(data_state.train.sampler is None),
+            persistent_workers=True,
+            collate_fn=ragged_collate,
+        )
+    )
+
+    valid_loader = cycle_loader(
+        DataLoader(
+            data_state.test.dataset,
+            batch_size=per_gpu_valid_bs,
+            sampler=data_state.test.sampler,
+            num_workers=config.data.num_workers,
+            pin_memory=True,
+            shuffle=(data_state.test.sampler is None),
+            collate_fn=ragged_collate,
+        )
+    )
+
+    return iter(train_loader), iter(valid_loader)   

cas9/data/match_cas9_length_distribution.py

@@ -0,0 +1,297 @@
+#!/usr/bin/env python3
+"""
+Script to construct a dataset of uniref50 sequences that matches the 
+sequence length distribution found in a cas9 FASTA file.
+
+The script prioritizes matching the length distribution over matching 
+a specified dataset size exactly.
+"""
+
+import argparse
+import sys
+from collections import Counter, defaultdict
+from Bio import SeqIO
+import numpy as np
+from tqdm import tqdm
+
+
+def load_sequences(fasta_path, description=""):
+    """Load sequences from a FASTA file."""
+    sequences = []
+    lengths = []
+    print(f"Loading {description} sequences from {fasta_path}...")
+    
+    for record in tqdm(SeqIO.parse(fasta_path, "fasta"), desc=f"Reading {description}"):
+        seq = str(record.seq).upper()
+        # Skip empty sequences
+        if not seq:
+            continue
+        # Filter out sequences with non-natural amino acids
+        if all(aa in "ACDEFGHIKLMNPQRSTVWY" for aa in seq):
+            sequences.append(seq)
+            lengths.append(len(seq))
+    
+    print(f"Loaded {len(sequences)} {description} sequences")
+    if sequences:
+        print(f"  Length range: {min(lengths)} - {max(lengths)}")
+        print(f"  Average length: {sum(lengths) / len(lengths):.1f}")
+    
+    return sequences, lengths
+
+
+def compute_length_distribution(lengths, bins=None):
+    """
+    Compute the distribution of sequence lengths.
+    
+    Args:
+        lengths: List of sequence lengths
+        bins: Optional list of bin edges. If None, uses histogram bins.
+    
+    Returns:
+        bin_edges: List of bin edges
+        counts: List of counts for each bin
+        normalized_dist: Normalized distribution (probabilities)
+    """
+    if bins is None:
+        # Use histogram to automatically determine bins
+        counts, bin_edges = np.histogram(lengths, bins='auto')
+    else:
+        counts, bin_edges = np.histogram(lengths, bins=bins)
+    
+    # Normalize to get probability distribution
+    total = sum(counts)
+    normalized_dist = counts / total if total > 0 else counts
+    
+    return bin_edges, counts, normalized_dist
+
+
+def group_sequences_by_length(sequences):
+    """Group sequences by their length."""
+    length_groups = defaultdict(list)
+    for seq in sequences:
+        length_groups[len(seq)].append(seq)
+    return length_groups
+
+
+def sample_to_match_distribution(
+    source_sequences, 
+    source_lengths,
+    target_lengths,
+    target_size=None,
+    bin_edges=None
+):
+    """
+    Sample sequences from source_sequences to match the length distribution 
+    of target_lengths.
+    
+    Args:
+        source_sequences: List of sequences to sample from
+        source_lengths: List of lengths corresponding to source_sequences
+        target_lengths: List of lengths to match the distribution of
+        target_size: Optional target number of sequences (approximate)
+        bin_edges: Optional bin edges for histogram (if None, computed from target)
+    
+    Returns:
+        sampled_sequences: List of sampled sequences
+        sampled_lengths: List of lengths of sampled sequences
+    """
+    # Compute target distribution
+    if bin_edges is None:
+        target_bin_edges, target_counts, target_dist = compute_length_distribution(target_lengths)
+    else:
+        _, target_counts, target_dist = compute_length_distribution(target_lengths, bins=bin_edges)
+        target_bin_edges = bin_edges
+    
+    print(f"\nTarget distribution:")
+    print(f"  Number of bins: {len(target_dist)}")
+    print(f"  Bin edges: {target_bin_edges[:5]}...{target_bin_edges[-5:]}")
+    
+    # Group source sequences by length
+    print("\nGrouping source sequences by length...")
+    length_groups = group_sequences_by_length(source_sequences)
+    print(f"  Found {len(length_groups)} unique lengths")
+    
+    # Determine how many sequences to sample per bin
+    if target_size is None:
+        # Use the same counts as target
+        target_counts_scaled = target_counts.copy()
+    else:
+        # Scale the distribution to approximately match target_size
+        total_target = sum(target_counts)
+        scale_factor = target_size / total_target if total_target > 0 else 1.0
+        target_counts_scaled = (target_counts * scale_factor).astype(int)
+        # Ensure we have at least 1 sequence per non-zero bin
+        target_counts_scaled = np.maximum(target_counts_scaled, (target_counts > 0).astype(int))
+    
+    print(f"\nTarget sample size: {sum(target_counts_scaled)} sequences")
+    
+    # Sample sequences for each bin
+    sampled_sequences = []
+    sampled_lengths = []
+    
+    print("\nSampling sequences to match distribution...")
+    for i in tqdm(range(len(target_dist)), desc="Processing bins"):
+        bin_start = int(target_bin_edges[i])
+        bin_end = int(target_bin_edges[i + 1])
+        n_samples = int(target_counts_scaled[i])
+        
+        if n_samples == 0:
+            continue
+        
+        # Find all source sequences in this length range
+        candidates = []
+        for length in range(bin_start, bin_end + 1):
+            if length in length_groups:
+                candidates.extend(length_groups[length])
+        
+        if not candidates:
+            # If no exact match, try to find sequences close to this bin
+            # Find the closest available lengths
+            available_lengths = sorted(length_groups.keys())
+            for length in available_lengths:
+                if bin_start <= length <= bin_end:
+                    candidates.extend(length_groups[length])
+                    break
+        
+        if not candidates:
+            # Still no candidates, try to find sequences in nearby bins
+            bin_center = (bin_start + bin_end) / 2
+            closest_length = min(available_lengths, key=lambda x: abs(x - bin_center))
+            candidates.extend(length_groups[closest_length])
+        
+        # Sample from candidates
+        if len(candidates) >= n_samples:
+            sampled = np.random.choice(candidates, size=n_samples, replace=False).tolist()
+        else:
+            # Not enough candidates, sample with replacement
+            sampled = np.random.choice(candidates, size=n_samples, replace=True).tolist()
+        
+        sampled_sequences.extend(sampled)
+        sampled_lengths.extend([len(seq) for seq in sampled])
+    
+    return sampled_sequences, sampled_lengths
+
+
+def write_fasta(sequences, output_path, description="sequences"):
+    """Write sequences to a FASTA file."""
+    print(f"\nWriting {len(sequences)} {description} to {output_path}...")
+    with open(output_path, 'w') as f:
+        for i, seq in enumerate(sequences):
+            f.write(f">sequence_{i+1}\n")
+            f.write(f"{seq}\n")
+    print(f"Saved to {output_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Match uniref50 sequence length distribution to cas9 distribution"
+    )
+    parser.add_argument(
+        'cas9_fasta',
+        type=str,
+        help='Path to FASTA file containing cas9 sequences'
+    )
+    parser.add_argument(
+        'uniref50_fasta',
+        type=str,
+        help='Path to FASTA file containing uniref50 sequences'
+    )
+    parser.add_argument(
+        'output_fasta',
+        type=str,
+        help='Path to output FASTA file'
+    )
+    parser.add_argument(
+        '--target-size',
+        type=int,
+        default=None,
+        help='Approximate target dataset size (distribution matching takes precedence)'
+    )
+    parser.add_argument(
+        '--seed',
+        type=int,
+        default=42,
+        help='Random seed for reproducibility (default: 42)'
+    )
+    
+    args = parser.parse_args()
+    
+    # Set random seed
+    np.random.seed(args.seed)
+    
+    # Load cas9 sequences
+    cas9_sequences, cas9_lengths = load_sequences(args.cas9_fasta, "cas9")
+    if not cas9_sequences:
+        print("Error: No valid cas9 sequences found!", file=sys.stderr)
+        sys.exit(1)
+    
+    # Load uniref50 sequences
+    uniref50_sequences, uniref50_lengths = load_sequences(args.uniref50_fasta, "uniref50")
+    if not uniref50_sequences:
+        print("Error: No valid uniref50 sequences found!", file=sys.stderr)
+        sys.exit(1)
+    
+    # Compute cas9 length distribution
+    print("\n" + "="*60)
+    print("Analyzing cas9 length distribution...")
+    print("="*60)
+    cas9_bin_edges, cas9_counts, cas9_dist = compute_length_distribution(cas9_lengths)
+    print(f"Cas9 distribution statistics:")
+    print(f"  Total sequences: {len(cas9_sequences)}")
+    print(f"  Number of bins: {len(cas9_dist)}")
+    print(f"  Length range: {min(cas9_lengths)} - {max(cas9_lengths)}")
+    
+    # Sample uniref50 sequences to match cas9 distribution
+    print("\n" + "="*60)
+    print("Sampling uniref50 sequences to match cas9 distribution...")
+    print("="*60)
+    sampled_sequences, sampled_lengths = sample_to_match_distribution(
+        uniref50_sequences,
+        uniref50_lengths,
+        cas9_lengths,
+        target_size=args.target_size,
+        bin_edges=cas9_bin_edges
+    )
+    
+    # Compare distributions
+    print("\n" + "="*60)
+    print("Distribution comparison:")
+    print("="*60)
+    sampled_bin_edges, sampled_counts, sampled_dist = compute_length_distribution(
+        sampled_lengths, bins=cas9_bin_edges
+    )
+    
+    print(f"\nCas9 distribution:")
+    print(f"  Total sequences: {len(cas9_sequences)}")
+    print(f"  Mean length: {np.mean(cas9_lengths):.1f}")
+    print(f"  Std length: {np.std(cas9_lengths):.1f}")
+    
+    print(f"\nSampled uniref50 distribution:")
+    print(f"  Total sequences: {len(sampled_sequences)}")
+    print(f"  Mean length: {np.mean(sampled_lengths):.1f}")
+    print(f"  Std length: {np.std(sampled_lengths):.1f}")
+    
+    # Compute distribution similarity (KL divergence or correlation)
+    # Remove zero bins for comparison
+    mask = (cas9_counts > 0) & (sampled_counts > 0)
+    if mask.sum() > 0:
+        cas9_dist_filtered = cas9_dist[mask]
+        sampled_dist_filtered = sampled_dist[mask]
+        # Normalize filtered distributions
+        cas9_dist_filtered = cas9_dist_filtered / cas9_dist_filtered.sum()
+        sampled_dist_filtered = sampled_dist_filtered / sampled_dist_filtered.sum()
+        
+        # Compute correlation
+        correlation = np.corrcoef(cas9_dist_filtered, sampled_dist_filtered)[0, 1]
+        print(f"\nDistribution correlation: {correlation:.4f}")
+    
+    # Write output
+    write_fasta(sampled_sequences, args.output_fasta, "sampled uniref50")
+    
+    print("\n" + "="*60)
+    print("Done!")
+    print("="*60)
+
+
+if __name__ == "__main__":
+    main()

cas9/data/match_cas9_length_distribution_filtered.py

@@ -0,0 +1,506 @@
+#!/usr/bin/env python3
+"""
+Script to construct a dataset of uniref50 sequences that matches the 
+sequence length distribution found in a cas9 FASTA file.
+
+This variation:
+1. Filters uniref sequences to only include those between cas9 min and max length
+2. Removes redundancy within the filtered sequences at 80% identity using mmseqs2
+3. Uses the remaining sequences to match the cas9 length distribution
+
+The script prioritizes matching the length distribution over matching 
+a specified dataset size exactly.
+"""
+
+import argparse
+import sys
+import os
+import subprocess
+import tempfile
+import shutil
+from collections import Counter, defaultdict
+from Bio import SeqIO
+import numpy as np
+from tqdm import tqdm
+
+
+def load_sequences(fasta_path, description=""):
+    """Load sequences from a FASTA file."""
+    sequences = []
+    lengths = []
+    print(f"Loading {description} sequences from {fasta_path}...")
+    
+    for record in tqdm(SeqIO.parse(fasta_path, "fasta"), desc=f"Reading {description}"):
+        seq = str(record.seq).upper()
+        # Skip empty sequences
+        if not seq:
+            continue
+        # Filter out sequences with non-natural amino acids
+        if all(aa in "ACDEFGHIKLMNPQRSTVWY" for aa in seq):
+            sequences.append(seq)
+            lengths.append(len(seq))
+    
+    print(f"Loaded {len(sequences)} {description} sequences")
+    if sequences:
+        print(f"  Length range: {min(lengths)} - {max(lengths)}")
+        print(f"  Average length: {sum(lengths) / len(lengths):.1f}")
+    
+    return sequences, lengths
+
+
+def filter_and_write_fasta(input_fasta, output_fasta, min_length, max_length, description=""):
+    """
+    Filter sequences by length range and write to output FASTA file.
+    
+    Args:
+        input_fasta: Input FASTA file path
+        output_fasta: Output FASTA file path
+        min_length: Minimum sequence length (inclusive)
+        max_length: Maximum sequence length (inclusive)
+        description: Description for progress messages
+    
+    Returns:
+        Number of sequences written
+    """
+    print(f"\nFiltering {description} sequences to length range [{min_length}, {max_length}]...")
+    
+    filtered_count = 0
+    with open(output_fasta, 'w') as out_handle:
+        for record in tqdm(SeqIO.parse(input_fasta, "fasta"), desc=f"Filtering {description}"):
+            seq = str(record.seq).upper()
+            # Skip empty sequences
+            if not seq:
+                continue
+            # Filter out sequences with non-natural amino acids
+            if not all(aa in "ACDEFGHIKLMNPQRSTVWY" for aa in seq):
+                continue
+            # Filter by length
+            seq_len = len(seq)
+            if min_length <= seq_len <= max_length:
+                SeqIO.write(record, out_handle, "fasta")
+                filtered_count += 1
+    
+    print(f"Filtered {filtered_count} sequences to length range [{min_length}, {max_length}]")
+    return filtered_count
+
+
+def remove_redundancy_mmseqs(input_fasta, output_fasta, identity_threshold=0.8, tmp_dir=None):
+    """
+    Remove redundancy using mmseqs2 clustering.
+    
+    Args:
+        input_fasta: Input FASTA file
+        output_fasta: Output FASTA file (representative sequences)
+        identity_threshold: Sequence identity threshold (0.0-1.0)
+        tmp_dir: Temporary directory for mmseqs2 files (if None, uses tempfile)
+    
+    Returns:
+        Number of sequences before and after clustering
+    """
+    print(f"\n{'='*60}")
+    print(f"Removing redundancy using mmseqs2 at {identity_threshold*100:.0f}% identity...")
+    print(f"{'='*60}")
+    
+    # Count input sequences
+    input_count = sum(1 for _ in SeqIO.parse(input_fasta, "fasta"))
+    print(f"Input sequences: {input_count}")
+    
+    # Create temporary directory if not provided
+    if tmp_dir is None:
+        tmp_dir = tempfile.mkdtemp(prefix="mmseqs_")
+        cleanup_tmp = True
+    else:
+        cleanup_tmp = False
+        os.makedirs(tmp_dir, exist_ok=True)
+    
+    try:
+        db_path = os.path.join(tmp_dir, "db")
+        cluster_path = os.path.join(tmp_dir, "cluster")
+        cluster_rep_path = os.path.join(tmp_dir, "cluster_rep")
+        
+        # Step 1: Create database
+        print("\nStep 1: Creating mmseqs2 database...")
+        cmd = ['mmseqs', 'createdb', input_fasta, db_path]
+        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+        
+        # Step 2: Cluster sequences
+        print(f"\nStep 2: Clustering at {identity_threshold*100:.0f}% identity...")
+        cmd = [
+            'mmseqs', 'cluster',
+            db_path,
+            cluster_path,
+            tmp_dir,
+            '--min-seq-id', str(identity_threshold)
+            # Note: mmseqs2 uses all available threads by default
+        ]
+        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+        
+        # Step 3: Extract representative sequences
+        print("\nStep 3: Extracting representative sequences...")
+        cmd = ['mmseqs', 'createsubdb', cluster_path, db_path, cluster_rep_path]
+        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+        
+        # Step 4: Convert to FASTA
+        print("\nStep 4: Converting to FASTA...")
+        cmd = ['mmseqs', 'convert2fasta', cluster_rep_path, output_fasta]
+        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+        
+        # Count output sequences
+        output_count = sum(1 for _ in SeqIO.parse(output_fasta, "fasta"))
+        print(f"\nOutput sequences: {output_count}")
+        print(f"Redundancy removed: {input_count - output_count} sequences ({100*(input_count-output_count)/input_count:.2f}%)")
+        
+        return input_count, output_count
+        
+    except subprocess.CalledProcessError as e:
+        print(f"Error running mmseqs2: {e}", file=sys.stderr)
+        print(f"stdout: {e.stdout}", file=sys.stderr)
+        print(f"stderr: {e.stderr}", file=sys.stderr)
+        raise
+    except FileNotFoundError:
+        print("Error: mmseqs2 not found in PATH. Please install mmseqs2.", file=sys.stderr)
+        raise
+    finally:
+        # Cleanup temporary directory
+        if cleanup_tmp and os.path.exists(tmp_dir):
+            shutil.rmtree(tmp_dir)
+
+
+def compute_length_distribution(lengths, bins=None):
+    """
+    Compute the distribution of sequence lengths.
+    
+    Args:
+        lengths: List of sequence lengths
+        bins: Optional list of bin edges. If None, uses histogram bins.
+    
+    Returns:
+        bin_edges: List of bin edges
+        counts: List of counts for each bin
+        normalized_dist: Normalized distribution (probabilities)
+    """
+    if bins is None:
+        # Use histogram to automatically determine bins
+        counts, bin_edges = np.histogram(lengths, bins='auto')
+    else:
+        counts, bin_edges = np.histogram(lengths, bins=bins)
+    
+    # Normalize to get probability distribution
+    total = sum(counts)
+    normalized_dist = counts / total if total > 0 else counts
+    
+    return bin_edges, counts, normalized_dist
+
+
+def group_sequences_by_length(sequences):
+    """Group sequences by their length."""
+    length_groups = defaultdict(list)
+    for seq in sequences:
+        length_groups[len(seq)].append(seq)
+    return length_groups
+
+
+def sample_to_match_distribution(
+    source_sequences, 
+    source_lengths,
+    target_lengths,
+    target_size=None,
+    bin_edges=None
+):
+    """
+    Sample sequences from source_sequences to match the length distribution 
+    of target_lengths.
+    
+    Args:
+        source_sequences: List of sequences to sample from
+        source_lengths: List of lengths corresponding to source_sequences
+        target_lengths: List of lengths to match the distribution of
+        target_size: Optional target number of sequences (approximate)
+        bin_edges: Optional bin edges for histogram (if None, computed from target)
+    
+    Returns:
+        sampled_sequences: List of sampled sequences
+        sampled_lengths: List of lengths of sampled sequences
+    """
+    # Compute target distribution
+    if bin_edges is None:
+        target_bin_edges, target_counts, target_dist = compute_length_distribution(target_lengths)
+    else:
+        _, target_counts, target_dist = compute_length_distribution(target_lengths, bins=bin_edges)
+        target_bin_edges = bin_edges
+    
+    print(f"\nTarget distribution:")
+    print(f"  Number of bins: {len(target_dist)}")
+    print(f"  Bin edges: {target_bin_edges[:5]}...{target_bin_edges[-5:]}")
+    
+    # Group source sequences by length
+    print("\nGrouping source sequences by length...")
+    length_groups = group_sequences_by_length(source_sequences)
+    print(f"  Found {len(length_groups)} unique lengths")
+    
+    # Determine how many sequences to sample per bin
+    if target_size is None:
+        # Use the same counts as target
+        target_counts_scaled = target_counts.copy()
+    else:
+        # Scale the distribution to approximately match target_size
+        total_target = sum(target_counts)
+        scale_factor = target_size / total_target if total_target > 0 else 1.0
+        target_counts_scaled = (target_counts * scale_factor).astype(int)
+        # Ensure we have at least 1 sequence per non-zero bin
+        target_counts_scaled = np.maximum(target_counts_scaled, (target_counts > 0).astype(int))
+    
+    print(f"\nTarget sample size: {sum(target_counts_scaled)} sequences")
+    
+    # Sample sequences for each bin
+    sampled_sequences = []
+    sampled_lengths = []
+    
+    print("\nSampling sequences to match distribution...")
+    for i in tqdm(range(len(target_dist)), desc="Processing bins"):
+        bin_start = int(target_bin_edges[i])
+        bin_end = int(target_bin_edges[i + 1])
+        n_samples = int(target_counts_scaled[i])
+        
+        if n_samples == 0:
+            continue
+        
+        # Find all source sequences in this length range
+        candidates = []
+        for length in range(bin_start, bin_end + 1):
+            if length in length_groups:
+                candidates.extend(length_groups[length])
+        
+        if not candidates:
+            # If no exact match, try to find sequences close to this bin
+            # Find the closest available lengths
+            available_lengths = sorted(length_groups.keys())
+            for length in available_lengths:
+                if bin_start <= length <= bin_end:
+                    candidates.extend(length_groups[length])
+                    break
+        
+        if not candidates:
+            # Still no candidates, try to find sequences in nearby bins
+            bin_center = (bin_start + bin_end) / 2
+            closest_length = min(available_lengths, key=lambda x: abs(x - bin_center))
+            candidates.extend(length_groups[closest_length])
+        
+        # Sample from candidates
+        if len(candidates) >= n_samples:
+            sampled = np.random.choice(candidates, size=n_samples, replace=False).tolist()
+        else:
+            # Not enough candidates, sample with replacement
+            sampled = np.random.choice(candidates, size=n_samples, replace=True).tolist()
+        
+        sampled_sequences.extend(sampled)
+        sampled_lengths.extend([len(seq) for seq in sampled])
+    
+    return sampled_sequences, sampled_lengths
+
+
+def write_fasta(sequences, output_path, description="sequences"):
+    """Write sequences to a FASTA file."""
+    print(f"\nWriting {len(sequences)} {description} to {output_path}...")
+    with open(output_path, 'w') as f:
+        for i, seq in enumerate(sequences):
+            f.write(f">sequence_{i+1}\n")
+            f.write(f"{seq}\n")
+    print(f"Saved to {output_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Match uniref50 sequence length distribution to cas9 distribution (with filtering and redundancy removal)"
+    )
+    parser.add_argument(
+        'cas9_fasta',
+        type=str,
+        help='Path to FASTA file containing cas9 sequences'
+    )
+    parser.add_argument(
+        'uniref50_fasta',
+        type=str,
+        help='Path to FASTA file containing uniref50 sequences'
+    )
+    parser.add_argument(
+        'output_fasta',
+        type=str,
+        help='Path to output FASTA file'
+    )
+    parser.add_argument(
+        '--target-size',
+        type=int,
+        default=None,
+        help='Approximate target dataset size (distribution matching takes precedence)'
+    )
+    parser.add_argument(
+        '--identity',
+        type=float,
+        default=0.8,
+        help='Sequence identity threshold for redundancy removal (default: 0.8)'
+    )
+    parser.add_argument(
+        '--tmp-dir',
+        type=str,
+        default=None,
+        help='Temporary directory for intermediate files (default: auto)'
+    )
+    parser.add_argument(
+        '--seed',
+        type=int,
+        default=42,
+        help='Random seed for reproducibility (default: 42)'
+    )
+    
+    args = parser.parse_args()
+    
+    # Set random seed
+    np.random.seed(args.seed)
+    
+    # Create temporary directory for intermediate files
+    if args.tmp_dir is None:
+        tmp_dir = tempfile.mkdtemp(prefix="match_cas9_")
+        cleanup_tmp = True
+    else:
+        tmp_dir = args.tmp_dir
+        cleanup_tmp = False
+        os.makedirs(tmp_dir, exist_ok=True)
+    
+    try:
+        # Load cas9 sequences
+        print("\n" + "="*60)
+        print("Step 1: Loading cas9 sequences...")
+        print("="*60)
+        cas9_sequences, cas9_lengths = load_sequences(args.cas9_fasta, "cas9")
+        if not cas9_sequences:
+            print("Error: No valid cas9 sequences found!", file=sys.stderr)
+            sys.exit(1)
+        
+        # Get cas9 length range
+        cas9_min_length = min(cas9_lengths)
+        cas9_max_length = max(cas9_lengths)
+        print(f"\nCas9 length range: {cas9_min_length} - {cas9_max_length}")
+        
+        # Filter uniref sequences by length range
+        print("\n" + "="*60)
+        print("Step 2: Filtering uniref sequences by length range...")
+        print("="*60)
+        filtered_fasta = os.path.join(tmp_dir, "uniref_filtered.fasta")
+        filtered_count = filter_and_write_fasta(
+            args.uniref50_fasta,
+            filtered_fasta,
+            cas9_min_length,
+            cas9_max_length,
+            "uniref50"
+        )
+        
+        if filtered_count == 0:
+            print("Error: No uniref sequences found in the specified length range!", file=sys.stderr)
+            sys.exit(1)
+        
+        # Remove redundancy from filtered sequences
+        print("\n" + "="*60)
+        print("Step 3: Removing redundancy from filtered sequences...")
+        print("="*60)
+        dedup_fasta = os.path.join(tmp_dir, "uniref_dedup.fasta")
+        input_count, output_count = remove_redundancy_mmseqs(
+            filtered_fasta,
+            dedup_fasta,
+            identity_threshold=args.identity,
+            tmp_dir=os.path.join(tmp_dir, "mmseqs_tmp")
+        )
+        
+        # Load deduplicated sequences
+        print("\n" + "="*60)
+        print("Step 4: Loading deduplicated sequences...")
+        print("="*60)
+        uniref50_sequences, uniref50_lengths = load_sequences(dedup_fasta, "uniref50 (deduplicated)")
+        if not uniref50_sequences:
+            print("Error: No valid deduplicated uniref50 sequences found!", file=sys.stderr)
+            sys.exit(1)
+        
+        # Compute cas9 length distribution
+        print("\n" + "="*60)
+        print("Step 5: Analyzing cas9 length distribution...")
+        print("="*60)
+        cas9_bin_edges, cas9_counts, cas9_dist = compute_length_distribution(cas9_lengths)
+        print(f"Cas9 distribution statistics:")
+        print(f"  Total sequences: {len(cas9_sequences)}")
+        print(f"  Number of bins: {len(cas9_dist)}")
+        print(f"  Length range: {min(cas9_lengths)} - {max(cas9_lengths)}")
+        
+        # Sample uniref50 sequences to match cas9 distribution
+        print("\n" + "="*60)
+        print("Step 6: Sampling uniref50 sequences to match cas9 distribution...")
+        print("="*60)
+        sampled_sequences, sampled_lengths = sample_to_match_distribution(
+            uniref50_sequences,
+            uniref50_lengths,
+            cas9_lengths,
+            target_size=args.target_size,
+            bin_edges=cas9_bin_edges
+        )
+        
+        # Compare distributions
+        print("\n" + "="*60)
+        print("Distribution comparison:")
+        print("="*60)
+        sampled_bin_edges, sampled_counts, sampled_dist = compute_length_distribution(
+            sampled_lengths, bins=cas9_bin_edges
+        )
+        
+        print(f"\nCas9 distribution:")
+        print(f"  Total sequences: {len(cas9_sequences)}")
+        print(f"  Mean length: {np.mean(cas9_lengths):.1f}")
+        print(f"  Std length: {np.std(cas9_lengths):.1f}")
+        
+        print(f"\nSampled uniref50 distribution:")
+        print(f"  Total sequences: {len(sampled_sequences)}")
+        print(f"  Mean length: {np.mean(sampled_lengths):.1f}")
+        print(f"  Std length: {np.std(sampled_lengths):.1f}")
+        
+        # Compute distribution similarity (correlation)
+        # Remove zero bins for comparison
+        mask = (cas9_counts > 0) & (sampled_counts > 0)
+        if mask.sum() > 0:
+            cas9_dist_filtered = cas9_dist[mask]
+            sampled_dist_filtered = sampled_dist[mask]
+            # Normalize filtered distributions
+            cas9_dist_filtered = cas9_dist_filtered / cas9_dist_filtered.sum()
+            sampled_dist_filtered = sampled_dist_filtered / sampled_dist_filtered.sum()
+            
+            # Compute correlation
+            correlation = np.corrcoef(cas9_dist_filtered, sampled_dist_filtered)[0, 1]
+            print(f"\nDistribution correlation: {correlation:.4f}")
+        
+        # Write output
+        print("\n" + "="*60)
+        print("Step 7: Writing final output...")
+        print("="*60)
+        write_fasta(sampled_sequences, args.output_fasta, "sampled uniref50")
+        
+        # Summary
+        print("\n" + "="*60)
+        print("SUMMARY")
+        print("="*60)
+        print(f"Cas9 sequences: {len(cas9_sequences)}")
+        print(f"Cas9 length range: {cas9_min_length} - {cas9_max_length}")
+        print(f"Uniref sequences after length filtering: {filtered_count}")
+        print(f"Uniref sequences after redundancy removal ({args.identity*100:.0f}%): {output_count}")
+        print(f"Final sampled sequences: {len(sampled_sequences)}")
+        print("="*60)
+        
+    finally:
+        # Cleanup temporary directory
+        if cleanup_tmp and os.path.exists(tmp_dir):
+            print(f"\nCleaning up temporary directory: {tmp_dir}")
+            shutil.rmtree(tmp_dir)
+    
+    print("\n" + "="*60)
+    print("Done!")
+    print("="*60)
+
+
+if __name__ == "__main__":
+    main()

cas9/data/orthologs.fasta

@@ -0,0 +1,218 @@
+>sacas9 len=1053
+MKRNYILGLDIGITSVGYGIIDYETRDVIDAGVRLFKEANVENNEGRRSKRGARRLKRRR
+RHRIQRVKKLLFDYNLLTDHSELSGINPYEARVKGLSQKLSEEEFSAALLHLAKRRGVHN
+VNEVEEDTGNELSTKEQISRNSKALEEKYVAELQLERLKKDGEVRGSINRFKTSDYVKEA
+KQLLKVQKAYHQLDQSFIDTYIDLLETRRTYYEGPGEGSPFGWKDIKEWYEMLMGHCTYF
+PEELRSVKYAYNADLYNALNDLNNLVITRDENEKLEYYEKFQIIENVFKQKKKPTLKQIA
+KEILVNEEDIKGYRVTSTGKPEFTNLKVYHDIKDITARKEIIENAELLDQIAKILTIYQS
+SEDIQEELTNLNSELTQEEIEQISNLKGYTGTHNLSLKAINLILDELWHTNDNQIAIFNR
+LKLVPKKVDLSQQKEIPTTLVDDFILSPVVKRSFIQSIKVINAIIKKYGLPNDIIIELAR
+EKNSKDAQKMINEMQKRNRQTNERIEEIIRTTGKENAKYLIEKIKLHDMQEGKCLYSLEA
+IPLEDLLNNPFNYEVDHIIPRSVSFDNSFNNKVLVKQEENSKKGNRTPFQYLSSSDSKIS
+YETFKKHILNLAKGKGRISKTKKEYLLEERDINRFSVQKDFINRNLVDTRYATRGLMNLL
+RSYFRVNNLDVKVKSINGGFTSFLRRKWKFKKERNKGYKHHAEDALIIANADFIFKEWKK
+LDKAKKVMENQMFEEKQAESMPEIETEQEYKEIFITPHQIKHIKDFKDYKYSHRVDKKPN
+RELINDTLYSTRKDDKGNTLIVNNLNGLYDKDNDKLKKLINKSPEKLLMYHHDPQTYQKL
+KLIMEQYGDEKNPLYKYYEETGNYLTKYSKKDNGPVIKKIKYYGNKLNAHLDITDDYPNS
+RNKVVKLSLKPYRFDVYLDNGVYKFVTVKNLDVIKKENYYEVNSKCYEEAKKLKKISNQA
+EFIASFYNNDLIKINGELYRVIGVNNDLLNRIEVNMIDITYREYLENMNDKRPPRIIKTI
+ASKTQSIKKYSTDILGNLYEVKSKKHPQIIKKG
+>cjcas9 len=984
+MARILAFDIGISSIGWAFSENDELKDCGVRIFTKVENPKTGESLALPRRLARSARKRLAR
+RKARLNHLKHLIANEFKLNYEDYQSFDESLAKAYKGSLISPYELRFRALNELLSKQDFAR
+VILHIAKRRGYDDIKNSDDKEKGAILKAIKQNEEKLANYQSVGEYLYKEYFQKFKENSKE
+FTNVRNKKESYERCIAQSFLKDELKLIFKKQREFGFSFSKKFEEEVLSVAFYKRALKDFS
+HLVGNCSFFTDEKRAPKNSPLAFMFVALTRIINLLNNLKNTEGILYTKDDLNALLNEVLK
+NGTLTYKQTKKLLGLSDDYEFKGEKGTYFIEFKKYKEFIKALGEHNLSQDDLNEIAKDIT
+LIKDEIKLKKALAKYDLNQNQIDSLSKLEFKDHLNISFKALKLVTPLMLEGKKYDEACNE
+LNLKVAINEDKKDFLPAFNETYYKDEVTNPVVLRAIKEYRKVLNALLKKYGKVHKINIEL
+AREVGKNHSQRAKIEKEQNENYKAKKDAELECEKLGLKINSKNILKLRLFKEQKEFCAYS
+GEKIKISDLQDEKMLEIDHIYPYSRSFDDSYMNKVLVFTKQNQEKLNQTPFEAFGNDSAK
+WQKIEVLAKNLPTKKQKRILDKNYKDKEQKNFKDRNLNDTRYIARLVLNYTKDYLDFLPL
+SDDENTKLNDTQKGSKVHVEAKSGMLTSALRHTWGFSAKDRNNHLHHAIDAVIIAYANNS
+IVKAFSDFKKEQESNSAELYAKKISELDYKNKRKFFEPFSGFRQKVLDKIDEIFVSKPER
+KKPSGALHEETFRKEEEFYQSYGGKEGVLKALELGKIRKVNGKIVKNGDMFRVDIFKHKK
+TNKFYAVPIYTMDFALKVLPNKAVARSKKGEIKDWILMDENYEFCFSLYKDSLILIQTKD
+MQEPEFVYYNAFTSSTVSLIVSKHDNKFETLSKNQKILFKNANEKEVIAKSIGIQNLKVF
+EKYIVSALGEVTKAEFRQREDFKK
+>spcas9 len=1368
+MDKKYSIGLDIGTNSVGWAVITDEYKVPSKKFKVLGNTDRHSIKKNLIGALLFDSGETAE
+ATRLKRTARRRYTRRKNRICYLQEIFSNEMAKVDDSFFHRLEESFLVEEDKKHERHPIFG
+NIVDEVAYHEKYPTIYHLRKKLVDSTDKADLRLIYLALAHMIKFRGHFLIEGDLNPDNSD
+VDKLFIQLVQTYNQLFEENPINASGVDAKAILSARLSKSRRLENLIAQLPGEKKNGLFGN
+LIALSLGLTPNFKSNFDLAEDAKLQLSKDTYDDDLDNLLAQIGDQYADLFLAAKNLSDAI
+LLSDILRVNTEITKAPLSASMIKRYDEHHQDLTLLKALVRQQLPEKYKEIFFDQSKNGYA
+GYIDGGASQEEFYKFIKPILEKMDGTEELLVKLNREDLLRKQRTFDNGSIPHQIHLGELH
+AILRRQEDFYPFLKDNREKIEKILTFRIPYYVGPLARGNSRFAWMTRKSEETITPWNFEE
+VVDKGASAQSFIERMTNFDKNLPNEKVLPKHSLLYEYFTVYNELTKVKYVTEGMRKPAFL
+SGEQKKAIVDLLFKTNRKVTVKQLKEDYFKKIECFDSVEISGVEDRFNASLGTYHDLLKI
+IKDKDFLDNEENEDILEDIVLTLTLFEDREMIEERLKTYAHLFDDKVMKQLKRRRYTGWG
+RLSRKLINGIRDKQSGKTILDFLKSDGFANRNFMQLIHDDSLTFKEDIQKAQVSGQGDSL
+HEHIANLAGSPAIKKGILQTVKVVDELVKVMGRHKPENIVIEMARENQTTQKGQKNSRER
+MKRIEEGIKELGSQILKEHPVENTQLQNEKLYLYYLQNGRDMYVDQELDINRLSDYDVDH
+IVPQSFLKDDSIDNKVLTRSDKNRGKSDNVPSEEVVKKMKNYWRQLLNAKLITQRKFDNL
+TKAERGGLSELDKAGFIKRQLVETRQITKHVAQILDSRMNTKYDENDKLIREVKVITLKS
+KLVSDFRKDFQFYKVREINNYHHAHDAYLNAVVGTALIKKYPKLESEFVYGDYKVYDVRK
+MIAKSEQEIGKATAKYFFYSNIMNFFKTEITLANGEIRKRPLIETNGETGEIVWDKGRDF
+ATVRKVLSMPQVNIVKKTEVQTGGFSKESILPKRNSDKLIARKKDWDPKKYGGFDSPTVA
+YSVLVVAKVEKGKSKKLKSVKELLGITIMERSSFEKNPIDFLEAKGYKEVKKDLIIKLPK
+YSLFELENGRKRMLASAGELQKGNELALPSKYVNFLYLASHYEKLKGSPEDNEQKQLFVE
+QHKHYLDEIIEQISEFSKRVILADANLDKVLSAYNKHRDKPIREQAENIIHLFTLTNLGA
+PAAFKYFDTTIDRKRYTSTKEVLDATLIHQSITGLYETRIDLSQLGGD
+>st3cas9 len=1388
+MTKPYSIGLDIGTNSVGWAVTTDNYKVPSKKMKVLGNTSKKYIKKNLLGVLLFDSGITAE
+GRRLKRTARRRYTRRRNRILYLQEIFSTEMATLDDAFFQRLDDSFLVPDDKRDSKYPIFG
+NLVEEKAYHDEFPTIYHLRKYLADSTKKADLRLVYLALAHMIKYRGHFLIEGEFNSKNND
+IQKNFQDFLDTYNAIFESDLSLENSKQLEEIVKDKISKLEKKDRILKLFPGEKNSGIFSE
+FLKLIVGNQADFRKCFNLDEKASLHFSKESYDEDLETLLGYIGDDYSDVFLKAKKLYDAI
+LLSGFLTVTDNETEAPLSSAMIKRYNEHKEDLALLKEYIRNISLKTYNEVFKDDTKNGYA
+GYIDGKTNQEDFYVYLKKLLAEFEGADYFLEKIDREDFLRKQRTFDNGSIPYQIHLQEMR
+AILDKQAKFYPFLAKNKERIEKILTFRIPYYVGPLARGNSDFAWSIRKRNEKITPWNFED
+VIDKESSAEAFINRMTSFDLYLPEEKVLPKHSLLYETFNVYNELTKVRFIAESMRDYQFL
+DSKQKKDIVRLYFKDKRKVTDKDIIEYLHAIYGYDGIELKGIEKQFNSSLSTYHDLLNII
+NDKEFLDDSSNEAIIEEIIHTLTIFEDREMIKQRLSKFENIFDKSVLKKLSRRHYTGWGK
+LSAKLINGIRDEKSGNTILDYLIDDGISNRNFMQLIHDDALSFKKKIQKAQIIGDEDKGN
+IKEVVKSLPGSPAIKKGILQSIKIVDELVKVMGGRKPESIVVEMARENQYTNQGKSNSQQ
+RLKRLEKSLKELGSKILKENIPAKLSKIDNNALQNDRLYLYYLQNGKDMYTGDDLDIDRL
+SNYDIDHIIPQAFLKDNSIDNKVLVSSASNRGKSDDVPSLEVVKKRKTFWYQLLKSKLIS
+QRKFDNLTKAERGGLSPEDKAGFIQRQLVETRQITKHVARLLDEKFNNKKDENNRAVRTV
+KIITLKSTLVSQFRKDFELYKVREINDFHHAHDAYLNAVVASALLKKYPKLEPEFVYGDY
+PKYNSFRERKSATEKVYFYSNIMNIFKKSISLADGRVIERPLIEVNEETGESVWNKESDL
+ATVRRVLSYPQVNVVKKVEEQNHGLDRGKPKGLFNANLSSKPKPNSNENLVGAKEYLDPK
+KYGGYAGISNSFTVLVKGTIEKGAKKKITNVLEFQGISILDRINYRKDKLNFLLEKGYKD
+IELIIELPKYSLFELSDGSRRMLASILSTNNKRGEIHKGNQIFLSQKFVKLLYHAKRISN
+TINENHRKYVENHKKEFEELFYYILEFNENYVGAKKNGKLLNSAFQSWQNHSIDELCSSF
+IGPTGSERKGLFELTSRGSAADFEFLGVKIPRYRDYTPSSLLKDATLIHQSVTGLYETRI
+DLAKLGEG
+>geocas9 len=1087
+MRYKIGLDIGITSVGWAVMNLDIPRIEDLGVRIFDRAENPQTGESLALPRRLARSARRRL
+RRRKHRLERIRRLVIREGILTKEELDKLFEEKHEIDVWQLRVEALDRKLNNDELARVLLH
+LAKRRGFKSNRKSERSNKENSTMLKHIEENRAILSSYRTVGEMIVKDPKFALHKRNKGEN
+YTNTIARDDLEREIRLIFSKQREFGNMSCTEEFENEYITIWASQRPVASKDDIEKKVGFC
+TFEPKEKRAPKATYTFQSFIAWEHINKLRLISPSGARGLTDEERRLLYEQAFQKNKITYH
+DIRTLLHLPDDTYFKGIVYDRGESRKQNENIRFLELDAYHQIRKAVDKVYGKGKSSSFLP
+IDFDTFGYALTLFKDDADIHSYLRNEYEQNGKRMPNLANKVYDNELIEELLNLSFTKFGH
+LSLKALRSILPYMEQGEVYSSACERAGYTFTGPKKKQKTMLLPNIPPIANPVVMRALTQA
+RKVVNAIIKKYGSPVSIHIELARDLSQTFDERRKTKKEQDENRKKNETAIRQLMEYGLTL
+NPTGHDIVKFKLWSEQNGRCAYSLQPIEIERLLEPGYVEVDHVIPYSRSLDDSYTNKVLV
+LTRENREKGNRIPAEYLGVGTERWQQFETFVLTNKQFSKKKRDRLLRLHYDENEETEFKN
+RNLNDTRYISRFFANFIREHLKFAESDDKQKVYTVNGRVTAHLRSRWEFNKNREESDLHH
+AVDAVIVACTTPSDIAKVTAFYQRREQNKELAKKTEPHFPQPWPHFADELRARLSKHPKE
+SIKALNLGNYDDQKLESLQPVFVSRMPKRSVTGAAHQETLRRYVGIDERSGKIQTVVKTK
+LSEIKLDASGHFPMYGKESDPRTYEAIRQRLLEHNNDPKKAFQEPLYKPKKNGEPGPVIR
+TVKIIDTKNQVIPLNDGKTVAYNSNIVRVDVFEKDGKYYCVPVYTMDIMKGILPNKAIEP
+NKPYSEWKEMTEDYTFRFSLYPNDLIRIELPREKTVKTAAGEEINVKDVFVYYKTIDSAN
+GGLELISHDHRFSLRGVGSRTLKRFEKYQVDVLGNIYKVRGEKRVGLASSAHSKPGKTIR
+PLQSTRD
+>nmcas9 len=1082
+MAAFKPNPINYILGLDIGIASVGWAMVEIDEEENPIRLIDLGVRVFERAEVPKTGDSLAM
+VRRLARSVRRLTRRRAHRLLRARRLLKREGVLQAADFDENGLIKSLPNTPWQLRAAALDR
+KLTPLEWSAVLLHLIKHRGYLSQRKNEGETADKELGALLKGVADNAHALQTGDFRTPAEL
+ALNKFEKESGHIRNQRGDYSHTFSRKDLQAELILLFEKQKEFGNPHISGGLKEGIETLLM
+TQRPALSGDAVQKMLGHCTFEPAEPKAAKNTYTAERFIWLTKLNNLRILEQGSERPLTDT
+ERATLMDEPYRKSKLTYAQARKLLGLEDTAFFKGLRYGKDNAEASTLMEMKAYHAISRAL
+EKEGLKDKKSPLNLSPELQDEIGTAFSLFKTDEDITGRLKDRIQPEILEALLKHISFDKF
+VQISLKALRRIVPLMEQGKRYDEACAEIYGDHYGKKNTEEKIYLPPIPADEIRNPVVLRA
+LSQARKVINGVVRRYGSPARIHIETAREVGKSFKDRKEIEKRQEENRKDREKAAAKFREY
+FPNFVGEPKSKDILKLRLYEQQHGKCLYSGKEINLGRLNEKGYVEIDHALPFSRTWDDSF
+NNKVLVLGSENQNKGNQTPYEYFNGKDNSREWQEFKARVETSRFPRSKKQRILLQKFDED
+GFKERNLNDTRYVNRFLCQFVADRMRLTGKGKKRVFASNGQITNLLRGFWGLRKVRAEND
+RHHALDAVVVACSTVAMQQKITRFVRYKEMNAFDGKTIDKETGEVLHQKTHFPQPWEFFA
+QEVMIRVFGKPDGKPEFEEADTPEKLRTLLAEKLSSRPEAVHEYVTPLFVSRAPNRKMSG
+QGHMETVKSAKRLDEGVSVLRVPLTQLKLKDLEKMVNREREPKLYEALKARLEAHKDDPA
+KAFAEPFYKYDKAGNRTQQVKAVRVEQVQKTGVWVRNHNGIADNATMVRVDVFEKGDKYY
+LVPIYSWQVAKGILPDRAVVQGKDEEDWQLIDDSFNFKFSLHPNDLVEVITKKARMFGYF
+ASCHRGTGNINIRIHDLDHKIGKNGILEGIGVKTALSFQKYQIDELGKEIRPCRLKKRPP
+VR
+>st1cas9 len=1121
+MSDLVLGLDIGIGSVGVGILNKVTGEIIHKNSRIFPAAQAENNLVRRTNRQGRRLARRKK
+HRRVRLNRLFEESGLITDFTKISINLNPYQLRVKGLTDELSNEELFIALKNMVKHRGISY
+LDDASDDGNSSVGDYAQIVKENSKQLETKTPGQIQLERYQTYGQLRGDFTVEKDGKKHRL
+INVFPTSAYRSEALRILQTQQEFNPQITDEFINRYLEILTGKRKYYHGPGNEKSRTDYGR
+YRTSGETLDNIFGILIGKCTFYPDEFRAAKASYTAQEFNLLNDLNNLTVPTETKKLSKEQ
+KNQIINYVKNEKAMGPAKLFKYIAKLLSCDVADIKGYRIDKSGKAEIHTFEAYRKMKTLE
+TLDIEQMDRETLDKLAYVLTLNTEREGIQEALEHEFADGSFSQKQVDELVQFRKANSSIF
+GKGWHNFSVKLMMELIPELYETSEEQMTILTRLGKQKTTSSSNKTKYIDEKLLTEEIYNP
+VVAKSVRQAIKIVNAAIKEYGDFDNIVIEMARETNEDDEKKAIQKIQKANKDEKDAAMLK
+AANQYNGKAELPHSVFHGHKQLATKIRLWHQQGERCLYTGKTISIHDLINNSNQFEVDHI
+LPLSITFDDSLANKVLVYATANQEKGQRTPYQALDSMDDAWSFRELKAFVRESKTLSNKK
+KEYLLTEEDISKFDVRKKFIERNLVDTRYASRVVLNALQEHFRAHKIDTKVSVVRGQFTS
+QLRRHWGIEKTRDTYHHHAVDALIIAASSQLNLWKKQKNTLVSYSEDQLLDIETGELISD
+DEYKESVFKAPYQHFVDTLKSKEFEDSILFSYQVDSKFNRKISDATIYATRQAKVGKDKA
+DETYVLGKIKDIYTQDGYDAFMKIYKKDKSKFLMYRHDPQTFEKVIEPILENYPNKQINE
+KGKEVPCNPFLKYKEEHGYIRKYSKKGNGPEIKSLKYYDSKLGNHIDITPKDSNNKVVLQ
+SVSPWRADVYFNKTTGKYEILGLKYADLQFEKGTGTYKISQEKYNDIKKKEGVDSDSEFK
+FTLYKNDLLLVKDTETKEQQLFRFLSRTMPKQKHYVELKPYDKQKFEGGEALIKVLGNVA
+NSGQCKKGLGKSNISIYKVRTDVLGNQHIIKNEGDKPKLDF
+>iSpyMac len=1360
+MDKKYSIGLDIGTNSVGWAVITDEYKVPSKKFKVLGNTDRHSIKKNLIGALLFDSGETAE
+ATRLKRTARRRYTRRKNRICYLQEIFSNEMAKVDDSFFHRLEESFLVEEDKKHERHPIFG
+NIVDEVAYHEKYPTIYHLRKKLVDSTDKADLRLIYLALAHMIKFRGHFLIEGDLNPDNSD
+VDKLFIQLVQTYNQLFEENPINASGVDAKAILSARLSKSRKLENLIAQLPGEKKNGLFGN
+LIALSLGLTPNFKSNFDLAEDAKLQLSKDTYDDDLDNLLAQIGDQYADLFLAAKNLSDAI
+LLSDILRVNTEITKAPLSASMIKRYDEHHQDLTLLKALVRQQLPEKYKEIFFDQSKNGYA
+GYIDGGASQEEFYKFIKPILEKMDGTEELLVKLKREDLLRKQRTFDNGSIPHQIHLGELH
+AILRRQEDFYPFLKDNREKIEKILTFRIPYYVGPLARGNSRFAWMTRKSEETITPWNFEE
+VVDKGASAQSFIERMTNFDKNLPNEKVLPKHSLLYEYFTVYNELTKVKYVTEGMRKPAFL
+SGEQKKAIVDLLFKTNRKVTVKQLKEDYFKKIECFDSVEISGVEDRFNASLGTYHDLLKI
+IKDKDFLDNEENEDILEDIVLTLTLFEDREMIEERLKTYAHLFDDKVMKQLKRRRYTGWG
+RLSRKLINGIRDKQSGKTILDFLKSDGFANANFMQLIHDDSLTFKEDIQKAQVSGQGDSL
+HEHIANLAGSPAIKKGILQTVKVVDELVKVMGRHKPENIVIEMARENQTTQKGQKNSRER
+MKRIEEGIKELGSQILKEHPVENTQLQNEKLYLYYLQNGRDMYVDQELDINRLSDYDVDH
+IVPQSFLKDDSIDNKVLTRSDKNRGKSDNVPSEEVVKKMKNYWRQLLNAKLITQRKFDNL
+TKAERGGLSELDKAGFIKRQLVETRQITKHVAQILDSRMNTKYDENDKLIREVKVITLKS
+KLVSDFRKDFQFYKVREINNYHHAHDAYLNAVVGTALIKKYPKLESEFVYGDYKVYDVRK
+MIAKSEQEIGKATAKYFFYSNIMNFFKTEITLANGEIRKRPLIETNGETGEIVWDKGRDF
+ATVRKVLSMPQVNIVKKTEIQTVGQNGGLFDDNPKSPLEVTPSKLVPLKKELNPKKYGGY
+QKPTTAYPVLLITDTKQLIPISVMNKKQFEQNPVKFLRDRGYQQVGKNDFIKLPKYTLVD
+IGDGIKRLWASSKEIHKGNQLVVSKKSQILLYHAHHLDSDLSNDYLQNHNQQFDVLFNEI
+ISFSKKCKLGKEHIQKIENVYSNKKNSASIEELAESFIKLLGFTQLGATSPFNFLGVKLN
+QKQYKGKKDYILPCTEGTLIRQSITGLYETRVDLSKIGED
+>Scplusplus len=1375
+MEKKYSIGLDIGTNSVGWAVITDDYKVPSKKFKVLGNTNRKSIKKNLMGALLFDSGETAE
+ATRLKRTARRRYTRRKNRIRYLQEIFANEMAKLDDSFFQRLEESFLVEEDKKNERHPIFG
+NLADEVAYHRNYPTIYHLRKKLADSPEKADLRLIYLALAHIIKFRGHFLIEGKLNAENSD
+VAKLFYQLIQTYNQLFEESPLDEIEVDAKGILSARLSKSKRLEKLIAVFPNEKKNGLFGN
+IIALALGLTPNFKSNFDLTEDAKLQLSKDTYDDDLDELLGQIGDQYADLFSAAKNLSDAI
+LLSDILRSNSEVTKAPLSASMVKRYDEHHQDLALLKTLVRQQFPEKYAEIFKDDTKNGYA
+GYVGADKKLRKRSGKLATEEEFYKFIKPILEKMDGAEELLAKLNRDDLLRKQRTFDNGSI
+PHQIHLKELHAILRRQEEFYPFLKENREKIEKILTFRIPYYVGPLARGNSRFAWLTRKSE
+EAITPWNFEEVVDKGASAQSFIERMTNFDEQLPNKKVLPKHSLLYEYFTVYNELTKVKYV
+TERMRKPEFLSGEQKKAIVDLLFKTNRKVTVKQLKEDYFKKIECFDSVEIIGVEDRFNAS
+LGTYHDLLKIIKDKDFLDNEENEDILEDIVLTLTLFEDREMIEERLKTYAHLFDDKVMKQ
+LKRRHYTGWGRLSRKMINGIRDKQSGKTILDFLKSDGFSNRNFMQLIHDDSLTFKEEIEK
+AQVSGQGDSLHEQIADLAGSPAIKKGILQTVKIVDELVKVMGHKPENIVIEMARENQTTT
+KGLQQSRERKKRIEEGIKELESQILKENPVENTQLQNEKLYLYYLQNGRDMYVDQELDIN
+RLSDYDVDHIVPQSFIKDDSIDNKVLTRSVENRGKSDNVPSEEVVKKMKNYWRQLLNAKL
+ITQRKFDNLTKAERGGLSEADKAGFIKRQLVETRQITKHVARILDSRMNTKRDKNDKPIR
+EVKVITLKSKLVSDFRKDFQLYKVRDINNYHHAHDAYLNAVVGTALIKKYPKLESEFVYG
+DYKVYDVRKMIAKSEQEIGKATAKRFFYSNIMNFFKTEVKLANGEIRKRPLIETNGETGE
+VVWNKEKDFATVRKVLAMPQVNIVKKTEVQTGGFSKESILSKRESAKLIPRKKGWDTRKY
+GGFGSPTVAYSILVVAKVEKGKAKKLKSVKVLVGITIMEKGSYEKDPIGFLEAKGYKDIK
+KELIFKLPKYSLFELENGRRRMLASAKELQKANELVLPQHLVRLLYYTQNISATTGSNNL
+GYIEQHREEFKEIFEKIIDFSEKYILKNKVNSNLKSSFDEQFAVSDSILLSNSFVSLLKY
+TSFGASGGFTFLDLDVKQGRLRYQTVTEVLDATLIYQSITGLYETRTDLSQLGGD
+>SpRYc len=1377
+MEKKYSIGLDIGTNSVGWAVITDDYKVPSKKFKVLGNTNRKSIKKNLMGALLFDSGETAE
+ATRLKRTARRRYTRRKNRIRYLQEIFANEMAKLDDSFFQRLEESFLVEEDKKNERHPIFG
+NLADEVAYHRNYPTIYHLRKKLADSPEKADLRLIYLALAHIIKFRGHFLIEGKLNAENSD
+VAKLFYQLIQTYNQLFEESPLDEIEVDAKGILSARLSKSKRLEKLIAVFPNEKKNGLFGN
+IIALALGLTPNFKSNFDLTEDAKLQLSKDTYDDDLDELLGQIGDQYADLFSAAKNLSDAI
+LLSDILRSNSEVTKAPLSASMVKRYDEHHQDLALLKTLVRQQFPEKYAEIFKDDTKNGYA
+GYVGADKKLRKRSGKLATEEEFYKFIKPILEKMDGAEELLAKLNRDDLLRKQRTFDNGSI
+PHQIHLKELHAILRRQEEFYPFLKENREKIEKILTFRIPYYVGPLARGNSRFAWLTRKSE
+EAITPWNFEEVVDKGASAQSFIERMTNFDEQLPNKKVLPKHSLLYEYFTVYNELTKVKYV
+TERMRKPEFLSGEQKKAIVDLLFKTNRKVTVKQLKEDYFKKIECFDSVEIIGVEDRFNAS
+LGTYHDLLKIIKDKDFLDNEENEDILEDIVLTLTLFEDREMIEERLKTYAHLFDDKVMKQ
+LKRRHYTGWGRLSRKMINGIRDKQSGKTILDFLKSDGFSNRNFMQLIHDDSLTFKEEIEK
+AQVSGQGDSLHEQIADLAGSPAIKKGILQTVKIVDELVKVMGHKPENIVIEMARENQTTT
+KGLQQSRERKKRIEEGIKELESQILKENPVENTQLQNEKLYLYYLQNGRDMYVDQELDIN
+RLSDYDVDHIVPQSFIKDDSIDNKVLTRSVENRGKSDNVPSEEVVKKMKNYWRQLLNAKL
+ITQRKFDNLTKAERGGLSEADKAGFIKRQLVETRQITKHVARILDSRMNTKRDKNDKPIR
+EVKVITLKSKLVSDFRKDFQLYKVRDINNYHHAHDAYLNAVVGTALIKKYPKLESEFVYG
+DYKVYDVRKMIAKSEQEIGKATAKRFFYSNIMNFFKTEVKLANGEIRKRPLIETNGETGE
+VVWNKEKDFATVRKVLAMPQVNIVKKTEVQTGGFSKESIRPKRNSDKLIARKKDWDPKKY
+GGFLWPTVAYSVLVVAKVEKGKSKKLKSVKELLGITIMERSSFEKNPIDFLEAKGYKEVK
+KDLIIKLPKYSLFELENGRKRMLASAKQLQKGNELALPSKYVNFLYLASHYEKLKGSPED
+NEQKQLFVEQHKHYLDEIIEQISEFSKRVILADANLDKVLSAYNKHRDKPIREQAENIIH
+LFTLTRLGAPRAFKYFDTTIDPKQYRSTKEVLDATLIHQSITGLYETRIDLSQLGGD

cas9/data/tokenizer.py

@@ -0,0 +1,42 @@
+# 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.
+# This implementation is adapted from https://github.com/louaaron/Score-Entropy-Discrete-Diffusion
+# which is released under MIT license
+
+import re
+
+
+def wt_detokenizer(string):
+    # contractions
+    string = string.replace("s '", "s'")
+    string = re.sub(r"/' [0-9]/", r"/'[0-9]/", string)
+    # number separators
+    string = string.replace(" @-@ ", "-")
+    string = string.replace(" @,@ ", ",")
+    string = string.replace(" @.@ ", ".")
+    # punctuation
+    string = string.replace(" : ", ": ")
+    string = string.replace(" ; ", "; ")
+    string = string.replace(" . ", ". ")
+    string = string.replace(" ! ", "! ")
+    string = string.replace(" ? ", "? ")
+    string = string.replace(" , ", ", ")
+    # double brackets
+    string = re.sub(r"\(\s*([^\)]*?)\s*\)", r"(\1)", string)
+    string = re.sub(r"\[\s*([^\]]*?)\s*\]", r"[\1]", string)
+    string = re.sub(r"{\s*([^}]*?)\s*}", r"{\1}", string)
+    string = re.sub(r"\"\s*([^\"]*?)\s*\"", r'"\1"', string)
+    string = re.sub(r"'\s*([^']*?)\s*'", r"'\1'", string)
+    # miscellaneous
+    string = string.replace("= = = =", "====")
+    string = string.replace("= = =", "===")
+    string = string.replace("= =", "==")
+    string = string.replace(" " + chr(176) + " ", chr(176))
+    string = string.replace(" \n", "\n")
+    string = string.replace("\n ", "\n")
+    string = string.replace(" N ", " 1 ")
+    string = string.replace(" 's", "'s")
+    return string

cas9/data/uniref_pre_batching.py

@@ -0,0 +1,220 @@
+import os
+import random
+import pandas as pd
+from datasets import Dataset, DatasetDict
+
+# import selfies as sf
+# from rdkit import Chem
+# from rdkit import RDLogger
+# RDLogger.DisableLog('rdApp.*')
+
+import sys
+
+from transformers import EsmTokenizer
+
+# ---------------------------------------------------------------------
+# CONFIG
+# ---------------------------------------------------------------------
+
+MAX_TOKENS_PER_BATCH = 4096
+SOURCE_DATASET = 'uniprot_sprot'
+TRAIN_RATIO, VAL_RATIO, TEST_RATIO = 0.8, 0.1, 0.1
+MIN_LENGTH = 800    
+MAX_LENGTH = 1600
+MAX_DATA_SIZE = None
+NUM_REPEATS = None
+
+# Set PROJECT_ROOT or edit paths below to your data locations
+PROJECT_ROOT = os.environ.get('PROJECT_ROOT', 'path/to/project')
+if SOURCE_DATASET == 'uniprot_sprot':
+    FASTA_PATH = os.path.join(PROJECT_ROOT, 'data/uniref/uniref50_filtered_matched_30k.fasta')
+elif SOURCE_DATASET == 'crispr_cas_atlas':
+    FASTA_PATH = os.path.join(PROJECT_ROOT, f'data/crispr_cas_atlas/cas9_sequences_reduced_30k.fasta')
+elif SOURCE_DATASET == 'gfp':
+    FASTA_PATH = os.path.join(PROJECT_ROOT, 'data/gfp/fpbase_valid_gfps.fasta')
+else:
+    raise ValueError(f"Invalid source dataset: {SOURCE_DATASET}")
+
+if MAX_DATA_SIZE is not None:
+    if MAX_DATA_SIZE // 1000 > 0:
+        if SOURCE_DATASET == 'uniprot_sprot':
+            OUTPUT_DIR = os.path.join(PROJECT_ROOT, f'data/uniref/uniref_dataset_esm2_tokenized_mt{MAX_TOKENS_PER_BATCH}_leq{MAX_LENGTH}_n{MAX_DATA_SIZE//1000}k_30k')
+        elif SOURCE_DATASET == 'crispr_cas_atlas':
+            OUTPUT_DIR = os.path.join(PROJECT_ROOT, f'data/crispr_cas_atlas/dataset/{SOURCE_DATASET}_dataset_esm2_tokenized_mt{MAX_TOKENS_PER_BATCH}_leq{MAX_LENGTH}_n{MAX_DATA_SIZE//1000}k')
+        else:
+            raise ValueError(f"Invalid source dataset: {SOURCE_DATASET}")
+    else:
+        if SOURCE_DATASET == 'uniprot_sprot':
+            OUTPUT_DIR = os.path.join(PROJECT_ROOT, f'data/uniref/uniref_dataset_esm2_tokenized_mt{MAX_TOKENS_PER_BATCH}_leq{MAX_LENGTH}_n{MAX_DATA_SIZE}_30k')
+        elif SOURCE_DATASET == 'crispr_cas_atlas':
+            OUTPUT_DIR = os.path.join(PROJECT_ROOT, f'data/crispr_cas_atlas/dataset/{SOURCE_DATASET}_dataset_esm2_tokenized_mt{MAX_TOKENS_PER_BATCH}_leq{MAX_LENGTH}_n{MAX_DATA_SIZE}')
+        else:
+            raise ValueError(f"Invalid source dataset: {SOURCE_DATASET}")
+else:
+    if SOURCE_DATASET == 'uniprot_sprot':
+        OUTPUT_DIR = os.path.join(PROJECT_ROOT, f'data/uniref/uniref_dataset_esm2_tokenized_mt{MAX_TOKENS_PER_BATCH}_leq{MAX_LENGTH}_30k')
+    elif SOURCE_DATASET == 'crispr_cas_atlas':
+        OUTPUT_DIR = os.path.join(PROJECT_ROOT, f'data/crispr_cas_atlas/dataset/{SOURCE_DATASET}_dataset_esm2_tokenized_mt{MAX_TOKENS_PER_BATCH}_leq{MAX_LENGTH}')
+    else:
+        raise ValueError(f"Invalid source dataset: {SOURCE_DATASET}")
+# ------------------------------------------------------------
+# 1) Load sequences from FASTA file
+# ------------------------------------------------------------
+from Bio import SeqIO
+
+sequences = []
+print("Loading sequences from FASTA and validating…")
+for record in SeqIO.parse(FASTA_PATH, "fasta"):
+    s = str(record.seq)
+    # skip empty sequences
+    if not s:
+        continue
+    sequences.append(s)
+
+# remove all duplicates 
+sequences = list(set(sequences))
+
+## filter out any sequence that contains non-natural amino acids
+print(f"len(sequences) before filtering: {len(sequences)}")
+sequences = [
+    seq for seq in sequences
+    if all(aa in "ACDEFGHIKLMNPQRSTVWY" for aa in seq)
+    and (MAX_LENGTH is None or len(seq) <= MAX_LENGTH)
+    and (MIN_LENGTH is None or len(seq) >= MIN_LENGTH)
+]
+print(f"len(sequences) after filtering: {len(sequences)}")
+print(f"Average length of sequences: {sum(len(seq) for seq in sequences) / len(sequences)}")
+
+# ---- Initialize ESM tokenizer ----
+tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
+print(f"Vocabulary Size: {tokenizer.vocab_size}")
+
+# shuffle before splitting
+random.shuffle(sequences)
+
+if MAX_DATA_SIZE is not None:
+    print(f"Truncating to {MAX_DATA_SIZE} sequences")
+    sequences = sequences[:MAX_DATA_SIZE]
+    print(f"len(sequences) after truncation: {len(sequences)}")
+
+n = len(sequences)
+n_train = int(n * TRAIN_RATIO)
+n_val = int(n * VAL_RATIO)
+
+train_sequences = sequences[:n_train]
+val_sequences   = sequences[n_train:n_train + n_val]
+test_sequences  = sequences[n_train + n_val:]
+
+print(f"len(train_sequences) before repetition: {len(train_sequences)}")
+print(f"len(val_sequences) before repetition: {len(val_sequences)}")
+print(f"len(test_sequences) before repetition: {len(test_sequences)}")
+
+if NUM_REPEATS is not None:
+    train_sequences = train_sequences * NUM_REPEATS
+    val_sequences = val_sequences * NUM_REPEATS
+    test_sequences = test_sequences * NUM_REPEATS
+
+print(f"len(train_sequences) after repetition: {len(train_sequences)}")
+print(f"len(val_sequences) after repetition: {len(val_sequences)}")
+print(f"len(test_sequences) after repetition: {len(test_sequences)}")
+# ------------------------------------------------------------
+# 2) helper: tokenize sequences using the ESM tokenizer
+# ------------------------------------------------------------
+def tokenize_sequences(sequences):
+    enc = {"input_ids": [], "attention_mask": []}
+    for s in sequences:
+        res = tokenizer(s, add_special_tokens=True)
+        enc["input_ids"].append(res["input_ids"])
+        enc["attention_mask"].append(res["attention_mask"])
+    return enc
+
+
+# ------------------------------------------------------------
+# 3) build batched dataset (token-based batching: groups sequences by length, packs up to max_tokens)
+# ------------------------------------------------------------
+from tqdm import tqdm
+
+def build_batched_dataset(sequences, max_tokens=8192):
+    # tokenize all first
+    toks = tokenize_sequences(sequences)
+    input_ids_list = toks["input_ids"]
+    attn_mask_list = toks["attention_mask"]
+
+    # add lengths
+    items = []
+    for ids, mask in tqdm(zip(input_ids_list, attn_mask_list), total=len(input_ids_list), desc="Collecting items with length"):
+        items.append({
+            "input_ids": ids,
+            "attention_mask": mask,
+            "length": len(ids),
+        })
+
+    # sort by length
+    items.sort(key=lambda x: x["length"])
+
+    batched_input_ids = []
+    batched_attention_masks = []
+    batched_lengths = []
+    batched_batch_sizes = []
+
+    i = 0
+    n = len(items)
+    while i < n:
+        L = items[i]["length"]
+        # how many of length L can we pack?
+        max_bs = max_tokens // L
+        if max_bs < 1:
+            max_bs = 1
+
+        # collect up to max_bs items with same length
+        cur_ids = []
+        cur_masks = []
+        taken = 0
+        j = i
+        while j < n and items[j]["length"] == L and taken < max_bs:
+            cur_ids.append(items[j]["input_ids"])         # length L
+            cur_masks.append(items[j]["attention_mask"])  # length L
+            j += 1
+            taken += 1
+
+        # now this is one batch: shape (B, L)
+        batched_input_ids.append(cur_ids)
+        batched_attention_masks.append(cur_masks)
+        batched_lengths.append(L)
+        batched_batch_sizes.append(taken)
+
+        i = j
+
+    ds = Dataset.from_dict({
+        "input_ids": batched_input_ids,           # list of (B, L)
+        "attention_mask": batched_attention_masks,
+        "seq_length": batched_lengths,
+        "batch_size": batched_batch_sizes,
+    })
+    return ds
+
+
+# ------------------------------------------------------------
+# 4) build batched datasets for each split
+# ------------------------------------------------------------
+train_ds = build_batched_dataset(train_sequences, max_tokens=MAX_TOKENS_PER_BATCH)
+val_ds   = build_batched_dataset(val_sequences,   max_tokens=MAX_TOKENS_PER_BATCH)
+test_ds  = build_batched_dataset(test_sequences,  max_tokens=MAX_TOKENS_PER_BATCH)
+
+for ds in [train_ds, val_ds, test_ds]:
+    num_seqs = 0
+    for item in ds:
+        num_seqs += len(item["input_ids"])
+    print(f"Number of sequences in {ds}: {num_seqs}")
+
+dsdict = DatasetDict(
+    {
+        "train": train_ds,
+        "validation": val_ds,
+        "test": test_ds,
+    }
+)
+
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+dsdict.save_to_disk(OUTPUT_DIR)
+print(f"saved to {OUTPUT_DIR}")

cas9/data/utils.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.
+# This implementation is adapted from https://github.com/pytorch/data/blob/main/torchdata/stateful_dataloader/sampler.py#L132
+# which is released under BSD-3 license
+
+import itertools
+from typing import Any, Dict, Optional
+
+import numpy as np
+import torch
+from torch import Tensor
+from torch.utils.data import DataLoader, Dataset, Sampler
+
+
+def cycle_loader(dataloader: DataLoader, sampler: Sampler = None) -> Tensor:
+    while 1:
+        if sampler is not None:
+            sampler.set_epoch(np.random.randint(0, 100000))
+        for data in dataloader:
+            yield data
+
+
+class StatefulDistributedSampler(torch.utils.data.distributed.DistributedSampler):
+    """
+    From: https://github.com/pytorch/data/blob/main/torchdata/stateful_dataloader/sampler.py#L132
+    """
+
+    _YIELDED = "yielded"
+
+    def __init__(
+        self,
+        dataset: Dataset,
+        num_replicas: Optional[int] = None,
+        rank: Optional[int] = None,
+        shuffle: bool = True,
+        seed: int = 0,
+        drop_last: bool = False,
+    ) -> None:
+        super().__init__(dataset, num_replicas, rank, shuffle, seed, drop_last)
+        self.yielded = 0
+        self.next_yielded = None
+
+    def __iter__(self):
+        self.yielded = 0
+        if self.next_yielded is not None:
+            self.yielded = self.next_yielded
+            self.next_yielded = None
+        it = super().__iter__()
+        for idx in itertools.islice(it, self.yielded, None):
+            self.yielded += 1
+            yield idx
+
+    def state_dict(self) -> Dict[str, Any]:
+        return {self._YIELDED: self.yielded}
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        if self._YIELDED not in state_dict:
+            raise ValueError("Invalid state_dict")
+        if state_dict[self._YIELDED] < 0:
+            raise ValueError("Cannot load state_dict with negative yielded value")
+        self.next_yielded = state_dict[self._YIELDED]

cas9/domain_detector.py

@@ -0,0 +1,226 @@
+# domain_detector.py
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+import tempfile
+from dataclasses import dataclass
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple
+
+
+@dataclass(frozen=True)
+class HitCriteria:
+    ievalue_max: float
+    min_dom_len: int
+
+
+@dataclass(frozen=True)
+class DomainDetectorConfig:
+    hmm_path: str
+    cpu: int = 8
+
+    # Stage1 acceptance
+    stage1: HitCriteria = HitCriteria(ievalue_max=1e-2, min_dom_len=35)
+
+    # Stage2 acceptance (recall-first)
+    stage2: HitCriteria = HitCriteria(ievalue_max=10.0, min_dom_len=20)
+
+    # Stage2 hmmscan reporting knobs (big so rescue sees weak hits)
+    stage2_report_E: float = 1e6
+    stage2_report_domE: float = 1e6
+    stage2_incE: float = 1e6
+    stage2_incdomE: float = 1e6
+
+    # Optional Stage1 sensitivity knobs (usually leave None)
+    stage1_report_E: Optional[float] = None
+    stage1_report_domE: Optional[float] = None
+    stage1_F1: Optional[float] = None
+    stage1_F2: Optional[float] = None
+    stage1_F3: Optional[float] = None
+
+    # Model name prefixes in your HMM library
+    hnh_prefix: str = "HNH_"
+    ruvc_prefix: str = "RuvC_"
+
+
+class DomainDetector:
+    """
+    Cas9 nuclease domain presence detector (HNH + RuvC), Option-B (recall-first).
+
+    Uses hmmscan:
+      - Stage1: fast scan over batch
+      - Stage2: --max scan only on Stage1 misses (rescue)
+    """
+
+    def __init__(self, config: DomainDetectorConfig):
+        self.cfg = config
+        self._hmmscan = self._check_exe("hmmscan")
+        self._validate_hmm(self.cfg.hmm_path)
+
+    @staticmethod
+    def _check_exe(name: str) -> str:
+        exe = shutil.which(name)
+        if not exe:
+            raise RuntimeError(f"Required executable not found on PATH: {name}")
+        return exe
+
+    @staticmethod
+    def _validate_hmm(hmm_path: str) -> None:
+        if not os.path.exists(hmm_path):
+            raise FileNotFoundError(f"HMM file not found: {hmm_path}")
+        if os.path.getsize(hmm_path) == 0:
+            raise ValueError(f"HMM file is empty: {hmm_path}")
+
+    @staticmethod
+    def _write_fasta_from_pairs(pairs: Sequence[Tuple[str, str]], path: str) -> None:
+        with open(path, "w") as out:
+            for sid, seq in pairs:
+                out.write(f">{sid}\n")
+                seq = seq.strip().replace(" ", "").replace("\n", "")
+                for i in range(0, len(seq), 80):
+                    out.write(seq[i:i+80] + "\n")
+
+    def _run_hmmscan(
+        self,
+        fasta_path: str,
+        domtblout_path: str,
+        *,
+        max_mode: bool,
+        report_E: Optional[float] = None,
+        report_domE: Optional[float] = None,
+        incE: Optional[float] = None,
+        incdomE: Optional[float] = None,
+        F1: Optional[float] = None,
+        F2: Optional[float] = None,
+        F3: Optional[float] = None,
+    ) -> None:
+        cmd: List[str] = [
+            self._hmmscan,
+            "--cpu", str(self.cfg.cpu),
+            "--domtblout", domtblout_path,
+        ]
+        if max_mode:
+            cmd.append("--max")
+
+        if F1 is not None:
+            cmd += ["--F1", str(F1)]
+        if F2 is not None:
+            cmd += ["--F2", str(F2)]
+        if F3 is not None:
+            cmd += ["--F3", str(F3)]
+
+        if report_E is not None:
+            cmd += ["-E", str(report_E)]
+        if report_domE is not None:
+            cmd += ["--domE", str(report_domE)]
+        if incE is not None:
+            cmd += ["--incE", str(incE)]
+        if incdomE is not None:
+            cmd += ["--incdomE", str(incdomE)]
+
+        cmd += [self.cfg.hmm_path, fasta_path]
+
+        # discard stdout; errors go to CalledProcessError
+        with open(os.devnull, "w") as devnull:
+            subprocess.run(cmd, check=True, stdout=devnull, stderr=subprocess.DEVNULL)
+
+    def _parse_domtblout_presence(self, domtblout_path: str, crit: HitCriteria) -> Tuple[Dict[str, bool], Dict[str, bool]]:
+        """
+        Returns (has_hnh, has_ruvc) dicts keyed by query ID.
+        domtblout columns (1-based):
+          1=model, 4=query, 13=i-Evalue, 20=env_from, 21=env_to
+        """
+        has_hnh: Dict[str, bool] = {}
+        has_ruvc: Dict[str, bool] = {}
+
+        with open(domtblout_path, "r") as f:
+            for line in f:
+                if not line.strip() or line.startswith("#"):
+                    continue
+                parts = line.split()
+                if len(parts) < 21:
+                    continue
+                model = parts[0]
+                qid = parts[3]
+                try:
+                    ie = float(parts[12])
+                    env_from = int(parts[19])
+                    env_to = int(parts[20])
+                except ValueError:
+                    continue
+                dom_len = env_to - env_from + 1
+                if ie > crit.ievalue_max or dom_len < crit.min_dom_len:
+                    continue
+
+                if model.startswith(self.cfg.hnh_prefix):
+                    has_hnh[qid] = True
+                elif model.startswith(self.cfg.ruvc_prefix):
+                    has_ruvc[qid] = True
+
+        return has_hnh, has_ruvc
+
+    def predict_batch(
+        self,
+        seqs: Sequence[str],
+        ids: Optional[Sequence[str]] = None,
+    ) -> List[bool]:
+        """
+        Classify sequences as pass/fail (domains present: HNH AND RuvC).
+        - seqs: list of amino acid sequences
+        - ids: optional list of IDs (same length). If omitted, uses seq0..seqN.
+
+        Returns: List[bool] aligned with input order.
+        """
+        if ids is None:
+            ids = [f"seq{i}" for i in range(len(seqs))]
+        if len(ids) != len(seqs):
+            raise ValueError("ids and seqs must have same length")
+
+        pairs = list(zip(ids, seqs))
+
+        with tempfile.TemporaryDirectory() as td:
+            fa1 = os.path.join(td, "stage1.faa")
+            dom1 = os.path.join(td, "stage1.domtblout")
+
+            self._write_fasta_from_pairs(pairs, fa1)
+            self._run_hmmscan(
+                fasta_path=fa1,
+                domtblout_path=dom1,
+                max_mode=False,
+                report_E=self.cfg.stage1_report_E,
+                report_domE=self.cfg.stage1_report_domE,
+                F1=self.cfg.stage1_F1,
+                F2=self.cfg.stage1_F2,
+                F3=self.cfg.stage1_F3,
+            )
+
+            h1, r1 = self._parse_domtblout_presence(dom1, self.cfg.stage1)
+
+            stage1_pass = {sid for sid in ids if h1.get(sid, False) and r1.get(sid, False)}
+            miss = [sid for sid in ids if sid not in stage1_pass]
+
+            stage2_pass: set[str] = set()
+            if miss:
+                miss_pairs = [(sid, seqs[ids.index(sid)]) for sid in miss]  # stable small sizes
+                fa2 = os.path.join(td, "stage2.faa")
+                dom2 = os.path.join(td, "stage2.domtblout")
+
+                self._write_fasta_from_pairs(miss_pairs, fa2)
+                self._run_hmmscan(
+                    fasta_path=fa2,
+                    domtblout_path=dom2,
+                    max_mode=True,
+                    report_E=self.cfg.stage2_report_E,
+                    report_domE=self.cfg.stage2_report_domE,
+                    incE=self.cfg.stage2_incE,
+                    incdomE=self.cfg.stage2_incdomE,
+                )
+                h2, r2 = self._parse_domtblout_presence(dom2, self.cfg.stage2)
+                stage2_pass = {sid for sid in miss if h2.get(sid, False) and r2.get(sid, False)}
+
+            final_pass = stage1_pass | stage2_pass
+            return [sid in final_pass for sid in ids]
+
+    def predict_one(self, seq: str, id_: str = "seq0") -> bool:
+        return self.predict_batch([seq], ids=[id_])[0]

cas9/evaluate_val.py

@@ -0,0 +1,1066 @@
+#!/usr/bin/env python3
+"""
+Script to evaluate Editflows models by generating sequences and evaluating them.
+1. Decodes test set into actual string sequences
+2. Calculates diversity, sequence loss (val_unweighted_total_loss), and optionally Cas9 scores for decoded sequences
+3. Uses multi-step generation (generate_from_x0_multi_edit) to sample new sequences
+4. Evaluates generated sequences on diversity, sequence loss, and optionally Cas9 scores
+5. Calculates and plots pLDDT scores for both decoded and generated sequences (as separate subplots)
+
+If Cas9 classifier is not provided, only diversity and sequence loss metrics will be calculated.
+"""
+
+import os
+import argparse
+import torch
+from datasets import load_from_disk
+import yaml
+from easydict import EasyDict as edict
+from tqdm import tqdm
+import numpy as np
+import random
+from collections import Counter
+import matplotlib.pyplot as plt
+from transformers import AutoTokenizer, EsmForProteinFolding
+
+from cas9.generate import build_model_and_stuff, tokenize_input_str, detokenize_output
+from cas9.model.utils import generate_from_x0_multi_edit
+from cas9.objectives import Cas9Classification
+
+# Try to import rapidfuzz for Levenshtein diversity (optional)
+HAS_RAPIDFUZZ = False
+HAS_LEVENSHTEIN = False
+try:
+    from rapidfuzz.distance import Levenshtein as RapidLevenshtein
+    HAS_RAPIDFUZZ = True
+except ImportError:
+    try:
+        from Levenshtein import distance as levenshtein_distance_fallback
+        HAS_LEVENSHTEIN = True
+    except ImportError:
+        print("Note: rapidfuzz and python-Levenshtein not available. Levenshtein diversity will be skipped.")
+
+
+def decode_validation_set(val_dataset, tokenizer, pad_id, bos_id, eos_id, num_samples=1000):
+    """
+    Decode validation dataset into actual string sequences.
+    
+    Returns:
+        Tuple of (all_sequences, sampled_sequences)
+    """
+    all_sequences = []
+    
+    print(f"Decoding sequences from validation dataset...")
+    for batch_item in tqdm(val_dataset, desc="Decoding"):
+        input_ids_batch = batch_item["input_ids"]
+        
+        # Convert to tensor if needed
+        if isinstance(input_ids_batch[0], list):
+            batch_tensor = torch.tensor(input_ids_batch, dtype=torch.long)
+        elif isinstance(input_ids_batch, torch.Tensor):
+            batch_tensor = input_ids_batch
+        else:
+            batch_tensor = torch.tensor(input_ids_batch, dtype=torch.long)
+        
+        # Decode each sequence to string
+        for seq_tensor in batch_tensor:
+            # Find valid length (non-padding tokens)
+            seq_list = seq_tensor.tolist()
+            valid_length = 0
+            for i, tok in enumerate(seq_list):
+                if tok == pad_id:
+                    break
+                valid_length = i + 1
+            
+            seq_list = seq_list[:valid_length]
+            
+            # Decode to string - tokenizer will handle BOS/EOS automatically
+            if len(seq_list) > 0:
+                try:
+                    # For ESM tokenizer (protein), use batch_decode
+                    if hasattr(tokenizer, 'batch_decode'):
+                        seq_str = tokenizer.batch_decode([seq_list], skip_special_tokens=True)[0]
+                    else:
+                        seq_str = tokenizer.decode(seq_list, skip_special_tokens=True)
+                    
+                    # Remove spaces (ESM tokenizer adds spaces between tokens)
+                    seq_str = seq_str.replace(" ", "")
+                    
+                    # Verify we got a valid sequence string
+                    if seq_str and len(seq_str) > 0:
+                        # Check if it contains valid amino acids
+                        valid_chars = sum(1 for c in seq_str if c in "ACDEFGHIKLMNPQRSTVWY")
+                        if valid_chars > len(seq_str) * 0.9:  # At least 90% valid amino acids
+                            all_sequences.append(seq_str)
+                except Exception as e:
+                    pass
+    
+    # Sample random subset if needed
+    if len(all_sequences) > num_samples:
+        sampled = random.sample(all_sequences, num_samples)
+        print(f"Sampled {num_samples} sequences from {len(all_sequences)} total")
+    else:
+        sampled = all_sequences
+        print(f"Using all {len(all_sequences)} sequences from validation set")
+    
+    # Debug: Print first few sequences to verify they're decoded correctly
+    if len(sampled) > 0:
+        print(f"\nSample of decoded sequences (first 3):")
+        for i, seq in enumerate(sampled[:3]):
+            print(f"  Sequence {i+1}: Length={len(seq)}, First 50 chars: {seq[:50]}")
+    
+    return all_sequences, sampled
+
+
+def calculate_plddt_from_sequence_string(sequence_string, esmfold_tokenizer, esm_model, device):
+    """
+    Calculate pLDDT score for a sequence string using ESMFold.
+    Based on generate_and_analyze_plddt.py
+    """
+    try:
+        tok = esmfold_tokenizer([sequence_string], return_tensors="pt", add_special_tokens=False).to(device)
+        with torch.no_grad():
+            out = esm_model(**tok)
+            plddt = out.plddt.mean(-1).mean(-1)  # Average across both confidence and sequence length
+            # Handle scalar or tensor output
+            if plddt.dim() > 0:
+                plddt = plddt[0]  # Take first element if batch dimension exists
+            return plddt.cpu().item()
+    except Exception as e:
+        print(f"Error calculating pLDDT for sequence (length {len(sequence_string)}): {e}")
+        return None
+
+
+def calculate_plddt_scores(sequences, esmfold_tokenizer, esm_model, device, batch_size=1):
+    """
+    Calculate pLDDT scores for a list of sequences.
+    
+    Args:
+        sequences: List of sequence strings
+        esmfold_tokenizer: ESMFold tokenizer
+        esm_model: ESMFold model
+        device: torch device
+        batch_size: Batch size for processing (default 1 for ESMFold)
+    
+    Returns:
+        List of pLDDT scores (None for failed calculations)
+    """
+    plddt_scores = []
+    
+    for i in tqdm(range(0, len(sequences), batch_size), desc="Computing pLDDT scores"):
+        batch_seqs = sequences[i:i+batch_size]
+        for seq in batch_seqs:
+            score = calculate_plddt_from_sequence_string(seq, esmfold_tokenizer, esm_model, device)
+            plddt_scores.append(score)
+    
+    return plddt_scores
+
+
+def plot_plddt_histogram(decoded_plddt, generated_plddt, output_dir, num_bins=50, dataset_type="test"):
+    """
+    Plot histogram comparing pLDDT scores between decoded and generated sequences.
+    Uses two subplots (one above the other) instead of overlapping distributions.
+    
+    Args:
+        decoded_plddt: List of pLDDT scores for decoded sequences
+        generated_plddt: List of pLDDT scores for generated sequences
+        output_dir: Directory to save the plot
+        num_bins: Number of bins for histogram
+        dataset_type: Type of dataset ("test", "validation", etc.) for labeling
+    """
+    # Filter out None values
+    decoded_plddt_valid = [s for s in decoded_plddt if s is not None]
+    generated_plddt_valid = [g for g in generated_plddt if g is not None]
+    
+    if len(decoded_plddt_valid) == 0 or len(generated_plddt_valid) == 0:
+        print("Warning: No valid pLDDT scores to plot. Skipping histogram.")
+        return
+    
+    # Create output directory if it doesn't exist
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Professional color scheme - clear, distinct colors
+    color_decoded = '#2E86AB'  # Professional blue
+    color_generated = '#E63946'  # Clear red/coral
+    
+    # Determine bin edges based on all scores
+    all_scores = decoded_plddt_valid + generated_plddt_valid
+    min_score = min(all_scores)
+    max_score = max(all_scores)
+    bin_edges = np.linspace(min_score, max_score, num_bins + 1)
+    
+    # Create figure with two subplots stacked vertically
+    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8), sharex=True)
+    
+    # Top subplot: Decoded sequences (true data)
+    ax1.hist(decoded_plddt_valid, bins=bin_edges, alpha=0.7, color=color_decoded, 
+             density=True, edgecolor=color_decoded, linewidth=1.2)
+    ax1.set_ylabel('Density', fontsize=14, fontweight='medium')
+    ax1.set_title(f'Decoded sequences ({dataset_type} set)', fontsize=13, fontweight='medium', pad=10)
+    ax1.grid(True, alpha=0.2, linestyle='--', linewidth=0.5)
+    ax1.spines['top'].set_visible(False)
+    ax1.spines['right'].set_visible(False)
+    ax1.spines['left'].set_linewidth(0.8)
+    ax1.spines['bottom'].set_linewidth(0.8)
+    ax1.tick_params(axis='both', which='major', labelsize=12, length=4, width=0.8)
+    
+    # Bottom subplot: Generated sequences
+    ax2.hist(generated_plddt_valid, bins=bin_edges, alpha=0.7, color=color_generated, 
+             density=True, edgecolor=color_generated, linewidth=1.2)
+    ax2.set_xlabel('pLDDT Score', fontsize=14, fontweight='medium')
+    ax2.set_ylabel('Density', fontsize=14, fontweight='medium')
+    ax2.set_title('Generated sequences', fontsize=13, fontweight='medium', pad=10)
+    ax2.grid(True, alpha=0.2, linestyle='--', linewidth=0.5)
+    ax2.spines['top'].set_visible(False)
+    ax2.spines['right'].set_visible(False)
+    ax2.spines['left'].set_linewidth(0.8)
+    ax2.spines['bottom'].set_linewidth(0.8)
+    ax2.tick_params(axis='both', which='major', labelsize=12, length=4, width=0.8)
+    
+    # Adjust spacing between subplots
+    plt.tight_layout()
+    
+    # Save plot
+    filename = os.path.join(output_dir, 'plddt_histogram_comparison.png')
+    plt.savefig(filename, dpi=300, bbox_inches='tight', facecolor='white')
+    print(f"\nSaved pLDDT histogram to {filename}")
+    plt.close()
+    
+
+
+def calculate_sequence_loss(editflow, sequences, tokenizer, pad_id, bos_id, eos_id, device, cfg, batch_size=32):
+    """
+    Calculate sequence loss (val_unweighted_total_loss) for sequences.
+    Uses the same loss calculation as validation_step in base_models.py.
+    
+    Args:
+        editflow: EditFlow LightningModule
+        sequences: List of sequence strings
+        tokenizer: Tokenizer
+        pad_id, bos_id, eos_id: Special token IDs
+        device: torch device
+        cfg: Config object
+        batch_size: Batch size for processing
+    
+    Returns:
+        List of loss values (one per sequence)
+    """
+    editflow.eval()
+    all_losses = []
+    
+    print(f"Calculating sequence loss for {len(sequences)} sequences...")
+    
+    # Process in batches
+    for i in tqdm(range(0, len(sequences), batch_size), desc="Computing losses"):
+        batch_seqs = sequences[i:i+batch_size]
+        
+        # Tokenize batch
+        x1_batch = []
+        for seq_str in batch_seqs:
+            x1 = tokenize_input_str(seq_str, cfg, tokenizer, bos_id, eos_id, pad_id, device)
+            x1_batch.append(x1.squeeze(0))
+        
+        # Pad to same length
+        max_len = max(x.shape[0] for x in x1_batch)
+        padded_batch = []
+        for x in x1_batch:
+            padding = torch.full((max_len - x.shape[0],), pad_id, dtype=torch.long, device=device)
+            padded_batch.append(torch.cat([x, padding]))
+        x1_tensor = torch.stack(padded_batch).to(device)
+        
+        # Calculate loss for this batch
+        with torch.no_grad():
+            # Call preparation to get necessary tensors
+            if editflow.reparameterize:
+                lam_total, logits_type, logits_ins, logits_sub, z_t, z_1, x_t, mask, weight, M_t = editflow.preparation(x1_tensor)
+                
+                if editflow.loc_prop_path:
+                    loss, loss_components = editflow.loss_fn.reparameterized_forward_localized(
+                        lam_total, logits_type, logits_ins, logits_sub, 
+                        z_t, z_1, x_t, mask, weight, M_t, editflow.lam_prop, 
+                        editflow.eps_id, editflow.bos_id, editflow.eos_id, 
+                        editflow.gamma_rate, editflow.gamma_edit, 
+                        editflow.use_aux_ce, editflow.aux_ce_weight
+                    )
+                else:
+                    loss, loss_components = editflow.loss_fn.reparameterized_forward(
+                        lam_total, logits_type, logits_ins, logits_sub,
+                        z_t, z_1, x_t, mask, weight, editflow.eps_id, 
+                        editflow.bos_id, editflow.eos_id, editflow.gamma_rate, 
+                        editflow.gamma_edit, editflow.use_aux_ce, editflow.aux_ce_weight
+                    )
+            else:
+                lam_ins, logits_ins, lam_del, lam_sub, logits_sub, z_t, z_1, x_t, mask, weight, M_t = editflow.preparation(x1_tensor)
+                
+                if editflow.loc_prop_path:
+                    loss, loss_components = editflow.loss_fn.forward_localized(
+                        lam_ins, logits_ins, lam_del, lam_sub, logits_sub, 
+                        z_t, z_1, x_t, mask, weight, M_t, editflow.lam_prop, 
+                        editflow.eps_id, editflow.bos_id, editflow.eos_id, 
+                        editflow.use_aux_ce, editflow.aux_ce_weight
+                    )
+                else:
+                    loss, loss_components = editflow.loss_fn.forward(
+                        lam_ins, logits_ins, lam_del, lam_sub, logits_sub, 
+                        z_t, z_1, x_t, mask, weight, editflow.eps_id, 
+                        editflow.bos_id, editflow.eos_id, 
+                        editflow.use_aux_ce, editflow.aux_ce_weight
+                    )
+            
+            # Extract unweighted total loss (matching val_unweighted_total_loss)
+            if "loss_total_unweighted" in loss_components:
+                # For reparameterized models
+                unweighted_loss = loss_components["loss_total_unweighted"]
+            elif "loss_base" in loss_components:
+                # For non-reparameterized models, loss_base is rate + edit (unweighted)
+                unweighted_loss = loss_components["loss_base"]
+            else:
+                # Fallback to total loss
+                unweighted_loss = loss
+            
+            # Get per-sequence losses (loss is averaged over batch, so we need to compute per-sequence)
+            # Since loss is batch-averaged, we'll use the batch loss for all sequences in the batch
+            # For more accurate per-sequence loss, we'd need to compute individually
+            batch_loss_value = unweighted_loss.item()
+            all_losses.extend([batch_loss_value] * len(batch_seqs))
+    
+    return all_losses
+
+
+def kgrams(s: str, k: int):
+    """Extract k-grams from a string."""
+    s = s.strip()
+    if len(s) < k:
+        return {s} if s else set()
+    return {s[i:i+k] for i in range(len(s) - k + 1)}
+
+
+def jaccard(a: set, b: set) -> float:
+    """Calculate Jaccard similarity between two sets."""
+    if not a and not b:
+        return 1.0
+    inter = len(a & b)
+    union = len(a | b)
+    return inter / union if union else 1.0
+
+
+def diversity_kmer_jaccard(seqs, k=3, pairs=50000, seed=0):
+    """
+    Diversity = 1 - average Jaccard similarity over random pairs.
+    Works for variable-length strings.
+    
+    Returns:
+        diversity (1 - avg_sim), avg_similarity
+    """
+    if len(seqs) < 2:
+        return 0.0, 1.0
+    
+    rng = random.Random(seed)
+    grams = [kgrams(s, k) for s in seqs]
+    n = len(seqs)
+    
+    # Limit pairs to avoid excessive computation
+    max_pairs = min(pairs, n * (n - 1) // 2)
+    if max_pairs == 0:
+        return 0.0, 1.0
+    
+    total_sim = 0.0
+    for _ in range(max_pairs):
+        i = rng.randrange(n)
+        j = rng.randrange(n - 1)
+        if j >= i:
+            j += 1
+        total_sim += jaccard(grams[i], grams[j])
+    
+    avg_sim = total_sim / max_pairs
+    return 1.0 - avg_sim, avg_sim
+
+
+def diversity_levenshtein(seqs, pairs=20000, seed=0):
+    """
+    Diversity = 1 - average normalized Levenshtein similarity over random pairs.
+    
+    Returns:
+        diversity (1 - avg_sim), avg_similarity
+    """
+    if len(seqs) < 2:
+        return 0.0, 1.0
+    
+    if not (HAS_RAPIDFUZZ or HAS_LEVENSHTEIN):
+        # No Levenshtein implementation available
+        return 0.0, 1.0
+    
+    rng = random.Random(seed)
+    n = len(seqs)
+    
+    # Limit pairs to avoid excessive computation
+    max_pairs = min(pairs, n * (n - 1) // 2)
+    if max_pairs == 0:
+        return 0.0, 1.0
+    
+    total_sim = 0.0
+    valid_pairs = 0
+    for _ in range(max_pairs):
+        i = rng.randrange(n)
+        j = rng.randrange(n - 1)
+        if j >= i:
+            j += 1
+        
+        if HAS_RAPIDFUZZ:
+            # Use rapidfuzz for fast normalized similarity
+            sim = RapidLevenshtein.normalized_similarity(seqs[i], seqs[j])
+        elif HAS_LEVENSHTEIN:
+            # Fallback: use python-Levenshtein distance and normalize
+            # Simple normalization: 1 - (distance / max_length)
+            dist = levenshtein_distance_fallback(seqs[i], seqs[j])
+            max_len = max(len(seqs[i]), len(seqs[j]))
+            sim = 1.0 - (dist / max_len) if max_len > 0 else 1.0
+        else:
+            # Should not reach here, but skip if somehow we do
+            continue
+        
+        total_sim += sim
+        valid_pairs += 1
+    
+    if valid_pairs == 0:
+        return 0.0, 1.0
+    
+    avg_sim = total_sim / valid_pairs
+    return 1.0 - avg_sim, avg_sim
+
+
+def calculate_diversity(sequences):
+    """
+    Calculate diversity metrics for a set of sequences.
+    Uses k-mer Jaccard diversity and optionally Levenshtein diversity.
+    
+    Args:
+        sequences: List of sequence strings
+    
+    Returns:
+        Dictionary with diversity metrics:
+        - unique_count: Number of unique sequences
+        - uniqueness_ratio: Fraction of unique sequences
+        - kmer_diversity: k-mer Jaccard diversity (1 - avg_similarity)
+        - kmer_avg_similarity: Average k-mer Jaccard similarity
+        - levenshtein_diversity: Levenshtein diversity (if available)
+        - levenshtein_avg_similarity: Average Levenshtein similarity (if available)
+    """
+    if len(sequences) == 0:
+        return {
+            'unique_count': 0,
+            'uniqueness_ratio': 0.0,
+            'kmer_diversity': 0.0,
+            'kmer_avg_similarity': 1.0,
+            'levenshtein_diversity': 0.0,
+            'levenshtein_avg_similarity': 1.0
+        }
+    
+    # Unique fraction
+    unique_count = len(set(sequences))
+    uniqueness_ratio = unique_count / len(sequences) if len(sequences) > 0 else 0.0
+    
+    # k-mer Jaccard diversity
+    kmer_div, kmer_sim = diversity_kmer_jaccard(sequences, k=3, pairs=50000, seed=0)
+    
+    # Levenshtein diversity (if available)
+    if HAS_RAPIDFUZZ or HAS_LEVENSHTEIN:
+        lev_div, lev_sim = diversity_levenshtein(sequences, pairs=20000, seed=0)
+    else:
+        lev_div, lev_sim = 0.0, 1.0
+    
+    return {
+        'unique_count': unique_count,
+        'uniqueness_ratio': uniqueness_ratio,
+        'kmer_diversity': kmer_div,
+        'kmer_avg_similarity': kmer_sim,
+        'levenshtein_diversity': lev_div,
+        'levenshtein_avg_similarity': lev_sim
+    }
+
+
+def evaluate_cas9_scores(sequences, cas9_classifier, threshold=0.5):
+    """
+    Evaluate Cas9 scores for sequences.
+    
+    Args:
+        sequences: List of sequence strings
+        cas9_classifier: Cas9Classification object
+        threshold: Score threshold for validity (default 0.5)
+    
+    Returns:
+        validity_rate, average_score, list of scores
+    """
+    if len(sequences) == 0:
+        return 0.0, 0.0, []
+    
+    # Get scores in batches
+    scores = []
+    batch_size = 32
+    
+    for i in tqdm(range(0, len(sequences), batch_size), desc="Computing Cas9 scores"):
+        batch_seqs = sequences[i:i+batch_size]
+        batch_scores = cas9_classifier.get_scores(batch_seqs)
+        scores.extend(batch_scores)
+    
+    scores = np.array(scores)
+    validity_rate = np.mean(scores > threshold)
+    avg_score = np.mean(scores)
+    
+    return validity_rate, avg_score, scores.tolist()
+
+
+def generate_sequences_multi_edit(model, source_dist, tokenizer, pad_id, bos_id, eos_id, eps_id,
+                                  input_sequences, device, cfg, num_steps=20, batch_size=32, num_generations_per_sequence=1):
+    """
+    Generate sequences from input sequences using multi-edit generation.
+    
+    Args:
+        model: Editflows model
+        input_sequences: List of input sequence strings
+        device: torch device
+        cfg: Config object
+        num_steps: Number of generation steps
+        batch_size: Batch size for generation
+        num_generations_per_sequence: Number of sequences to generate per input sequence
+    
+    Returns:
+        List of generated sequence strings
+    """
+    model.eval()
+    generated_sequences = []
+    
+    # Get allowed tokens
+    allowed_tokens = torch.tensor(
+        [tok for tok in source_dist._allowed_tokens if tok not in (eps_id,)],
+        device=device,
+        dtype=torch.long,
+    )
+    
+    print(f"Generating sequences using multi-edit generation (num_steps={num_steps}, batch_size={batch_size}, num_generations_per_sequence={num_generations_per_sequence})...")
+    print(f"Total sequences to generate: {len(input_sequences)} input sequences × {num_generations_per_sequence} = {len(input_sequences) * num_generations_per_sequence}")
+    
+    # Process in batches
+    for i in tqdm(range(0, len(input_sequences), batch_size), desc="Generating"):
+        batch_seqs = input_sequences[i:i+batch_size]
+        
+        # Generate num_generations_per_sequence for each sequence in the batch
+        for gen_idx in range(num_generations_per_sequence):
+            # Tokenize batch
+            x0_batch = []
+            for seq_str in batch_seqs:
+                x0 = tokenize_input_str(seq_str, cfg, tokenizer, bos_id, eos_id, pad_id, device)
+                x0_batch.append(x0.squeeze(0))
+            
+            # Pad to same length
+            max_len = max(x.shape[0] for x in x0_batch)
+            padded_batch = []
+            for x in x0_batch:
+                padding = torch.full((max_len - x.shape[0],), pad_id, dtype=torch.long, device=device)
+                padded_batch.append(torch.cat([x, padding]))
+            x0_tensor = torch.stack(padded_batch).to(device)
+            
+            # Generate using multi-edit generation
+            with torch.no_grad():
+                x_gen = generate_from_x0_multi_edit(
+                    model,
+                    x0_tensor,
+                    pad_id=pad_id,
+                    bos_id=bos_id,
+                    eos_id=eos_id,
+                    allowed_tokens=allowed_tokens,
+                    num_steps=num_steps,
+                    device=device,
+                )
+            
+            # Decode generated sequences
+            for j in range(x_gen.shape[0]):
+                gen_seq = detokenize_output(x_gen[j:j+1], cfg, tokenizer, bos_id, eos_id, pad_id)
+                generated_sequences.append(gen_seq)
+    
+    return generated_sequences
+
+
+def get_evaluate_val_argument_parser():
+    parser = argparse.ArgumentParser(description='Evaluate Editflows model by generating sequences')
+    parser.add_argument('--config', type=str, required=True,
+                        help='Path to config YAML file')
+    parser.add_argument('--ckpt', type=str, required=True,
+                        help='Path to model checkpoint (.ckpt file)')
+    parser.add_argument('--val_data', type=str, default=None,
+                        help='Path to validation dataset (output from uniref_pre_batching.py). Deprecated: use --test_data instead.')
+    parser.add_argument('--test_data', type=str, default=None,
+                        help='Path to test dataset (output from uniref_pre_batching.py)')
+    parser.add_argument('--cas9_classifier_ckpt', type=str, default=None,
+                        help='Path to Cas9 classifier checkpoint (optional, if not provided, validity rate will not be calculated)')
+    parser.add_argument('--cas9_classifier_config', type=str, default=None,
+                        help='Path to Cas9 classifier config (optional, required if --cas9_classifier_ckpt is provided)')
+    parser.add_argument('--num_samples', type=int, default=1000,
+                        help='Number of sequences to sample from test set (default: 1000)')
+    parser.add_argument('--num_generations_per_sequence', type=int, default=1,
+                        help='Number of sequences to generate per sampled sequence (default: 1). Total generated = num_samples × num_generations_per_sequence')
+    parser.add_argument('--num_steps', type=int, default=20,
+                        help='Number of generation steps (default: 20)')
+    parser.add_argument('--batch_size', type=int, default=32,
+                        help='Batch size for generation (default: 32)')
+    parser.add_argument('--validity_threshold', type=float, default=0.5,
+                        help='Cas9 score threshold for validity (default: 0.5)')
+    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu',
+                        help='Device to use (cuda/cpu)')
+    parser.add_argument('--seed', type=int, default=42,
+                        help='Random seed (default: 42)')
+    parser.add_argument('--output_dir', type=str, default='./evaluation_output',
+                        help='Directory to save pLDDT plots (default: ./evaluation_output)')
+    parser.add_argument('--output_fasta', type=str, default=None,
+                        help='Path to save generated sequences as FASTA file (default: {output_dir}/generated_sequences.fasta)')
+    parser.add_argument('--num_bins', type=int, default=50,
+                        help='Number of bins for pLDDT histogram (default: 50)')
+    parser.add_argument('--calculate_plddt', action='store_true',
+                        help='Enable pLDDT calculation and plotting (default: False, pLDDT calculation is disabled by default)')
+    parser.add_argument('--include_train', action='store_true',
+                        help='Also compute losses for the train set (default: False)')
+    return parser
+
+
+def _set_eval_run_seeds(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def run_evaluate_val(
+    args,
+    *,
+    run_seed=None,
+    save_fasta=True,
+    fasta_tag=None,
+):
+    """
+    Run the full evaluate_val pipeline and return metrics as a nested dict of JSON-serializable values.
+
+    Args:
+        args: Namespace from get_evaluate_val_argument_parser().parse_args() (or compatible).
+        run_seed: If int, seeds random/numpy/torch before stochastic steps. If None, behavior matches
+            legacy evaluate_val (no seeding).
+        save_fasta: If False, skip writing generated sequences to FASTA.
+        fasta_tag: If set (and save_fasta), write to output_dir/generated_sequences_{fasta_tag}.fasta
+            unless output_fasta is explicitly set on args.
+    """
+    if run_seed is not None:
+        _set_eval_run_seeds(int(run_seed))
+
+    device = torch.device(args.device)
+    
+    # Load config
+    with open(args.config, 'r') as f:
+        cfg = edict(yaml.safe_load(f))
+    
+    # Build model
+    print("Building model...")
+    editflow, source_dist, tokenizer, pad_id, bos_id, eos_id, eps_id = build_model_and_stuff(cfg, device)
+    
+    # Load checkpoint
+    print(f"Loading checkpoint from {args.ckpt}...")
+    ckpt = torch.load(args.ckpt, map_location=device, weights_only=False)
+    editflow.load_state_dict(ckpt["state_dict"], strict=False)
+    model = editflow.model.to(device)
+    model.eval()
+    
+    # Initialize Cas9 classifier (optional)
+    cas9_classifier = None
+    if args.cas9_classifier_ckpt is not None:
+        if args.cas9_classifier_config is None:
+            raise ValueError("--cas9_classifier_config is required when --cas9_classifier_ckpt is provided")
+        
+        print("Initializing Cas9 classifier...")
+        cas9_classifier = Cas9Classification(
+            device=device,
+            checkpoint_path=args.cas9_classifier_ckpt,
+            config_path=args.cas9_classifier_config,
+            shared_esm_model=model.esm_emb  # Share ESM model to save memory
+        )
+        
+        # Test classifier on a known Cas9 sequence to verify it works
+        test_cas9_seq = "MDKKYSIGLDIGTNSVGWAVITDEYKVPSKKFKVLGNTDRHSIKKNLIGALLFDSGETAEATRLKRTARRRYTRRKNRICYLQEIFSNEMAKVDDSFFHRLEESFLVEEDKKHERHPIFGNIVDEVAYHEKYPTIYHLRKKLVDSTDKADLRLIYLALAHMIKFRGHFLIEGDLNPDNSDVDKLFIQLVQTYNQLFEENPINASGVDAKAILSARLSKSRRLENLIAQLPGEKKNGLFGNLIALSLGLTPNFKSNFDLAEDAKLQLSKDTYDDDLDNLLAQIGDQYADLFLAAKNLSDAILLSDILRVNTEITKAPLSASMIKRYDEHHQDLTLLKALVRQQLPEKYKEIFFDQSKNGYAGYIDGGASQEEFYKFIKPILEKMDGTEELLVKLNREDLLRKQRTFDNGSIPHQIHLGELHAILRRQEDFYPFLKDNREKIEKILTFRIPYYVGPLARGNSRFAWMTRKSEETITPWNFEEVVDKGASAQSFIERMTNFDKNLPNEKVLPKHSLLYEYFTVYNELTKVKYVTEGMRKPAFLSGEQKKAIVDLLFKTNRKVTVKQLKEDYFKKIECFDSVEISGVEDRFNASLGTYHDLLKIIKDKDFLDNEENEDILEDIVLTLTLFEDREMIEERLKTYAHLFDDKVMKQLKRRRYTGWGRLSRKLINGIRDKQSGKTILDFLKSDGFANRNFMQLIHDDSLTFKEDIQKAQVSGQGDSLHEHIANLAGSPAIKKGILQTVKVVDELVKVMGRHKPENIVIEMARENQTTQKGQKNSRERMKRIEEGIKELGSQILKEHPVENTQLQNEKLYLYYLQNGRDMYVDQELDINRLSDYDVDHIVPQSFLKDDSIDNKVLTRSDKNRGKSDNVPSEEVVKKMKNYWRQLLNAKLITQRKFDNLTKAERGGLSELDKAGFIKRQLVETRQITKHVAQILDSRMNTKYDENDKLIREVKVITLKSKLVSDFRKDFQFYKVREINNYHHAHDAYLNAVVGTALIKKYPKLESEFVYGDYKVYDVRKMIAKSEQEIGKATAKYFFYSNIMNFFKTEITLANGEIRKRPLIETNGETGEIVWDKGRDFATVRKVLSMPQVNIVKKTEVQTGGFSKESILPKRNSDKLIARKKDWDPKKYGGFDSPTVAYSVLVVAKVEKGKSKKLKSVKELLGITIMERSSFEKNPIDFLEAKGYKEVKKDLIIKLPKYSLFELENGRKRMLASAGELQKGNELALPSKYVNFLYLASHYEKLKGSPEDNEQKQLFVEQHKHYLDEIIEQISEFSKRVILADANLDKVLSAYNKHRDKPIREQAENIIHLFTLTNLGAPAAFKYFDTTIDRKRYTSTKEVLDATLIHQSITGLYETRIDLSQLGGD"
+        test_score = cas9_classifier.get_scores([test_cas9_seq])[0]
+        print(f"Test Cas9 classifier on known Cas9 sequence: score = {test_score:.4f}")
+        if test_score < 0.5:
+            print(f"WARNING: Known Cas9 sequence got low score! This suggests a problem with the classifier or sequence format.")
+        else:
+            print(f"Classifier working correctly (score > 0.5 for known Cas9 sequence)")
+    else:
+        print("No Cas9 classifier provided. Validity rate will not be calculated.")
+    
+    # Determine which dataset path to use (test_data takes precedence)
+    if args.test_data is not None:
+        dataset_path = args.test_data
+        split_name = "test"
+        dataset_type = "test"
+    elif args.val_data is not None:
+        dataset_path = args.val_data
+        # Try "validation" or "val" first when --val_data is used
+        split_name = None
+        dataset_type = "validation"
+    else:
+        raise ValueError("Either --test_data or --val_data must be provided")
+    
+    # Load dataset
+    print(f"\nLoading {dataset_type} dataset from {dataset_path}...")
+    dataset_dict = load_from_disk(dataset_path)
+    
+    # Determine which split to use
+    if split_name is None:
+        # When --val_data is used, prioritize validation splits
+        if "validation" in dataset_dict:
+            test_dataset = dataset_dict["validation"]
+            split_name = "validation"
+        elif "val" in dataset_dict:
+            test_dataset = dataset_dict["val"]
+            split_name = "val"
+        elif "test" in dataset_dict:
+            # Fallback to test if validation splits not found
+            test_dataset = dataset_dict["test"]
+            split_name = "test"
+            dataset_type = "test"
+        else:
+            raise ValueError(f"Dataset must have 'validation', 'val', or 'test' split. Found splits: {list(dataset_dict.keys())}")
+    else:
+        if split_name in dataset_dict:
+            test_dataset = dataset_dict[split_name]
+        else:
+            raise ValueError(f"Dataset must have '{split_name}' split. Found splits: {list(dataset_dict.keys())}")
+    
+    # Load train dataset if --include_train flag is set
+    train_dataset = None
+    all_train_sequences = None
+    train_avg_loss = None
+    if args.include_train:
+        if "train" in dataset_dict:
+            train_dataset = dataset_dict["train"]
+            print(f"\nLoading train dataset from {dataset_path}...")
+            print(f"Train dataset loaded: {len(train_dataset)} items")
+        else:
+            print(f"Warning: --include_train flag set but 'train' split not found in dataset. Found splits: {list(dataset_dict.keys())}")
+            print("Skipping train set evaluation.")
+    
+    # Step 1: Decode test set into string sequences
+    print(f"\n{'='*70}")
+    print("STEP 1: Decoding test set")
+    print(f"{'='*70}")
+    all_decoded_sequences, sampled_decoded_sequences = decode_validation_set(
+        test_dataset,
+        tokenizer,
+        pad_id,
+        bos_id,
+        eos_id,
+        num_samples=args.num_samples
+    )
+    
+    # Step 2: Calculate metrics on decoded sequences
+    print(f"\n{'='*70}")
+    print("STEP 2: Evaluating decoded sequences")
+    print(f"{'='*70}")
+    
+    # Calculate sequence loss on FULL test set
+    print(f"\nCalculating sequence loss for FULL test set ({len(all_decoded_sequences)} sequences)...")
+    orig_losses_full = calculate_sequence_loss(
+        editflow, all_decoded_sequences, tokenizer, pad_id, bos_id, eos_id, device, cfg, batch_size=args.batch_size
+    )
+    orig_avg_loss_full = np.mean(orig_losses_full)
+    print(f"  Average sequence loss (full test set): {orig_avg_loss_full:.4f}")
+    
+    # Calculate other metrics on sampled sequences only
+    print(f"\nCalculating diversity for sampled decoded sequences ({len(sampled_decoded_sequences)} sequences)...")
+    orig_diversity = calculate_diversity(sampled_decoded_sequences)
+    
+    # Calculate sequence loss for sampled sequences (for comparison)
+    print(f"\nCalculating sequence loss for sampled decoded sequences ({len(sampled_decoded_sequences)} sequences)...")
+    orig_losses = calculate_sequence_loss(
+        editflow, sampled_decoded_sequences, tokenizer, pad_id, bos_id, eos_id, device, cfg, batch_size=args.batch_size
+    )
+    orig_avg_loss = np.mean(orig_losses)
+    print(f"  Average sequence loss (sampled): {orig_avg_loss:.4f}")
+    
+    # Calculate sequence loss for train set if --include_train flag is set
+    if args.include_train and train_dataset is not None:
+        print(f"\n{'='*70}")
+        print("TRAIN SET EVALUATION")
+        print(f"{'='*70}")
+        print(f"\nDecoding train set into string sequences...")
+        all_train_sequences, _ = decode_validation_set(
+            train_dataset,
+            tokenizer,
+            pad_id,
+            bos_id,
+            eos_id,
+            num_samples=10**9  # Decode all sequences, don't sample (use very large number)
+        )
+        print(f"\nCalculating sequence loss for FULL train set ({len(all_train_sequences)} sequences)...")
+        train_losses = calculate_sequence_loss(
+            editflow, all_train_sequences, tokenizer, pad_id, bos_id, eos_id, device, cfg, batch_size=args.batch_size
+        )
+        train_avg_loss = np.mean(train_losses)
+        print(f"  Average sequence loss (full train set): {train_avg_loss:.4f}")
+    
+    # Calculate Cas9 scores only if classifier is provided (on sampled sequences)
+    orig_validity_rate, orig_avg_score, orig_scores = None, None, None
+    if cas9_classifier is not None:
+        print(f"\nCalculating Cas9 scores for sampled decoded sequences...")
+        orig_validity_rate, orig_avg_score, orig_scores = evaluate_cas9_scores(
+            sampled_decoded_sequences, cas9_classifier, threshold=args.validity_threshold
+        )
+    
+    # Step 3: Generate new sequences using multi-edit generation
+    print(f"\n{'='*70}")
+    print("STEP 3: Generating new sequences using multi-edit generation")
+    print(f"{'='*70}")
+    generated_sequences = generate_sequences_multi_edit(
+        model, source_dist, tokenizer, pad_id, bos_id, eos_id, eps_id,
+        sampled_decoded_sequences, device, cfg, num_steps=args.num_steps, batch_size=args.batch_size,
+        num_generations_per_sequence=args.num_generations_per_sequence
+    )
+    print(f"Generated {len(generated_sequences)} sequences")
+    
+    # Save generated sequences as FASTA file
+    fasta_path_written = None
+    if save_fasta:
+        if args.output_fasta is None:
+            os.makedirs(args.output_dir, exist_ok=True)
+            if fasta_tag:
+                fasta_path = os.path.join(args.output_dir, f"generated_sequences_{fasta_tag}.fasta")
+            else:
+                fasta_path = os.path.join(args.output_dir, "generated_sequences.fasta")
+        else:
+            fasta_path = args.output_fasta
+            fasta_dir = os.path.dirname(fasta_path)
+            if fasta_dir:
+                os.makedirs(fasta_dir, exist_ok=True)
+
+        print(f"\nSaving generated sequences to FASTA file: {fasta_path}")
+        with open(fasta_path, 'w') as f:
+            for i, seq in enumerate(generated_sequences):
+                f.write(f">generated_sequence_{i+1}\n")
+                for j in range(0, len(seq), 80):
+                    f.write(seq[j:j+80] + "\n")
+        print(f"Saved {len(generated_sequences)} sequences to {fasta_path}")
+        fasta_path_written = fasta_path
+    
+    # Step 4: Calculate metrics on generated sequences
+    print(f"\n{'='*70}")
+    print("STEP 4: Evaluating generated sequences")
+    print(f"{'='*70}")
+    
+    print(f"\nCalculating diversity for generated sequences...")
+    gen_diversity = calculate_diversity(generated_sequences)
+    
+    # Calculate sequence loss for generated sequences
+    print(f"\nCalculating sequence loss for generated sequences...")
+    gen_losses = calculate_sequence_loss(
+        editflow, generated_sequences, tokenizer, pad_id, bos_id, eos_id, device, cfg, batch_size=args.batch_size
+    )
+    gen_avg_loss = np.mean(gen_losses)
+    print(f"  Average sequence loss: {gen_avg_loss:.4f}")
+    
+    # Calculate Cas9 scores only if classifier is provided
+    gen_validity_rate, gen_avg_score, gen_scores = None, None, None
+    if cas9_classifier is not None:
+        print(f"\nCalculating Cas9 scores for generated sequences...")
+        gen_validity_rate, gen_avg_score, gen_scores = evaluate_cas9_scores(
+            generated_sequences, cas9_classifier, threshold=args.validity_threshold
+        )
+    
+    def _to_float(x):
+        if x is None:
+            return None
+        return float(x)
+
+    def _diversity_plain(d):
+        return {k: int(v) if k == "unique_count" else float(v) for k, v in d.items()}
+
+    metrics = {
+        "run_seed": int(run_seed) if run_seed is not None else None,
+        "dataset_type": dataset_type,
+        "fasta_path": fasta_path_written,
+        "counts": {
+            "full_train": len(all_train_sequences) if all_train_sequences is not None else None,
+            "full_test": len(all_decoded_sequences),
+            "sampled_decoded": len(sampled_decoded_sequences),
+            "generated": len(generated_sequences),
+        },
+        "loss": {
+            "train_avg": _to_float(train_avg_loss),
+            "decoded_full_avg": _to_float(orig_avg_loss_full),
+            "decoded_sampled_avg": _to_float(orig_avg_loss),
+            "generated_avg": _to_float(gen_avg_loss),
+        },
+        "cas9": None,
+        "diversity_decoded": _diversity_plain(orig_diversity),
+        "diversity_generated": _diversity_plain(gen_diversity),
+        "plddt": None,
+    }
+    if cas9_classifier is not None:
+        metrics["cas9"] = {
+            "decoded_validity_rate": _to_float(orig_validity_rate),
+            "decoded_avg_score": _to_float(orig_avg_score),
+            "generated_validity_rate": _to_float(gen_validity_rate),
+            "generated_avg_score": _to_float(gen_avg_score),
+        }
+
+    # Step 5: Calculate and plot pLDDT scores (optional)
+    if args.calculate_plddt:
+        print(f"\n{'='*70}")
+        print("STEP 5: Calculating pLDDT scores")
+        print(f"{'='*70}")
+
+        print("\nClearing GPU memory...")
+        del model
+        del editflow
+        if cas9_classifier is not None:
+            del cas9_classifier
+        if device.type == 'cuda':
+            torch.cuda.empty_cache()
+            print("GPU memory cleared.")
+
+        print("Loading ESMFold model for pLDDT calculation...")
+        esmfold_tokenizer_path = "facebook/esmfold_v1"
+        esmfold_tokenizer = AutoTokenizer.from_pretrained(esmfold_tokenizer_path)
+        esmfold_device = device
+        esm_model = EsmForProteinFolding.from_pretrained(
+            esmfold_tokenizer_path,
+            torch_dtype=torch.bfloat16
+        ).to(esmfold_device).eval()
+        print("ESMFold model loaded successfully!")
+
+        print(f"\nCalculating pLDDT scores for sampled decoded sequences...")
+        decoded_plddt = calculate_plddt_scores(
+            sampled_decoded_sequences, esmfold_tokenizer, esm_model, esmfold_device, batch_size=1
+        )
+        decoded_plddt_valid = [s for s in decoded_plddt if s is not None]
+        if len(decoded_plddt_valid) > 0:
+            print(f"  Valid scores: {len(decoded_plddt_valid)}/{len(decoded_plddt)}")
+            print(f"  Mean pLDDT: {np.mean(decoded_plddt_valid):.2f}, Std: {np.std(decoded_plddt_valid):.2f}")
+
+        print(f"\nCalculating pLDDT scores for generated sequences...")
+        generated_plddt = calculate_plddt_scores(
+            generated_sequences, esmfold_tokenizer, esm_model, esmfold_device, batch_size=1
+        )
+        generated_plddt_valid = [g for g in generated_plddt if g is not None]
+        if len(generated_plddt_valid) > 0:
+            print(f"  Valid scores: {len(generated_plddt_valid)}/{len(generated_plddt)}")
+            print(f"  Mean pLDDT: {np.mean(generated_plddt_valid):.2f}, Std: {np.std(generated_plddt_valid):.2f}")
+
+        if len(decoded_plddt_valid) > 0 and len(generated_plddt_valid) > 0:
+            print(f"\nPlotting pLDDT histograms...")
+            plot_plddt_histogram(decoded_plddt_valid, generated_plddt_valid, args.output_dir, args.num_bins, dataset_type)
+        else:
+            print("Warning: Not enough valid pLDDT scores to plot.")
+
+        metrics["plddt"] = {
+            "decoded_mean": float(np.mean(decoded_plddt_valid)) if decoded_plddt_valid else None,
+            "decoded_std": float(np.std(decoded_plddt_valid)) if decoded_plddt_valid else None,
+            "decoded_valid_n": len(decoded_plddt_valid),
+            "decoded_total_n": len(decoded_plddt),
+            "generated_mean": float(np.mean(generated_plddt_valid)) if generated_plddt_valid else None,
+            "generated_std": float(np.std(generated_plddt_valid)) if generated_plddt_valid else None,
+            "generated_valid_n": len(generated_plddt_valid),
+            "generated_total_n": len(generated_plddt),
+        }
+
+    return metrics
+
+
+def print_evaluation_results(metrics, args):
+    """Print the summary block (same layout as historical evaluate_val.py)."""
+    dataset_type = metrics["dataset_type"]
+    counts = metrics["counts"]
+    loss_m = metrics["loss"]
+    orig_diversity = metrics["diversity_decoded"]
+    gen_diversity = metrics["diversity_generated"]
+
+    print("\n" + "="*70)
+    print("EVALUATION RESULTS")
+    print("="*70)
+
+    print(f"\nDataset Statistics:")
+    if counts["full_train"] is not None:
+        print(f"  Full train set sequences: {counts['full_train']}")
+    print(f"  Full test set sequences: {counts['full_test']}")
+    print(f"  Sampled decoded sequences: {counts['sampled_decoded']}")
+    print(f"  Generated sequences: {counts['generated']}")
+
+    print(f"\n" + "-"*70)
+    print("SEQUENCE LOSS (val_unweighted_total_loss)")
+    print("-"*70)
+    if loss_m["train_avg"] is not None:
+        print(f"Train set (FULL train set, {counts['full_train']} sequences):")
+        print(f"  Average loss: {loss_m['train_avg']:.4f}")
+    print(f"\nDecoded sequences (FULL {dataset_type} set, {counts['full_test']} sequences):")
+    print(f"  Average loss: {loss_m['decoded_full_avg']:.4f}")
+    print(f"\nDecoded sequences (sampled, {counts['sampled_decoded']} sequences):")
+    print(f"  Average loss: {loss_m['decoded_sampled_avg']:.4f}")
+    print(f"\nGenerated sequences ({counts['generated']} sequences):")
+    print(f"  Average loss: {loss_m['generated_avg']:.4f}")
+
+    cas9 = metrics["cas9"]
+    if cas9 is not None:
+        print(f"\n" + "-"*70)
+        print("CAS9 SCORES (validity threshold = {})".format(args.validity_threshold))
+        print("-"*70)
+        print(f"NOTE: If {dataset_type} set contains general proteins (not Cas9-specific),")
+        print("      low Cas9 scores are expected. Cas9 classifier is trained to identify Cas9 proteins.")
+        print(f"\nDecoded sequences (sampled from {dataset_type} set, {counts['sampled_decoded']} sequences):")
+        print(f"  Validity rate: {cas9['decoded_validity_rate']:.4f} ({cas9['decoded_validity_rate']*100:.2f}%)")
+        print(f"  Average Cas9 score: {cas9['decoded_avg_score']:.4f}")
+        print(f"\nGenerated sequences ({counts['generated']} sequences):")
+        print(f"  Validity rate: {cas9['generated_validity_rate']:.4f} ({cas9['generated_validity_rate']*100:.2f}%)")
+        print(f"  Average Cas9 score: {cas9['generated_avg_score']:.4f}")
+
+    print(f"\n" + "-"*70)
+    print("DIVERSITY METRICS")
+    print("-"*70)
+    print(f"Decoded sequences (sampled from {dataset_type} set, {counts['sampled_decoded']} sequences):")
+    print(f"  Unique sequences: {orig_diversity['unique_count']} / {counts['sampled_decoded']} ({orig_diversity['uniqueness_ratio']*100:.2f}%)")
+    print(f"  k-mer Jaccard diversity (k=3): {orig_diversity['kmer_diversity']:.4f}")
+    print(f"  k-mer Jaccard avg similarity: {orig_diversity['kmer_avg_similarity']:.4f}")
+    if HAS_RAPIDFUZZ or HAS_LEVENSHTEIN:
+        print(f"  Levenshtein diversity: {orig_diversity['levenshtein_diversity']:.4f}")
+        print(f"  Levenshtein avg similarity: {orig_diversity['levenshtein_avg_similarity']:.4f}")
+    print(f"\nGenerated sequences ({counts['generated']} sequences):")
+    print(f"  Unique sequences: {gen_diversity['unique_count']} / {counts['generated']} ({gen_diversity['uniqueness_ratio']*100:.2f}%)")
+    print(f"  k-mer Jaccard diversity (k=3): {gen_diversity['kmer_diversity']:.4f}")
+    print(f"  k-mer Jaccard avg similarity: {gen_diversity['kmer_avg_similarity']:.4f}")
+    if HAS_RAPIDFUZZ or HAS_LEVENSHTEIN:
+        print(f"  Levenshtein diversity: {gen_diversity['levenshtein_diversity']:.4f}")
+        print(f"  Levenshtein avg similarity: {gen_diversity['levenshtein_avg_similarity']:.4f}")
+
+    plddt = metrics.get("plddt")
+    if plddt is not None:
+        print(f"\n" + "-"*70)
+        print("PLDDT METRICS")
+        print("-"*70)
+        print(f"Decoded sequences (sampled from {dataset_type} set, {counts['sampled_decoded']} sequences):")
+        if plddt["decoded_mean"] is not None:
+            print(f"  Mean pLDDT: {plddt['decoded_mean']:.2f}")
+            print(f"  Std pLDDT: {plddt['decoded_std']:.2f}")
+        print(f"  Valid scores: {plddt['decoded_valid_n']}/{plddt['decoded_total_n']}")
+        print(f"\nGenerated sequences ({counts['generated']} sequences):")
+        if plddt["generated_mean"] is not None:
+            print(f"  Mean pLDDT: {plddt['generated_mean']:.2f}")
+            print(f"  Std pLDDT: {plddt['generated_std']:.2f}")
+        print(f"  Valid scores: {plddt['generated_valid_n']}/{plddt['generated_total_n']}")
+
+    print("="*70)
+
+
+def main():
+    parser = get_evaluate_val_argument_parser()
+    args = parser.parse_args()
+    # Legacy behavior: do not fix RNG (matches previous commented-out seed lines).
+    metrics = run_evaluate_val(args)
+    print_evaluation_results(metrics, args)
+
+
+if __name__ == "__main__":
+    main()

cas9/evaluate_val_multi_sample.py

@@ -0,0 +1,187 @@
+#!/usr/bin/env python3
+"""
+Run evaluate_val-style evaluation multiple times and report per-run metric lists.
+
+Stochastic steps (sampling from the test set, flow time / noise, generation) can make
+scalar metrics differ between runs. This script calls run_evaluate_val() repeatedly and
+prints every value plus simple mean / std / min / max summaries.
+
+Seeding:
+  --base_seed unset: same as legacy evaluate_val (no explicit RNG seeding; variation across runs).
+  --base_seed K: run i uses seed K + i before each replicate (reproducible multi-run bracket).
+"""
+
+import argparse
+import json
+import numpy as np
+
+from cas9.evaluate_val import get_evaluate_val_argument_parser, run_evaluate_val, print_evaluation_results
+
+
+def _get_nested(d, *keys):
+    for k in keys:
+        d = d[k]
+    return d
+
+
+def _stats(vals):
+    arr = np.array([v for v in vals if v is not None], dtype=float)
+    if arr.size == 0:
+        return None
+    return {
+        "mean": float(arr.mean()),
+        "std": float(arr.std()),
+        "min": float(arr.min()),
+        "max": float(arr.max()),
+    }
+
+
+def _fmt_list(vals, prec=4):
+    parts = []
+    for v in vals:
+        if v is None:
+            parts.append("None")
+        elif isinstance(v, float):
+            parts.append(f"{v:.{prec}f}")
+        elif isinstance(v, int):
+            parts.append(str(v))
+        else:
+            parts.append(str(v))
+    return "[" + ", ".join(parts) + "]"
+
+
+def _print_scalar_block(title, rows, all_metrics):
+    print(f"\n{title}")
+    print("-" * len(title))
+    for label, path in rows:
+        vals = [_get_nested(m, *path) for m in all_metrics]
+        print(f"  {label}")
+        print(f"    values: {_fmt_list(vals)}")
+        st = _stats(vals)
+        if st:
+            print(f"    mean={st['mean']:.6f}  std={st['std']:.6f}  min={st['min']:.6f}  max={st['max']:.6f}")
+        else:
+            print("    (no numeric values)")
+
+
+def print_multi_run_summary(all_metrics, args):
+    n = len(all_metrics)
+    print("\n" + "=" * 70)
+    print(f"MULTI-RUN SUMMARY ({n} runs)")
+    print("=" * 70)
+    seeds = [m["run_seed"] for m in all_metrics]
+    print(f"Per-run seeds: {seeds}")
+
+    loss_rows = [
+        ("Train avg (full train; None if not computed)", ("loss", "train_avg")),
+        ("Decoded avg (full test/val)", ("loss", "decoded_full_avg")),
+        ("Decoded avg (sampled)", ("loss", "decoded_sampled_avg")),
+        ("Generated avg", ("loss", "generated_avg")),
+    ]
+    _print_scalar_block("SEQUENCE LOSS (val_unweighted_total_loss)", loss_rows, all_metrics)
+
+    if all_metrics[0]["cas9"] is not None:
+        cas9_rows = [
+            ("Decoded validity rate", ("cas9", "decoded_validity_rate")),
+            ("Decoded avg Cas9 score", ("cas9", "decoded_avg_score")),
+            ("Generated validity rate", ("cas9", "generated_validity_rate")),
+            ("Generated avg Cas9 score", ("cas9", "generated_avg_score")),
+        ]
+        _print_scalar_block("CAS9 SCORES", cas9_rows, all_metrics)
+
+    div_dec = [
+        ("unique_count", ("diversity_decoded", "unique_count")),
+        ("uniqueness_ratio", ("diversity_decoded", "uniqueness_ratio")),
+        ("kmer_diversity", ("diversity_decoded", "kmer_diversity")),
+        ("kmer_avg_similarity", ("diversity_decoded", "kmer_avg_similarity")),
+        ("levenshtein_diversity", ("diversity_decoded", "levenshtein_diversity")),
+        ("levenshtein_avg_similarity", ("diversity_decoded", "levenshtein_avg_similarity")),
+    ]
+    _print_scalar_block("DIVERSITY — decoded (sampled)", div_dec, all_metrics)
+
+    div_gen = [
+        ("unique_count", ("diversity_generated", "unique_count")),
+        ("uniqueness_ratio", ("diversity_generated", "uniqueness_ratio")),
+        ("kmer_diversity", ("diversity_generated", "kmer_diversity")),
+        ("kmer_avg_similarity", ("diversity_generated", "kmer_avg_similarity")),
+        ("levenshtein_diversity", ("diversity_generated", "levenshtein_diversity")),
+        ("levenshtein_avg_similarity", ("diversity_generated", "levenshtein_avg_similarity")),
+    ]
+    _print_scalar_block("DIVERSITY — generated", div_gen, all_metrics)
+
+    if any(m.get("plddt") is not None for m in all_metrics):
+        p_rows = [
+            ("Decoded mean pLDDT", ("plddt", "decoded_mean")),
+            ("Decoded std pLDDT", ("plddt", "decoded_std")),
+            ("Generated mean pLDDT", ("plddt", "generated_mean")),
+            ("Generated std pLDDT", ("plddt", "generated_std")),
+        ]
+        _print_scalar_block("PLDDT", p_rows, all_metrics)
+
+    print("=" * 70)
+
+
+def main():
+    parser = get_evaluate_val_argument_parser()
+    parser.add_argument(
+        "--n_runs",
+        type=int,
+        default=3,
+        help="Number of independent evaluate_val runs (default: 3)",
+    )
+    parser.add_argument(
+        "--base_seed",
+        type=int,
+        default=None,
+        help="If set, run i uses RNG seed (base_seed + i). If unset, do not seed (like default evaluate_val).",
+    )
+    parser.add_argument(
+        "--save_fasta",
+        action="store_true",
+        help="Write generated_sequences_run{i}.fasta per run under output_dir",
+    )
+    parser.add_argument(
+        "--print_each_run",
+        action="store_true",
+        help="Print the full EVALUATION RESULTS block after every replicate",
+    )
+    parser.add_argument(
+        "--json_out",
+        type=str,
+        default=None,
+        help="Optional path to write a JSON list of per-run metric dicts",
+    )
+    args = parser.parse_args()
+
+    if args.n_runs < 1:
+        raise ValueError("--n_runs must be >= 1")
+
+    all_metrics = []
+    for i in range(args.n_runs):
+        print("\n" + "#" * 70)
+        print(f"RUN {i + 1} / {args.n_runs}")
+        print("#" * 70)
+        run_seed = (args.base_seed + i) if args.base_seed is not None else None
+        if run_seed is not None:
+            print(f"(run_seed={run_seed})")
+        fasta_tag = f"run{i + 1}" if args.save_fasta else None
+        metrics = run_evaluate_val(
+            args,
+            run_seed=run_seed,
+            save_fasta=args.save_fasta,
+            fasta_tag=fasta_tag,
+        )
+        all_metrics.append(metrics)
+        if args.print_each_run:
+            print_evaluation_results(metrics, args)
+
+    print_multi_run_summary(all_metrics, args)
+
+    if args.json_out:
+        with open(args.json_out, "w") as f:
+            json.dump(all_metrics, f, indent=2)
+        print(f"\nWrote per-run metrics to {args.json_out}")
+
+
+if __name__ == "__main__":
+    main()

cas9/generate.py

@@ -0,0 +1,78 @@
+import torch
+from transformers import EsmTokenizer
+
+# Cas9 stack uses model/base_models.py (not reparam_models.py, which is the GFP/root variant).
+from cas9.model.base_models import EditFlow, ProteinEditFlowModel, ReparameterizedProteinEditFlowModel
+from cas9.model.utils import generate_from_x0, generate_from_x0_multi_edit
+from cas9.logic import flow
+
+
+def build_model_and_stuff(cfg, device):
+    """
+    Rebuild exactly what train.py builds, but we won't set up lightning Trainer.
+    Returns:
+      editflow, source_distribution, tokenizer, pad_id, bos_id, eos_id, eps_id
+    """
+    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 = 1
+    bos_id = 0
+    eos_id = 2
+    # Match train.py: reparameterized (8-value) vs plain (5-value) model per config flag.
+    if getattr(cfg.training, "reparameterize", False):
+        model = ReparameterizedProteinEditFlowModel(vocab_size=vocab_size, pad_id=pad_id, config=cfg.model)
+    else:
+        model = ProteinEditFlowModel(vocab_size=vocab_size, pad_id=pad_id, config=cfg.model)
+
+    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,
+    ).to(device)
+
+    return editflow, source_distribution, tokenizer, pad_id, bos_id, eos_id, eps_id
+
+
+def tokenize_input_str(input_str, cfg, tokenizer, bos_id, eos_id, pad_id, device):
+    # Signature matches the Cas9 pcomol_cas.py call site (cfg/pad_id accepted for
+    # compatibility; the operation is tokenize input + ensure BOS/EOS -> (1, L)).
+    toks = tokenizer(input_str, return_tensors='pt')
+    ids = toks["input_ids"][0].to(device)
+    if ids[0].item() != bos_id:
+        ids = torch.cat([torch.tensor([bos_id], device=device), ids], dim=0)
+    if ids[-1].item() != eos_id:
+        ids = torch.cat([ids, torch.tensor([eos_id], device=device)], dim=0)
+    x0 = ids.unsqueeze(0)  # (1, L)
+    return x0
+
+
+def detokenize_output(x, tokenizer, bos_id, eos_id, pad_id):
+    """Convert a single generated sequence (1, L) back to string."""
+    seq = x[0].tolist()
+    seq = [tok for tok in seq if tok != pad_id]           # strip padding
+    if len(seq) > 0 and seq[0] == bos_id:                 # strip BOS
+        seq = seq[1:]
+    if len(seq) > 0 and seq[-1] == eos_id:                # strip EOS
+        seq = seq[:-1]
+    return tokenizer.batch_decode([seq], skip_special_tokens=True)[0]

cas9/generate_cas9_csvs.py

@@ -0,0 +1,348 @@
+#!/usr/bin/env python3
+"""
+Generate CSV outputs for one or more Cas9 inputs using the same rollout sampler
+as `pcomol_cas.py`.
+
+This script:
+1. Parses Cas9 name -> sequence mappings from `run_pcomol.sh`
+2. Runs input-conditioned generation via `short_rollout_batch`
+3. Writes one CSV per requested Cas9 name with the same column layout commonly
+   produced by `pcomol_cas.py` for Cas9 runs
+
+Objective-specific columns are intentionally left empty.
+"""
+
+import argparse
+import csv
+import time
+from pathlib import Path
+from typing import Dict, List
+
+import torch
+import yaml
+from easydict import EasyDict as edict
+
+from cas9.generate import build_model_and_stuff, detokenize_output, tokenize_input_str
+from cas9.pcomole import short_rollout_batch
+
+
+CSV_COLUMNS = [
+    "final_len",
+    "length_diff",
+    "cas9_score",
+    "deletion_count_score",
+    "pam_matching_score",
+    "cas9_weight",
+    "deletion_count_weight",
+    "pam_matching_weight",
+    "predicted_pam",
+    "pam_matching_ce_loss_score",
+    "pam_matching_probability",
+    "deletion_rate_scale",
+    "pam_min_confidence",
+    "generation_time",
+    "sequence",
+]
+
+
+def parse_fasta_sequences(path: Path) -> Dict[str, str]:
+    """Parse a FASTA of Cas9 orthologs into {name: sequence}."""
+    seqs: Dict[str, str] = {}
+    name, buf = None, []
+    for line in Path(path).read_text().splitlines():
+        line = line.strip()
+        if line.startswith(">"):
+            if name:
+                seqs[name] = "".join(buf)
+            name, buf = line[1:].split()[0], []
+        elif line:
+            buf.append(line)
+    if name:
+        seqs[name] = "".join(buf)
+    return seqs
+
+
+def load_sequences(path: Path) -> Dict[str, str]:
+    """Dispatch on suffix: .sh uses the legacy run_pcomol.sh parser."""
+    return (parse_run_pcomol_sequences(path) if Path(path).suffix == ".sh"
+            else parse_fasta_sequences(path))
+
+
+def parse_run_pcomol_sequences(path: Path) -> Dict[str, str]:
+    sequences: Dict[str, str] = {}
+    with open(path, "r") as handle:
+        for raw_line in handle:
+            line = raw_line.strip()
+            if not line or line.startswith("#") or "=" not in line:
+                continue
+            name, value = line.split("=", 1)
+            name = name.strip()
+            value = value.strip()
+            if not value.startswith('"') or not value.endswith('"'):
+                continue
+            seq = value[1:-1].strip()
+            if not seq:
+                continue
+            if all(ch.isalpha() for ch in seq):
+                sequences[name] = seq
+    return sequences
+
+
+def normalize_cas9_names(requested_names: List[str], available: Dict[str, str]) -> List[str]:
+    normalized = []
+    missing = []
+    lower_to_name = {name.lower(): name for name in available}
+    for name in requested_names:
+        resolved = available.get(name)
+        if resolved is None:
+            canonical = lower_to_name.get(name.lower())
+            if canonical is None:
+                missing.append(name)
+                continue
+            resolved_name = canonical
+        else:
+            resolved_name = name
+        normalized.append(resolved_name)
+
+    if missing:
+        available_names = ", ".join(sorted(available))
+        raise ValueError(
+            f"Unknown Cas9 name(s): {', '.join(missing)}. Available names: {available_names}"
+        )
+    return normalized
+
+
+def build_empty_row(final_len: int, length_diff: int, deletion_rate_scale: float, generation_time: float, sequence: str) -> Dict[str, object]:
+    row = {column: "" for column in CSV_COLUMNS}
+    row["final_len"] = final_len
+    row["length_diff"] = length_diff
+    row["deletion_rate_scale"] = deletion_rate_scale
+    row["generation_time"] = generation_time
+    row["sequence"] = sequence
+    return row
+
+
+def write_csv_rows(output_path: Path, rows: List[Dict[str, object]]) -> None:
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(output_path, "w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=CSV_COLUMNS)
+        writer.writeheader()
+        writer.writerows(rows)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Generate n sequences per requested Cas9 input using the same "
+            "sampling rollout as pcomol_cas.py."
+        )
+    )
+    parser.add_argument(
+        "--pcomol_config",
+        type=str,
+        required=True,
+        help="Path to pcomol/config YAML (e.g. configs/config_uniprot_leq_250.yaml)",
+    )
+    parser.add_argument(
+        "--ckpt",
+        type=str,
+        required=True,
+        help="Path to model checkpoint (.ckpt)",
+    )
+    parser.add_argument(
+        "--cas9_names",
+        nargs="+",
+        required=True,
+        help="Cas9 ortholog names, e.g. spcas9 st1cas9",
+    )
+    parser.add_argument(
+        "--n",
+        type=int,
+        nargs="+",
+        required=True,
+        help=(
+            "Number of generated sequences to produce. Pass one value to use for "
+            "all Cas9 inputs, or one value per Cas9 name in the same order."
+        ),
+    )
+    parser.add_argument(
+        "--sequences", "--run_pcomol_sh",
+        dest="run_pcomol_sh",
+        type=str,
+        default=str(Path(__file__).with_name("data") / "orthologs.fasta"),
+        help="FASTA of Cas9 orthologs (or a legacy run_pcomol.sh)",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        required=True,
+        help="Directory where one CSV per Cas9 name will be written",
+    )
+    parser.add_argument(
+        "--output_suffix",
+        type=str,
+        default="_generations.csv",
+        help="Suffix appended to each Cas9 name for its CSV filename",
+    )
+    parser.add_argument(
+        "--num_steps",
+        type=int,
+        default=32,
+        help="Number of flow steps (default: 32)",
+    )
+    parser.add_argument(
+        "--max_len_cap",
+        type=int,
+        default=None,
+        help="Maximum sequence length cap (default: None)",
+    )
+    parser.add_argument(
+        "--deletion_rate_scale",
+        type=float,
+        default=1.0,
+        help="Scaling factor for deletion rate (default: 1.0)",
+    )
+    parser.add_argument(
+        "--zero_lam_ins",
+        action="store_true",
+        help="Set insertion rate to 0 (deletions and substitutions only)",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="Optional random seed for reproducibility",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default=None,
+        help="Device to use (cuda/cpu). Default: auto-detect",
+    )
+    args = parser.parse_args()
+
+    if args.seed is not None:
+        torch.manual_seed(args.seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(args.seed)
+
+    if args.device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(args.device)
+
+    with open(args.pcomol_config, "r") as handle:
+        cfg = edict(yaml.safe_load(handle))
+
+    sequence_map = load_sequences(Path(args.run_pcomol_sh))
+    if not sequence_map:
+        raise ValueError(f"No Cas9 sequences found in {args.run_pcomol_sh}")
+
+    cas9_names = normalize_cas9_names(args.cas9_names, sequence_map)
+
+    if any(n <= 0 for n in args.n):
+        raise ValueError("All values passed to --n must be positive")
+
+    if len(args.n) == 1:
+        n_per_cas9 = {name: args.n[0] for name in cas9_names}
+    elif len(args.n) == len(cas9_names):
+        n_per_cas9 = {name: n for name, n in zip(cas9_names, args.n)}
+    else:
+        raise ValueError(
+            f"--n must have either 1 value or exactly {len(cas9_names)} values "
+            f"to match --cas9_names, but got {len(args.n)} values."
+        )
+
+    editflow, source_dist, tokenizer, pad_id, bos_id, eos_id, eps_id = build_model_and_stuff(
+        cfg, device
+    )
+
+    ckpt = torch.load(args.ckpt, map_location=device)
+    editflow.load_state_dict(ckpt["state_dict"], strict=False)
+    model = editflow.model.to(device)
+    model.eval()
+
+    allowed_tokens = torch.tensor(
+        [tok for tok in source_dist._allowed_tokens if tok not in (eps_id,)],
+        device=device,
+        dtype=torch.long,
+    )
+    time_grid = torch.linspace(0.0, 1.0, steps=args.num_steps, device=device)
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    for cas9_name in cas9_names:
+        input_sequence = sequence_map[cas9_name].replace(" ", "").strip()
+        target_n = n_per_cas9[cas9_name]
+        x0 = tokenize_input_str(input_sequence, cfg, tokenizer, bos_id, eos_id, pad_id, device)
+
+        rows: List[Dict[str, object]] = []
+        rows.append(
+            build_empty_row(
+                final_len=len(input_sequence),
+                length_diff=0,
+                deletion_rate_scale=args.deletion_rate_scale,
+                generation_time=0,
+                sequence=input_sequence,
+            )
+        )
+
+        generated = 0
+        attempts = 0
+        max_attempts = max(target_n * 10, 50)
+
+        while generated < target_n and attempts < max_attempts:
+            attempts += 1
+            start_time = time.time()
+            try:
+                with torch.no_grad():
+                    x_t = short_rollout_batch(
+                        model,
+                        x0,
+                        time_grid,
+                        start_idx=0,
+                        pad_id=pad_id,
+                        bos_id=bos_id,
+                        eos_id=eos_id,
+                        allowed_tokens=allowed_tokens,
+                        max_len_cap=args.max_len_cap,
+                        num_rollouts=1,
+                        num_steps=args.num_steps,
+                        protected_mask=None,
+                        deletion_rate_scale=args.deletion_rate_scale,
+                        zero_lam_ins=args.zero_lam_ins,
+                    )
+            except torch.cuda.OutOfMemoryError:
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+                continue
+
+            generation_time = time.time() - start_time
+            out_str = detokenize_output(x_t, cfg, tokenizer, bos_id, eos_id, pad_id)
+            out_str = out_str.replace(" ", "").strip()
+
+            rows.append(
+                build_empty_row(
+                    final_len=len(out_str),
+                    length_diff=len(input_sequence) - len(out_str),
+                    deletion_rate_scale=args.deletion_rate_scale,
+                    generation_time=round(generation_time, 4),
+                    sequence=out_str,
+                )
+            )
+            generated += 1
+
+        if generated < target_n:
+            raise RuntimeError(
+                f"Only generated {generated}/{target_n} sequences for {cas9_name} "
+                f"after {attempts} attempts."
+            )
+
+        output_path = output_dir / f"{cas9_name}{args.output_suffix}"
+        write_csv_rows(output_path, rows)
+        print(f"Wrote {generated} generated sequences for {cas9_name} to {output_path}")
+
+
+if __name__ == "__main__":
+    main()

cas9/hmm/README.md

@@ -0,0 +1,3 @@
+Tier B artifact slot — not yet populated.
+
+See the repository README, section 4, for what belongs here.

cas9/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.

cas9/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."
+    )

cas9/logic/flow.py

@@ -0,0 +1,437 @@
+# 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
+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 = 2, bos_id = 0, eos_id = 2, sample_type='regular', eps_id = -1):
+        """
+        For each sequence in x1, sample an x0 whose *core* length (excluding BOS/EOS)
+        is in [max(0, (1-scale_size)*len_valid(x1)), (1+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
+            # sample core length in [max(0, (1-scale_size)*len_valid(x1)), (1+scale_size)*len_valid(x1)]
+            min_core_len = max(0, int((1 - scale_size) * valid_len[b].item()))
+            # print(f"x1 length: {valid_len[b].item()}, min_core_len: {min_core_len}, max_core_len: {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")
+    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")

cas9/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

cas9/logic/state.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.
+
+import logging
+from pathlib import Path
+
+import torch
+from cas9.data import DataState
+
+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)

cas9/logic/training.py

@@ -0,0 +1,141 @@
+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).clamp(max=0.9999)
+
+        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_components = 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,
+        )
+        # Log component losses if logger is available
+        if logger is not None:
+            logger.log_loss(value=loss_components["loss_base"].item(), step=state.step, name="loss_base")
+            logger.log_loss(value=loss_components["loss_rate"].item(), step=state.step, name="loss_rate")
+            logger.log_loss(value=loss_components["loss_edit"].item(), step=state.step, name="loss_edit")
+            logger.log_loss(value=loss_components["loss_aux_unweighted"].item(), step=state.step, name="loss_aux_unweighted")
+            logger.log_loss(value=loss_components["loss_aux_weighted"].item(), step=state.step, name="loss_aux_weighted")
+
+    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()

cas9/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",
+]

cas9/model/base_models.py

@@ -0,0 +1,935 @@
+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(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
+        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)
+
+        # print(f"mask: {mask.tolist()}")
+
+        # --- 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
+
+    def get_esm_embeddings(self, x_t, mask):
+        """
+        Extract ESM embeddings for given tokens.
+        
+        This method allows external code to reuse the ESM model for other purposes
+        (e.g., sharing with Cas9Classifier to save memory).
+        
+        Args:
+            x_t: (B, L) LongTensor - tokenized protein sequences
+            mask: (B, L) BoolTensor - True for valid tokens, False for padding
+        
+        Returns:
+            embeddings: (B, L, esm_hidden_size) - ESM embeddings
+        """
+        return self.esm_emb(x_t, attention_mask=mask).last_hidden_state
+
+class ReparameterizedProteinEditFlowModel(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
+        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)
+
+        # print(f"mask: {mask.tolist()}")
+
+        # --- 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)
+        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):
+            # For lambdas: force to 0 on pads
+            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
+
+            # kill logits on pads
+            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, logits_type, pi_type
+
+class ReparameterizedSMILESEditFlowModel(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_total: (B, L)         >= 0
+      logits_type: (B, L, 3)    over {ins, del, sub}
+      logits_ins: (B, L, V)
+      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_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
+        t: (B,) or scalar in [0,1]
+        """
+        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_total = self.softplus(self.lam_total_head(h)).squeeze(-1)  # (B, L)
+        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)
+
+        # --- 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_total = lam_total * pad_mask_f
+
+            # kill logits on pads
+            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_total, logits_type, logits_ins, 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
+        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.loc_prop_path = getattr(config.training, "loc_prop_path", False)
+
+        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.lam_prop = getattr(self.cfg.training, "lam_prop", 0.0)
+        self.use_aux_ce = getattr(self.cfg.training, "use_aux_ce", False)
+        self.aux_ce_weight = getattr(self.cfg.training, "aux_ce_weight", 0.0)
+        self.reparameterize = getattr(self.cfg.training, "reparameterize", False)
+        self.gamma_rate = getattr(self.cfg.training, "gamma_rate", 1)
+        self.gamma_edit = getattr(self.cfg.training, "gamma_edit", 1)
+        # print(f"use_aux_ce: {self.use_aux_ce}, aux_ce_weight: {self.aux_ce_weight}")
+
+        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)
+            # print(f"allowed_tokens: {allowed_tokens}")
+            
+            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).clamp(max=0.9999)
+
+            # Use lambda_indep if available (for convex schedulers), otherwise fallback to d_alpha_t / sigma_t
+            if hasattr(self.path, 'lambda_indep'):
+                weight = self.path.lambda_indep(t)  # (B,)
+            else:
+                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)
+
+            if self.loc_prop_path:
+                z_t, M_t, m_t = self.path.sample_localized(
+                    z0=z_0, z1=z_1, t=t, lambda_prop=self.lam_prop, return_M=True
+                )
+            else:
+                z_t = self.path.sample(z_0, z_1, t=t)
+                M_t = None
+                m_t = None
+            
+            x_t, mask = remove_eps(z_t, self.eps_id, self.pad_id)
+
+        if self.reparameterize:
+            lam_ins, logits_ins, lam_del, lam_sub, logits_sub, lam_total, logits_type, pi_type = self.model(x_t=x_t, mask=mask,t=t)
+            
+            return lam_total, logits_type, logits_ins, logits_sub, z_t, z_1, x_t, mask, weight, M_t
+
+        else:
+            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, M_t
+    
+
+    def training_step(self, batch, batch_idx):
+        x_1 = torch.tensor(batch["input_ids"]).to(self.device)
+        B = x_1.shape[0]
+        
+        if self.reparameterize:
+            lam_total, logits_type, logits_ins, logits_sub, z_t, z_1, x_t, mask, weight, M_t = self.preparation(x_1)
+
+            if self.loc_prop_path:
+                loss, loss_components = self.loss_fn.reparameterized_forward_localized(lam_total, logits_type, logits_ins, logits_sub, 
+                                z_t, z_1, x_t, mask, weight, M_t, self.lam_prop, self.eps_id, self.bos_id, self.eos_id, self.gamma_rate, self.gamma_edit, self.use_aux_ce, self.aux_ce_weight)        
+            else:
+                loss, loss_components = self.loss_fn.reparameterized_forward(lam_total, logits_type, logits_ins, logits_sub,
+                                z_t, z_1, x_t, mask, weight, self.eps_id, self.bos_id, self.eos_id, self.gamma_rate, self.gamma_edit, self.use_aux_ce, self.aux_ce_weight)
+            
+        else:
+            lam_ins, logits_ins, lam_del, lam_sub, logits_sub, z_t, z_1, x_t, mask, weight, M_t = self.preparation(x_1)
+
+            if self.loc_prop_path:
+                loss, loss_components = self.loss_fn.forward_localized(lam_ins, logits_ins, lam_del, lam_sub, logits_sub, 
+                                    z_t, z_1, x_t, mask, weight, M_t, self.lam_prop, self.eps_id, self.bos_id, self.eos_id, self.use_aux_ce, self.aux_ce_weight)
+            else:
+                loss, loss_components = self.loss_fn.forward(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.use_aux_ce, self.aux_ce_weight)
+        
+        # Log total loss (same metric name for both reparameterized and non-reparameterized models)
+        # This ensures both model types appear on the same WandB chart for comparison
+        self.log("train_loss", loss, prog_bar=True, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+        # Log full non-aux loss
+        # Log individual non-aux loss components
+        self.log("train_loss_rate", loss_components["loss_rate"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+        self.log("train_loss_edit", loss_components["loss_edit"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+        # Log unweighted and weighted aux losses (only if use_aux_ce is True)
+        if self.use_aux_ce and self.aux_ce_weight > 0.0:
+            self.log("train_loss_base", loss_components["loss_base"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+            self.log("train_loss_aux_unweighted", loss_components["loss_aux_unweighted"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+            self.log("train_loss_aux_weighted", loss_components["loss_aux_weighted"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+        # Log new weighted/unweighted components (only available for reparameterized methods)
+        if "loss_rate_unweighted" in loss_components:
+            self.log("train_loss_rate_unweighted", loss_components["loss_rate_unweighted"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+            self.log("train_loss_rate_weighted", loss_components["loss_rate_weighted"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+            self.log("train_loss_edit_unweighted", loss_components["loss_edit_unweighted"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+            self.log("train_loss_edit_weighted", loss_components["loss_edit_weighted"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+            self.log("train_loss_total_weighted", loss_components["loss_total_weighted"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+            self.log("train_loss_total_unweighted", loss_components["loss_total_unweighted"], prog_bar=False, 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]
+        
+        if self.reparameterize:
+            lam_total, logits_type, logits_ins, logits_sub, z_t, z_1, x_t, mask, weight, M_t = self.preparation(x_1)
+            
+            if self.loc_prop_path:
+                loss, loss_components = self.loss_fn.reparameterized_forward_localized(lam_total, logits_type, logits_ins, logits_sub, 
+                                z_t, z_1, x_t, mask, weight, M_t, self.lam_prop, self.eps_id, self.bos_id, self.eos_id, self.gamma_rate, self.gamma_edit, self.use_aux_ce, self.aux_ce_weight)        
+            else:
+                loss, loss_components = self.loss_fn.reparameterized_forward(lam_total, logits_type, logits_ins, logits_sub,
+                                z_t, z_1, x_t, mask, weight, self.eps_id, self.bos_id, self.eos_id, self.gamma_rate, self.gamma_edit, self.use_aux_ce, self.aux_ce_weight)
+        else:
+            lam_ins, logits_ins, lam_del, lam_sub, logits_sub, z_t, z_1, x_t, mask, weight, M_t = self.preparation(x_1)
+            
+            if self.loc_prop_path:
+                loss, loss_components = self.loss_fn.forward_localized(lam_ins, logits_ins, lam_del, lam_sub, logits_sub, 
+                                    z_t, z_1, x_t, mask, weight, M_t, self.lam_prop, self.eps_id, self.bos_id, self.eos_id, self.use_aux_ce, self.aux_ce_weight)
+            else:
+                loss, loss_components = self.loss_fn.forward(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.use_aux_ce, self.aux_ce_weight)
+        
+        # Log total loss (same metric name for both reparameterized and non-reparameterized models)
+        # This ensures both model types appear on the same WandB chart for comparison
+        self.log("val_loss", loss, prog_bar=True, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+        # Log full non-aux loss
+        # Log individual non-aux loss components
+        self.log("val_loss_rate", loss_components["loss_rate"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+        self.log("val_loss_edit", loss_components["loss_edit"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+        # Log unweighted and weighted aux losses (only if use_aux_ce is True)
+        if self.use_aux_ce and self.aux_ce_weight > 0.0:
+            self.log("val_loss_base", loss_components["loss_base"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+            self.log("val_loss_aux_unweighted", loss_components["loss_aux_unweighted"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+            self.log("val_loss_aux_weighted", loss_components["loss_aux_weighted"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+        # Log new weighted/unweighted components (only available for reparameterized methods)
+        if "loss_rate_unweighted" in loss_components:
+            self.log("val_loss_rate_unweighted", loss_components["loss_rate_unweighted"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+            self.log("val_loss_rate_weighted", loss_components["loss_rate_weighted"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+            self.log("val_loss_edit_unweighted", loss_components["loss_edit_unweighted"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+            self.log("val_loss_edit_weighted", loss_components["loss_edit_weighted"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+            self.log("val_loss_total_weighted", loss_components["loss_total_weighted"], prog_bar=False, on_step=True, on_epoch=True, batch_size=B, sync_dist=True)
+            self.log("val_loss_total_unweighted", loss_components["loss_total_unweighted"], prog_bar=False, 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
+
+    def on_after_backward(self):
+        """
+        Hook called after backward() but before gradient clipping.
+        Used to log gradient information before clipping occurs.
+        """
+        # Get gradient clipping value from trainer (default is 1.0)
+        clip_val = getattr(self.trainer, 'gradient_clip_val', 1.0)
+        if clip_val is None:
+            clip_val = float('inf')
+        
+        # Collect all parameter gradients
+        total_grad_norm = 0.0
+        param_grad_norms = {}
+        num_clipped = 0
+        total_params = 0
+        
+        for name, param in self.named_parameters():
+            if param.grad is not None:
+                # Compute gradient norm for this parameter
+                grad_norm = param.grad.data.norm(2).item()
+                param_grad_norms[name] = grad_norm
+                total_grad_norm += grad_norm ** 2
+                total_params += 1
+                
+                # Check if this gradient would be clipped
+                if grad_norm > clip_val:
+                    num_clipped += 1
+        
+        # Compute total gradient norm (L2 norm across all parameters)
+        total_grad_norm = total_grad_norm ** 0.5
+        
+        # Log total gradient norm
+        self.log("train_grad_norm", total_grad_norm, on_step=True, on_epoch=False, sync_dist=True)
+        
+        # Log ratio of clipped gradients
+        if total_params > 0:
+            clip_ratio = num_clipped / total_params
+            self.log("train_grad_clip_ratio", clip_ratio, on_step=True, on_epoch=False, sync_dist=True)
+        else:
+            self.log("train_grad_clip_ratio", 0.0, on_step=True, on_epoch=False, sync_dist=True)
+        
+        # Log individual parameter gradient norms (as a dictionary/table in wandb)
+        # We'll log the max, min, mean, and std of parameter gradients
+        if param_grad_norms:
+            grad_norms_list = list(param_grad_norms.values())
+            self.log("train_grad_norm_max", max(grad_norms_list), on_step=True, on_epoch=False, sync_dist=True)
+            self.log("train_grad_norm_min", min(grad_norms_list), on_step=True, on_epoch=False, sync_dist=True)
+            self.log("train_grad_norm_mean", sum(grad_norms_list) / len(grad_norms_list), on_step=True, on_epoch=False, sync_dist=True)
+            
+            # Compute std
+            mean_grad = sum(grad_norms_list) / len(grad_norms_list)
+            variance = sum((x - mean_grad) ** 2 for x in grad_norms_list) / len(grad_norms_list)
+            std_grad = variance ** 0.5
+            self.log("train_grad_norm_std", std_grad, on_step=True, on_epoch=False, sync_dist=True)
+            
+            # Log full parameter gradients as a wandb table (every N steps to avoid too much data)
+            # For now, we'll log a summary. To log full details, we can use wandb.Table
+            if hasattr(self.logger, 'experiment') and hasattr(self.logger.experiment, 'log'):
+                # Log detailed gradient info every 100 steps
+                if self.global_step % 100 == 0:
+                    import wandb
+                    # Create a table with parameter names and their gradient norms
+                    table_data = [[name, norm] for name, norm in param_grad_norms.items()]
+                    table = wandb.Table(columns=["parameter", "grad_norm"], data=table_data)
+                    self.logger.experiment.log({"train_grad_norms_table": table}, step=self.global_step)
+
+    # 2
+
+                

cas9/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
+

cas9/model/transformer.py

@@ -0,0 +1,478 @@
+# 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 sys
+import os
+# Add flow_matching to path if needed for imports
+_flow_matching_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'flow_matching')
+if _flow_matching_dir not in sys.path:
+    sys.path.insert(0, _flow_matching_dir)
+
+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)

cas9/model/utils.py

@@ -0,0 +1,949 @@
+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
+
+
+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,
+    sample_type: str = 'regular',
+):
+    """
+    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
+    pos_temperature: float = 1.0,        # temperature for sampling position (reparameterized models only)
+    device: torch.device = None,
+    is_reparameterized: bool = None,      # If None, will auto-detect from model output
+    convert_to_vanilla_outputs: bool = False,  # If False, use direct sampling for reparameterized models
+):
+    """
+    Discrete edit sampler for Edit Flows with temperature on:
+      - operation choice (insert/delete/sub)
+      - token choice (for insert/sub)
+      - position choice (for reparameterized models when convert_to_vanilla_outputs=False)
+
+    At each step we apply at most ONE edit per sequence.
+    
+    Supports both base and reparameterized models:
+    - Base: outputs (lam_ins, logits_ins, lam_del, lam_sub, logits_sub)
+    - Reparameterized: outputs (lam_total, logits_type, logits_ins, logits_sub)
+    
+    For reparameterized models:
+    - If convert_to_vanilla_outputs=True (default): converts to base format and uses
+      best-position-per-operation approach
+    - If convert_to_vanilla_outputs=False: uses direct sampling:
+      1. Samples position from lam_total using pos_temperature
+      2. Samples edit type from logits_type at sampled position using op_temperature
+      3. Samples token if needed (insert/sub) using token_temperature
+    """
+    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())
+
+    # Auto-detect model type if not specified
+    if is_reparameterized is None:
+        # Try to detect from model class name first (more efficient)
+        model_class_name = model.__class__.__name__
+        if "Reparameterized" in model_class_name:
+            is_reparameterized = True
+        else:
+            # Fall back to test forward pass to detect model type
+            test_t = torch.zeros(1, device=device)
+            test_mask = torch.ones(1, x.size(1), dtype=torch.bool, device=device)
+            test_output = model(x_t=x[:1], mask=test_mask, t=test_t)
+            # Reparameterized models return 4 (SMILES) or 8 (Protein) values
+            is_reparameterized = len(test_output) in (4, 8)
+
+    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
+        model_output = model(x_t=x, mask=mask, t=t)
+        
+        if is_reparameterized:
+            if len(model_output) == 8:
+                # ReparameterizedProteinEditFlowModel: returns all decomposed rates + total + type info
+                lam_ins, logits_ins, lam_del, lam_sub, logits_sub, lam_total, logits_type, pi_type = model_output
+                # pi_type is already computed, no need to recompute
+            elif len(model_output) == 4:
+                # ReparameterizedSMILESEditFlowModel: (lam_total, logits_type, logits_ins, logits_sub)
+                lam_total, logits_type, logits_ins, logits_sub = model_output
+                pi_type = F.softmax(logits_type, dim=-1)  # (B, L, 3) over {ins, del, sub}
+                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)
+            else:
+                raise ValueError(f"Unexpected reparameterized model output length: {len(model_output)}. Expected 4 or 8.")
+            
+            if convert_to_vanilla_outputs:
+                # For ReparameterizedProteinEditFlowModel, we already have lam_ins/del/sub
+                # For ReparameterizedSMILESEditFlowModel, we computed them above
+                pass  # lam_ins, lam_del, lam_sub are already set
+            else:
+                # Use direct sampling approach: keep reparameterized outputs as-is
+                # We'll sample position and edit type separately below
+                lam_ins = None  # Not used in direct sampling mode
+                lam_del = None
+                lam_sub = None
+        else:
+            # Base model: (lam_ins, logits_ins, lam_del, lam_sub, logits_sub)
+            lam_ins, logits_ins, lam_del, lam_sub, logits_sub = model_output
+
+        # 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()
+
+            if is_reparameterized and not convert_to_vanilla_outputs:
+                # Direct sampling approach for reparameterized models
+                lam_total_b = lam_total[b]  # (L,)
+                logits_type_b = logits_type[b]  # (L, 3)
+                logits_ins_b = logits_ins[b]  # (L, V)
+                logits_sub_b = logits_sub[b]  # (L, V)
+                
+                # 1. Sample position from lam_total with pos_temperature
+                # Mask out invalid positions (after EOS, or BOS/EOS for certain operations)
+                # For now, we'll allow sampling from all valid positions, then filter based on edit type
+                lam_total_valid = lam_total_b[:length_b].clone()  # Only consider valid positions
+                
+                # Apply temperature to position distribution
+                if pos_temperature is not None and pos_temperature > 0.0:
+                    pos_logits = lam_total_valid / pos_temperature
+                    pos_probs = F.softmax(pos_logits, dim=-1)
+                    sampled_pos = int(torch.multinomial(pos_probs, 1).item())
+                elif pos_temperature == 0.0:
+                    # Greedy sampling: choose position with highest lam_total
+                    sampled_pos = int(torch.argmax(lam_total_valid).item())
+                else:
+                    # Default behavior when pos_temperature is None: use softmax without temperature scaling
+                    pos_probs = F.softmax(lam_total_valid, dim=-1)
+                    sampled_pos = int(torch.multinomial(pos_probs, 1).item())
+                
+                # 2. Sample edit type from logits_type at the sampled position with op_temperature
+                edit_type_logits = logits_type_b[sampled_pos]  # (3,) for {ins, del, sub}
+                
+                if op_temperature is not None and op_temperature > 0.0:
+                    edit_type_logits_scaled = edit_type_logits / op_temperature
+                else:
+                    edit_type_logits_scaled = edit_type_logits
+                
+                edit_type_probs = F.softmax(edit_type_logits_scaled, dim=-1)
+                op_idx = int(torch.multinomial(edit_type_probs, 1).item())
+                
+                # 3. Apply the sampled edit
+                # 0 -> insert, 1 -> delete, 2 -> sub
+                if op_idx == 0:
+                    # insertion: can insert at any position, but skip after EOS
+                    if sampled_pos < eos_pos:
+                        ins_tok = sample_token_from_logits(logits_ins_b[sampled_pos])
+                        tokens = tokens[:sampled_pos + 1] + [ins_tok] + tokens[sampled_pos + 1:]
+                    # else: skip insertion if position is at or after EOS
+                elif op_idx == 1:
+                    # deletion: skip BOS/EOS
+                    if tokens[sampled_pos] != bos_id and tokens[sampled_pos] != eos_id:
+                        tokens = tokens[:sampled_pos] + tokens[sampled_pos + 1:]
+                    # else: skip deletion if position is BOS/EOS
+                else:  # op_idx == 2
+                    # substitution: skip BOS/EOS
+                    if tokens[sampled_pos] != bos_id and tokens[sampled_pos] != eos_id:
+                        sub_tok = sample_token_from_logits(logits_sub_b[sampled_pos])
+                        tokens = tokens[:sampled_pos] + [sub_tok] + tokens[sampled_pos + 1:]
+                    # else: skip substitution if position is BOS/EOS
+            else:
+                # Original approach: convert to vanilla outputs or use base model outputs
+                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
+    
+@torch.no_grad()
+def generate_from_x0_ctmc(
+    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: Optional[torch.Tensor] = None,  # 1D tensor of vocab ids we can generate
+    num_steps: int = 32,
+    max_len_cap: Optional[int] = None,
+    op_temperature: float = 1.0,          # accepted but unused (for API compat)
+    token_temperature: float = 1.0,
+    pos_temperature: float = 1.0,         # accepted but unused (for API compat)
+    is_reparameterized: Optional[bool] = None,
+    convert_to_vanilla_outputs: bool = False,
+    device: Optional[torch.device] = None,
+):
+    """
+    CTMC-style discrete-time sampler for Edit Flows / DFM.
+
+    At each step:
+      - For each position j, we sample independent Bernoulli events:
+          insert with prob h * λ_ins[t,j]
+          delete/sub with prob h * (λ_del[t,j] + λ_sub[t,j])
+        and, if a del/sub event occurs, choose delete vs sub proportional to λ_del vs λ_sub.
+      - We apply all resulting edit operations simultaneously (left-to-right).
+
+    Supports:
+      - Base model:         (lam_ins, logits_ins, lam_del, lam_sub, logits_sub)
+      - Reparameterized:    (lam_total, logits_type, logits_ins, logits_sub)
+        * If convert_to_vanilla_outputs=True:
+            lam_ins/del/sub = lam_total * softmax(logits_type)[..., k]
+        * If convert_to_vanilla_outputs=False:
+            probabilities are computed directly from lam_total and π_type.
+    """
+
+    if device is None:
+        device = x0.device
+
+    x = x0.clone().to(device)
+    B = x.size(0)
+
+    if num_steps <= 0:
+        return x
+
+    def sample_token_from_logits(logits_row: torch.Tensor) -> int:
+        """
+        logits_row: (V,). Apply allowed_tokens mask + temperature, 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, -1e9)  # effectively remove disallowed tokens
+
+        if token_temperature is not None and token_temperature > 0.0 and token_temperature != 1.0:
+            logit = logit / token_temperature
+
+        probs = F.softmax(logit, dim=-1)
+        idx = torch.multinomial(probs, num_samples=1)
+        return int(idx.item())
+
+    # Auto-detect reparameterized vs base model if not specified
+    if is_reparameterized is None:
+        model_class_name = model.__class__.__name__
+        if "Reparameterized" in model_class_name:
+            is_reparameterized = True
+        else:
+            # Fallback: look at forward output length
+            with torch.no_grad():
+                t_test = torch.zeros(1, device=device)
+                mask_test = torch.ones(1, x.size(1), dtype=torch.bool, device=device)
+                test_out = model(x_t=x[:1], mask=mask_test, t=t_test)
+            # Reparameterized models return 4 (SMILES) or 8 (Protein) values
+            is_reparameterized = (len(test_out) in (4, 8))
+
+    # Time step size h; t_k = k * h, k = 0..num_steps-1
+    if num_steps == 1:
+        h = 1.0
+    else:
+        h = 1.0 / float(num_steps - 1)
+
+    for step in range(num_steps):
+        t_scalar = step * h
+        t = torch.full((B,), t_scalar, device=device, dtype=torch.float32)
+
+        mask = (x != pad_id)
+        model_out = model(x_t=x, mask=mask, t=t)
+
+        if is_reparameterized:
+            if len(model_out) == 8:
+                # ReparameterizedProteinEditFlowModel: returns all decomposed rates + total + type info
+                lam_ins, logits_ins, lam_del, lam_sub, logits_sub, lam_total, logits_type, pi_type = model_out
+                # pi_type is already computed, no need to recompute
+            elif len(model_out) == 4:
+                # ReparameterizedSMILESEditFlowModel: (lam_total, logits_type, logits_ins, logits_sub)
+                lam_total, logits_type, logits_ins, logits_sub = model_out
+                pi_type = F.softmax(logits_type, dim=-1)  # (B, L, 3) over {ins, del, sub}
+            else:
+                raise ValueError(f"Unexpected reparameterized model output length: {len(model_out)}. Expected 4 or 8.")
+
+            if convert_to_vanilla_outputs:
+                if len(model_out) == 8:
+                    # For ReparameterizedProteinEditFlowModel, we already have lam_ins/del/sub
+                    pass  # lam_ins, lam_del, lam_sub are already set
+                else:
+                    # Convert to "vanilla" λ_ins/λ_del/λ_sub, then reuse base logic
+                    lam_ins = lam_total * pi_type[..., 0]
+                    lam_del = lam_total * pi_type[..., 1]
+                    lam_sub = lam_total * pi_type[..., 2]
+            else:
+                # We'll use lam_total + pi_type directly in the loop
+                lam_ins = lam_del = lam_sub = None  # not used in this branch
+        else:
+            # Base model: (lam_ins, logits_ins, lam_del, lam_sub, logits_sub)
+            lam_ins, logits_ins, lam_del, lam_sub, logits_sub = model_out
+            pi_type = None  # not used for base
+
+        new_batch = []
+        max_len_this_round = 0
+
+        for b in range(B):
+            seq = x[b]
+            valid = (seq != pad_id)
+            tokens = seq[valid].tolist()
+
+            # If sequence somehow became empty, reinsert BOS/EOS
+            if len(tokens) == 0:
+                tokens = [bos_id, eos_id]
+
+            # Find EOS position (default to last if missing)
+            try:
+                eos_pos = tokens.index(eos_id)
+            except ValueError:
+                eos_pos = len(tokens) - 1
+
+            Lb = len(tokens)
+
+            delete_mask = [False] * Lb
+            sub_tokens = [None] * Lb
+            ins_tokens = [None] * Lb
+
+            if is_reparameterized and not convert_to_vanilla_outputs:
+                # ----- Reparameterized CTMC branch: use lam_total + π_type directly -----
+                lam_total_b = lam_total[b, :Lb]          # (Lb,)
+                pi_b = pi_type[b, :Lb, :]                # (Lb, 3)
+                logits_ins_b = logits_ins[b, :Lb, :]     # (Lb, V)
+                logits_sub_b = logits_sub[b, :Lb, :]     # (Lb, V)
+
+                for j in range(Lb):
+                    tok_j = tokens[j]
+
+                    # π_type components
+                    pi_ins = float(pi_b[j, 0].item())
+                    pi_del = float(pi_b[j, 1].item())
+                    pi_sub = float(pi_b[j, 2].item())
+                    lam_tot_ij = float(lam_total_b[j].item())
+
+                    # -------- Insertion event at position j --------
+                    if j < eos_pos and lam_tot_ij > 0.0 and pi_ins > 0.0:
+                        p_ins = h * lam_tot_ij * pi_ins
+                        p_ins = min(p_ins, 1.0)
+                        if p_ins > 0.0 and torch.rand(1, device=device).item() < p_ins:
+                            ins_tok = sample_token_from_logits(logits_ins_b[j])
+                            ins_tokens[j] = ins_tok
+
+                    # -------- Delete/substitute event at position j --------
+                    if tok_j == bos_id or tok_j == eos_id:
+                        continue  # never delete/sub BOS/EOS
+
+                    pi_ds = pi_del + pi_sub
+                    if lam_tot_ij <= 0.0 or pi_ds <= 0.0:
+                        continue
+
+                    lam_ds = lam_tot_ij * pi_ds
+                    p_ds = h * lam_ds
+                    p_ds = min(p_ds, 1.0)
+                    if p_ds <= 0.0:
+                        continue
+
+                    if torch.rand(1, device=device).item() < p_ds:
+                        # A delete/sub event occurs; choose which
+                        p_del_given = pi_del / pi_ds
+                        choose_del = (torch.rand(1, device=device).item() < p_del_given)
+
+                        if choose_del:
+                            delete_mask[j] = True
+                            ins_tokens[j] = None
+                            sub_tokens[j] = None
+                        else:
+                            sub_tok = sample_token_from_logits(logits_sub_b[j])
+                            sub_tokens[j] = sub_tok
+
+            else:
+                # ----- Base CTMC branch (or reparam+vanilla with lam_ins/lam_del/lam_sub) -----
+                lam_ins_b = lam_ins[b, :Lb]
+                lam_del_b = lam_del[b, :Lb]
+                lam_sub_b = lam_sub[b, :Lb]
+                logits_ins_b = logits_ins[b, :Lb, :]
+                logits_sub_b = logits_sub[b, :Lb, :]
+
+                for j in range(Lb):
+                    tok_j = tokens[j]
+
+                    # -------- Insertion event at position j --------
+                    if j < eos_pos:
+                        lam_ij = float(lam_ins_b[j].item())
+                        if lam_ij > 0.0:
+                            p_ins = h * lam_ij
+                            p_ins = min(p_ins, 1.0)
+                            if p_ins > 0.0 and torch.rand(1, device=device).item() < p_ins:
+                                ins_tok = sample_token_from_logits(logits_ins_b[j])
+                                ins_tokens[j] = ins_tok
+
+                    # -------- Delete/substitute event at position j --------
+                    if tok_j == bos_id or tok_j == eos_id:
+                        continue
+
+                    lam_del_ij = float(lam_del_b[j].item())
+                    lam_sub_ij = float(lam_sub_b[j].item())
+                    lam_ds = lam_del_ij + lam_sub_ij
+                    if lam_ds <= 0.0:
+                        continue
+
+                    p_ds = h * lam_ds
+                    p_ds = min(p_ds, 1.0)
+                    if p_ds <= 0.0:
+                        continue
+
+                    if torch.rand(1, device=device).item() < p_ds:
+                        if lam_del_ij == 0.0:
+                            choose_del = False
+                        elif lam_sub_ij == 0.0:
+                            choose_del = True
+                        else:
+                            p_del_given = lam_del_ij / lam_ds
+                            choose_del = (torch.rand(1, device=device).item() < p_del_given)
+
+                        if choose_del:
+                            delete_mask[j] = True
+                            ins_tokens[j] = None
+                            sub_tokens[j] = None
+                        else:
+                            sub_tok = sample_token_from_logits(logits_sub_b[j])
+                            sub_tokens[j] = sub_tok
+
+            # -------- Apply all edits simultaneously (left-to-right) --------
+            new_tokens = []
+            for j in range(Lb):
+                tok_j = tokens[j]
+
+                if delete_mask[j]:
+                    pass
+                elif sub_tokens[j] is not None:
+                    new_tokens.append(sub_tokens[j])
+                else:
+                    new_tokens.append(tok_j)
+
+                if ins_tokens[j] is not None:
+                    new_tokens.append(ins_tokens[j])
+
+            # Ensure EOS is present
+            if eos_id not in new_tokens:
+                new_tokens.append(eos_id)
+
+            # Enforce max length cap
+            if max_len_cap is not None and len(new_tokens) > max_len_cap:
+                new_tokens = new_tokens[:max_len_cap]
+                if new_tokens[-1] != eos_id:
+                    new_tokens[-1] = eos_id
+
+            new_seq = torch.tensor(new_tokens, device=device, dtype=torch.long)
+            new_batch.append(new_seq)
+            max_len_this_round = max(max_len_this_round, new_seq.size(0))
+
+        x_next = x.new_full((B, max_len_this_round), pad_id)
+        for b, seq_b in enumerate(new_batch):
+            x_next[b, :seq_b.size(0)] = seq_b
+
+        x = x_next
+
+    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
+        model_output = model(x_t=x, mask=mask, t=t)
+        
+        # Handle both base models (5 values) and ReparameterizedProteinEditFlowModel (8 values)
+        # ReparameterizedSMILESEditFlowModel returns 4 values, but we don't use it here
+        if len(model_output) == 8:
+            # ReparameterizedProteinEditFlowModel: returns all decomposed rates + total + type info
+            lam_ins, logits_ins, lam_del, lam_sub, logits_sub, lam_total, logits_type, pi_type = model_output
+        elif len(model_output) == 5:
+            # Base model: (lam_ins, logits_ins, lam_del, lam_sub, logits_sub)
+            lam_ins, logits_ins, lam_del, lam_sub, logits_sub = model_output
+            lam_total = None  # Not used in base model path
+            logits_type = None
+            pi_type = None
+        else:
+            raise ValueError(f"Unexpected model output length: {len(model_output)}. Expected 5 (base) or 8 (ReparameterizedProteinEditFlowModel)")
+
+        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

cas9/objectives.py

@@ -0,0 +1,1248 @@
+# import warnings
+# import logging
+# warnings.filterwarnings("ignore", message="to-Python converter.*already registered", category=RuntimeWarning)
+# warnings.filterwarnings("ignore", message="pkg_resources is deprecated", category=UserWarning)
+# warnings.filterwarnings("ignore", category=FutureWarning)
+# warnings.filterwarnings("ignore", category=DeprecationWarning)
+# warnings.filterwarnings("ignore", category=UserWarning)
+# from rdkit import RDLogger, rdBase
+# RDLogger.DisableLog("rdApp.*")
+# logging.basicConfig(level=logging.ERROR)
+# logging.getLogger("chemprop").setLevel(logging.ERROR)
+# logging.getLogger("hyperopt").setLevel(logging.ERROR)
+# rdBase.DisableLog('rdApp.error')
+
+import torch
+import torch.nn.functional as F
+import logging
+from rdkit import Chem
+# Make chemprop imports conditional to avoid import errors when not needed
+try:
+    from chemprop import data, models
+    # Try to import featurizers separately as it may not exist in all versions
+    try:
+        from chemprop import featurizers
+    except (ImportError, AttributeError):
+        featurizers = None
+    CHEMPROP_AVAILABLE = True
+except (ImportError, AttributeError) as e:
+    CHEMPROP_AVAILABLE = False
+    # Create dummy objects to avoid NameError if someone tries to use them
+    data = None
+    featurizers = None
+    models = None
+from lightning import pytorch as pl
+
+# ---- QUIET MODE (put these lines at the top of your script) ----
+import os, warnings, logging
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+os.environ["TRANSFORMERS_VERBOSITY"] = "error"
+os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
+
+# Make PyTorch stop suggesting Tensor Core settings
+torch.set_float32_matmul_precision("high")
+
+# Silence Python warnings (fine-tune as needed)
+warnings.filterwarnings("ignore", category=FutureWarning)
+warnings.filterwarnings("ignore", category=UserWarning, message=r".*predict_dataloader.*many workers.*")
+warnings.filterwarnings("ignore", message=r"Dropping last batch of size .*")
+
+# Quiet RDKit
+from rdkit import RDLogger
+RDLogger.DisableLog("rdApp.*")
+
+# Quiet common loggers (Lightning, Chemprop, etc.)
+logging.basicConfig(level=logging.ERROR, force=True)
+for name in [
+    "lightning", "pytorch_lightning", "lightning.pytorch",
+    "chemprop", "rdkit", "urllib3", "torch"
+]:
+    logging.getLogger(name).setLevel(logging.ERROR)
+# ---------------------------------------------------------------
+
+
+# from admet_ai import ADMETModel
+import sys
+import os
+import numpy as np
+from transformers import AutoModelForMaskedLM
+import warnings
+import numpy as np
+import torch.nn as nn
+from rdkit import Chem
+from collections import defaultdict
+import pdb
+
+import math
+import sys
+import os
+from pathlib import Path
+
+# Resolve cas_predictor / protein2pam through the repo-root path helpers.
+# Both are optional: if absent, the dependent objectives raise a clear error
+# at construction time rather than at import time.
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from paths import cas_predictor_root, protein2pam_root  # noqa: E402
+
+try:
+    cas_predictor_path = cas_predictor_root()
+    if str(cas_predictor_path) not in sys.path:
+        sys.path.insert(0, str(cas_predictor_path))
+except FileNotFoundError:
+    cas_predictor_path = Path("/nonexistent/cas_predictor")
+
+try:
+    protein2pam_path = protein2pam_root()
+    if str(protein2pam_path) not in sys.path:
+        sys.path.insert(0, str(protein2pam_path))
+except FileNotFoundError:
+    protein2pam_path = Path("/nonexistent/protein2pam")
+
+try:
+    # Import using importlib to avoid conflicts with editflows/model package
+    import importlib.util
+    import importlib
+    
+    # Temporarily add cas_predictor to sys.path at the front to ensure it takes precedence
+    cas_predictor_str = str(cas_predictor_path)
+    original_path = sys.path.copy()
+    if cas_predictor_str not in sys.path:
+        sys.path.insert(0, cas_predictor_str)
+    
+    try:
+        # Load model.py directly - use cas_predictor.model as the module name
+        model_file = cas_predictor_path / "model.py"
+        spec = importlib.util.spec_from_file_location("cas_predictor.model", model_file)
+        cas_model = importlib.util.module_from_spec(spec)
+        # Set __file__ to help with relative imports
+        cas_model.__file__ = str(model_file)
+        # Register the module in sys.modules before exec_module so imports work
+        sys.modules['cas_predictor.model'] = cas_model
+        sys.modules['model'] = cas_model  # Also register as 'model' for lightning_module imports
+        spec.loader.exec_module(cas_model)
+        Cas9Classifier = cas_model.Cas9Classifier
+        
+        # Load lightning_module.py - it will import from 'model' which should now resolve correctly
+        lightning_file = cas_predictor_path / "lightning_module.py"
+        spec_lightning = importlib.util.spec_from_file_location("cas_predictor.lightning_module", lightning_file)
+        cas_lightning = importlib.util.module_from_spec(spec_lightning)
+        # Set __file__ to help with relative imports
+        cas_lightning.__file__ = str(lightning_file)
+        # Register the module in sys.modules before exec_module
+        sys.modules['cas_predictor.lightning_module'] = cas_lightning
+        spec_lightning.loader.exec_module(cas_lightning)
+        Cas9ClassifierModule = cas_lightning.Cas9ClassifierModule
+        
+        import yaml
+        try:
+            from easydict import EasyDict as edict
+        except ImportError:
+            # Fallback: use regular dict if easydict not available
+            class EasyDict(dict):
+                def __init__(self, *args, **kwargs):
+                    super().__init__(*args, **kwargs)
+                def __getattr__(self, key):
+                    try:
+                        return self[key]
+                    except KeyError:
+                        raise AttributeError(key)
+                def __setattr__(self, key, value):
+                    self[key] = value
+            edict = EasyDict
+        CAS9_PREDICTOR_AVAILABLE = True
+    finally:
+        # Restore original sys.path
+        sys.path[:] = original_path
+except (ImportError, Exception) as e:
+    print(f"Warning: Could not import Cas9 predictor modules: {e}")
+    import traceback
+    traceback.print_exc()
+    CAS9_PREDICTOR_AVAILABLE = False
+
+# SMARTS patterns
+_AMIDE_SMARTS        = Chem.MolFromSmarts("[CX3](=[OX1])[NX3]")                     # C(=O)-N
+_CARBONYL_C_SMARTS   = Chem.MolFromSmarts("[CX3](=[OX1])")                          # carbonyl C
+_DIPEPTIDE_SMARTS    = Chem.MolFromSmarts("[CX3](=[OX1])N[#6X4][CX3](=[OX1])N")     # amide–C(sp3)–amide
+
+def _amide_bond_indices(mol, ignore_ring_amides=False):
+    ids = set()
+    for c_idx, _, n_idx in mol.GetSubstructMatches(_AMIDE_SMARTS):
+        b = mol.GetBondBetweenAtoms(c_idx, n_idx)
+        if b and b.GetBondType() == Chem.rdchem.BondType.SINGLE:
+            if ignore_ring_amides and b.IsInRing():
+                continue
+            ids.add(b.GetIdx())
+    return ids
+
+def _carbonyl_c_indices(mol):
+    return {m[0] for m in mol.GetSubstructMatches(_CARBONYL_C_SMARTS)}
+
+def _carbonyl_neighbor_stats(mol, c_indices):
+    stats = {"total": 0, "with_N": 0, "with_O": 0, "with_S": 0, "pure_amide": 0}
+    for c_idx in c_indices:
+        c = mol.GetAtomWithIdx(c_idx)
+        stats["total"] += 1
+        hasN = hasO = hasS = False
+        for b in c.GetBonds():
+            if b.GetBondType() != Chem.rdchem.BondType.SINGLE:
+                continue
+            z = b.GetOtherAtom(c).GetAtomicNum()
+            if   z == 7:  hasN = True
+            elif z == 8:  hasO = True
+            elif z == 16: hasS = True
+        stats["with_N"] += int(hasN)
+        stats["with_O"] += int(hasO)
+        stats["with_S"] += int(hasS)
+        if hasN and not (hasO or hasS):
+            stats["pure_amide"] += 1
+    return stats
+
+def _adjacent_amide_pairs(mol):
+    # Count distinct amide–C(sp3)–amide windows (dedup by central carbon)
+    centers = set()
+    for match in mol.GetSubstructMatches(_DIPEPTIDE_SMARTS):
+        centers.add(match[3])  # central sp3 carbon index
+    return len(centers)
+
+def analyze_peptide_likeness(smiles: str,
+                             ignore_ring_amides: bool = False,
+                             amide_density_target: float = 0.12):
+    """
+    Compute peptide-likeness metrics and a continuous score in [0,1].
+    - ignore_ring_amides=False: include macro/cyclic peptides by default.
+    - amide_density_target: amide-per-atom density to hit score~1 for peptides
+      (~0.10–0.15 works well; default 0.12 ≈ 1 amide per ~8 heavy atoms).
+    """
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        raise ValueError(f"Invalid SMILES: {smiles}")
+
+    n_heavy_atoms  = mol.GetNumAtoms()
+    n_heavy_bonds  = mol.GetNumBonds()
+
+    amide_bonds    = _amide_bond_indices(mol, ignore_ring_amides=ignore_ring_amides)
+    n_amide_bonds  = len(amide_bonds)
+
+    carbonyl_cs    = _carbonyl_c_indices(mol)
+    cstats         = _carbonyl_neighbor_stats(mol, carbonyl_cs)
+    total_carb     = max(1, cstats["total"])
+
+    # Core features
+    f1 = cstats["with_N"] / total_carb  # acyl-N fraction ∈ [0,1]
+    amide_per_atom = n_amide_bonds / max(1, n_heavy_atoms)
+    f2 = min(1.0, amide_per_atom / max(1e-8, amide_density_target))  # saturate at 1
+    n_adjacent = _adjacent_amide_pairs(mol)
+    f3 = 1.0 - math.exp(-n_adjacent)  # 0, 0.63, 0.86, 0.95, ... as pairs increase
+
+    # Penalty for non-peptidic carbonyls (carbamates/anhydrides/thioesters)
+    pure_amide_fraction = cstats["pure_amide"] / total_carb
+    penalty = 1.0 - pure_amide_fraction  # 0 (all pure amide) … 1 (no pure amide)
+
+    # Final heuristic score in [0,1]
+    score = 0.55 * f1 + 0.25 * f2 + 0.20 * f3 - 0.25 * penalty
+    score = max(0.0, min(1.0, score))
+
+    return {
+        "n_heavy_atoms": n_heavy_atoms,
+        "n_heavy_bonds": n_heavy_bonds,
+        "n_carbonyls": cstats["total"],
+        "n_amide_bonds": n_amide_bonds,
+        "amide_bond_ratio_all_bonds": n_amide_bonds / max(1, n_heavy_bonds),
+        "acyl_N_fraction": f1,
+        "pure_amide_fraction": pure_amide_fraction,
+        "amide_per_atom": amide_per_atom,
+        "n_adjacent_amide_pairs": n_adjacent,
+        "peptide_likeness": score,  # <<< continuous score in [0,1]
+    }
+
+
+def score_combination(ratios, scores, admet_scores):
+    high_mask = ratios > 0.6
+    low_mask  = ratios < 0.1
+    mid_mask  = ~(high_mask | low_mask)
+
+    # start with zeros
+    final_scores = torch.zeros_like(scores)
+
+    # high-peptide: use peptide scores
+    final_scores[high_mask] = scores[high_mask]
+
+    # low-peptide: use admet scores
+    final_scores[low_mask] = admet_scores[low_mask]
+
+    # middle band: linear blend
+    if mid_mask.any():
+        r_mid = ratios[mid_mask]
+        alpha = (r_mid - 0.1) / 0.5          # in [0, 1]
+        blended = alpha * scores[mid_mask] + (1 - alpha) * admet_scores[mid_mask]
+        final_scores[mid_mask] = blended
+
+def detokenize_output(x, cfg, tokenizer, bos_id, eos_id, pad_id):
+    """
+    Convert a single generated sequence (1, L) back to string.
+    """
+    seq = x[0].tolist()
+    # strip padding
+    seq = [tok for tok in seq if tok != pad_id]
+    # strip BOS/EOS
+    if len(seq) > 0 and seq[0] == bos_id:
+        seq = seq[1:]
+    if len(seq) > 0 and seq[-1] == eos_id:
+        seq = seq[:-1]
+
+    if cfg.task == 'protein':
+        # esm tokenizer has batch_decode
+        return tokenizer.batch_decode([seq], skip_special_tokens=True)[0]
+    elif cfg.task in ('smiles', 'selfies'):
+        return tokenizer.decode(seq)
+    else:
+        return " ".join(map(str, seq))
+
+
+
+class Cas9Classification:
+    """
+    Cas9 binary classification objective.
+    
+    This objective works with protein sequences (not SMILES).
+    It uses ESM-2 embeddings and a trained Cas9Classifier to predict
+    whether a protein sequence is Cas9-like.
+    
+    Note: This objective expects protein sequences, not SMILES sequences.
+    If used with SMILES-based generation, you'll need to convert SMILES to proteins first.
+    """
+    def __init__(self, device, checkpoint_path=None, config_path=None, shared_esm_model=None):
+        if not CAS9_PREDICTOR_AVAILABLE:
+            raise ImportError(
+                "Cas9 predictor modules not available. "
+                "Ensure cas_predictor is in the correct location."
+            )
+        
+        self.device = device
+        
+        # Default checkpoint path if not provided
+        if checkpoint_path is None:
+            from paths import cas9_classifier_ckpt
+            checkpoint_path = str(cas9_classifier_ckpt())
+        
+        # Load config
+        config = None
+        if config_path:
+            with open(config_path, "r") as f:
+                config_dict = yaml.safe_load(f)
+                config = edict(config_dict)
+        else:
+            # Try to load from checkpoint
+            try:
+                checkpoint = torch.load(checkpoint_path, map_location="cpu")
+                if "hyper_parameters" in checkpoint:
+                    config_dict = checkpoint["hyper_parameters"]
+                    if isinstance(config_dict, dict):
+                        config = edict(config_dict)
+            except Exception as e:
+                print(f"Warning: Could not load config from checkpoint: {e}")
+        
+        # If config still not found, create a default config with reasonable defaults
+        if config is None:
+            print("Warning: Could not load config from checkpoint or config_path. Using default config.")
+            config = edict({
+                "model": edict({
+                    "esm_model_name": "facebook/esm2_t33_650M_UR50D",
+                    "freeze_esm": True,
+                    "hidden_size": 512,
+                    "dropout": 0.1,
+                    "pooling_type": "mean"
+                }),
+                "optim": edict({
+                    "lr": 5e-5,
+                    "beta1": 0.9,
+                    "beta2": 0.999,
+                    "eps": 1e-8,
+                    "weight_decay": 0.01,
+                    "warmup_ratio": 0.1
+                })
+            })
+        
+        # Load model from checkpoint
+        # If shared_esm_model is provided, we need to manually instantiate and load
+        # to avoid creating a duplicate ESM model
+        if shared_esm_model is not None:
+            # Manually instantiate with shared ESM model
+            self.model = Cas9ClassifierModule(config, shared_esm_model=shared_esm_model)
+            # Load checkpoint state dict (excluding ESM weights since we're using shared model)
+            checkpoint = torch.load(checkpoint_path, map_location="cpu")
+            state_dict = checkpoint["state_dict"]
+            # Filter out ESM weights to avoid loading them (we're using shared model)
+            filtered_state_dict = {k: v for k, v in state_dict.items() 
+                                 if not k.startswith("model.esm_emb")}
+            self.model.load_state_dict(filtered_state_dict, strict=False)
+            print(f"Loaded Cas9Classifier checkpoint with shared ESM model. "
+                  f"Filtered out {len(state_dict) - len(filtered_state_dict)} ESM parameters.")
+        else:
+            # Original behavior: use load_from_checkpoint (creates its own ESM model)
+            self.model = Cas9ClassifierModule.load_from_checkpoint(
+                checkpoint_path,
+                config=config,
+                strict=False,
+            )
+        
+        self.model.eval()
+        self.model.to(device)
+        
+        # Load ESM tokenizer (same approach as BindingAffinity)
+        # We only need the alphabet for tokenization, not the full model
+        # If we have a shared ESM model, we can get the tokenizer from it
+        # Otherwise, load the ESM model just to get the alphabet
+        if shared_esm_model is not None:
+            # Use the shared ESM model's tokenizer if available
+            # For transformers.EsmModel, we need to use EsmTokenizer separately
+            from transformers import EsmTokenizer
+            esm_model_name = getattr(config.model, "esm_model_name", "facebook/esm2_t33_650M_UR50D")
+            tokenizer = EsmTokenizer.from_pretrained(esm_model_name)
+            # Create a simple batch converter function
+            def batch_converter(data):
+                # data is list of (name, sequence) tuples
+                sequences = [seq for _, seq in data]
+                encoded = tokenizer(sequences, padding=True, return_tensors="pt", add_special_tokens=True)
+                # Return format: (names, sequences, tokens)
+                names = [name for name, _ in data]
+                return names, sequences, encoded["input_ids"]
+            self.batch_converter = batch_converter
+        else:
+            # Use HuggingFace transformers API instead of old esm.pretrained API
+            from transformers import EsmTokenizer
+            esm_model_name = getattr(config.model, "esm_model_name", "facebook/esm2_t33_650M_UR50D")
+            tokenizer = EsmTokenizer.from_pretrained(esm_model_name)
+            # Create a simple batch converter function
+            def batch_converter(data):
+                # data is list of (name, sequence) tuples
+                sequences = [seq for _, seq in data]
+                encoded = tokenizer(sequences, padding=True, return_tensors="pt", add_special_tokens=True)
+                # Return format: (names, sequences, tokens)
+                names = [name for name, _ in data]
+                return names, sequences, encoded["input_ids"]
+            self.batch_converter = batch_converter
+    
+    def get_scores(self, protein_seqs):
+        """
+        Get Cas9 classification scores for protein sequences.
+        
+        Args:
+            protein_seqs: List of protein sequence strings (amino acid sequences) or single string
+        
+        Returns:
+            scores: List of probabilities (0-1) that each sequence is Cas9-like
+        """
+        # Handle single string input
+        if isinstance(protein_seqs, str):
+            protein_seqs = [protein_seqs]
+        
+        if not protein_seqs:
+            return []
+        
+        scores = []
+        
+        with torch.no_grad():
+            # Tokenize sequences using ESM batch converter
+            data = [(f"seq_{i}", seq) for i, seq in enumerate(protein_seqs)]
+            _, _, batch_tokens = self.batch_converter(data)
+            batch_tokens = batch_tokens.to(self.device)
+            
+            # Create attention mask (True for valid tokens, False for padding)
+            # ESM uses 1 for padding, so we check for non-padding tokens
+            # Convert to bool tensor explicitly (model expects BoolTensor)
+            attention_mask = (batch_tokens != 1).bool().to(self.device)  # 1 is ESM's padding token
+            
+            # Debug: Print shapes and values for first sequence if debugging
+            if len(protein_seqs) == 1 and len(protein_seqs[0]) < 2000:  # Only debug for single sequences and reasonable lengths
+                import os
+                if os.environ.get("DEBUG_CAS9", "0") == "1":
+                    print(f"[DEBUG] Cas9 classifier - sequence length: {len(protein_seqs[0])}")
+                    print(f"[DEBUG] Cas9 classifier - batch_tokens shape: {batch_tokens.shape}")
+                    print(f"[DEBUG] Cas9 classifier - attention_mask shape: {attention_mask.shape}")
+                    print(f"[DEBUG] Cas9 classifier - attention_mask sum (valid tokens): {attention_mask.sum().item()}")
+                    print(f"[DEBUG] Cas9 classifier - batch_tokens min/max: {batch_tokens.min().item()}/{batch_tokens.max().item()}")
+            
+            # Forward pass
+            logits = self.model(batch_tokens, attention_mask)
+            
+            # Debug: Print logits if debugging
+            if len(protein_seqs) == 1 and len(protein_seqs[0]) < 2000:
+                import os
+                if os.environ.get("DEBUG_CAS9", "0") == "1":
+                    print(f"[DEBUG] Cas9 classifier - logits shape: {logits.shape}")
+                    print(f"[DEBUG] Cas9 classifier - logits value: {logits.item()}")
+            
+            probs = torch.sigmoid(logits).cpu()
+            
+            # Debug: Print probabilities if debugging
+            if len(protein_seqs) == 1 and len(protein_seqs[0]) < 2000:
+                import os
+                if os.environ.get("DEBUG_CAS9", "0") == "1":
+                    print(f"[DEBUG] Cas9 classifier - probabilities: {probs.tolist()}")
+            
+            scores = probs.tolist()
+        
+        return scores
+    
+    def __call__(self, protein_tokens, protein_seqs):
+        """
+        Objective call interface.
+        
+        Args:
+            protein_tokens: Unused (kept for interface compatibility)
+            protein_seqs: List of protein sequence strings
+        
+        Returns:
+            Tuple of ('cas9', scores) where scores is a list of probabilities
+        """
+        scores = self.get_scores(protein_seqs)
+        return 'cas9', scores
+
+
+class DeletionCount:
+    """
+    Objective that maximizes the number of deletions (length reduction) in protein sequences.
+    
+    This objective computes the difference between the original sequence length and
+    the current sequence length, normalized by a maximum deletion percentage to produce
+    a value between 0 and 1.
+    
+    Args:
+        original_seq: The original protein sequence string (before any edits)
+        max_deletion_percentage: Maximum deletion percentage (0-1). The objective value will be computed as
+                                 (deletion_count / (original_length * max_deletion_percentage)), clamped to [0, 1].
+                                 If None, defaults to 1.0 (100% deletion possible).
+    """
+    def __init__(self, original_seq, max_deletion_percentage=None):
+        # Store original sequence length (excluding spaces)
+        self.original_length = len(original_seq.replace(' ', ''))
+        # Set max_deletion_percentage to 1.0 (100%) if not provided
+        if max_deletion_percentage is None:
+            self.max_deletion_percentage = 1.0
+        else:
+            if max_deletion_percentage <= 0 or max_deletion_percentage > 1.0:
+                raise ValueError(f"max_deletion_percentage must be in (0, 1], got {max_deletion_percentage}")
+            self.max_deletion_percentage = float(max_deletion_percentage)
+        
+        # Calculate max_deletion as absolute value for backward compatibility
+        self.max_deletion = self.original_length * self.max_deletion_percentage
+    
+    def __call__(self, protein_tokens, protein_seqs):
+        """
+        Objective call interface.
+        
+        Args:
+            protein_tokens: Unused (kept for interface compatibility)
+            protein_seqs: List of protein sequence strings
+        
+        Returns:
+            Tuple of ('deletion_count', scores) where scores is a list of normalized deletion
+            percentages (deletion_count / (original_length * max_deletion_percentage)), clamped to [0, 1]
+        """
+        # Handle single string input
+        if isinstance(protein_seqs, str):
+            protein_seqs = [protein_seqs]
+        
+        if not protein_seqs:
+            return 'deletion_count', []
+        
+        scores = []
+        for seq in protein_seqs:
+            # Remove spaces and compute current length
+            current_length = len(seq.replace(' ', ''))
+            # Number of deletions = original_length - current_length
+            deletion_count = self.original_length - current_length
+            # Normalize by max_deletion and clamp to [0, 1]
+            normalized_score = max(0.0, min(1.0, deletion_count / self.max_deletion))
+            scores.append(float(normalized_score))
+        
+        return 'deletion_count', scores
+
+
+class PAMMatching:
+    """
+    PAM matching objective using temperature-scaled log probabilities.
+    
+    This objective uses the HuggingFace protein2pam model (cas9_full) to predict PAM logits
+    for protein sequences. The score is computed using temperature-scaled log probabilities:
+    at each target position (non-N), we compute log P(target_nucleotide | position) after
+    applying temperature scaling to the logits. The final score is the geometric mean of
+    probabilities (exp of mean log prob) over target positions.
+    
+    This approach amplifies small improvements, making them more visible in the weighted sum
+    during optimization, which helps when candidates show only marginal improvements.
+    
+    Args:
+        device: torch device
+        target_pam: Target PAM sequence string of length 10 (e.g., "NGGNNNNNNN")
+        model_name: HuggingFace model identifier (default: "Profluent-Bio/protein2pam-cas9_full")
+        sigmoid_temperature: Temperature for scaling logits before softmax. Lower values (0.1-0.3)
+                            sharpen the distribution, making small improvements more visible.
+                            Default 0.2. (Note: parameter name kept for backward compatibility)
+        pam_prediction_temperature: Temperature for PAM prediction (in predict_pam method).
+                                    Values < 1.0 make distributions sharper (more confident),
+                                    > 1.0 make them softer. Default 1.0 (no temperature scaling).
+                                    Note: This only affects prediction, not scoring.
+        shared_esm_model: Optional shared ESM model (transformers.EsmModel) to use instead of
+                         loading a separate ESM backbone. This reduces memory usage when multiple
+                         models need ESM embeddings. Default None (loads its own ESM model).
+    """
+    def __init__(self, device, target_pam, model_name="Profluent-Bio/protein2pam-cas9_full", sigmoid_temperature=0.2, use_entropy_for_n_positions=True, use_ce_loss=False, pam_min_confidence=0.55, pam_prediction_temperature=1.0, shared_esm_model=None):
+        self.device = device
+        self.target_pam = target_pam.upper()
+        
+        # Validate target PAM length
+        if len(self.target_pam) != 10:
+            raise ValueError(f"Target PAM must be exactly 10 nucleotides, got {len(self.target_pam)}")
+        
+        # Validate nucleotides
+        valid_nucleotides = set('ACGTN')
+        if not all(nuc in valid_nucleotides for nuc in self.target_pam):
+            raise ValueError(f"Target PAM contains invalid nucleotides. Only ACGTN allowed.")
+        
+        # Convert target PAM to class indices (ACGT = 0,1,2,3, N = -1 for masking)
+        nucleotides = ['A', 'C', 'G', 'T']
+        self.target_indices = []
+        self.target_mask = []
+        for nuc in self.target_pam:
+            if nuc == 'N':
+                # N means any nucleotide - we'll mask this position in loss calculation
+                self.target_indices.append(-1)  # Placeholder, will be masked
+                self.target_mask.append(False)
+            else:
+                self.target_indices.append(nucleotides.index(nuc))
+                self.target_mask.append(True)
+        
+        self.target_indices = torch.tensor(self.target_indices, device=device, dtype=torch.long)
+        self.target_mask = torch.tensor(self.target_mask, device=device, dtype=torch.bool)
+        self.sigmoid_temperature = float(sigmoid_temperature)
+        self.use_entropy_for_n_positions = bool(use_entropy_for_n_positions)
+        self.use_ce_loss = bool(use_ce_loss)
+        self.pam_min_confidence = float(pam_min_confidence)
+        self.pam_prediction_temperature = float(pam_prediction_temperature)
+        
+        # Load HuggingFace model. The protein2pam package root is resolved by
+        # paths.protein2pam_root() (honours $PROTEIN2PAM_ROOT).
+        
+        try:
+            protein2pam_path_resolved = protein2pam_root()
+        except FileNotFoundError as exc:
+            raise ImportError(str(exc)) from exc
+        
+        # Ensure the path is in sys.path
+        protein2pam_path_str = str(protein2pam_path_resolved)
+        if protein2pam_path_str not in sys.path:
+            sys.path.insert(0, protein2pam_path_str)
+            print(f"[DEBUG] Added protein2pam path to sys.path: {protein2pam_path_str}")
+        
+        try:
+            # Try importing directly from huggingface submodule to avoid protein2pam/__init__.py
+            # which imports PAMOracle (requires torch at import time)
+            import importlib.util
+            
+            # Import the modules directly without going through __init__.py
+            huggingface_dir = protein2pam_path_resolved / "protein2pam" / "huggingface"
+            
+            # Load configuration_esm
+            config_file = huggingface_dir / "configuration_esm.py"
+            if config_file.exists():
+                spec_config = importlib.util.spec_from_file_location("protein2pam.huggingface.configuration_esm", config_file)
+                if spec_config and spec_config.loader:
+                    # Create package structure
+                    if 'protein2pam' not in sys.modules:
+                        sys.modules['protein2pam'] = type(sys)('protein2pam')
+                    if 'protein2pam.huggingface' not in sys.modules:
+                        sys.modules['protein2pam.huggingface'] = type(sys)('protein2pam.huggingface')
+                    
+                    config_module = importlib.util.module_from_spec(spec_config)
+                    sys.modules['protein2pam.huggingface.configuration_esm'] = config_module
+                    spec_config.loader.exec_module(config_module)
+            
+            # Load modeling_esm (depends on configuration_esm)
+            modeling_file = huggingface_dir / "modeling_esm.py"
+            if modeling_file.exists():
+                spec_modeling = importlib.util.spec_from_file_location("protein2pam.huggingface.modeling_esm", modeling_file)
+                if spec_modeling and spec_modeling.loader:
+                    modeling_module = importlib.util.module_from_spec(spec_modeling)
+                    sys.modules['protein2pam.huggingface.modeling_esm'] = modeling_module
+                    spec_modeling.loader.exec_module(modeling_module)
+            
+            # Load tokenizer
+            tokenizer_file = huggingface_dir / "tokenizer.py"
+            if tokenizer_file.exists():
+                spec_tokenizer = importlib.util.spec_from_file_location("protein2pam.huggingface.tokenizer", tokenizer_file)
+                if spec_tokenizer and spec_tokenizer.loader:
+                    tokenizer_module = importlib.util.module_from_spec(spec_tokenizer)
+                    sys.modules['protein2pam.huggingface.tokenizer'] = tokenizer_module
+                    spec_tokenizer.loader.exec_module(tokenizer_module)
+            
+            # Now import from the loaded modules
+            EsmForSequenceClassification = sys.modules['protein2pam.huggingface.modeling_esm'].EsmForSequenceClassification
+            get_tokenizer = sys.modules['protein2pam.huggingface.tokenizer'].get_tokenizer
+            
+            self.tokenizer = get_tokenizer()
+            print(f"Loading PAM prediction model: {model_name}...")
+            # Load model without device_map (device_map expects string/dict, not torch.device)
+            # Then move to device manually for compatibility with different transformers versions
+            self.model = EsmForSequenceClassification.from_pretrained(model_name)
+            
+            # If shared_esm_model is provided, replace the ESM backbone to share parameters
+            if shared_esm_model is not None:
+                print(f"Replacing PAM model's ESM backbone with shared ESM model...")
+                # EsmForSequenceClassification typically has an 'esm' attribute containing the backbone
+                if hasattr(self.model, 'esm'):
+                    self.model.esm = shared_esm_model
+                    print(f"Successfully replaced PAM model's ESM backbone with shared model.")
+                elif hasattr(self.model, 'model') and hasattr(self.model.model, 'esm'):
+                    # Some models wrap ESM in a 'model' attribute
+                    self.model.model.esm = shared_esm_model
+                    print(f"Successfully replaced PAM model's ESM backbone (via .model.esm) with shared model.")
+                else:
+                    print(f"Warning: Could not find ESM backbone in PAM model structure. Available attributes: {[attr for attr in dir(self.model) if not attr.startswith('_')]}")
+                    # Try to find it recursively
+                    for attr_name in ['esm', 'backbone', 'encoder']:
+                        if hasattr(self.model, attr_name):
+                            setattr(self.model, attr_name, shared_esm_model)
+                            print(f"Replaced {attr_name} with shared ESM model.")
+                            break
+            
+            self.model = self.model.to(device)
+            self.model.eval()
+            print(f"PAM prediction model loaded successfully.")
+        except ImportError as e:
+            # Show more details about the import error
+            import traceback
+            print(f"[DEBUG] protein2pam import error: {e}")
+            print(f"[DEBUG] sys.path entries containing 'protein2pam': {[p for p in sys.path if 'protein2pam' in p]}")
+            print(f"[DEBUG] Expected protein2pam path: {protein2pam_path_resolved}")
+            print(f"[DEBUG] Path exists: {protein2pam_path_resolved.exists() if protein2pam_path_resolved else False}")
+            traceback.print_exc()
+            raise ImportError(
+                f"protein2pam package not available: {e}. "
+                "Please ensure protein2pam is installed and in the correct location."
+            )
+        except Exception as e:
+            import traceback
+            print(f"[DEBUG] Failed to load PAM prediction model: {e}")
+            traceback.print_exc()
+            raise RuntimeError(f"Failed to load PAM prediction model: {e}")
+    
+    def get_pam_probability_distributions(self, protein_seqs):
+        """
+        Get detailed probability distributions for PAM predictions.
+        
+        Args:
+            protein_seqs: List of protein sequence strings
+        
+        Returns:
+            List of dictionaries, each containing:
+            - 'probabilities': (10, 4) tensor of probabilities for each position and nucleotide
+            - 'predicted_pam': Predicted PAM string
+            - 'per_position': List of dicts with position info (nucleotide probs, entropy, etc.)
+        """
+        if not protein_seqs:
+            return []
+        
+        # Handle single string input
+        if isinstance(protein_seqs, str):
+            protein_seqs = [protein_seqs]
+        
+        nucleotides = ['A', 'C', 'G', 'T']
+        results = []
+        
+        with torch.no_grad():
+            # Tokenize sequences
+            encodings = self.tokenizer.encode_batch(protein_seqs)
+            input_batch = dict(
+                input_ids=torch.tensor([encoding.ids for encoding in encodings], device=self.device),
+                attention_mask=torch.tensor([encoding.attention_mask for encoding in encodings], device=self.device),
+            )
+            
+            # Get PAM predictions
+            output = self.model(**input_batch)
+            logits = output.logits  # (batch_size, 10, 4)
+            probabilities = F.softmax(logits, dim=-1)  # (batch_size, 10, 4)
+            log_probs = F.log_softmax(logits, dim=-1)  # (batch_size, 10, 4)
+            
+            # Process each sequence
+            for i in range(probabilities.shape[0]):
+                probs = probabilities[i]  # (10, 4)
+                log_probs_seq = log_probs[i]  # (10, 4)
+                pam_seq = []
+                per_position = []
+                
+                for pos_idx in range(probs.shape[0]):
+                    pos_probs = probs[pos_idx, :]  # (4,)
+                    pos_log_probs = log_probs_seq[pos_idx, :]  # (4,)
+                    
+                    # Compute entropy for this position
+                    entropy = -(pos_probs * pos_log_probs).sum().item()
+                    max_entropy = math.log(4)  # Maximum entropy for uniform distribution
+                    normalized_entropy = entropy / max_entropy
+                    
+                    max_prob = pos_probs.max().item()
+                    max_idx = pos_probs.argmax().item()
+                    predicted_nuc = nucleotides[max_idx]
+                    
+                    # Create position info
+                    pos_info = {
+                        'position': pos_idx,
+                        'probabilities': {nuc: pos_probs[j].item() for j, nuc in enumerate(nucleotides)},
+                        'predicted': predicted_nuc,
+                        'max_probability': max_prob,
+                        'entropy': entropy,
+                        'normalized_entropy': normalized_entropy,
+                    }
+                    per_position.append(pos_info)
+                    
+                    # Predict nucleotide (use instance min_confidence)
+                    if max_prob < self.pam_min_confidence:
+                        pam_seq.append('N')
+                    else:
+                        pam_seq.append(predicted_nuc)
+                
+                results.append({
+                    'probabilities': probs.cpu(),
+                    'predicted_pam': ''.join(pam_seq),
+                    'per_position': per_position,
+                })
+        
+        return results
+    
+    def predict_pam(self, protein_seqs, min_confidence=None):
+        """
+        Predict PAM sequences for protein sequences, with support for 'N' when confidence is low.
+        
+        Uses temperature scaling (pam_prediction_temperature) to control distribution sharpness.
+        Lower temperatures (< 1.0) produce sharper distributions and more confident predictions.
+        
+        Args:
+            protein_seqs: List of protein sequence strings
+            min_confidence: Minimum probability threshold for predicting a specific nucleotide.
+                           If max probability < min_confidence, predict 'N'.
+                           If None, uses self.pam_min_confidence (default: None)
+        
+        Returns:
+            predicted_pams: List of predicted PAM strings (10 nucleotides each)
+        """
+        if min_confidence is None:
+            min_confidence = self.pam_min_confidence
+        if not protein_seqs:
+            return []
+        
+        # Handle single string input
+        if isinstance(protein_seqs, str):
+            protein_seqs = [protein_seqs]
+        
+        # Check if model and tokenizer are initialized
+        if not hasattr(self, 'model') or self.model is None:
+            print(f"[ERROR] PAMMatching.predict_pam: model is not initialized. Cannot predict PAM.")
+            return []
+        
+        if not hasattr(self, 'tokenizer') or self.tokenizer is None:
+            print(f"[ERROR] PAMMatching.predict_pam: tokenizer is not initialized. Cannot predict PAM.")
+            return []
+        
+        predicted_pams = []
+        nucleotides = ['A', 'C', 'G', 'T']
+        
+        try:
+            with torch.no_grad():
+                # Tokenize sequences
+                encodings = self.tokenizer.encode_batch(protein_seqs)
+                input_batch = dict(
+                    input_ids=torch.tensor([encoding.ids for encoding in encodings], device=self.device),
+                    attention_mask=torch.tensor([encoding.attention_mask for encoding in encodings], device=self.device),
+                )
+                
+                # Get PAM predictions
+                output = self.model(**input_batch)
+                logits = output.logits  # (batch_size, 10, 4) - 10 positions, 4 nucleotides (ACGT)
+                # Apply temperature scaling for prediction (sharper distributions when < 1.0)
+                scaled_logits = logits / self.pam_prediction_temperature
+                probabilities = F.softmax(scaled_logits, dim=-1)  # (batch_size, 10, 4)
+                
+                # Predict PAM for each sequence
+                for i in range(probabilities.shape[0]):
+                    probs = probabilities[i]  # (10, 4)
+                    pam_seq = []
+                    
+                    for pos_idx in range(probs.shape[0]):
+                        pos_probs = probs[pos_idx, :]  # (4,)
+                        max_prob = pos_probs.max().item()
+                        max_idx = pos_probs.argmax().item()
+                        
+                        # If confidence is low, predict 'N'
+                        if max_prob < min_confidence:
+                            pam_seq.append('N')
+                        else:
+                            pam_seq.append(nucleotides[max_idx])
+                    
+                    predicted_pams.append(''.join(pam_seq))
+        except Exception as e:
+            import traceback
+            print(f"[ERROR] PAMMatching.predict_pam failed: {e}")
+            traceback.print_exc()
+            return []
+        
+        return predicted_pams
+    
+    def get_scores(self, protein_seqs):
+        """
+        Get PAM matching scores using either cross-entropy loss or log probability approach.
+        
+        If use_ce_loss=True:
+            Uses cross-entropy loss between predicted logits and target PAM indices.
+            Score = exp(-mean_ce_loss) to convert to [0, 1] range (higher is better).
+        
+        If use_ce_loss=False (default):
+            Uses raw probabilities (no temperature scaling) to maintain high probability values.
+            For non-N positions: score is the geometric mean of probabilities (exp of mean log prob).
+            For N positions: if use_entropy_for_n_positions=True, score is based on entropy.
+        
+        Args:
+            protein_seqs: List of protein sequence strings
+        
+        Returns:
+            scores: List of scores in [0, 1] (higher = better match to target PAM)
+        """
+        if not protein_seqs:
+            return []
+        
+        # Handle single string input
+        if isinstance(protein_seqs, str):
+            protein_seqs = [protein_seqs]
+        
+        scores = []
+        max_entropy = math.log(4)  # Maximum entropy for uniform distribution over 4 nucleotides
+        
+        with torch.no_grad():
+            # Tokenize sequences
+            encodings = self.tokenizer.encode_batch(protein_seqs)
+            input_batch = dict(
+                input_ids=torch.tensor([encoding.ids for encoding in encodings], device=self.device),
+                attention_mask=torch.tensor([encoding.attention_mask for encoding in encodings], device=self.device),
+            )
+            
+            # Get PAM predictions (logits)
+            output = self.model(**input_batch)
+            logits = output.logits  # (batch_size, 10, 4) - 10 positions, 4 nucleotides (ACGT)
+            
+            for i in range(logits.shape[0]):
+                seq_logits = logits[i]  # (10, 4)
+                
+                if self.use_ce_loss:
+                    # Cross-entropy loss approach
+                    # Separate non-N and N positions
+                    non_n_mask = self.target_mask  # True for non-N positions
+                    n_mask = ~non_n_mask  # True for N positions
+                    
+                    score_components = []
+                    
+                    # For non-N positions: compute cross-entropy loss against target nucleotides
+                    if non_n_mask.any():
+                        masked_logits = seq_logits[non_n_mask]  # (K, 4)
+                        masked_targets = self.target_indices[non_n_mask]  # (K,) in {0,1,2,3}
+                        
+                        # Compute cross-entropy loss per position (reduction='none')
+                        ce_loss_per_pos = F.cross_entropy(
+                            masked_logits, masked_targets, reduction='none'
+                        )  # (K,)
+                        
+                        # Mean cross-entropy loss over non-N positions
+                        mean_ce_loss = ce_loss_per_pos.mean().item()
+                        
+                        # Convert to score in [0, 1] range using exp(-ce_loss)
+                        # Lower CE loss = higher score (better match)
+                        non_n_score = math.exp(-mean_ce_loss)
+                        score_components.append(non_n_score)
+                    
+                    # For N positions: handle based on entropy flag
+                    if n_mask.any():
+                        if self.use_entropy_for_n_positions:
+                            # Use entropy scoring (encourage uniform distribution)
+                            log_probs = F.log_softmax(seq_logits, dim=-1)  # (10, 4)
+                            probs = torch.exp(log_probs)  # (10, 4)
+                            n_probs = probs[n_mask]  # (M, 4) where M is number of N positions
+                            n_log_probs = log_probs[n_mask]  # (M, 4)
+                            entropy_per_pos = -(n_probs * n_log_probs).sum(dim=-1)  # (M,)
+                            mean_entropy = entropy_per_pos.mean().item()
+                            n_score = mean_entropy / max_entropy
+                            score_components.append(n_score)
+                        else:
+                            # Use CE loss against uniform distribution (consistent with CE loss approach)
+                            # This encourages uniform distribution over all 4 nucleotides
+                            n_logits = seq_logits[n_mask]  # (M, 4) where M is number of N positions
+                            
+                            # Compute cross-entropy loss against uniform distribution [0.25, 0.25, 0.25, 0.25]
+                            # CE_uniform = -sum(0.25 * log(p_i)) = -0.25 * sum(log(p_i))
+                            # Lower CE loss (closer to uniform) = higher score
+                            log_probs = F.log_softmax(n_logits, dim=-1)  # (M, 4)
+                            uniform_target = 0.25  # Uniform probability for each nucleotide
+                            ce_loss_uniform_per_pos = -(uniform_target * log_probs).sum(dim=-1)  # (M,)
+                            
+                            # Mean cross-entropy loss over N positions
+                            mean_ce_loss_uniform = ce_loss_uniform_per_pos.mean().item()
+                            
+                            # Convert to score in [0, 1] range using exp(-ce_loss)
+                            # Lower CE loss (more uniform) = higher score
+                            n_score = math.exp(-mean_ce_loss_uniform)
+                            score_components.append(n_score)
+                    
+                    # Combine scores: if both non-N and N positions exist, take weighted average
+                    # Weight by number of positions of each type
+                    if len(score_components) == 2:
+                        num_non_n = non_n_mask.sum().item()
+                        num_n = n_mask.sum().item()
+                        total_positions = num_non_n + num_n
+                        score = (score_components[0] * num_non_n + score_components[1] * num_n) / total_positions
+                    elif len(score_components) == 1:
+                        score = score_components[0]
+                    else:
+                        score = 0.0
+                    
+                    scores.append(score)
+                else:
+                    # Original log probability approach (default)
+                    # Use raw probabilities (no temperature scaling) to maintain high probabilities
+                    # This treats all probability values equally, appropriate for maintaining
+                    # high probability rather than optimizing low probabilities upward
+                    log_probs = F.log_softmax(seq_logits, dim=-1)  # (10, 4)
+                    probs = torch.exp(log_probs)  # (10, 4) - needed for entropy calculation
+                    
+                    # Separate non-N and N positions
+                    non_n_mask = self.target_mask  # True for non-N positions
+                    n_mask = ~non_n_mask  # True for N positions
+                    
+                    score_components = []
+                    
+                    # For non-N positions: maximize log probability of target nucleotide
+                    if non_n_mask.any():
+                        masked_log_probs = log_probs[non_n_mask]  # (K, 4)
+                        masked_targets = self.target_indices[non_n_mask]  # (K,) in {0,1,2,3}
+                        
+                        # Get log probability of target nucleotide at each position
+                        target_log_probs = masked_log_probs.gather(
+                            1, masked_targets.unsqueeze(1)
+                        ).squeeze(1)  # (K,)
+                        
+                        # Mean log probability (geometric mean in probability space)
+                        mean_log_prob = target_log_probs.mean().item()
+                        
+                        # Convert to score in [0, 1] range using geometric mean probability
+                        # This gives equal weight to maintaining high probabilities
+                        non_n_score = math.exp(mean_log_prob)
+                        score_components.append(non_n_score)
+                    
+                    # For N positions: maximize entropy (encourage uniform distribution) if enabled
+                    if n_mask.any() and self.use_entropy_for_n_positions:
+                        n_probs = probs[n_mask]  # (M, 4) where M is number of N positions
+                        
+                        # Compute entropy for each N position: H = -sum(p_i * log(p_i))
+                        # Use log_probs for numerical stability: H = -sum(p_i * log_p_i)
+                        n_log_probs = log_probs[n_mask]  # (M, 4)
+                        entropy_per_pos = -(n_probs * n_log_probs).sum(dim=-1)  # (M,)
+                        
+                        # Mean entropy across N positions
+                        mean_entropy = entropy_per_pos.mean().item()
+                        
+                        # Normalize by max entropy to get score in [0, 1]
+                        # Higher entropy (more uniform) = higher score
+                        n_score = mean_entropy / max_entropy
+                        score_components.append(n_score)
+                    
+                    # Combine scores: if both non-N and N positions exist, take weighted average
+                    # Weight by number of positions of each type
+                    # If entropy is disabled, only use non-N score
+                    if len(score_components) == 2:
+                        # Weighted average based on number of positions
+                        num_non_n = non_n_mask.sum().item()
+                        num_n = n_mask.sum().item()
+                        total_positions = num_non_n + num_n
+                        score = (score_components[0] * num_non_n + score_components[1] * num_n) / total_positions
+                    elif len(score_components) == 1:
+                        score = score_components[0]
+                    else:
+                        # No positions to evaluate (shouldn't happen, but handle gracefully)
+                        score = 0.0
+                    
+                    scores.append(score)
+        
+        return scores
+    
+    def get_score_for_pam(self, protein_seqs, pam_sequence, use_temperature_scaling=False):
+        """
+        Get probability score for a specific PAM sequence (not necessarily the target).
+        
+        This computes the log probability score for an arbitrary PAM sequence.
+        By default, uses raw probabilities (no temperature scaling) for more
+        interpretable results. Set use_temperature_scaling=True to match the
+        optimization objective.
+        
+        Args:
+            protein_seqs: List of protein sequence strings
+            pam_sequence: PAM sequence string of length 10 (e.g., "NGGNNNNNNN")
+            use_temperature_scaling: If True, apply temperature scaling (default: False for display)
+        
+        Returns:
+            scores: List of scores in [0, 1] (probability of the given PAM sequence)
+        """
+        if not protein_seqs:
+            return []
+        
+        # Handle single string input
+        if isinstance(protein_seqs, str):
+            protein_seqs = [protein_seqs]
+        
+        # Validate and convert PAM sequence to indices
+        pam_sequence = pam_sequence.upper()
+        if len(pam_sequence) != 10:
+            raise ValueError(f"PAM sequence must be exactly 10 nucleotides, got {len(pam_sequence)}")
+        
+        nucleotides = ['A', 'C', 'G', 'T']
+        pam_indices = []
+        pam_mask = []
+        for nuc in pam_sequence:
+            if nuc == 'N':
+                pam_indices.append(-1)
+                pam_mask.append(False)
+            else:
+                if nuc not in nucleotides:
+                    raise ValueError(f"Invalid nucleotide in PAM sequence: {nuc}")
+                pam_indices.append(nucleotides.index(nuc))
+                pam_mask.append(True)
+        
+        pam_indices = torch.tensor(pam_indices, device=self.device, dtype=torch.long)
+        pam_mask = torch.tensor(pam_mask, device=self.device, dtype=torch.bool)
+        
+        scores = []
+        
+        with torch.no_grad():
+            # Tokenize sequences
+            encodings = self.tokenizer.encode_batch(protein_seqs)
+            input_batch = dict(
+                input_ids=torch.tensor([encoding.ids for encoding in encodings], device=self.device),
+                attention_mask=torch.tensor([encoding.attention_mask for encoding in encodings], device=self.device),
+            )
+            
+            # Get PAM predictions (logits)
+            output = self.model(**input_batch)
+            logits = output.logits  # (batch_size, 10, 4) - 10 positions, 4 nucleotides (ACGT)
+            
+            for i in range(logits.shape[0]):
+                seq_logits = logits[i]  # (10, 4)
+                
+                # Apply temperature scaling only if requested (for display, use raw probabilities)
+                if use_temperature_scaling:
+                    scaled_logits = seq_logits / self.sigmoid_temperature
+                    log_probs = F.log_softmax(scaled_logits, dim=-1)  # (10, 4)
+                else:
+                    # Use raw probabilities (no temperature scaling) for more interpretable results
+                    log_probs = F.log_softmax(seq_logits, dim=-1)  # (10, 4)
+                
+                if pam_mask.any():
+                    masked_log_probs = log_probs[pam_mask]  # (K, 4)
+                    masked_pam_indices = pam_indices[pam_mask]  # (K,) in {0,1,2,3}
+                    
+                    # Get log probability of PAM nucleotide at each position
+                    pam_log_probs = masked_log_probs.gather(
+                        1, masked_pam_indices.unsqueeze(1)
+                    ).squeeze(1)  # (K,)
+                    
+                    # Mean log probability (linear in log space)
+                    mean_log_prob = pam_log_probs.mean().item()
+                    
+                    # Convert to score in [0, 1] range using geometric mean probability
+                    score = math.exp(mean_log_prob)
+                else:
+                    # All positions are N: use neutral score
+                    score = 0.0
+                
+                scores.append(score)
+        
+        return scores
+        
+        # ========================================================================
+        # OLD IMPLEMENTATION (logit-margin-based with sigmoid) - COMMENTED OUT
+        # ========================================================================
+        # This was the previous implementation that struggled to show small improvements
+        # because sigmoid saturates quickly. Kept for reference.
+        #
+        # def get_scores_old(self, protein_seqs):
+        #     """
+        #     Get PAM matching scores (logit-margin-based, in [0, 1]) for protein sequences.
+        #     
+        #     At each target (non-N) position: margin = (target PAM logit) - (largest non-target logit),
+        #     then score_pos = sigmoid(margin / sigmoid_temperature). The sequence score is the mean
+        #     over those positions. Higher is better.
+        #     
+        #     Args:
+        #         protein_seqs: List of protein sequence strings
+        #     
+        #     Returns:
+        #         scores: List of scores in [0, 1] (higher = target PAM preferred over alternatives)
+        #     """
+        #     if not protein_seqs:
+        #         return []
+        #     
+        #     # Handle single string input
+        #     if isinstance(protein_seqs, str):
+        #         protein_seqs = [protein_seqs]
+        #     
+        #     scores = []
+        #     
+        #     with torch.no_grad():
+        #         # Tokenize sequences
+        #         encodings = self.tokenizer.encode_batch(protein_seqs)
+        #         input_batch = dict(
+        #             input_ids=torch.tensor([encoding.ids for encoding in encodings], device=self.device),
+        #             attention_mask=torch.tensor([encoding.attention_mask for encoding in encodings], device=self.device),
+        #         )
+        #         
+        #         # Get PAM predictions (logits)
+        #         output = self.model(**input_batch)
+        #         logits = output.logits  # (batch_size, 10, 4) - 10 positions, 4 nucleotides (ACGT)
+        #         
+        #         for i in range(logits.shape[0]):
+        #             seq_logits = logits[i]  # (10, 4)
+        #             masked_positions = self.target_mask
+        #             if masked_positions.any():
+        #                 masked_logits = seq_logits[masked_positions]  # (K, 4)
+        #                 masked_targets = self.target_indices[masked_positions]  # (K,) in {0,1,2,3}
+        #                 # Target logit at each position
+        #                 target_logits = masked_logits.gather(1, masked_targets.unsqueeze(1)).squeeze(1)  # (K,)
+        #                 # Largest non-target logit: mask out target class then max over dim=-1
+        #                 logits_copy = masked_logits.clone()
+        #                 logits_copy.scatter_(1, masked_targets.unsqueeze(1), -1e9)
+        #                 max_non_target = logits_copy.max(dim=1).values  # (K,)
+        #                 margin = target_logits - max_non_target  # (K,)
+        #                 margin = margin / self.sigmoid_temperature
+        #                 score_per_pos = torch.sigmoid(margin)
+        #                 score = score_per_pos.mean().item()
+        #             else:
+        #                 # All positions are N: no target to match, use neutral score
+        #                 score = 0.5
+        #             scores.append(score)
+        #     
+        #     return scores
+    
+    def __call__(self, protein_tokens, protein_seqs):
+        """
+        Objective call interface.
+        
+        Args:
+            protein_tokens: Unused (kept for interface compatibility)
+            protein_seqs: List of protein sequence strings
+        
+        Returns:
+            Tuple of ('pam_matching', scores) where scores is a list of values in [0, 1]:
+            - For non-N positions: temperature-scaled log probability (geometric mean over positions)
+            - For N positions: if use_entropy_for_n_positions=True, normalized entropy (encouraging uniform distribution)
+            - If use_entropy_for_n_positions=False, N positions are ignored (score only considers non-N positions)
+            - Combined score is weighted average if both types exist and entropy is enabled
+            Higher scores indicate better match to target PAM.
+        """
+        scores = self.get_scores(protein_seqs)
+        return 'pam_matching', scores
+

cas9/pam_detector_hmm.py

@@ -0,0 +1,456 @@
+# cas9_pam_hmm.py
+from __future__ import annotations
+
+import os
+import re
+import hashlib
+import subprocess
+import tempfile
+from dataclasses import dataclass
+from typing import Dict, Iterable, List, Optional, Tuple, Union
+
+import torch
+
+
+# ----------------------------
+# Data structures
+# ----------------------------
+@dataclass(frozen=True)
+class HMMHit:
+    """
+    Represents one domain hit line from HMMER domtblout.
+    Coordinates are 1-based (HMMER convention).
+    """
+    model: str
+    accession: str
+    query_id: str
+    query_len: int
+
+    # Per-domain statistics
+    i_evalue: float
+    bitscore: float
+
+    # Alignment coordinates
+    hmm_from: int
+    hmm_to: int
+    ali_from: int
+    ali_to: int
+    env_from: int
+    env_to: int
+
+    # Optional
+    hmm_len: Optional[int] = None
+    bias: Optional[float] = None
+
+
+# ----------------------------
+# domtblout parsing
+# ----------------------------
+class DomTbloutParser:
+    """
+    Robust parser for HMMER3 --domtblout files.
+    Assumes standard HMMER3 domtblout layout (space-delimited; description is trailing).
+    """
+
+    @staticmethod
+    def parse(domtbl_path: str) -> List[HMMHit]:
+        hits: List[HMMHit] = []
+        with open(domtbl_path, "r") as f:
+            for line in f:
+                if not line.strip() or line.startswith("#"):
+                    continue
+                cols = line.strip().split()
+
+                # Standard domtblout columns (HMMER3):
+                # 1 target name        cols[0]
+                # 2 target accession   cols[1]
+                # 3 tlen               cols[2]
+                # 4 query name         cols[3]
+                # 5 query accession    cols[4]
+                # 6 qlen               cols[5]
+                # 7 E-value (full)     cols[6]
+                # 8 score (full)       cols[7]
+                # 9 bias (full)        cols[8]
+                # 10 #                 cols[9]
+                # 11 of                cols[10]
+                # 12 c-Evalue          cols[11]
+                # 13 i-Evalue          cols[12]
+                # 14 score (domain)    cols[13]
+                # 15 bias (domain)     cols[14]
+                # 16 hmm_from          cols[15]
+                # 17 hmm_to            cols[16]
+                # 18 ali_from          cols[17]
+                # 19 ali_to            cols[18]
+                # 20 env_from          cols[19]
+                # 21 env_to            cols[20]
+                # 22 acc               cols[21]
+                # 23+ description      cols[22:]
+
+                try:
+                    model = cols[0]
+                    accession = cols[1]
+                    hmm_len = int(cols[2])
+                    query_id = cols[3]
+                    query_len = int(cols[5])
+
+                    i_evalue = float(cols[12])
+                    bitscore = float(cols[13])
+                    bias = float(cols[14])
+
+                    hmm_from = int(cols[15])
+                    hmm_to = int(cols[16])
+                    ali_from = int(cols[17])
+                    ali_to = int(cols[18])
+                    env_from = int(cols[19])
+                    env_to = int(cols[20])
+                except (IndexError, ValueError) as e:
+                    # If your HMMER output is slightly different, raise with context.
+                    raise RuntimeError(f"Failed parsing domtblout line:\n{line}\nError: {e}") from e
+
+                hits.append(
+                    HMMHit(
+                        model=model,
+                        accession=accession,
+                        query_id=query_id,
+                        query_len=query_len,
+                        i_evalue=i_evalue,
+                        bitscore=bitscore,
+                        hmm_from=hmm_from,
+                        hmm_to=hmm_to,
+                        ali_from=ali_from,
+                        ali_to=ali_to,
+                        env_from=env_from,
+                        env_to=env_to,
+                        hmm_len=hmm_len,
+                        bias=bias,
+                    )
+                )
+        return hits
+
+
+# ----------------------------
+# hmmscan runner
+# ----------------------------
+class HMMSCAN:
+    """
+    Thin wrapper around hmmscan for 1 or many sequences.
+    Uses --domtblout for machine-parseable results.
+    """
+
+    def __init__(
+        self,
+        hmm_db_path: str,
+        hmmscan_bin: str = "hmmscan",
+        cpu: int = 1,
+        extra_args: Optional[List[str]] = None,
+        env: Optional[Dict[str, str]] = None,
+    ):
+        self.hmm_db_path = hmm_db_path
+        self.hmmscan_bin = hmmscan_bin
+        self.cpu = int(cpu)
+        self.extra_args = extra_args or []
+        self.env = env
+
+        if not os.path.exists(self.hmm_db_path):
+            raise FileNotFoundError(f"HMM DB not found: {self.hmm_db_path}")
+
+    def scan_fasta(self, fasta_path: str, domtblout_path: str) -> None:
+        cmd = [
+            self.hmmscan_bin,
+            "--cpu",
+            str(self.cpu),
+            "--domtblout",
+            domtblout_path,
+            self.hmm_db_path,
+            fasta_path,
+        ] + self.extra_args
+
+        try:
+            result = subprocess.run(
+                cmd,
+                check=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                env=self.env,
+                text=True,
+            )
+        except subprocess.CalledProcessError as e:
+            error_msg = e.stderr.decode('utf-8') if isinstance(e.stderr, bytes) else e.stderr
+            stdout_msg = e.stdout.decode('utf-8') if isinstance(e.stdout, bytes) else e.stdout
+            raise RuntimeError(
+                f"hmmscan failed with exit code {e.returncode}.\n"
+                f"Command: {' '.join(cmd)}\n"
+                f"STDOUT:\n{stdout_msg}\n"
+                f"STDERR:\n{error_msg}\n"
+                f"HMM DB path: {self.hmm_db_path}\n"
+                f"FASTA path: {fasta_path}\n"
+                f"FASTA exists: {os.path.exists(fasta_path)}"
+            ) from e
+
+    def scan_sequences(
+        self,
+        seqs: Union[str, List[str]],
+        ids: Optional[Union[str, List[str]]] = None,
+    ) -> Tuple[str, str]:
+        """
+        Writes temporary FASTA and runs hmmscan.
+        Returns: (domtblout_path, fasta_path) inside a tempdir.
+        Caller is responsible for tempdir cleanup if they want persistence.
+        """
+        if isinstance(seqs, str):
+            seq_list = [seqs]
+        else:
+            seq_list = list(seqs)
+
+        if ids is None:
+            id_list = [f"q{i}" for i in range(len(seq_list))]
+        elif isinstance(ids, str):
+            id_list = [ids]
+        else:
+            id_list = list(ids)
+
+        if len(id_list) != len(seq_list):
+            raise ValueError("ids must have same length as seqs")
+
+        td = tempfile.TemporaryDirectory()
+        fasta_path = os.path.join(td.name, "queries.faa")
+        domtbl_path = os.path.join(td.name, "out.domtbl")
+
+        with open(fasta_path, "w") as f:
+            for qid, seq in zip(id_list, seq_list):
+                seq = seq.replace(" ", "").strip()
+                f.write(f">{qid}\n{seq}\n")
+
+        self.scan_fasta(fasta_path, domtbl_path)
+
+        # Return temp paths; also return tempdir object by attaching it for lifecycle
+        # (so it doesn't get GC'ed early).
+        setattr(self, "_last_tempdir", td)
+        return domtbl_path, fasta_path
+
+
+# ----------------------------
+# PI-domain masking policy ("Option 1")
+# ----------------------------
+class Cas9PIMasker:
+    """
+    Builds a compact 'no deletion' mask for the PAM-interacting (PI) region.
+
+    Option 1 implemented:
+      - choose highest-confidence PI-family hit
+      - then mask a capped window around the alignment core (ali_from..ali_to)
+        using max_mask_len
+
+    Notes:
+      - HMMER coords are 1-based AA positions.
+      - Assumes your token tensor x is [BOS] + AAs + [EOS] (+ pad),
+        so AA position p corresponds to token index p.
+    """
+
+    DEFAULT_MODEL_PREFERENCE = ["Cas9_PI", "Cas9_PI_C", "Cas9_PI2", "CjCas9_PI_CTD"]
+
+    def __init__(
+        self,
+        hmmscan: HMMSCAN,
+        *,
+        model_preference: Optional[List[str]] = None,
+        use_env_coords: bool = False,      # False = ali coords (smaller)
+        max_mask_len: int = 200,           # cap masked window length
+        evalue_cutoff: float = 1e-5,       # per-domain i-evalue cutoff
+        min_ali_len: int = 30,             # ignore tiny alignments
+        fallback_last_n: Optional[int] = 250,  # if no hit, optionally mask last N AAs
+        cache_size: int = 2048,
+    ):
+        self.hmmscan = hmmscan
+        self.model_preference = model_preference or self.DEFAULT_MODEL_PREFERENCE
+        self.model_rank = {m: i for i, m in enumerate(self.model_preference)}
+        self.use_env_coords = use_env_coords
+        self.max_mask_len = int(max_mask_len)
+        self.evalue_cutoff = float(evalue_cutoff)
+        self.min_ali_len = int(min_ali_len)
+        self.fallback_last_n = fallback_last_n
+        self.cache_size = int(cache_size)
+
+        # Simple LRU-ish cache keyed by sha1(sequence)
+        self._cache: Dict[str, Optional[Tuple[int, int]]] = {}
+        self._cache_order: List[str] = []
+
+    # ---- selection logic ----
+    def _hit_interval(self, hit: HMMHit) -> Tuple[int, int]:
+        if self.use_env_coords:
+            s, t = hit.env_from, hit.env_to
+        else:
+            s, t = hit.ali_from, hit.ali_to
+        if s > t:
+            s, t = t, s
+        return s, t
+
+    def _hit_ali_len(self, hit: HMMHit) -> int:
+        s, t = self._hit_interval(hit)
+        return t - s + 1
+
+    def _is_candidate(self, hit: HMMHit) -> bool:
+        if hit.model not in self.model_rank:
+            return False
+        if hit.i_evalue > self.evalue_cutoff:
+            return False
+        if self._hit_ali_len(hit) < self.min_ali_len:
+            return False
+        return True
+
+    def _choose_best_hit(self, hits: List[HMMHit]) -> Optional[HMMHit]:
+        """
+        Pick the highest-confidence PI-family hit.
+
+        Ranking:
+          1) model preference (Cas9_PI preferred over PI_C over PI2 over CTD)
+          2) lowest i-evalue
+          3) highest bitscore
+        """
+        candidates = [h for h in hits if self._is_candidate(h)]
+        if not candidates:
+            return None
+
+        def key(h: HMMHit):
+            return (self.model_rank[h.model], h.i_evalue, -h.bitscore)
+
+        return min(candidates, key=key)
+
+    def _cap_interval(self, s: int, t: int, seq_len: int) -> Tuple[int, int]:
+        """
+        Cap [s,t] to length max_mask_len around the alignment core midpoint.
+        """
+        if s > t:
+            s, t = t, s
+
+        s = max(1, min(s, seq_len))
+        t = max(1, min(t, seq_len))
+
+        L = t - s + 1
+        if L <= self.max_mask_len:
+            return s, t
+
+        mid = (s + t) // 2
+        half = self.max_mask_len // 2
+        s2 = mid - half
+        t2 = s2 + self.max_mask_len - 1
+
+        if s2 < 1:
+            s2 = 1
+            t2 = min(seq_len, self.max_mask_len)
+        if t2 > seq_len:
+            t2 = seq_len
+            s2 = max(1, seq_len - self.max_mask_len + 1)
+
+        return s2, t2
+
+    # ---- caching ----
+    def _seq_key(self, seq: str) -> str:
+        return hashlib.sha1(seq.encode("utf-8")).hexdigest()
+
+    def _cache_get(self, seq: str) -> Optional[Tuple[int, int]]:
+        k = self._seq_key(seq)
+        if k in self._cache:
+            return self._cache[k]
+        return None
+
+    def _cache_put(self, seq: str, interval: Optional[Tuple[int, int]]) -> None:
+        k = self._seq_key(seq)
+        if k in self._cache:
+            self._cache[k] = interval
+            return
+        self._cache[k] = interval
+        self._cache_order.append(k)
+        if len(self._cache_order) > self.cache_size:
+            old = self._cache_order.pop(0)
+            self._cache.pop(old, None)
+
+    # ---- public API ----
+    def pi_core_interval(self, seq: str) -> Optional[Tuple[int, int]]:
+        """
+        Returns (start,end) 1-based AA coordinates of the compact PI-core region,
+        after applying the max_mask_len cap.
+
+        If no PI hit is found and fallback_last_n is set, returns last N AAs.
+        """
+        seq = seq.replace(" ", "").strip()
+        if not seq:
+            return None
+
+        cached = self._cache_get(seq)
+        if cached is not None or (self._seq_key(seq) in self._cache):
+            # cached could be None meaning "no hit"
+            return cached
+
+        domtbl, _ = self.hmmscan.scan_sequences(seq, ids="q")
+        hits = DomTbloutParser.parse(domtbl)
+        best = self._choose_best_hit(hits)
+
+        if best is None:
+            if self.fallback_last_n is None:
+                self._cache_put(seq, None)
+                return None
+            n = int(self.fallback_last_n)
+            L = len(seq)
+            s = max(1, L - n + 1)
+            t = L
+            interval = (s, t)
+            self._cache_put(seq, interval)
+            return interval
+
+        s, t = self._hit_interval(best)
+        interval = self._cap_interval(s, t, seq_len=len(seq))
+        self._cache_put(seq, interval)
+        return interval
+
+    def build_no_del_mask(
+        self,
+        x: torch.Tensor,
+        aa_seq: str,
+        pad_id: int,
+        *,
+        bos_at_index0: bool = True,
+    ) -> torch.Tensor:
+        """
+        Build (B,Lmax) bool mask where True => deletion disallowed.
+        Currently supports B=1 (single accepted sequence), which matches your use.
+        """
+        if x.dim() != 2 or x.shape[0] != 1:
+            raise ValueError(f"Expected x shape (1,Lmax), got {tuple(x.shape)}")
+
+        Lmax = x.shape[1]
+        device = x.device
+        mask = torch.zeros((1, Lmax), device=device, dtype=torch.bool)
+
+        aa_seq = aa_seq.replace(" ", "").strip()
+        seq_len = len(aa_seq)
+
+        interval = self.pi_core_interval(aa_seq)
+        if interval is None:
+            return mask
+
+        s, t = interval  # 1-based AA positions
+        s = max(1, min(s, seq_len))
+        t = max(1, min(t, seq_len))
+        if s > t:
+            s, t = t, s
+
+        # Token mapping:
+        # if BOS at index 0, AA position p is at token index p
+        # else AA position p is at token index p-1
+        offset = 0 if bos_at_index0 else -1
+        ts = s + offset
+        tt = t + offset
+
+        ts = max(0, min(ts, Lmax - 1))
+        tt = max(0, min(tt, Lmax - 1))
+        if ts > tt:
+            ts, tt = tt, ts
+
+        mask[0, ts:tt + 1] = True
+
+        # Don't bother masking PAD; your deletion rates are already masked on PAD tokens.
+        # But if you want, you can explicitly clear PAD positions:
+        mask[0, x[0] == pad_id] = False
+
+        return mask

cas9/predict_pam_distribution.py

@@ -0,0 +1,206 @@
+#!/usr/bin/env python3
+"""
+Script to predict PAM sequences for protein sequences and print the PAM distribution.
+
+Edit the SEQUENCES dictionary below with your protein sequences.
+This script shows the full softmaxed probability distribution for each position in the PAM.
+"""
+
+from collections import Counter
+import torch
+
+# Import PAMMatching from cas9.objectives
+from cas9.objectives import PAMMatching
+
+# ============================================================================
+# EDIT THIS DICTIONARY WITH YOUR SEQUENCES
+# Key = sequence name, Value = protein sequence string
+# ============================================================================
+
+SEQUENCES = {
+    "sacas9": "MKRNYILGLDIGITSVGYGIIDYETRDVIDAGVRLFKEANVENNEGRRSKRGARRLKRRRRHRIQRVKKLLFDYNLLTDHSELSGINPYEARVKGLSQKLSEEEFSAALLHLAKRRGVHNVNEVEEDTGNELSTKEQISRNSKALEEKYVAELQLERLKKDGEVRGSINRFKTSDYVKEAKQLLKVQKAYHQLDQSFIDTYIDLLETRRTYYEGPGEGSPFGWKDIKEWYEMLMGHCTYFPEELRSVKYAYNADLYNALNDLNNLVITRDENEKLEYYEKFQIIENVFKQKKKPTLKQIAKEILVNEEDIKGYRVTSTGKPEFTNLKVYHDIKDITARKEIIENAELLDQIAKILTIYQSSEDIQEELTNLNSELTQEEIEQISNLKGYTGTHNLSLKAINLILDELWHTNDNQIAIFNRLKLVPKKVDLSQQKEIPTTLVDDFILSPVVKRSFIQSIKVINAIIKKYGLPNDIIIELAREKNSKDAQKMINEMQKRNRQTNERIEEIIRTTGKENAKYLIEKIKLHDMQEGKCLYSLEAIPLEDLLNNPFNYEVDHIIPRSVSFDNSFNNKVLVKQEENSKKGNRTPFQYLSSSDSKISYETFKKHILNLAKGKGRISKTKKEYLLEERDINRFSVQKDFINRNLVDTRYATRGLMNLLRSYFRVNNLDVKVKSINGGFTSFLRRKWKFKKERNKGYKHHAEDALIIANADFIFKEWKKLDKAKKVMENQMFEEKQAESMPEIETEQEYKEIFITPHQIKHIKDFKDYKYSHRVDKKPNRELINDTLYSTRKDDKGNTLIVNNLNGLYDKDNDKLKKLINKSPEKLLMYHHDPQTYQKLKLIMEQYGDEKNPLYKYYEETGNYLTKYSKKDNGPVIKKIKYYGNKLNAHLDITDDYPNSRNKVVKLSLKPYRFDVYLDNGVYKFVTVKNLDVIKKENYYEVNSKCYEEAKKLKKISNQAEFIASFYNNDLIKINGELYRVIGVNNDLLNRIEVNMIDITYREYLENMNDKRPPRIIKTIASKTQSIKKYSTDILGNLYEVKSKKHPQIIKKG",
+    
+    "cjcas9": "MARILAFDIGISSIGWAFSENDELKDCGVRIFTKVENPKTGESLALPRRLARSARKRLARRKARLNHLKHLIANEFKLNYEDYQSFDESLAKAYKGSLISPYELRFRALNELLSKQDFARVILHIAKRRGYDDIKNSDDKEKGAILKAIKQNEEKLANYQSVGEYLYKEYFQKFKENSKEFTNVRNKKESYERCIAQSFLKDELKLIFKKQREFGFSFSKKFEEEVLSVAFYKRALKDFSHLVGNCSFFTDEKRAPKNSPLAFMFVALTRIINLLNNLKNTEGILYTKDDLNALLNEVLKNGTLTYKQTKKLLGLSDDYEFKGEKGTYFIEFKKYKEFIKALGEHNLSQDDLNEIAKDITLIKDEIKLKKALAKYDLNQNQIDSLSKLEFKDHLNISFKALKLVTPLMLEGKKYDEACNELNLKVAINEDKKDFLPAFNETYYKDEVTNPVVLRAIKEYRKVLNALLKKYGKVHKINIELAREVGKNHSQRAKIEKEQNENYKAKKDAELECEKLGLKINSKNILKLRLFKEQKEFCAYSGEKIKISDLQDEKMLEIDHIYPYSRSFDDSYMNKVLVFTKQNQEKLNQTPFEAFGNDSAKWQKIEVLAKNLPTKKQKRILDKNYKDKEQKNFKDRNLNDTRYIARLVLNYTKDYLDFLPLSDDENTKLNDTQKGSKVHVEAKSGMLTSALRHTWGFSAKDRNNHLHHAIDAVIIAYANNSIVKAFSDFKKEQESNSAELYAKKISELDYKNKRKFFEPFSGFRQKVLDKIDEIFVSKPERKKPSGALHEETFRKEEEFYQSYGGKEGVLKALELGKIRKVNGKIVKNGDMFRVDIFKHKKTNKFYAVPIYTMDFALKVLPNKAVARSKKGEIKDWILMDENYEFCFSLYKDSLILIQTKDMQEPEFVYYNAFTSSTVSLIVSKHDNKFETLSKNQKILFKNANEKEVIAKSIGIQNLKVFEKYIVSALGEVTKAEFRQREDFKK",
+    
+    "spcas9": "MDKKYSIGLDIGTNSVGWAVITDEYKVPSKKFKVLGNTDRHSIKKNLIGALLFDSGETAEATRLKRTARRRYTRRKNRICYLQEIFSNEMAKVDDSFFHRLEESFLVEEDKKHERHPIFGNIVDEVAYHEKYPTIYHLRKKLVDSTDKADLRLIYLALAHMIKFRGHFLIEGDLNPDNSDVDKLFIQLVQTYNQLFEENPINASGVDAKAILSARLSKSRRLENLIAQLPGEKKNGLFGNLIALSLGLTPNFKSNFDLAEDAKLQLSKDTYDDDLDNLLAQIGDQYADLFLAAKNLSDAILLSDILRVNTEITKAPLSASMIKRYDEHHQDLTLLKALVRQQLPEKYKEIFFDQSKNGYAGYIDGGASQEEFYKFIKPILEKMDGTEELLVKLNREDLLRKQRTFDNGSIPHQIHLGELHAILRRQEDFYPFLKDNREKIEKILTFRIPYYVGPLARGNSRFAWMTRKSEETITPWNFEEVVDKGASAQSFIERMTNFDKNLPNEKVLPKHSLLYEYFTVYNELTKVKYVTEGMRKPAFLSGEQKKAIVDLLFKTNRKVTVKQLKEDYFKKIECFDSVEISGVEDRFNASLGTYHDLLKIIKDKDFLDNEENEDILEDIVLTLTLFEDREMIEERLKTYAHLFDDKVMKQLKRRRYTGWGRLSRKLINGIRDKQSGKTILDFLKSDGFANRNFMQLIHDDSLTFKEDIQKAQVSGQGDSLHEHIANLAGSPAIKKGILQTVKVVDELVKVMGRHKPENIVIEMARENQTTQKGQKNSRERMKRIEEGIKELGSQILKEHPVENTQLQNEKLYLYYLQNGRDMYVDQELDINRLSDYDVDHIVPQSFLKDDSIDNKVLTRSDKNRGKSDNVPSEEVVKKMKNYWRQLLNAKLITQRKFDNLTKAERGGLSELDKAGFIKRQLVETRQITKHVAQILDSRMNTKYDENDKLIREVKVITLKSKLVSDFRKDFQFYKVREINNYHHAHDAYLNAVVGTALIKKYPKLESEFVYGDYKVYDVRKMIAKSEQEIGKATAKYFFYSNIMNFFKTEITLANGEIRKRPLIETNGETGEIVWDKGRDFATVRKVLSMPQVNIVKKTEVQTGGFSKESILPKRNSDKLIARKKDWDPKKYGGFDSPTVAYSVLVVAKVEKGKSKKLKSVKELLGITIMERSSFEKNPIDFLEAKGYKEVKKDLIIKLPKYSLFELENGRKRMLASAGELQKGNELALPSKYVNFLYLASHYEKLKGSPEDNEQKQLFVEQHKHYLDEIIEQISEFSKRVILADANLDKVLSAYNKHRDKPIREQAENIIHLFTLTNLGAPAAFKYFDTTIDRKRYTSTKEVLDATLIHQSITGLYETRIDLSQLGGD",
+    
+    "nmcas9": "MAAFKPNPINYILGLDIGIASVGWAMVEIDEEENPIRLIDLGVRVFERAEVPKTGDSLAMVRRLARSVRRLTRRRAHRLLRARRLLKREGVLQAADFDENGLIKSLPNTPWQLRAAALDRKLTPLEWSAVLLHLIKHRGYLSQRKNEGETADKELGALLKGVADNAHALQTGDFRTPAELALNKFEKESGHIRNQRGDYSHTFSRKDLQAELILLFEKQKEFGNPHISGGLKEGIETLLMTQRPALSGDAVQKMLGHCTFEPAEPKAAKNTYTAERFIWLTKLNNLRILEQGSERPLTDTERATLMDEPYRKSKLTYAQARKLLGLEDTAFFKGLRYGKDNAEASTLMEMKAYHAISRALEKEGLKDKKSPLNLSPELQDEIGTAFSLFKTDEDITGRLKDRIQPEILEALLKHISFDKFVQISLKALRRIVPLMEQGKRYDEACAEIYGDHYGKKNTEEKIYLPPIPADEIRNPVVLRALSQARKVINGVVRRYGSPARIHIETAREVGKSFKDRKEIEKRQEENRKDREKAAAKFREYFPNFVGEPKSKDILKLRLYEQQHGKCLYSGKEINLGRLNEKGYVEIDHALPFSRTWDDSFNNKVLVLGSENQNKGNQTPYEYFNGKDNSREWQEFKARVETSRFPRSKKQRILLQKFDEDGFKERNLNDTRYVNRFLCQFVADRMRLTGKGKKRVFASNGQITNLLRGFWGLRKVRAENDRHHALDAVVVACSTVAMQQKITRFVRYKEMNAFDGKTIDKETGEVLHQKTHFPQPWEFFAQEVMIRVFGKPDGKPEFEEADTPEKLRTLLAEKLSSRPEAVHEYVTPLFVSRAPNRKMSGQGHMETVKSAKRLDEGVSVLRVPLTQLKLKDLEKMVNREREPKLYEALKARLEAHKDDPAKAFAEPFYKYDKAGNRTQQVKAVRVEQVQKTGVWVRNHNGIADNATMVRVDVFEKGDKYYLVPIYSWQVAKGILPDRAVVQGKDEEDWQLIDDSFNFKFSLHPNDLVEVITKKARMFGYFASCHRGTGNINIRIHDLDHKIGKNGILEGIGVKTALSFQKYQIDELGKEIRPCRLKKRPPVR",
+    
+    "st1cas9": "MSDLVLGLDIGIGSVGVGILNKVTGEIIHKNSRIFPAAQAENNLVRRTNRQGRRLARRKKHRRVRLNRLFEESGLITDFTKISINLNPYQLRVKGLTDELSNEELFIALKNMVKHRGISYLDDASDDGNSSVGDYAQIVKENSKQLETKTPGQIQLERYQTYGQLRGDFTVEKDGKKHRLINVFPTSAYRSEALRILQTQQEFNPQITDEFINRYLEILTGKRKYYHGPGNEKSRTDYGRYRTSGETLDNIFGILIGKCTFYPDEFRAAKASYTAQEFNLLNDLNNLTVPTETKKLSKEQKNQIINYVKNEKAMGPAKLFKYIAKLLSCDVADIKGYRIDKSGKAEIHTFEAYRKMKTLETLDIEQMDRETLDKLAYVLTLNTEREGIQEALEHEFADGSFSQKQVDELVQFRKANSSIFGKGWHNFSVKLMMELIPELYETSEEQMTILTRLGKQKTTSSSNKTKYIDEKLLTEEIYNPVVAKSVRQAIKIVNAAIKEYGDFDNIVIEMARETNEDDEKKAIQKIQKANKDEKDAAMLKAANQYNGKAELPHSVFHGHKQLATKIRLWHQQGERCLYTGKTISIHDLINNSNQFEVDHILPLSITFDDSLANKVLVYATANQEKGQRTPYQALDSMDDAWSFRELKAFVRESKTLSNKKKEYLLTEEDISKFDVRKKFIERNLVDTRYASRVVLNALQEHFRAHKIDTKVSVVRGQFTSQLRRHWGIEKTRDTYHHHAVDALIIAASSQLNLWKKQKNTLVSYSEDQLLDIETGELISDDEYKESVFKAPYQHFVDTLKSKEFEDSILFSYQVDSKFNRKISDATIYATRQAKVGKDKADETYVLGKIKDIYTQDGYDAFMKIYKKDKSKFLMYRHDPQTFEKVIEPILENYPNKQINEKGKEVPCNPFLKYKEEHGYIRKYSKKGNGPEIKSLKYYDSKLGNHIDITPKDSNNKVVLQSVSPWRADVYFNKTTGKYEILGLKYADLQFEKGTGTYKISQEKYNDIKKKEGVDSDSEFKFTLYKNDLLLVKDTETKEQQLFRFLSRTMPKQKHYVELKPYDKQKFEGGEALIKVLGNVANSGQCKKGLGKSNISIYKVRTDVLGNQHIIKNEGDKPKLDF",
+    "nmecas9":
+    "MAAFKPNPINYILGLDIGIASVGWAMVEIDEEENPIRLIDLGVRVFERAEVPKTGDSLAMVRRLARSVRRLTRRRAHRLLRARRLLKREGVLQAADFDENGLIKSLPNTPWQLRAAALDRKLTPLEWSAVLLHLIKHRGYLSQRKNEGETADKELGALLKGVADNAHALQTGDFRTPAELALNKFEKESGHIRNQRGDYSHTFSRKDLQAELILLFEKQKEFGNPHISGGLKEGIETLLMTQRPALSGDAVQKMLGHCTFEPAEPKAAKNTYTAERFIWLTKLNNLRILEQGSERPLTDTERATLMDEPYRKSKLTYAQARKLLGLEDTAFFKGLRYGKDNAEASTLMEMKAYHAISRALEKEGLKDKKSPLNLSPELQDEIGTAFSLFKTDEDITGRLKDRIQPEILEALLKHISFDKFVQISLKALRRIVPLMEQGKRYDEACAEIYGDHYGKKNTEEKIYLPPIPADEIRNPVVLRALSQARKVINGVVRRYGSPARIHIETAREVGKSFKDRKEIEKRQEENRKDREKAAAKFREYFPNFVGEPKSKDILKLRLYEQQHGKCLYSGKEINLGRLNEKGYVEIDHALPFSRTWDDSFNNKVLVLGSENQNKGNQTPYEYFNGKDNSREWQEFKARVETSRFPRSKKQRILLQKFDEDGFKERNLNDTRYVNRFLCQFVADRMRLTGKGKKRVFASNGQITNLLRGFWGLRKVRAENDRHHALDAVVVACSTVAMQQKITRFVRYKEMNAFDGKTIDKETGEVLHQKTHFPQPWEFFAQEVMIRVFGKPDGKPEFEEADTPEKLRTLLAEKLSSRPEAVHEYVTPLFVSRAPNRKMSGQGHMETVKSAKRLDEGVSVLRVPLTQLKLKDLEKMVNREREPKLYEALKARLEAHKDDPAKAFAEPFYKYDKAGNRTQQVKAVRVEQVQKTGVWVRNHNGIADNATMVRVDVFEKGDKYYLVPIYSWQVAKGILPDRAVVQGKDEEDWQLIDDSFNFKFSLHPNDLVEVITKKARMFGYFASCHRGTGNINIRIHDLDHKIGKNGILEGIGVKTALSFQKYQIDELGKEIRPCRLKKRPPVR"
+}
+
+# Model configuration
+MODEL_NAME = "Profluent-Bio/protein2pam-cas9_full"  # HuggingFace model name
+DEVICE = None  # None = auto-detect (cuda if available, else cpu)
+MIN_CONFIDENCE = 0.7  # Minimum probability threshold for predicting a specific nucleotide (vs. 'N')
+                       # If max probability < MIN_CONFIDENCE, predict 'N'
+
+
+def get_pam_probability_distributions(sequences_list, device=None, model_name="Profluent-Bio/protein2pam-cas9_full", min_confidence=0.55):
+    """
+    Get detailed probability distributions for PAM predictions.
+    
+    Args:
+        sequences_list: List of protein sequence strings
+        device: torch device (default: auto-detect)
+        model_name: HuggingFace model name for PAM prediction
+        min_confidence: Minimum probability threshold for predicting a specific nucleotide
+    
+    Returns:
+        List of dictionaries containing probability distributions and predicted PAMs
+    """
+    if not sequences_list:
+        print("Error: No sequences provided.")
+        return []
+    
+    # Auto-detect device if not provided
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    print(f"Using device: {device}")
+    print(f"Loading PAM prediction model: {model_name}...")
+    
+    # Create a dummy PAMMatching object (target_pam doesn't matter for prediction)
+    pam_predictor = PAMMatching(
+        device=device,
+        target_pam="NNNNNNNNNN",  # Dummy target, we only use get_pam_probability_distributions()
+        model_name=model_name,
+        use_entropy_for_n_positions=True,
+        use_ce_loss=False
+    )
+    
+    print(f"Computing PAM probability distributions for {len(sequences_list)} sequences...")
+    
+    # Get probability distributions
+    prob_distributions = pam_predictor.get_pam_probability_distributions(sequences_list)
+    
+    # Recompute predicted PAMs with custom min_confidence
+    nucleotides = ['A', 'C', 'G', 'T']
+    for dist in prob_distributions:
+        predicted_pam = []
+        for pos_info in dist['per_position']:
+            max_prob = pos_info['max_probability']
+            if max_prob < min_confidence:
+                predicted_pam.append('N')
+            else:
+                predicted_pam.append(pos_info['predicted'])
+        dist['predicted_pam_custom'] = ''.join(predicted_pam)
+    
+    return prob_distributions
+
+
+def print_sequence_pam_distribution(seq_name, prob_dist, min_confidence):
+    """Print the PAM probability distribution for a single sequence."""
+    print("\n" + "="*80)
+    print(f"Sequence: {seq_name}")
+    print("="*80)
+    
+    print(f"\nPredicted PAM (min_confidence={min_confidence:.2f}): {prob_dist['predicted_pam_custom']}")
+    print(f"Predicted PAM (default min_confidence=0.55): {prob_dist['predicted_pam']}")
+    
+    print("\nPosition-wise Probability Distribution:")
+    print("-" * 80)
+    print(f"{'Pos':<5} {'A':<10} {'C':<10} {'G':<10} {'T':<10} {'Pred':<6} {'Max Prob':<10} {'Entropy':<10}")
+    print("-" * 80)
+    
+    nucleotides = ['A', 'C', 'G', 'T']
+    for pos_info in prob_dist['per_position']:
+        pos = pos_info['position']
+        probs = pos_info['probabilities']
+        pred = pos_info['predicted']
+        max_prob = pos_info['max_probability']
+        entropy = pos_info['entropy']
+        
+        # Determine if this position would be 'N' with custom threshold
+        final_pred = 'N' if max_prob < min_confidence else pred
+        
+        print(f"{pos+1:<5} "
+              f"{probs['A']:<10.4f} "
+              f"{probs['C']:<10.4f} "
+              f"{probs['G']:<10.4f} "
+              f"{probs['T']:<10.4f} "
+              f"{final_pred:<6} "
+              f"{max_prob:<10.4f} "
+              f"{entropy:<10.4f}")
+    
+    print("-" * 80)
+
+
+def print_summary_distribution(pam_distribution_dict):
+    """Print summary distribution across all sequences."""
+    print("\n" + "="*80)
+    print("SUMMARY: PAM DISTRIBUTION ACROSS ALL SEQUENCES")
+    print("="*80)
+    
+    total = len(pam_distribution_dict)
+    pam_counter = Counter(pam_distribution_dict.values())
+    
+    print(f"\nTotal sequences: {total}")
+    print(f"Unique PAMs: {len(pam_counter)}")
+    
+    # Sort by frequency (descending)
+    sorted_pams = sorted(pam_counter.items(), key=lambda x: x[1], reverse=True)
+    
+    print("\nPAM Distribution (sorted by frequency):")
+    print("-" * 80)
+    print(f"{'PAM Sequence':<15} {'Count':<10} {'Percentage':<10}")
+    print("-" * 80)
+    
+    for pam, count in sorted_pams:
+        percentage = (count / total) * 100
+        print(f"{pam:<15} {count:<10} {percentage:>6.2f}%")
+    
+    print("-" * 80)
+    print(f"{'Total':<15} {total:<10} {'100.00%':<10}")
+    print("="*80)
+
+
+def main():
+    # Get sequences from the dictionary at the top
+    sequences_dict = SEQUENCES
+    
+    if not sequences_dict:
+        print("Error: SEQUENCES dictionary is empty. Please add sequences to the dictionary at the top of the script.")
+        return
+    
+    print(f"Processing {len(sequences_dict)} sequences...")
+    print(f"Using min_confidence threshold: {MIN_CONFIDENCE}")
+    
+    # Extract sequence names and values
+    sequence_names = list(sequences_dict.keys())
+    sequences_list = list(sequences_dict.values())
+    
+    # Set device
+    device = DEVICE
+    if device is not None:
+        device = torch.device(device)
+    
+    # Get probability distributions
+    prob_distributions = get_pam_probability_distributions(
+        sequences_list,
+        device=device,
+        model_name=MODEL_NAME,
+        min_confidence=MIN_CONFIDENCE
+    )
+    
+    # Create dictionary mapping sequence names to predicted PAMs
+    predicted_pams_dict = {
+        name: dist['predicted_pam_custom'] 
+        for name, dist in zip(sequence_names, prob_distributions)
+    }
+    
+    # Print detailed distributions for each sequence
+    for seq_name, prob_dist in zip(sequence_names, prob_distributions):
+        print_sequence_pam_distribution(seq_name, prob_dist, MIN_CONFIDENCE)
+    
+    # Print summary distribution
+    print_summary_distribution(predicted_pams_dict)
+    
+    # Return dictionaries for potential programmatic use
+    return predicted_pams_dict, prob_distributions
+
+
+if __name__ == "__main__":
+    main()

cas9/scripts/evaluate.sh

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+# Validation metrics over N sampling runs (mean/std/min/max).
+set -euo pipefail
+ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
+cd "$ROOT"
+python -u -m cas9.evaluate_val_multi_sample \
+  --n_runs "${N_RUNS:-10}" \
+  --config "$ROOT/configs/config_cas9.yaml" \
+  --ckpt "$ROOT/cas9/ckpt/last.ckpt" \
+  --output_dir "$ROOT/outputs/cas9_eval" \
+  --json_out "$ROOT/outputs/cas9_eval/multi_run_metrics.json" \
+  "$@"

cas9/scripts/generate.sh

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+# Unguided / seeded Cas9 rollouts (no objectives) -> one CSV per ortholog.
+set -euo pipefail
+ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
+cd "$ROOT"
+python -u -m cas9.generate_cas9_csvs \
+  --config "$ROOT/configs/config_cas9.yaml" \
+  --ckpt "$ROOT/cas9/ckpt/last.ckpt" \
+  --sequences "$ROOT/cas9/data/orthologs.fasta" \
+  --num_steps 32 \
+  --output_dir "$ROOT/outputs" \
+  "$@"

cas9/scripts/pcomole.sh

@@ -0,0 +1,36 @@
+#!/usr/bin/env bash
+# Official Cas9 pCoMole editor: cas9-likeness + deletion-count + PAM-matching,
+# under Cas9-score, PAM-matching and length constraints.
+set -euo pipefail
+ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
+cd "$ROOT"
+
+ORTHOLOG="${ORTHOLOG:-cjcas9}"
+INPUT_SEQUENCE="$(awk -v n=">$ORTHOLOG" '
+  $1==n {f=1; next} /^>/{f=0} f {printf "%s", $0}' cas9/data/orthologs.fasta)"
+if [ -z "$INPUT_SEQUENCE" ]; then
+  echo "Ortholog '$ORTHOLOG' not found in cas9/data/orthologs.fasta" >&2; exit 1
+fi
+
+python -u -m cas9.pcomole \
+  --pcomol_config "$ROOT/configs/config_cas9.yaml" \
+  --ckpt "$ROOT/cas9/ckpt/last.ckpt" \
+  --input "$INPUT_SEQUENCE" \
+  --cas9_objective --deletion_count_obj --pam_distr_objective \
+  --objective_weights 1 1 1 \
+  --target_pam matching \
+  --num_steps 5 \
+  --num_rollouts 5 \
+  --num_candidates 10 \
+  --num_sequences 50 \
+  --max_deletion_percentage 0.2 \
+  --deletion_rate_scale 2500 \
+  --pam_min_confidence 0.7 \
+  --pam_prediction_temperature 1 \
+  --pam_probability_threshold 0.0 \
+  --cas9_score_threshold 0.5 \
+  --max_target_length 930 \
+  --min_target_length 890 \
+  --use_ce_loss --allow_equal_length --zero_lam_ins \
+  --output_file "$ROOT/outputs/cas9_${ORTHOLOG}_pcomole.csv" \
+  "$@"

cas9/scripts/train.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+set -euo pipefail
+ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
+cd "$ROOT"
+python -u -m cas9.train --config configs/config_cas9.yaml "$@"

cas9/train.py

@@ -0,0 +1,275 @@
+import os
+import argparse
+from datasets import load_from_disk
+import torch
+from torch.utils.data import DataLoader
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
+from pytorch_lightning.loggers import WandbLogger
+
+from cas9.data import data
+from cas9.model.base_models import EditFlow, ProteinEditFlowModel, SMILESEditFlowModel, ReparameterizedProteinEditFlowModel, ReparameterizedSMILESEditFlowModel
+from smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
+from smiles_tokenizer.selfies_tokenizers import SelfiesTokenizer
+
+from cas9.logic import flow
+from transformers import EsmTokenizer
+import datetime
+
+import yaml
+from easydict import EasyDict as edict
+
+import pdb
+
+DEFAULT_CONFIG_PATH = 'configs/config_test.yaml'
+
+def main(config_path=None):
+    # Use provided config path or default to hardcoded path
+    if config_path is None:
+        config_path = DEFAULT_CONFIG_PATH
+    
+    with open(config_path, 'r') as f:
+        config_dict = yaml.safe_load(f)
+    cfg = edict(config_dict)
+
+    print(f"cfg: {cfg}")
+
+    run_name = f"lr{cfg.optim.lr}_epoch{cfg.optim.n_epochs}_scale{cfg.model.scale_size}_optimal{cfg.model.p_optimal}_{cfg.logging.run_name}"
+    workdir = os.path.join(cfg.work_dir, run_name)
+    os.makedirs(workdir, exist_ok=True)
+    
+    pl.seed_everything(cfg.training.seed, workers=True)
+
+    # Data  
+    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 = 1
+        bos_id = 0
+        eos_id = 2
+        if cfg.training.reparameterize:
+            model = ReparameterizedProteinEditFlowModel(vocab_size=vocab_size, pad_id=pad_id, config=cfg.model)
+        else:
+            model = ProteinEditFlowModel(vocab_size=vocab_size, pad_id=pad_id, config=cfg.model)
+            # model = SMILESEditFlowModel(vocab_size=vocab_size, pad_id=pad_id, config=cfg.model)
+
+        if cfg.training.no_training:
+            print(f"No training mode: freezing all parameters")
+            for param in model.parameters():
+                param.requires_grad = False
+            
+            first_param = next(model.parameters())
+            first_param.requires_grad = True
+            print(f"Keeping parameter with shape {first_param.shape} with requires_grad=True for backward compatibility")
+
+    # if cfg.task == 'protein':
+    #     tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
+    #     vocab_size = tokenizer.vocab_size
+    #     source_distribution = flow.get_source_distribution(
+    #         source_distribution=cfg.flow.source_distribution, vocab_size=vocab_size, special_token_ids=[0,1,2,3, 24, 25, 26, 27, 28, 29, 30, 31]
+    #     )
+    #     pad_id = 1
+    #     bos_id = 0
+    #     eos_id = 2
+    #     if cfg.training.reparameterize:
+    #         model = ReparameterizedProteinEditFlowModel(vocab_size=vocab_size, pad_id=pad_id, config=cfg.model)
+    #     else:
+    #         model = ProteinEditFlowModel(vocab_size=vocab_size, pad_id=pad_id, config=cfg.model)
+    #         # model = SMILESEditFlowModel(vocab_size=vocab_size, pad_id=pad_id, config=cfg.model)
+
+    #     if cfg.training.no_training:
+    #         print(f"No training mode: freezing all parameters")
+    #         for param in model.parameters():
+    #             param.requires_grad = False
+            
+    #         first_param = next(model.parameters())
+    #         first_param.requires_grad = True
+    #         print(f"Keeping parameter with shape {first_param.shape} with requires_grad=True for backward compatibility")
+
+            
+    # 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 = 1
+    #     bos_id = 0
+    #     eos_id = 2
+    #     model = ProteinEditFlowModel(vocab_size=vocab_size, pad_id=pad_id, config=cfg.model)
+    elif cfg.task == 'smiles':
+        vocab_size = 587
+        tokenizer = SMILES_SPE_Tokenizer('smiles_tokenizer/new_vocab.txt',
+                                         'smiles_tokenizer/new_splits.txt')
+        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 = 0
+        bos_id = 2
+        eos_id = 3
+
+        if cfg.training.reparameterize:
+            model = ReparameterizedSMILESEditFlowModel(vocab_size=vocab_size, pad_id=pad_id, config=cfg.model)
+        else:
+            model = SMILESEditFlowModel(vocab_size=vocab_size, pad_id=pad_id, config=cfg.model)
+    elif cfg.task == 'selfies':
+        vocab_size = 44
+        tokenizer = SelfiesTokenizer.load(cfg.data.get('tokenizer_path', 'path/to/tokenizer/vocab.json'))
+        source_distribution = flow.get_source_distribution(
+            source_distribution=cfg.flow.source_distribution, vocab_size=vocab_size, special_token_ids=[0,1,2]
+        )
+        pad_id = 0
+        bos_id = 1
+        eos_id = 2
+        if cfg.training.reparameterize:
+            print(f"Using reparameterized SMILES edit flow model")
+            model = ReparameterizedSMILESEditFlowModel(vocab_size=vocab_size, pad_id=pad_id, config=cfg.model)
+        else:
+            model = SMILESEditFlowModel(vocab_size=vocab_size, pad_id=pad_id, config=cfg.model)
+    else:
+        raise NotImplementedError
+
+    num_parameters = sum(p.numel() for p in model.parameters())
+    print(f"NUM PARAMETERS: {num_parameters}")
+    
+    # Print accumulation batch size
+    grad_accumulation_steps = getattr(cfg.training, 'grad_accumulation_steps', 1)
+    batch_size = getattr(cfg.training, 'batch_size', 1)
+    num_gpus = getattr(cfg.compute, 'ngpus', 1)
+    effective_batch_size = batch_size * grad_accumulation_steps * num_gpus
+    print(f"GRADIENT ACCUMULATION STEPS: {grad_accumulation_steps}")
+    print(f"EFFECTIVE BATCH SIZE (batch_size * grad_accum_steps * num_gpus): {effective_batch_size} (batch_size={batch_size}, grad_accum_steps={grad_accumulation_steps}, num_gpus={num_gpus})")
+    print(f"LAM_PROP: {getattr(cfg.training, 'lam_prop', None)}")
+
+    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
+    )
+
+    
+    print("--------------------------------")
+    print(f"run_name: {run_name}")
+    print(f"config.training.loc_prop_path: {getattr(cfg.training, 'loc_prop_path', False)}")
+    print(f"config.model.d_model: {cfg.model.d_model}")
+    print("--------------------------------")
+
+    # Dataloader
+    if cfg.task == 'protein':
+        # train_dataloader, val_dataloader = data.get_data_loaders(config=cfg, data_state=None)
+        train_dataset = load_from_disk(cfg.data.train_path)
+        val_dataset = load_from_disk(cfg.data.val_path)
+        
+        # Apply debug subset if specified
+        subset_ratio = getattr(cfg.data, 'debug_subset_ratio', None)
+        if subset_ratio is not None and subset_ratio < 1.0:
+            train_size = len(train_dataset)
+            train_subset_size = int(train_size * subset_ratio)
+            train_dataset = train_dataset.shuffle(seed=cfg.training.seed).select(range(train_subset_size))
+            
+            val_size = len(val_dataset)
+            val_subset_size = int(val_size * subset_ratio)
+            val_dataset = val_dataset.select(range(val_subset_size))
+            
+            print(f"DEBUG MODE: Using {train_subset_size}/{train_size} training samples ({subset_ratio*100:.1f}%) and {val_subset_size}/{val_size} validation samples ({subset_ratio*100:.1f}%)")
+        else:
+            train_dataset = train_dataset.shuffle(seed=cfg.training.seed)
+        
+        print(f"train_dataset: {train_dataset[0]['input_ids']}")
+        print(f"train_dataset length: {len(train_dataset[0]['input_ids'])}")
+        print(f"train_dataset keys: {train_dataset[0].keys()}")
+        train_dataloader = DataLoader(train_dataset, batch_size=None, shuffle=True, num_workers=4)
+        val_dataloader = DataLoader(val_dataset, batch_size=None, shuffle=False, num_workers=4)
+    elif cfg.task == 'smiles':
+        train_dataset = load_from_disk(cfg.data.train_path)
+        val_dataset = load_from_disk(cfg.data.val_path)
+        train_dataloader = DataLoader(train_dataset, batch_size=None, shuffle=True, num_workers=4)
+        val_dataloader = DataLoader(val_dataset, batch_size=None, shuffle=False, num_workers=4)
+    elif cfg.task == 'selfies':
+        train_dataset = load_from_disk(cfg.data.train_path)
+        val_dataset = load_from_disk(cfg.data.val_path)
+        
+        # Apply debug subset if specified
+        subset_ratio = getattr(cfg.data, 'debug_subset_ratio', None)
+        if subset_ratio is not None and subset_ratio < 1.0:
+            train_size = len(train_dataset)
+            train_subset_size = int(train_size * subset_ratio)
+            train_dataset = train_dataset.shuffle(seed=cfg.training.seed).select(range(train_subset_size))
+            
+            val_size = len(val_dataset)
+            val_subset_size = int(val_size * subset_ratio)
+            val_dataset = val_dataset.select(range(val_subset_size))
+            
+            print(f"DEBUG MODE: Using {train_subset_size}/{train_size} training samples ({subset_ratio*100:.1f}%) and {val_subset_size}/{val_size} validation samples ({subset_ratio*100:.1f}%)")
+        else:
+            train_dataset = train_dataset.shuffle(seed=cfg.training.seed)
+        
+        print(f"train_dataset: {train_dataset[0]['input_ids']}")
+        print(f"train_dataset length: {len(train_dataset[0]['input_ids'])}")
+        train_dataloader = DataLoader(train_dataset, batch_size=None, shuffle=True, num_workers=4)
+        val_dataloader = DataLoader(val_dataset, batch_size=None, shuffle=False, num_workers=4)
+    else:
+        raise NotImplementedError
+    
+    ckpt = ModelCheckpoint(
+        dirpath=os.path.join(workdir, "checkpoint"),
+        monitor="val_loss",          # the metric you log
+        mode="min",                  # lower is better
+        save_top_k=3,                # keep best 3
+        save_last=True,
+        filename="epoch{epoch:04d}-val{val_loss:.2f}",
+        auto_insert_metric_name=False,  # <- this stops the extra "val_loss=..."
+    )
+    lrmon = LearningRateMonitor(logging_interval="step")
+
+    # Get checkpoint path from config (if specified)
+    ckpt_path = getattr(cfg.training, 'ckpt_path', None) or getattr(cfg, 'ckpt_path', None)
+    
+    # Only resume wandb if we're actually resuming from a checkpoint
+    # Otherwise, start a fresh run to avoid step counter mismatches
+    # Setting resume=None ensures a fresh run, preventing "step less than current step" warnings
+    wandb_resume = "allow" if ckpt_path else None
+    
+    wandb_logger = WandbLogger(
+        project=cfg.logging.get('project', 'Gated proposal model'),
+        name=run_name,
+        entity=cfg.logging.get('entity') or None,
+        resume=wandb_resume,
+    )
+    
+    trainer = pl.Trainer(
+        default_root_dir=workdir,
+        accelerator="gpu" if torch.cuda.is_available() else "cpu",
+        devices=cfg.compute.ngpus,
+        strategy="ddp_find_unused_parameters_true" if cfg.compute.ngpus > 1 else "auto",
+        precision='bf16-mixed',
+        max_epochs=cfg.optim.n_epochs,
+        log_every_n_steps=10,
+        callbacks=[ckpt, lrmon],
+        enable_checkpointing=True,
+        gradient_clip_val=1.0,
+        deterministic=False,
+        logger=wandb_logger,
+    )
+
+    trainer.fit(editflow, train_dataloader, val_dataloader, ckpt_path=ckpt_path if ckpt_path else None)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Train EditFlow model')
+    parser.add_argument('--config', type=str, default=None,
+                        help='Path to config YAML file (default: uses hardcoded default config)')
+    args = parser.parse_args()
+    print(f"args.config: {args.config}")
+    main(config_path=args.config)

configs/config_cas9.yaml

@@ -0,0 +1,77 @@
+task: protein
+
+compute:
+  ngpus: 1
+  nodes: 1
+
+logging:
+  enable_wandb: False
+  project: pCoMole
+
+data:
+  # Cas9 training arrows are not shipped. Rebuild them with
+  #   python cas9/data/cas9_pre_batching.py
+  # (ESM-2 tokenized, length-bucketed) or point these at your own
+  # HuggingFace load_from_disk() dataset.
+  train_path: data/cas9/train
+  val_path: data/cas9/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.03
+  optimizer: AdamW
+  lr: 1.0e-5
+  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
+  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.25
+
+work_dir: outputs/cas9

paths.py

@@ -62,3 +62,67 @@ def smiles_vocab_files() -> tuple[Path, Path]:
 
 def selfies_vocab_path() -> Path:
     return REPO_ROOT / "smiles_tokenizer" / "selfies_vocab.json"
+
+
+# --- Cas9 ---------------------------------------------------------------
+
+def protein2pam_root() -> Path:
+    """Root of the protein2pam package (provides the custom ESM loader).
+
+    Weights come from the Hub ("Profluent-Bio/protein2pam-cas9_full"); this is
+    only the Python package that defines the model class.
+    """
+    env = os.environ.get("PROTEIN2PAM_ROOT")
+    root = Path(env).expanduser().resolve() if env else (REPO_ROOT / "cas9" / "protein2pam")
+    if not root.is_dir():
+        raise FileNotFoundError(
+            "protein2pam was not found. Install it next to the Cas9 code at "
+            f"{REPO_ROOT / 'cas9' / 'protein2pam'} or set PROTEIN2PAM_ROOT.\n"
+            "  pip install protein2pam==0.2.0"
+        )
+    return root
+
+
+def cas_predictor_root() -> Path:
+    """Root of the cas_predictor package (model.py + lightning_module.py)."""
+    env = os.environ.get("CAS_PREDICTOR_ROOT")
+    root = Path(env).expanduser().resolve() if env else (REPO_ROOT / "cas9" / "cas_predictor")
+    if not root.is_dir():
+        raise FileNotFoundError(
+            "cas_predictor was not found. Place it at "
+            f"{REPO_ROOT / 'cas9' / 'cas_predictor'} or set CAS_PREDICTOR_ROOT."
+        )
+    return root
+
+
+def cas9_classifier_ckpt() -> Path:
+    """Cas9-likeness classifier checkpoint (hard constraint + objective)."""
+    env = os.environ.get("CAS9_CLASSIFIER_CKPT")
+    p = resolve_path(env) if env else (REPO_ROOT / "cas9" / "classifier_ckpt" / "last.ckpt")
+    if not p.is_file():
+        raise FileNotFoundError(
+            f"Cas9 classifier checkpoint not found at {p}. Download it into "
+            "cas9/classifier_ckpt/ or set CAS9_CLASSIFIER_CKPT."
+        )
+    return p
+
+
+def cas9_classifier_config() -> Path | None:
+    """Optional classifier config; None falls back to the ckpt hyper_parameters."""
+    env = os.environ.get("CAS9_CLASSIFIER_CONFIG")
+    if env:
+        return resolve_path(env)
+    default = REPO_ROOT / "cas9" / "classifier_ckpt" / "config.yaml"
+    return default if default.is_file() else None
+
+
+def cas9_hmm_db(name: str) -> Path:
+    """HMM database under cas9/hmm/ (e.g. 'cas9_domain_detector_RuvCmix.hmm')."""
+    env = os.environ.get("CAS9_HMM_DIR")
+    root = Path(env).expanduser().resolve() if env else (REPO_ROOT / "cas9" / "hmm")
+    p = root / name
+    if not p.is_file():
+        raise FileNotFoundError(
+            f"HMM database {name} not found at {p}. Place it in cas9/hmm/ or set CAS9_HMM_DIR."
+        )
+    return p

requirements.txt

@@ -27,3 +27,12 @@ chemprop>=2.0.0
 
 # Optional logging
 # wandb
+
+# Cas9
+h5py                 # CICERO-format PAM datasets
+biopython            # FASTA / sequence handling
+rapidfuzz            # sequence similarity in Cas9 diagnostics
+# HMMER (`hmmscan`) must be on PATH for the Cas9 domain / PAM-interacting
+# detectors: conda install -c bioconda hmmer
+# protein2pam provides the PAM oracle's model class:
+#   pip install protein2pam==0.2.0