| """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() |
|
|
