pku-embl/metagenomics-skills/caduceus

import importlib.metadata as md
import json, os, sys
from pathlib import Path
import torch
from transformers import AutoModelForMaskedLM, AutoTokenizer
if len(sys.argv) != 2:
    raise SystemExit("usage: python smoke_caduceus.py FRESH_OUTPUT_DIR")
out = Path(sys.argv[1])
if out.exists():
    raise SystemExit(f"refusing to overwrite: {out}")
out.mkdir(parents=True)
log = (out / "run.log").open("w")
sys.stdout = log; sys.stderr = log
model_name = os.getenv(
    "CADUCEUS_MODEL",
    "kuleshov-group/caduceus-ph_seqlen-131k_d_model-256_n_layer-16")
revision = os.getenv(
    "CADUCEUS_REVISION", "b0477522ac5d044ad03578aa724ec8e4bdbd405b")
sequence = os.getenv("CADUCEUS_SMOKE_SEQUENCE", "ACGTACGTACGT")
if not sequence or set(sequence.upper()) - set("ACGTN"):
    raise ValueError("smoke sequence must contain only A/C/G/T/N")
tokenizer = AutoTokenizer.from_pretrained(
    model_name, revision=revision, trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained(
    model_name, revision=revision, trust_remote_code=True).eval()
batch = tokenizer(sequence, return_tensors="pt", truncation=False)
with torch.inference_mode():
    logits = model(**batch).logits
assert logits.shape[:2] == batch["input_ids"].shape
torch.save(logits.cpu(), out / "smoke-logits.pt")
metadata = {
    "model": model_name, "revision": revision, "sequence_length": len(sequence),
    "token_shape": list(batch["input_ids"].shape), "logit_shape": list(logits.shape),
    "torch": torch.__version__, "transformers": md.version("transformers")}
(out / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
assert (out / "smoke-logits.pt").stat().st_size > 0
assert (out / "metadata.json").stat().st_size > 0
print(json.dumps(metadata, sort_keys=True))
log.flush()
sys.stdout = sys.__stdout__; sys.stderr = sys.__stderr__
log.close()

Long DNA representation LM?
├─ RC-aware / long-range Mamba → caduceus
├─ Multi-species efficient transformer → dnabert2
├─ NTv3 / NT family → nucleotide-transformer
└─ Generative all-domains → evo2