recent changes

configs/data_task/clean/remap.yaml

@@ -1,5 +1,5 @@
 name: remap
-type: clean
+task_type: clean
 
 nr_raw_path: dpacman/data_files/raw/remap/remap2022_nr_macs2_hg38_v1_0.bed
 nr_processed_dir: dpacman/data_files/processed/remap

dpacman/classifier/loss.py

@@ -1,58 +1,170 @@
 """
-Define loss functions needed for training the model
+Define loss functions needed for training the model — padding safe (-1 sentinel)
 """
 
 import torch
-from torch.nn import functional as F
+import torch.nn.functional as F
 
+def _expand_like(mask: torch.Tensor, like: torch.Tensor):
+    # Make mask broadcastable to logits/targets (handles (B,L) vs (B,L,1))
+    while mask.dim() < like.dim():
+        mask = mask.unsqueeze(-1)
+    return mask.expand_as(like)
 
-def bce_loss_masked(logits, targets, nonpeak_mask, pos_weight=None):
+def bce_loss_masked(logits, targets, nonpeak_mask, pos_weight=None, eps=1e-8):
     """
-    Compute the masked Binary Cross Entropy, only on certain positions.
-    We will only compute BCE on positions whre nonpeak_mask == 1.0; the mask represents non-peak positions
+    Compute masked BCE with logits over non-peak positions only.
+    Expects nonpeak_mask already broadcastable to logits.
     """
+    # Clamp targets into [0,1] to be safe, even if pads slip through earlier
+    t = targets.clamp(0.0, 1.0)
     loss = F.binary_cross_entropy_with_logits(
-        logits, targets, reduction="none", pos_weight=pos_weight
+        logits, t, reduction="none", pos_weight=pos_weight
     )
-    denom = nonpeak_mask.sum().clamp_min(1.0)
-    return (loss * nonpeak_mask).sum() / denom
-
+    m = _expand_like(nonpeak_mask, loss).to(loss.dtype)
+    denom = m.sum().clamp_min(eps)
+    return (loss * m).sum() / denom
 
 def mse_peaks_only(logits, targets, peak_mask, eps=1e-8):
     """
-    Calculate MSE on peaks only.
+    Calculate MSE on peaks only (on probabilities), masking everything else.
     """
     probs = torch.sigmoid(logits)
-    mse_peaks = F.mse_loss(probs * peak_mask, targets * peak_mask, reduction="sum") / (
-        peak_mask.sum() + eps
-    )
-    return mse_peaks
-
+    per_elem = F.mse_loss(probs, targets, reduction="none")
+    m = _expand_like(peak_mask, per_elem).to(per_elem.dtype)
+    denom = m.sum().clamp_min(eps)
+    return (per_elem * m).sum() / denom
 
-def calculate_loss(logits, targets, eps=1e-8, alpha=1.0, gamma=1.0):
+def calculate_loss(
+    logits,
+    targets,
+    eps: float = 1e-8,
+    alpha: float = 1.0,
+    gamma: float = 1.0,
+    pos_weight=None,
+    pad_value: float = -1.0,
+):
     """
-    Combine masked-BCE + global-MSE to get a loss vlaue
+    Combine masked-BCE (non-peak) + masked-MSE on probs (peak), ignoring padding.
+    Assumes targets == -1 are pads; non-peak = 0; peak > 0.
     """
-    # Calculate peak and non-peak masks.
-    # Anything outside a peak will have a label equal to 0.
-    nonpeak_mask = (targets == 0).float()
-    peak_mask = (targets > 0).float()
+    valid = (targets != pad_value)
 
-    bce_nonpeak = bce_loss_masked(logits, targets, nonpeak_mask)
-    mse_peak = mse_peaks_only(logits, targets, peak_mask, eps=eps)
+    # Peak / non-peak masks that exclude pads
+    nonpeak_mask = valid & (targets == 0)
+    peak_mask    = valid & (targets > 0)
 
-    loss = alpha * bce_nonpeak + gamma * mse_peak
+    # For safety, zero-out targets at pad positions so they never feed into BCE/MSE
+    targets_safe = torch.where(valid, targets, torch.zeros_like(targets))
 
-    return loss
+    bce_nonpeak = bce_loss_masked(logits, targets_safe, nonpeak_mask, pos_weight=pos_weight, eps=eps)
+    mse_peak    = mse_peaks_only(logits, targets_safe, peak_mask, eps=eps)
 
+    return alpha * bce_nonpeak + gamma * mse_peak
 
-def accuracy_percentage(logits, targets, peak_thresh=0.5):
+def accuracy_percentage(
+    logits,
+    targets,
+    peak_thresh: float = 0.5,
+    eps: float = 1e-8,
+    pad_value: float = -1.0,
+):
     """
-    Compute accuracy in predicting high-confidence peaks (probability > 0.5)
+    Compute accuracy for predicting high-confidence peaks (prob >= 0.5), ignoring padding.
     """
+    valid = (targets != pad_value)
     probs = torch.sigmoid(logits)
-    preds_bin = (probs >= 0.5).float()
-    labels = (targets >= peak_thresh).float()
-    correct = (preds_bin == labels).float().sum()
-    total = torch.numel(labels)
-    return (correct / max(1, total)).item() * 100.0
+    preds_bin = (probs >= 0.5)
+    labels    = (targets >= peak_thresh)
+
+    v = _expand_like(valid, preds_bin)
+    correct = ((preds_bin == labels) & v).to(torch.float32).sum()
+    total   = v.to(torch.float32).sum().clamp_min(eps)
+    return (correct / total).item() * 100.0
+
+if __name__ == "__main__":
+    import torch
+
+    torch.manual_seed(0)
+    PAD = -1.0
+
+    def make_targets_BL(B=2, L=8, pad_positions=(6, 7)):
+        """Create (B,L) targets: 0=non-peak, >0=peak, -1=pad."""
+        t = torch.zeros(B, L)
+        # sprinkle a few peaks (values in [0.6, 1.0])
+        t[:, 1] = torch.rand(B) * 0.4 + 0.6
+        t[:, 3] = torch.rand(B) * 0.4 + 0.6
+        # pads
+        for p in pad_positions:
+            t[:, p] = PAD
+        return t
+
+    def make_targets_BLC(B=2, L=8, C=3, pad_positions=(6, 7)):
+        """
+        Create (B,L,C) targets by broadcasting a (B,L) base across channels
+        (so masking needs to expand correctly).
+        """
+        base = make_targets_BL(B, L, pad_positions)  # (B,L)
+        t = base.unsqueeze(-1).expand(-1, -1, C).clone()
+        # Make channel 1 slightly different to show per-channel variety
+        t[..., 1] = torch.where(t[..., 1] > 0, (t[..., 1] * 0.85).clamp(0, 1), t[..., 1])
+        return t
+
+    def mask_stats(name, logits, targets, pad_value=PAD):
+        valid = (targets != pad_value)
+        nonpeak_mask = valid & (targets == 0)
+        peak_mask    = valid & (targets > 0)
+
+        m_nonpeak = _expand_like(nonpeak_mask, logits)
+        m_peak    = _expand_like(peak_mask, logits)
+
+        print(f"\n[{name}]")
+        print(f"  logits.shape  = {tuple(logits.shape)}")
+        print(f"  targets.shape = {tuple(targets.shape)}")
+        # Previews (first batch)
+        if targets.dim() == 2:  # (B,L)
+            print(f"  targets[0,:] preview: {targets[0]}")
+        else:  # (B,L,C)
+            print(f"  targets[0,:,0] ch0 preview: {targets[0,:,0]}")
+            print(f"  targets[0,:,1] ch1 preview: {targets[0,:,1]}")
+        # Mask counts after EXPANSION (these define denominators)
+        print(f"  #non-peak elems used = {m_nonpeak.sum().item():.0f}")
+        print(f"  #peak elems used     = {m_peak.sum().item():.0f}")
+
+    # =========================
+    # Case A: (B, L)
+    # =========================
+    B, L = 2, 8
+    logits_BL  = torch.randn(B, L)                 # raw scores
+    targets_BL = make_targets_BL(B, L)             # 0, >0, and -1 pads
+
+    mask_stats("BL", logits_BL, targets_BL, pad_value=PAD)
+
+    loss_BL = calculate_loss(
+        logits_BL, targets_BL, pad_value=PAD, alpha=1.0, gamma=1.0
+    )
+    acc_BL = accuracy_percentage(
+        logits_BL, targets_BL, pad_value=PAD, peak_thresh=0.5
+    )
+
+    print(f"  loss_BL = {loss_BL.item():.6f}")
+    print(f"  acc_BL  = {acc_BL:.2f}%")
+
+    # =========================
+    # Case B: (B, L, C)
+    # =========================
+    B, L, C = 2, 8, 3
+    logits_BLC  = torch.randn(B, L, C)             # raw scores with channels
+    targets_BLC = make_targets_BLC(B, L, C)        # broadcasted targets + tweaks
+
+    mask_stats("BLC", logits_BLC, targets_BLC, pad_value=PAD)
+
+    loss_BLC = calculate_loss(
+        logits_BLC, targets_BLC, pad_value=PAD, alpha=1.0, gamma=1.0
+    )
+    acc_BLC = accuracy_percentage(
+        logits_BLC, targets_BLC, pad_value=PAD, peak_thresh=0.5
+    )
+
+    print(f"  loss_BLC = {loss_BLC.item():.6f}")
+    print(f"  acc_BLC  = {acc_BLC:.2f}%")

dpacman/data_modules/pair.py

@@ -314,7 +314,7 @@ class ShelfCollator:
         tr_key: str = "tr_sequence",
         dna_key: str = "dna_sequence",
         dtype: torch.dtype = torch.float32,
-        pad_value: float = 0.0,
+        pad_value: float = -1.0,
     ):
         self.tr_path = tr_shelf_path
         self.dna_path = dna_shelf_path

dpacman/data_tasks/embeddings/embedders.py

@@ -26,6 +26,8 @@ from sklearn.preprocessing import OneHotEncoder
 import math
 import rootutils
 from dpacman.utils import pylogger
+from tqdm import trange
+from tqdm.contrib.logging import logging_redirect_tqdm
 
 root = rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
 logger = pylogger.RankedLogger(__name__, rank_zero_only=True)
@@ -44,7 +46,7 @@ class CaduceusEmbedder:
             model_name, trust_remote_code=True
         )
         self.model = (
-            AutoModel.from_pretrained(model_name, trust_remote_code=True)
+            AutoModelForMaskedLM.from_pretrained(model_name, trust_remote_code=True)
             .to(device)
             .eval()
         )
@@ -52,42 +54,39 @@ class CaduceusEmbedder:
         self.chunk_size = chunk_size
         self.step = chunk_size - overlap
 
-    def embed(self, seqs, batch_size=1):
+    def embed(self, seqs, batch_size=1,  pooling=False):
         """
         seqs: List[str] of DNA sequences (each <= chunk_size for this test)
         returns: np.ndarray of shape (N, L, D), raw per‐token embeddings
         """
-        # outputs = []
-        # for seq in seqs:
-        #     # --- new: raw per‐token embeddings in one shot ---
-        #     toks = self.tokenizer(
-        #         seq,
-        #         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  # (1, L, D)
-        #     outputs.append(out.cpu().numpy()[0])             # (L, D)
-
-        # return np.stack(outputs, axis=0)  # (N, L, D)
-        outputs = []
-        for seq in tqdm(
-            seqs, total=len(seqs), desc="DNA: Caduceus", dynamic_ncols=True
-        ):
-            toks = self.tokenizer(
-                seq,
-                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  # (1, L, D)
-            outputs.append(out.cpu().numpy()[0])  # (L, D)
-        return outputs  # list of variable-length (L_i, D) arrays
+        n = len(seqs)
+        if n == 0:
+            return {}
 
+        # (Optional) quick info; uses logger if provided, else print
+        max_len = max(len(s) for s in seqs)
+        logger.info(f"Max length (will be padded/truncated to tokenizer setting): {max_len}")
+    
+        outputs = {}  # seq -> embedding
+
+        with logging_redirect_tqdm():
+            for i in range(0, n, batch_size):
+                batch_seqs = seqs[i : i + batch_size]
+                logger.info(f"Embedding batch {n//(batch_size*(i+1))}")
+        
+                for seq in tqdm(batch_seqs, total=len(batch_seqs), desc="DNA: Caduceus", dynamic_ncols=True):
+                    toks = self.tokenizer(  # note: the tokenization
+                        seq,
+                        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  # (1, L+1, D)
+                    outputs[seq] = out.cpu().numpy().squeeze(0)[0:-1,:]        # (L, D)
+            return outputs  # list of variable-length (L_i, D) arrays
+        
     def benchmark(self, lengths=None):
         """
         Time embedding on single-sequence of various lengths.

dpacman/scripts/run_embeddings.sh

@@ -8,10 +8,11 @@ timestamp=$(date "+%Y-%m-%d_%H-%M-%S")
 run_dir="$HOME/DPACMAN/logs/${main_task}/${data_task_type}/runs/${timestamp}"
 mkdir -p "$run_dir"
 
-nohup python -u -m scripts.preprocess \
+nohup python -s -u -m scripts.preprocess \
   hydra.run.dir="${run_dir}" \
-  data_task="${data_task_type}/protein" \
-  data_task.debug="false" \
+  data_task="${data_task_type}/dna" \
+  data_task.chrom_model="caduceus" \
+  data_task.debug="true" \
   > "${run_dir}/run.log" 2>&1 &
 
 echo $! > "${run_dir}/pid.txt"

dpacman/scripts/run_train.sh

@@ -16,6 +16,13 @@ fi
 
 nohup python -u -m scripts.train \
   hydra.run.dir="${run_dir}" \
+  data_module.train_file="data_files/processed/splits/by_dna/train.csv" \
+  data_module.val_file="data_files/processed/splits/by_dna/val.csv" \
+  data_module.test_file="data_files/processed/splits/by_dna/test.csv" \
+  data_module.tr_shelf_path="data_files/processed/embeddings/fimo_hits_only/trs_esm.shelf" \
+  data_module.dna_shelf_path="data_files/processed/embeddings/fimo_hits_only/peaks_caduceus.shelf" \
+  model.glm_input_dim=256 \
+  model.compressed_dim=1029 \
   > "${run_dir}/run.log" 2>&1 &
 
 echo $! > "${run_dir}/pid.txt"

h100_env.yaml

@@ -43,6 +43,7 @@ dependencies:
     - tqdm==4.67.1
     - matplotlib==3.10.3
     - transformers==4.55.2
+    - huggingface_hub==0.34.4
     - biopython==1.85
     - ortools==9.14.6206
     - fair-esm==2.0.0

h100_env2.yaml

@@ -0,0 +1,57 @@
+name: dnabind3
+channels:
+  - conda-forge
+  - defaults
+
+dependencies:
+  - python=3.10
+  - pip>=24
+  # compiled / heavy libs via conda-forge
+  - numpy>=2.0,<3.0
+  - scikit-learn>=1.5,<1.7
+  - pandas>=2.2,<2.3
+  - matplotlib>=3.8,<3.11
+  - lxml>=5.2,<6
+  - lightning=2.5.1
+  - torchmetrics>=1.3
+  - dask
+  - distributed
+  - dask-ml
+  # toolchain for JIT/building CUDA extensions (mamba-ssm, Triton kernels)
+  - cuda-toolkit=12.4
+  - cmake
+  - ninja
+
+  - pip:
+      # Force CUDA wheels and keep them from being overwritten by CPU builds
+      - --index-url=https://download.pytorch.org/whl/cu124
+      - torch==2.6.0+cu124
+
+      # HF stack + hard deps used at runtime
+      - transformers==4.53.0
+      - tokenizers>=0.21,<0.22
+      - safetensors>=0.4.3
+      - huggingface-hub==0.34.4
+      - regex
+
+      # Your libs
+      - rootutils==1.0.7
+      - hydra-core==1.3.2
+      - hydra-colorlog==1.2.0
+      - omegaconf==2.3.0
+      - pymex==0.9.31
+      - gitpython==3.1.44
+      - black==25.1.0
+      - tqdm==4.67.1
+      - biopython==1.85
+      - ortools==9.14.6206
+      - fair-esm==2.0.0
+      - rich==14.1.0
+      - wandb==0.21.1
+
+      # Mamba + Triton (for CUDA kernels)
+      - mamba-ssm==2.2.4
+      - triton>=3.0,<3.5
+
+      # your package
+      - -e .