"""bioai.training.train_vae -- train the DiscreteVAE on 200-nt dsRNA precursors. CLI:: python -m bioai.training.train_vae --epochs 10 Extracts 200-nt precursors from ``data/synthetic/pest_transcripts.fasta`` (or real data if available at ``data/external/pest_transcripts.fasta``), tiles them into 200-nt windows with 50% overlap, and trains the VAE for the given number of epochs. Saves the checkpoint to ``checkpoints/vae_best.pt`` under the project root. This is a STRETCH GOAL: even an untrained VAE is enough for the demo (we just need the ``sample()`` method to produce *some* 200-nt precursor that the downstream ranker can dice into siRNAs). But training does make the samples look more like real transcripts (GC content, k-mer frequencies). """ from __future__ import annotations import argparse import sys from pathlib import Path from typing import List, Tuple import torch import torch.nn.functional as F from torch.utils.data import DataLoader, Dataset from ..models.discrete_vae import DiscreteVAE from ..models.sirna_cnn import resolve_device from ..sequence_utils import fasta_iter, tile_sequence # Portable paths (resolved from bioai.paths) from bioai.paths import VAE_CHECKPOINT as CHECKPOINT_PATH # noqa: E402 from bioai.paths import EXTERNAL_DIR, SYNTHETIC_DIR # noqa: E402 # A/C/G/T -> {0,1,2,3} BASE_TO_IDX = {"A": 0, "C": 1, "G": 2, "T": 3, "U": 3} # --------------------------------------------------------------------------- # # Dataset # --------------------------------------------------------------------------- # class PrecursorDataset(Dataset): def __init__(self, fasta_paths: List[Path], seq_len: int = 200, step: int = 100, max_per_transcript: int = 5, max_total: int = 2000): self.seq_len = seq_len self.records: List[torch.Tensor] = [] for path in fasta_paths: if not path.exists(): continue for _, seq in fasta_iter(path): # skip overly short transcripts if len(seq) < seq_len: continue windows = tile_sequence(seq, window=seq_len, step=step, max_candidates=max_per_transcript) for _, _, sub in windows: if len(self.records) >= max_total: break toks = [BASE_TO_IDX.get(b, 0) for b in sub[:seq_len]] if len(toks) < seq_len: # pad with A toks = toks + [0] * (seq_len - len(toks)) self.records.append(torch.tensor(toks, dtype=torch.long)) if len(self.records) >= max_total: break if len(self.records) >= max_total: break if not self.records: # fallback: generate random sequences so training still runs print("[train_vae] no real transcripts found -- using random 200-mers") for _ in range(256): self.records.append(torch.randint(0, 4, (seq_len,))) def __len__(self) -> int: return len(self.records) def __getitem__(self, idx: int) -> torch.Tensor: return self.records[idx] # --------------------------------------------------------------------------- # # Training # --------------------------------------------------------------------------- # def train( epochs: int = 10, batch_size: int = 32, lr: float = 1e-3, device: str = "auto", seq_len: int = 200, fasta_paths: List[Path] | None = None, checkpoint_path: Path | None = None, ) -> str: device_t = resolve_device(device) print(f"[train_vae] device = {device_t}") if fasta_paths is None: candidates = [ EXTERNAL_DIR / "pest_transcripts.fasta", SYNTHETIC_DIR / "pest_transcripts.fasta", ] fasta_paths = candidates print(f"[train_vae] FASTA sources: {[str(p) for p in fasta_paths]}") ds = PrecursorDataset(fasta_paths, seq_len=seq_len) print(f"[train_vae] {len(ds)} precursors") loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False) model = DiscreteVAE( seq_len=seq_len, vocab_size=4, latent_dim=64, hidden_dim=128, num_heads=4, num_layers=2, ).to(device_t) optimizer = torch.optim.Adam(model.parameters(), lr=lr) checkpoint_path = checkpoint_path or CHECKPOINT_PATH checkpoint_path.parent.mkdir(parents=True, exist_ok=True) best_loss = float("inf") for epoch in range(1, epochs + 1): model.train() running_total = 0.0 running_recon = 0.0 running_kl = 0.0 n_batches = 0 for tokens in loader: tokens = tokens.to(device_t) optimizer.zero_grad() recon_logits, mu, logvar = model(tokens) loss, info = DiscreteVAE.elbo_loss(recon_logits, tokens, mu, logvar, beta=1.0) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() running_total += info["total"] running_recon += info["recon"] running_kl += info["kl"] n_batches += 1 avg_total = running_total / max(1, n_batches) avg_recon = running_recon / max(1, n_batches) avg_kl = running_kl / max(1, n_batches) print( f"Epoch {epoch:3d}/{epochs}: " f"loss={avg_total:.4f} recon={avg_recon:.4f} kl={avg_kl:.4f}" ) if avg_total < best_loss: best_loss = avg_total torch.save(model.state_dict(), checkpoint_path) print(f" -> saved checkpoint to {checkpoint_path}") # quick sanity: sample one precursor and check GC content model.eval() with torch.no_grad(): sample = model.sample(1, device=device_t).cpu().numpy()[0] gc = (sum(1 for t in sample if t in (1, 2)) / len(sample)) * 100 print(f"[train_vae] sample GC = {gc:.1f}% (target ~50%)") print(f"[train_vae] done. best_loss={best_loss:.4f}") return str(checkpoint_path) # --------------------------------------------------------------------------- # # CLI # --------------------------------------------------------------------------- # def main(argv: List[str] | None = None) -> int: p = argparse.ArgumentParser(description="Train the DiscreteVAE on 200-nt precursors.") p.add_argument("--epochs", type=int, default=10) p.add_argument("--batch-size", type=int, default=32) p.add_argument("--lr", type=float, default=1e-3) p.add_argument("--device", type=str, default="auto", choices=["auto", "cpu", "cuda"]) p.add_argument("--seq-len", type=int, default=200) p.add_argument("--checkpoint", type=str, default=str(CHECKPOINT_PATH)) args = p.parse_args(argv) ckpt = train( epochs=args.epochs, batch_size=args.batch_size, lr=args.lr, device=args.device, seq_len=args.seq_len, checkpoint_path=Path(args.checkpoint), ) print(f"[train_vae] checkpoint: {ckpt}") return 0 if __name__ == "__main__": sys.exit(main())