preliminary plug + play

dpacman/data/compute_embeddings.py

@@ -0,0 +1,243 @@
+"""
+Plug-and-play embedding extraction for:
+  • Chromosome sequences (from raw UCSC JSON)  
+  • TF sequences (transcription_factors.fasta)
+
+Usage example (DNA + protein in one go):
+  module load miniconda/24.7.1
+  conda activate dpacman
+  python dpacman/data/compute_embeddings.py \
+    --genome-json-dir ../data_files/raw/genomes/hg38 \
+    --tf-fasta         ../data_files/processed/tfclust/hg38_tf/transcription_factors.fasta \
+    --chrom-model      caduceus \
+    --tf-model         esm-dbp \
+    --out-dir          ../data_files/processed/tfclust/hg38_tf/embeddings \
+    --device           cuda
+"""
+import os
+import re
+import argparse
+import json
+import numpy as np
+from pathlib import Path
+import torch
+from transformers import AutoTokenizer, AutoModel, AutoModelForMaskedLM, pipeline
+import esm
+
+# ---- model wrappers ----
+
+class CaduceusEmbedder:
+    def __init__(self, device, chunk_size=131_072, overlap=0):
+        """
+        device: 'cpu' or 'cuda'
+        chunk_size: max bases (and thus tokens) to send in one forward pass
+        overlap: how many bases each window overlaps the previous; 0 = no overlap
+        """
+        model_name = "kuleshov-group/caduceus-ph_seqlen-131k_d_model-256_n_layer-16"
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name, trust_remote_code=True
+        )
+        self.model = AutoModel.from_pretrained(
+            model_name, trust_remote_code=True
+        ).to(device).eval()
+        self.device     = device
+        self.chunk_size = chunk_size
+        self.step       = chunk_size - overlap
+
+    def embed(self, seqs):
+        all_embs = []
+        for seq in seqs:
+            window_vecs = []
+            # slide windows of up to chunk_size bases
+            for i in range(0, len(seq), self.step):
+                chunk = seq[i : i + self.chunk_size]
+                if not chunk:
+                    break
+                # enforce truncation so tokens <= chunk_size
+                toks = self.tokenizer(
+                    chunk,
+                    return_tensors="pt",
+                    padding=False,
+                    truncation=True,
+                    max_length=self.chunk_size
+                ).to(self.device)
+                with torch.no_grad():
+                    out = self.model(**toks).last_hidden_state
+                # mean-pool tokens → (D,)
+                window_vecs.append(out.mean(dim=1).squeeze(0).cpu())
+            # average over windows → one (D,) vector per full sequence
+            seq_emb = torch.stack(window_vecs, dim=0).mean(dim=0).numpy()
+            all_embs.append(seq_emb)
+        return np.vstack(all_embs)  # shape (N, D)
+
+class DNABertEmbedder:
+    def __init__(self, device):
+        self.tokenizer = AutoTokenizer.from_pretrained("zhihan1996/DNA_bert_6", trust_remote_code=True)
+        self.model     = AutoModel.from_pretrained("zhihan1996/DNA_bert_6", trust_remote_code=True).to(device)
+        self.device    = device
+
+    def embed(self, seqs):
+        embs = []
+        for s in seqs:
+            tokens = self.tokenizer(s, return_tensors="pt", padding=True)["input_ids"].to(self.device)
+            with torch.no_grad():
+                out = self.model(tokens).last_hidden_state.mean(1)
+            embs.append(out.cpu().numpy())
+        return np.vstack(embs)
+
+class NucleotideTransformerEmbedder:
+    def __init__(self, device):
+        # HF “feature-extraction” returns a list of (L, D) arrays for each input
+        # device: “cpu” or “cuda”
+        self.pipe = pipeline(
+            "feature-extraction",
+            model="InstaDeepAI/nucleotide-transformer-500m-1000g",
+            device= -1 if device=="cpu" else 0    # HF uses -1 for CPU, 0 for GPU #:contentReference[oaicite:0]{index=0}
+        )
+
+    def embed(self, seqs):
+        """
+        seqs: List[str] of raw DNA sequences
+        returns: (N, D) array, one D-dim vector per sequence
+        """
+        all_embeddings = self.pipe(seqs, truncation=True, padding=True)
+        # all_embeddings is a List of shape (L, D) arrays
+        pooled = [ np.mean(x, axis=0) for x in all_embeddings ]
+        return np.vstack(pooled)       
+
+class ESMEmbedder:
+    def __init__(self, device):
+        self.model, self.alphabet = esm.pretrained.esm1b_t33_650M_UR50S()
+        self.batch_converter = self.alphabet.get_batch_converter()
+        self.model.to(device).eval()
+        self.device = device
+
+    def embed(self, seqs):
+        batch = [(str(i), seq) for i, seq in enumerate(seqs)]
+        _, _, toks = self.batch_converter(batch)
+        toks = toks.to(self.device)
+        with torch.no_grad():
+            results = self.model(toks, repr_layers=[33], return_contacts=False)
+        reps = results["representations"][33]
+        return reps[:, 1:-1].mean(1).cpu().numpy()
+
+class ESMDBPEmbedder:
+    def __init__(self, device):
+        # Load a local ESM-DBP model from pretrained directory
+        model_path = Path(__file__).resolve().parent.parent / 'pretrained'/ 'ESM-DBP'/ 'ESM-DBP.model'
+        self.model, self.alphabet = esm.pretrained.load_model_and_alphabet_and_params(str(model_path))
+        self.batch_converter = self.alphabet.get_batch_converter()
+        self.model.to(device).eval()
+        self.device = device
+
+    def embed(self, seqs):
+        batch = [(str(i), seq) for i, seq in enumerate(seqs)]
+        _, _, toks = self.batch_converter(batch)
+        toks = toks.to(self.device)
+        with torch.no_grad():
+            results = self.model(toks, repr_layers=[33], return_contacts=False)
+        reps = results["representations"][33]
+        return reps[:, 1:-1].mean(1).cpu().numpy()
+
+class GPNEmbedder:
+    def __init__(self, device):
+        model_name = "songlab/gpn-msa-sapiens"
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForMaskedLM.from_pretrained(model_name)
+        self.model.to(device)
+        self.model.eval()
+        self.device = device
+
+    def embed(self, seqs):
+        inputs = self.tokenizer(
+            seqs,
+            return_tensors="pt",
+            padding=True,
+            truncation=True
+        ).to(self.device)
+
+        with torch.no_grad():
+            last_hidden = self.model(**inputs).last_hidden_state
+        return last_hidden.mean(dim=1).cpu().numpy()
+
+class ProGenEmbedder:
+    def __init__(self, device):
+        model_name = "jinyuan22/ProGen2-base"
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModel.from_pretrained(model_name).to(device).eval()
+        self.device = device
+
+    def embed(self, seqs):
+        inputs = self.tokenizer(
+            seqs,
+            return_tensors="pt",
+            padding=True,
+            truncation=True
+        ).to(self.device)
+        with torch.no_grad():
+            last_hidden = self.model(**inputs).last_hidden_state
+        return last_hidden.mean(dim=1).cpu().numpy()
+
+# ---- main pipeline ----
+
+def get_embedder(name, device, for_dna=True):
+    name = name.lower()
+    if for_dna:
+        if name=="caduceus":   return CaduceusEmbedder(device)
+        if name=="dnabert":    return DNABertEmbedder(device)
+        if name=="nucleotide": return NucleotideTransformerEmbedder(device)
+        if name=="gpn":        return GPNEmbedder(device)
+    else:
+        if name in ("esm",):    return ESMEmbedder(device)
+        if name in ("esm-dbp","esm_dbp"): return ESMDBPEmbedder(device)
+        if name=="progen":      return ProGenEmbedder(device)
+    raise ValueError(f"Unknown model {name} (for_dna={for_dna})")
+
+
+def embed_and_save(seqs, ids, embedder, out_path):
+    embs = embedder.embed(seqs)
+    np.save(out_path, embs)
+    with open(out_path.with_suffix(".ids"), "w") as f:
+        f.write("\n".join(ids))
+
+if __name__=="__main__":
+    p = argparse.ArgumentParser()
+    p.add_argument("--genome-json-dir", default="data_files/raw/genomes/hg38", help="dir of UCSC JSONs")
+    p.add_argument("--tf-fasta",      required=True, help="input TF FASTA file")
+    p.add_argument("--chrom-model",   default="caduceus")
+    p.add_argument("--tf-model",      default="esm-dbp")
+    p.add_argument("--out-dir",       default="data_files/processed/tfclust/hg38_tf/embeddings")
+    p.add_argument("--device",        default="cpu")
+    args = p.parse_args()
+
+    os.makedirs(args.out_dir, exist_ok=True)
+    device = args.device
+
+    #Load only primary chromosome JSONs (chr1–22, X, Y, M)
+    genome_dir = Path(args.genome_json_dir)
+    chrom_seqs, chrom_ids = [], []
+    primary_pattern = re.compile(r"^hg38_chr(?:[1-9]|1[0-9]|2[0-2]|X|Y|M)\.json$")
+    for j in sorted(genome_dir.iterdir()):
+        if not primary_pattern.match(j.name):
+            continue
+        data = json.loads(j.read_text())
+        seq  = data.get("dna") or data.get("sequence")
+        chrom = data.get("chrom") or j.stem.split("_")[-1]
+        chrom_seqs.append(seq)
+        chrom_ids.append(chrom)
+    chrom_embedder = get_embedder(args.chrom_model, device, for_dna=True)
+    out_chrom = Path(args.out_dir)/f"chrom_{args.chrom_model}.npy"
+    embed_and_save(chrom_seqs, chrom_ids, chrom_embedder, out_chrom)
+
+    #Load TF sequences
+    tf_seqs, tf_ids = [], []
+    with open(args.tf_fasta) as f:
+        for header in f:
+            seq = next(f).strip()
+            tf_ids.append(header[1:].split()[0])
+            tf_seqs.append(seq)
+    tf_embedder = get_embedder(args.tf_model, device, for_dna=False)
+    out_tf = Path(args.out_dir)/f"tf_{args.tf_model}.npy"
+    embed_and_save(tf_seqs, tf_ids, tf_embedder, out_tf)
+
+    print("Done.")