File size: 5,253 Bytes
7c33ad4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | import argparse
import torch
import yaml
from easydict import EasyDict as edict
from pathlib import Path
import sys
REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from constraints import GFP
from model.reparam_models import EditFlow, ProteinEditFlowModel
from model.utils import generate_from_x0, generate_from_x0_multi_edit
from logic import flow
# tokenizers used in train.py
from transformers import EsmTokenizer
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 = 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)
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, tokenizer, bos_id, eos_id, device):
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()
# 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]
return tokenizer.batch_decode([seq], skip_special_tokens=True)[0]
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--config", type=str, default="../configs/config_gfp.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)
parser.add_argument("--num_samples", type=int, default=1)
parser.add_argument("--output_csv", type=str, default=None)
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)
editflow.load_state_dict(ckpt["state_dict"], strict=False)
model = editflow.model.to(device)
model.eval()
x0 = tokenize_input_str(args.input, tokenizer, bos_id, eos_id, device)
allowed_tokens = torch.tensor(
[tok for tok in source_dist._allowed_tokens if tok not in (eps_id,) and tok not in range(24,33)],
device=device,
dtype=torch.long,
)
samples = []
for _ in range(args.num_samples):
x_gen = generate_from_x0_multi_edit(
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 = detokenize_output(x_gen, tokenizer, bos_id, eos_id, pad_id)
out_str = out_str.replace(' ', '')
print(len(out_str))
print('----------------------------')
print(f"Input Sequence: {args.input}\n")
print(f"Designed Sequence: {out_str}\n")
gfp_classifier = GFP(device)
gfp_probs = gfp_classifier.get_scores(out_str, return_probs=True)
print(gfp_probs)
samples.append(out_str)
if args.output_csv:
with open(args.output_csv, 'a') as f:
for sample in samples:
f.write(sample + '\n')
if __name__ == "__main__":
main()
|