pCoMole/peptide/generate.py

#!/usr/bin/env python3
"""
EditFlows sequence generation script.

Takes checkpoint and sampling parameters, and either a single input sequence
or a FASTA file. For each input, runs the EditFlows model to generate one
new sequence and prints it as it is generated. When input is a FASTA, the
script outputs a FASTA of all generated sequences (to stdout and optionally
to a file via --output_fasta).
"""

import argparse
import sys

import torch
import yaml
from easydict import EasyDict as edict

import sys
sys.path.append('/scratch/pranamlab/tong/pCoMol')
from model.reparam_models import EditFlow, ProteinEditFlowModel
from model.utils import generate_from_x0, generate_from_x0_multi_edit
from transformers import AutoModel, AutoTokenizer
from logic import flow

# from pcomol_cas import short_rollout_batch
import pdb

def build_model_and_stuff(cfg, device):
    """
    Rebuild exactly what train.py builds, but we won't set up lightning Trainer.
    Returns:
      editflow_module  (LightningModule)
      source_dist
      (pad_id, bos_id, eos_id)
    """
    tokenizer = AutoTokenizer.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)

    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 main():
    # parser = argparse.ArgumentParser(
    #     description="Generate sequences with EditFlows model from one or more inputs."
    # )
    # parser.add_argument(
    #     "--pcomol_config",
    #     type=str,
    #     required=True,
    #     help="Path to pcomol/config YAML (e.g. configs/config_peptides.yaml)",
    # )
    # parser.add_argument(
    #     "--ckpt",
    #     type=str,
    #     required=True,
    #     help="Path to model checkpoint (.ckpt)",
    # )
    # group = parser.add_mutually_exclusive_group(required=True)
    # group.add_argument(
    #     "--input",
    #     type=str,
    #     help="Single input sequence string",
    # )
    # group.add_argument(
    #     "--input_fasta",
    #     type=str,
    #     help="Path to FASTA file of input sequences",
    # )
    # parser.add_argument(
    #     "--output_fasta",
    #     type=str,
    #     default=None,
    #     help="When using --input_fasta, write generated sequences to this FASTA file (default: only stdout)",
    # )

    # # Sampling parameters (aligned with pcomol_cas.py)
    # 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,
    #     help="Scaling factor for deletion rate (default: 1)",
    # )
    # parser.add_argument(
    #     "--zero_lam_ins",
    #     action="store_true",
    #     help="Set insertion rate to 0 (deletions and substitutions only)",
    # )
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str, default="./configs/config_test.yaml")
    parser.add_argument("--ckpt", type=str, required=True, help="path to lightning checkpoint (.ckpt)")
    parser.add_argument("--input", type=str, required=True, help="input x_0 as raw string (smiles/protein/selfies)")
    parser.add_argument("--num_steps", type=int, default=32)
    parser.add_argument("--max_len_cap", type=int, default=None)
    parser.add_argument("--op_temperature", type=float, default=1)
    parser.add_argument("--token_temperature", type=float, default=1)
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    with open(args.config, "r") as f:
        cfg = edict(yaml.safe_load(f))

    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, weights_only=False)
    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)

    x0 = tokenizer(args.input, return_tensors='pt')['input_ids'].to(device)

    with torch.no_grad():
        x_T = generate_from_x0(
            model,
            x0,
            pad_id=pad_id,
            bos_id=bos_id,
            eos_id=eos_id,
            allowed_tokens=allowed_tokens,
            num_steps=args.num_steps,
            max_len_cap=args.max_len_cap,
            op_temperature=args.op_temperature,      # soften op choice
            token_temperature=args.token_temperature,   # soften token choice
        )

    out_str = tokenizer.decode(x_T.squeeze(0))
    out_str = out_str.replace(" ", "")[5:-5].strip()

    print(out_str)

if __name__ == "__main__":
    main()