Upload train_splice_cnn_v4.py with huggingface_hub

train_splice_cnn_v4.py

@@ -0,0 +1,278 @@
+# ============================================================
+# MutationPredictorCNN_v4 Training Script (401 bp FASTA)
+# Proper sequence-based training
+# ============================================================
+
+import argparse
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from sklearn.metrics import roc_auc_score
+import pysam
+from tqdm import tqdm
+import os
+
+# ============================================================
+# Arguments
+# ============================================================
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument("--train_csv", required=True)
+parser.add_argument("--fasta", required=True)
+parser.add_argument("--output_model", required=True)
+
+parser.add_argument("--epochs", type=int, default=30)
+parser.add_argument("--batch_size", type=int, default=256)
+parser.add_argument("--num_workers", type=int, default=8)
+parser.add_argument("--lr", type=float, default=0.001)
+
+args = parser.parse_args()
+
+# ============================================================
+# Config
+# ============================================================
+
+WINDOW = 401
+HALF = WINDOW // 2
+SEQ_LEN = WINDOW - 2
+DEVICE = "cpu"
+
+print("Loading FASTA...")
+fasta = pysam.FastaFile(args.fasta)
+
+# ============================================================
+# Encoding
+# ============================================================
+
+BASE_MAP = {"A":0,"C":1,"G":2,"T":3}
+COMP = {"A":"T","T":"A","C":"G","G":"C","N":"N"}
+
+def fetch_seq(chrom, pos):
+
+    start = pos - HALF - 1
+    end = pos + HALF
+
+    try:
+        return fasta.fetch(str(chrom), start, end).upper()
+    except:
+        try:
+            return fasta.fetch("chr"+str(chrom), start, end).upper()
+        except:
+            return None
+
+
+def encode_seq(seq):
+
+    arr = np.zeros((11, SEQ_LEN), dtype=np.float32)
+
+    for i in range(SEQ_LEN):
+
+        j = i + 1
+        base = seq[j] if j < len(seq) else "N"
+
+        if base in BASE_MAP:
+            arr[BASE_MAP[base], i] = 1
+            comp = COMP[base]
+            if comp in BASE_MAP:
+                arr[4 + BASE_MAP[comp], i] = 1
+
+        arr[8, i] = (j - HALF) / HALF
+
+        if seq[j:j+2] == "GT":
+            arr[9, i] = 1
+
+        if seq[j:j+2] == "AG":
+            arr[10, i] = 1
+
+    return arr
+
+
+def mut_onehot(ref, alt):
+
+    types = [
+        "A>C","A>G","A>T",
+        "C>A","C>G","C>T",
+        "G>A","G>C","G>T",
+        "T>A","T>C","T>G"
+    ]
+
+    vec = np.zeros(12, dtype=np.float32)
+
+    key = f"{ref}>{alt}"
+
+    if key in types:
+        vec[types.index(key)] = 1
+
+    return vec
+
+
+# ============================================================
+# Dataset
+# ============================================================
+
+class SpliceDataset(Dataset):
+
+    def __init__(self, df):
+
+        self.df = df.reset_index(drop=True)
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+
+        row = self.df.iloc[idx]
+
+        seq = fetch_seq(row.chrom, int(row.pos))
+
+        if seq is None or len(seq) != WINDOW:
+            seq = "N" * WINDOW
+
+        seq_enc = encode_seq(seq)
+
+        mut = mut_onehot(row.ref, row.alt)
+
+        region = np.zeros(2, dtype=np.float32)
+        splice = np.zeros(3, dtype=np.float32)
+
+        label = float(row.label)
+
+        return (
+            torch.tensor(seq_enc),
+            torch.tensor(mut),
+            torch.tensor(region),
+            torch.tensor(splice),
+            torch.tensor(label)
+        )
+
+
+# ============================================================
+# Model
+# ============================================================
+
+class MutationPredictorCNN_v4(nn.Module):
+
+    def __init__(self):
+
+        super().__init__()
+
+        self.conv1 = nn.Conv1d(11, 64, 7, padding=3)
+        self.conv2 = nn.Conv1d(64, 128, 5, padding=2)
+        self.conv3 = nn.Conv1d(128, 256, 3, padding=1)
+
+        self.pool = nn.AdaptiveAvgPool1d(1)
+
+        self.mut_fc = nn.Linear(12, 32)
+
+        self.region_fc = nn.Linear(2, 8)
+        self.splice_fc = nn.Linear(3, 16)
+
+        self.fc1 = nn.Linear(312, 128)
+        self.fc2 = nn.Linear(128, 64)
+        self.fc3 = nn.Linear(64, 1)
+
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(0.3)
+
+    def forward(self, seq, mut, region, splice):
+
+        x = self.relu(self.conv1(seq))
+        x = self.relu(self.conv2(x))
+        x = self.relu(self.conv3(x))
+
+        x = self.pool(x).squeeze(-1)
+
+        m = self.relu(self.mut_fc(mut))
+        r = self.relu(self.region_fc(region))
+        s = self.relu(self.splice_fc(splice))
+
+        x = torch.cat([x,m,r,s], dim=1)
+
+        x = self.dropout(self.relu(self.fc1(x)))
+        x = self.relu(self.fc2(x))
+
+        return self.fc3(x)
+
+
+# ============================================================
+# Load dataset
+# ============================================================
+
+print("Loading dataset...")
+
+df = pd.read_csv(args.train_csv)
+
+train_ds = SpliceDataset(df)
+
+train_dl = DataLoader(
+    train_ds,
+    batch_size=args.batch_size,
+    shuffle=True,
+    num_workers=args.num_workers
+)
+
+# ============================================================
+# Train
+# ============================================================
+
+model = MutationPredictorCNN_v4().to(DEVICE)
+
+criterion = nn.BCEWithLogitsLoss()
+
+optimizer = torch.optim.Adam(
+    model.parameters(),
+    lr=args.lr
+)
+
+best_auc = 0
+
+for epoch in range(args.epochs):
+
+    model.train()
+
+    losses = []
+    probs = []
+    labels = []
+
+    for seq, mut, region, splice, label in train_dl:
+
+        seq = seq.to(DEVICE)
+        mut = mut.to(DEVICE)
+        region = region.to(DEVICE)
+        splice = splice.to(DEVICE)
+        label = label.to(DEVICE).unsqueeze(1)
+
+        optimizer.zero_grad()
+
+        logits = model(seq, mut, region, splice)
+
+        loss = criterion(logits, label)
+
+        loss.backward()
+
+        optimizer.step()
+
+        losses.append(loss.item())
+
+        probs.extend(torch.sigmoid(logits).detach().cpu().numpy())
+        labels.extend(label.cpu().numpy())
+
+    auc = roc_auc_score(labels, probs)
+
+    print(f"Epoch {epoch+1}/{args.epochs} Loss={np.mean(losses):.4f} AUC={auc:.4f}")
+
+    if auc > best_auc:
+
+        best_auc = auc
+
+        torch.save(
+            model.state_dict(),
+            args.output_model
+        )
+
+        print("Saved best model")
+
+print("Training complete.")