File size: 2,929 Bytes
12fea4a | 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 | 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]
|