File size: 2,535 Bytes
29899b4 bc0d37c 29899b4 bc0d37c 29899b4 bc0d37c 29899b4 bc0d37c 29899b4 bc0d37c 29899b4 bc0d37c 29899b4 bc0d37c 29899b4 bc0d37c 29899b4 bc0d37c 29899b4 bc0d37c 29899b4 bc0d37c 29899b4 bc0d37c 29899b4 | 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 | """
Embed DNA sequences from ReMap peaks.
"""
from .utils import pad_token_embeddings, embed_and_save
from dpacman.data_tasks.embeddings import get_embedder
import os
import torch
import json
import pandas as pd
from pathlib import Path
from omegaconf import DictConfig
import rootutils
from dpacman.utils import pylogger
root = rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
logger = pylogger.RankedLogger(__name__, rank_zero_only=True)
def main(cfg: DictConfig):
logger.info(
f"Making embeddings using {cfg.data_task.chrom_model} for dna sequences at {cfg.data_task.input_file}"
)
# make out dir if necessary
out_dir = Path(root) / cfg.data_task.out_dir
os.makedirs(out_dir, exist_ok=True)
# set device
device = "cpu"
if cfg.data_task.device == "gpu":
if torch.cuda.is_available():
device = "cuda"
logger.info(f"Using device: {device}")
# read the input file
input_file = Path(root) / cfg.data_task.input_file
if str(input_file).endswith(".json"):
# load the json and isolate the sequences and ids
with open(input_file, "r") as f:
d = json.load(f)
df = pd.DataFrame.from_dict(d, orient="index").reset_index()
df.columns = ["seq_id", "sequence"]
if cfg.data_task.debug:
logger.info(f"DEBUG MODE. Only embedding 5 sequences")
df = df.sample(n=5, random_state=42).reset_index(drop=True)
# crucial: sort by SIZE so that we are padding things in a way that will support pre-batching
df["sequence_length"] = df["sequence"].str.len()
df = df.sort_values(by="sequence_length", ascending=True).reset_index(drop=True)
# turn into list of sequences and IDs
peak_seqs = df["sequence"].tolist()
peak_ids = df["seq_id"].tolist()
logger.info(
f"Embedding {len(peak_seqs)} binding peak sequences from processed remap data"
)
# Get the DNA embedder
dna_embedder = get_embedder(cfg.data_task.chrom_model, device, for_dna=True)
logger.info(f"Device of embedding model: {dna_embedder.device}")
out_peaks = str(out_dir / f"peaks_{cfg.data_task.chrom_model}.pkl")
if cfg.data_task.debug:
out_peaks = out_peaks.replace(".pkl", "_debug.pkl")
embed_and_save(
peak_seqs,
peak_ids,
dna_embedder,
out_peaks,
batch_size=cfg.data_task.batch_size,
)
logger.info(f"Finished embedding DNA sequences. Saved to: {out_peaks}")
if __name__ == "__main__":
main()
|