File size: 3,731 Bytes
0f5ae29 | 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 | """Reconstruct DNA from raw mean-pooled embeddings saved as a NumPy array."""
import argparse
import json
from pathlib import Path
import numpy as np
import torch
from transformers import AutoTokenizer
from model import QueryDecoderReconstructor
class Decoder:
def __init__(self, model_dir, device="cpu"):
folder = Path(model_dir)
checkpoint = torch.load(folder / "model.pt", map_location="cpu", weights_only=True)
config = checkpoint["config"]["model"]
self.device = torch.device(device)
self.input_dim = checkpoint["input_dim"]
self.stats = json.loads((folder / "normalization.json").read_text())
assert self.stats["method"] == "standard"
assert np.isfinite([self.stats["mean"], self.stats["std"]]).all()
assert self.stats["std"] > 0
self.model = QueryDecoderReconstructor(
input_dim=self.input_dim, mode=checkpoint["mode"],
seq_length=checkpoint["effective_seq_length"], output_dim=checkpoint["output_dim"],
**{key: config[key] for key in ["d_model", "nhead", "num_layers", "dim_feedforward",
"dropout", "n_context_tokens", "aux_length_loss_weight"]},
)
self.model.load_state_dict(checkpoint["state_dict"], strict=True)
self.model.to(self.device).eval()
self.character = checkpoint["tokenizer_type"] == "char"
if self.character:
self.eos_id = 4
assert checkpoint["output_dim"] == 5
else:
self.tokenizer = AutoTokenizer.from_pretrained(folder / "tokenizer", local_files_only=True)
self.eos_id = json.loads((folder / "tokenizer_info.json").read_text())["eos_id"]
assert len(self.tokenizer) == checkpoint["output_dim"]
def decode_tokens(self, ids):
ids = list(ids)
if self.eos_id in ids:
ids = ids[:ids.index(self.eos_id)]
if self.character:
return "".join("ACGT"[i] for i in ids)
return self.tokenizer.decode(ids, skip_special_tokens=True).replace(" ", "")
@torch.inference_mode()
def predict(self, embeddings, batch_size=32):
raw = np.asarray(embeddings, dtype=np.float32)
assert raw.ndim == 2 and raw.shape[1] == self.input_dim, f"Expected [N, {self.input_dim}] embeddings"
assert np.isfinite(raw).all(), "Embeddings must be finite"
assert batch_size > 0
sequences = []
for start in range(0, len(raw), batch_size):
normalized = (raw[start:start + batch_size] - self.stats["mean"]) / self.stats["std"]
assert np.isfinite(normalized).all()
logits, _ = self.model(torch.from_numpy(normalized).to(self.device))
assert torch.isfinite(logits).all(), "Non-finite model output"
sequences.extend(self.decode_tokens(ids) for ids in logits.argmax(-1).cpu().tolist())
return sequences
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model", required=True, choices=["dnabert2", "evo2", "ntv2"])
parser.add_argument("--embeddings", required=True, type=Path)
parser.add_argument("--output", required=True, type=Path)
parser.add_argument("--device", default="cpu")
parser.add_argument("--batch-size", type=int, default=32)
args = parser.parse_args()
decoder = Decoder(Path(__file__).parent / args.model, args.device)
sequences = decoder.predict(np.load(args.embeddings, allow_pickle=False), args.batch_size)
args.output.write_text(json.dumps(sequences, indent=2) + "\n", encoding="utf-8")
print(f"Wrote {len(sequences)} reconstructions to {args.output}")
if __name__ == "__main__":
main()
|