File size: 2,835 Bytes
6dd9839 | 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 | #!/usr/bin/env python3
"""Encode an amino-acid sequence to a ProRiboGen-compatible VESM3B H5."""
from __future__ import annotations
import argparse
import os
import re
import subprocess
import sys
import tempfile
from pathlib import Path
from common import PKG_ROOT, device, load_env, resolve, vesm_paths
def write_fasta(path: Path, p_id: str, sequence: str) -> None:
seq = re.sub(r"\s+", "", sequence).upper()
if not seq or not re.fullmatch(r"[ACDEFGHIKLMNPQRSTVWY*BXZ]+", seq):
raise ValueError("Protein sequence is empty or contains invalid characters")
path.write_text(f">{p_id}\n{seq}\n", encoding="utf-8")
def encode_protein_to_h5(
protein_seq: str,
*,
p_id: str = "QUERY",
output_h5: Path | None = None,
fp16: bool = True,
) -> Path:
load_env()
base, weights = vesm_paths()
if not base.is_dir():
raise FileNotFoundError(f"VESM 基座不存在: {base}")
if not weights.is_file():
raise FileNotFoundError(f"VESM 权重不存在: {weights}")
work = resolve(os.environ.get("WORKSPACE", "workspace"))
work.mkdir(parents=True, exist_ok=True)
out = output_h5 or (work / f"{p_id}_vesm3b.h5")
with tempfile.TemporaryDirectory(dir=work) as tmp:
fa = Path(tmp) / f"{p_id}.fasta"
write_fasta(fa, p_id, protein_seq)
script = PKG_ROOT / "protein_encoder" / "build_vesm3b_protein_embeddings_h5.py"
cmd = [
sys.executable,
str(script),
"--base-model-dir",
str(base),
"--vesm-weights",
str(weights),
"--fasta",
str(fa),
"--output",
str(out),
"--device",
device(),
]
if fp16:
cmd.append("--fp16")
subprocess.run(cmd, check=True)
return out
def main() -> None:
ap = argparse.ArgumentParser(description="Protein AA sequence → VESM3B H5")
ap.add_argument("--protein", required=True, help="AA sequence, or @path/to.fasta")
ap.add_argument("--p-id", default="QUERY")
ap.add_argument("--output", type=Path, default=None)
ap.add_argument("--no-fp16", action="store_true")
args = ap.parse_args()
load_env()
protein = args.protein
if protein.startswith("@"):
text = Path(protein[1:]).read_text(encoding="utf-8")
lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
if lines and lines[0].startswith(">"):
protein = "".join(lines[1:])
if args.p_id == "QUERY":
args.p_id = lines[0][1:].split()[0]
else:
protein = "".join(lines)
out = encode_protein_to_h5(
protein, p_id=args.p_id, output_h5=args.output, fp16=not args.no_fp16
)
print(f"OK -> {out}")
if __name__ == "__main__":
main()
|