Imranyai/CodonFM-80M-mRNA-stability
0
1"""2data_setup.py — Download, preprocess, audit, and prepare all datasets for3CodonFM-80M mRNA stability fine-tuning and benchmarking.4 5Usage:6 # Download and preprocess everything7 python data_setup.py --all8 9 # Download only training datasets10 python data_setup.py --training11 12 # Download only benchmark datasets13 python data_setup.py --benchmark14 15 # Audit datasets (inspect stats, find issues)16 python data_setup.py --audit17 18 # Export preprocessed training data to local files19 python data_setup.py --training --export ./processed_data20 21 # Show codon vocabulary and tokenizer details22 python data_setup.py --vocab23"""24 25import argparse26import json27import os28import sys29import csv30import urllib.request31from collections import Counter32from pathlib import Path33 34import numpy as np35 36try:37 import pandas as pd38except ImportError:39 pd = None40 41try:42 from datasets import load_dataset43except ImportError:44 load_dataset = None45 46 47# ============================================================48# 1. CODON VOCABULARY & TOKENIZER49# ============================================================50 51RNA_BASES = ['A', 'U', 'G', 'C']52ALL_CODONS = [b1 + b2 + b3 for b1 in RNA_BASES for b2 in RNA_BASES for b3 in RNA_BASES]53 54# Biological codon table (RNA) → Amino Acid55CODON_TABLE = {56 'UUU': 'Phe', 'UUC': 'Phe', 'UUA': 'Leu', 'UUG': 'Leu',57 'CUU': 'Leu', 'CUC': 'Leu', 'CUA': 'Leu', 'CUG': 'Leu',58 'AUU': 'Ile', 'AUC': 'Ile', 'AUA': 'Ile', 'AUG': 'Met/Start',59 'GUU': 'Val', 'GUC': 'Val', 'GUA': 'Val', 'GUG': 'Val',60 'UCU': 'Ser', 'UCC': 'Ser', 'UCA': 'Ser', 'UCG': 'Ser',61 'CCU': 'Pro', 'CCC': 'Pro', 'CCA': 'Pro', 'CCG': 'Pro',62 'ACU': 'Thr', 'ACC': 'Thr', 'ACA': 'Thr', 'ACG': 'Thr',63 'GCU': 'Ala', 'GCC': 'Ala', 'GCA': 'Ala', 'GCG': 'Ala',64 'UAU': 'Tyr', 'UAC': 'Tyr', 'UAA': 'Stop', 'UAG': 'Stop',65 'CAU': 'His', 'CAC': 'His', 'CAA': 'Gln', 'CAG': 'Gln',66 'AAU': 'Asn', 'AAC': 'Asn', 'AAA': 'Lys', 'AAG': 'Lys',67 'GAU': 'Asp', 'GAC': 'Asp', 'GAA': 'Glu', 'GAG': 'Glu',68 'UGU': 'Cys', 'UGC': 'Cys', 'UGA': 'Stop', 'UGG': 'Trp',69 'CGU': 'Arg', 'CGC': 'Arg', 'CGA': 'Arg', 'CGG': 'Arg',70 'AGU': 'Ser', 'AGC': 'Ser', 'AGA': 'Arg', 'AGG': 'Arg',71 'GGU': 'Gly', 'GGC': 'Gly', 'GGA': 'Gly', 'GGG': 'Gly',72}73 74# Token vocabulary (matches CodonFM config: vocab_size=69, pad_token_id=3)75SPECIAL_TOKENS = {'[CLS]': 0, '[SEP]': 1, '[MASK]': 2, '[PAD]': 3, '[UNK]': 4}76CODON_TO_ID = {codon: i + 5 for i, codon in enumerate(ALL_CODONS)}77ID_TO_CODON = {v: k for k, v in CODON_TO_ID.items()}78ID_TO_CODON.update({v: k for k, v in SPECIAL_TOKENS.items()})79 80VOCAB_SIZE = len(SPECIAL_TOKENS) + len(ALL_CODONS) # 5 + 64 = 6981assert VOCAB_SIZE == 6982 83PAD_TOKEN_ID = 384CLS_TOKEN_ID = 085SEP_TOKEN_ID = 186MASK_TOKEN_ID = 287UNK_TOKEN_ID = 488 89 90def seq_to_codons(seq: str) -> list:91 """Split an mRNA/DNA sequence into codon triplets."""92 seq = seq.upper().replace('T', 'U').strip()93 return [seq[i:i+3] for i in range(0, len(seq) - len(seq) % 3, 3)]94 95 96def tokenize_mRNA(seq: str, max_length: int = 2046) -> dict:97 """Tokenize an mRNA/DNA sequence into CodonFM token IDs."""98 codons = seq_to_codons(seq)99 token_ids = [CLS_TOKEN_ID]100 for codon in codons[:max_length - 2]:101 token_ids.append(CODON_TO_ID.get(codon, UNK_TOKEN_ID))102 token_ids.append(SEP_TOKEN_ID)103 attention_mask = [1] * len(token_ids)104 return {'input_ids': token_ids, 'attention_mask': attention_mask}105 106 107def validate_sequence(seq: str) -> dict:108 """Validate an mRNA/DNA sequence for CodonFM compatibility."""109 seq_clean = seq.upper().replace('T', 'U').strip()110 issues = []111 112 if len(seq_clean) == 0:113 issues.append("Empty sequence")114 if len(seq_clean) % 3 != 0:115 issues.append(f"Length {len(seq_clean)} not divisible by 3 (truncated to {len(seq_clean) - len(seq_clean) % 3})")116 117 invalid_chars = set(seq_clean) - {'A', 'U', 'G', 'C'}118 if invalid_chars:119 issues.append(f"Invalid characters: {invalid_chars}")120 121 codons = seq_to_codons(seq_clean)122 n_codons = len(codons)123 124 # Check for start codon125 starts_with_aug = codons[0] == 'AUG' if codons else False126 127 # Check for stop codons128 stop_codons = {'UAA', 'UAG', 'UGA'}129 internal_stops = [i for i, c in enumerate(codons[:-1]) if c in stop_codons]130 ends_with_stop = codons[-1] in stop_codons if codons else False131 132 if internal_stops:133 issues.append(f"Internal stop codons at positions: {internal_stops}")134 135 # Unknown codons136 unk_codons = [c for c in codons if c not in CODON_TO_ID]137 if unk_codons:138 issues.append(f"Unknown codons: {set(unk_codons)}")139 140 return {141 'valid': len(issues) == 0,142 'issues': issues,143 'length_nt': len(seq_clean),144 'length_codons': n_codons,145 'starts_with_AUG': starts_with_aug,146 'ends_with_stop': ends_with_stop,147 'n_internal_stops': len(internal_stops),148 }149 150 151# ============================================================152# 2. TRAINING DATASETS153# ============================================================154 155TRAINING_DATASETS = {156 'mogam-ai/CDS-BART-mRNA-stability': {157 'description': 'iCodon vertebrate mRNA stability profiles (human, mouse, frog, fish)',158 'source_paper': 'Diez et al. 2022, Scientific Reports "iCodon customizes gene expression based on the codon composition"',159 'seq_col': 'seq',160 'label_col': 'y',161 'splits': {'train': 'train', 'val': 'val', 'test': 'test'},162 'label_meaning': 'mRNA half-life z-score (higher = more stable, mean≈0, std≈1)',163 'species': ['Human', 'Mouse', 'Xenopus (frog)', 'Zebrafish'],164 'notes': 'RNA sequences (A,U,G,C). All divisible by 3. Subset of GleghornLab dataset.',165 },166 'GleghornLab/mrna_stability_other': {167 'description': 'Extended multi-species mRNA stability data (superset of mogam-ai dataset)',168 'source_paper': 'Li et al. 2024, Genome Research "CodonBERT large language model for mRNA vaccines"',169 'seq_col': 'rna',170 'label_col': 'labels',171 'splits': {'train': 'train', 'val': 'valid', 'test': 'test'},172 'label_meaning': 'mRNA half-life z-score (higher = more stable)',173 'species': ['Multiple vertebrate species'],174 'notes': 'Has extra "seqs" column (protein-encoded, not used). Contains 1 outlier sequence of 3 nt. Superset of mogam-ai.',175 'extra_col': 'seqs',176 },177}178 179 180def download_training_data(export_dir=None):181 """Download and inspect training datasets from HuggingFace Hub."""182 if load_dataset is None:183 print("ERROR: `datasets` library required. Run: pip install datasets")184 return None185 186 all_data = {}187 188 for repo_id, info in TRAINING_DATASETS.items():189 print(f"\n{'='*60}")190 print(f"📦 {repo_id}")191 print(f" {info['description']}")192 print(f"={'='*60}")193 194 ds = load_dataset(repo_id)195 196 for split_name, hf_split in info['splits'].items():197 split_data = ds[hf_split]198 seqs = split_data[info['seq_col']]199 labels = split_data[info['label_col']]200 201 print(f"\n [{split_name}] {len(seqs)} samples")202 print(f" Seq lengths (nt): min={min(len(s) for s in seqs)}, "203 f"mean={np.mean([len(s) for s in seqs]):.0f}, "204 f"max={max(len(s) for s in seqs)}")205 print(f" Seq lengths (cod): min={min(len(s)//3 for s in seqs)}, "206 f"mean={np.mean([len(s)//3 for s in seqs]):.0f}, "207 f"max={max(len(s)//3 for s in seqs)}")208 labels_arr = np.array(labels)209 print(f" Labels: min={labels_arr.min():.3f}, mean={labels_arr.mean():.3f}, "210 f"std={labels_arr.std():.3f}, max={labels_arr.max():.3f}")211 212 all_data[repo_id] = ds213 214 if export_dir:215 export_training_data(all_data, export_dir)216 217 return all_data218 219 220def preprocess_training_data(use_both_datasets=True, min_codons=3, max_codons=2046,221 remove_duplicates=True, deduplicate_across_splits=True):222 """223 Preprocess training data: clean, filter, deduplicate, and combine.224 225 Steps:226 1. Load both HF datasets227 2. Use GleghornLab as primary (superset) OR combine both228 3. Filter: remove sequences < min_codons or > max_codons codons229 4. Filter: remove sequences with NaN labels230 5. Filter: remove sequences with invalid characters231 6. Deduplicate: remove exact sequence duplicates within each split232 7. Deduplicate: ensure no train sequences appear in val/test (data leakage check)233 8. Return clean {train, val, test} dictionaries234 235 Returns:236 dict with 'train', 'val', 'test' keys, each containing 'sequences' and 'labels' lists237 """238 if load_dataset is None:239 raise ImportError("datasets library required: pip install datasets")240 241 print("Loading datasets...")242 243 if use_both_datasets:244 # Use GleghornLab (superset) — it contains ALL of mogam-ai plus extra data245 ds = load_dataset("GleghornLab/mrna_stability_other")246 raw_data = {247 'train': {'seqs': ds['train']['rna'], 'labels': ds['train']['labels']},248 'val': {'seqs': ds['valid']['rna'], 'labels': ds['valid']['labels']},249 'test': {'seqs': ds['test']['rna'], 'labels': ds['test']['labels']},250 }251 print(f" Using GleghornLab/mrna_stability_other (superset)")252 else:253 # Use mogam-ai only (smaller, cleaner)254 ds = load_dataset("mogam-ai/CDS-BART-mRNA-stability")255 raw_data = {256 'train': {'seqs': ds['train']['seq'], 'labels': ds['train']['y']},257 'val': {'seqs': ds['val']['seq'], 'labels': ds['val']['y']},258 'test': {'seqs': ds['test']['seq'], 'labels': ds['test']['y']},259 }260 print(f" Using mogam-ai/CDS-BART-mRNA-stability only")261 262 clean_data = {}263 total_removed = {'short': 0, 'long': 0, 'nan': 0, 'invalid': 0, 'duplicate': 0}264 265 for split in ['train', 'val', 'test']:266 seqs = raw_data[split]['seqs']267 labels = raw_data[split]['labels']268 orig_count = len(seqs)269 270 clean_seqs = []271 clean_labels = []272 seen = set()273 274 for seq, label in zip(seqs, labels):275 # Skip None/empty276 if seq is None or len(seq) == 0:277 total_removed['invalid'] += 1278 continue279 280 # Normalize: uppercase, T→U281 seq = seq.upper().replace('T', 'U').strip()282 283 # Check NaN label284 if np.isnan(label):285 total_removed['nan'] += 1286 continue287 288 # Check invalid characters289 if set(seq) - {'A', 'U', 'G', 'C'}:290 total_removed['invalid'] += 1291 continue292 293 # Check length294 n_codons = len(seq) // 3295 if n_codons < min_codons:296 total_removed['short'] += 1297 continue298 if n_codons > max_codons:299 total_removed['long'] += 1300 continue301 302 # Deduplicate within split303 if remove_duplicates:304 if seq in seen:305 total_removed['duplicate'] += 1306 continue307 seen.add(seq)308 309 clean_seqs.append(seq)310 clean_labels.append(float(label))311 312 clean_data[split] = {'sequences': clean_seqs, 'labels': clean_labels}313 print(f" [{split}] {orig_count} → {len(clean_seqs)} samples "314 f"(removed {orig_count - len(clean_seqs)})")315 316 # Cross-split deduplication: check for train/test leakage317 if deduplicate_across_splits:318 test_seqs = set(clean_data['test']['sequences'])319 val_seqs = set(clean_data['val']['sequences'])320 321 leakage_test = sum(1 for s in clean_data['train']['sequences'] if s in test_seqs)322 leakage_val = sum(1 for s in clean_data['train']['sequences'] if s in val_seqs)323 val_test_overlap = len(val_seqs & test_seqs)324 325 print(f"\n Data leakage check:")326 print(f" Train→Test overlap: {leakage_test} sequences")327 print(f" Train→Val overlap: {leakage_val} sequences")328 print(f" Val→Test overlap: {val_test_overlap} sequences")329 330 if leakage_test > 0 or leakage_val > 0:331 print(f" ⚠️ WARNING: Data leakage detected! Removing leaked sequences from train...")332 eval_seqs = test_seqs | val_seqs333 filtered_train = [(s, l) for s, l in334 zip(clean_data['train']['sequences'], clean_data['train']['labels'])335 if s not in eval_seqs]336 clean_data['train']['sequences'] = [x[0] for x in filtered_train]337 clean_data['train']['labels'] = [x[1] for x in filtered_train]338 print(f" Train after dedup: {len(clean_data['train']['sequences'])} samples")339 340 print(f"\n Removal summary: {total_removed}")341 print(f" Final sizes: train={len(clean_data['train']['sequences'])}, "342 f"val={len(clean_data['val']['sequences'])}, "343 f"test={len(clean_data['test']['sequences'])}")344 345 return clean_data346 347 348def export_training_data(data, export_dir):349 """Export preprocessed data to CSV files."""350 os.makedirs(export_dir, exist_ok=True)351 352 if isinstance(data, dict) and 'train' in data and 'sequences' in data.get('train', {}):353 # Already preprocessed format354 for split in ['train', 'val', 'test']:355 if split not in data:356 continue357 filepath = os.path.join(export_dir, f'{split}.csv')358 with open(filepath, 'w', newline='') as f:359 writer = csv.writer(f)360 writer.writerow(['sequence', 'stability_score'])361 for seq, label in zip(data[split]['sequences'], data[split]['labels']):362 writer.writerow([seq, label])363 print(f" Exported {split}: {len(data[split]['sequences'])} rows → {filepath}")364 else:365 print(" Export requires preprocessed data. Run preprocess_training_data() first.")366 367 368# ============================================================369# 3. BENCHMARK DATASETS370# ============================================================371 372CODONBERT_BASE_URL = "https://raw.githubusercontent.com/Sanofi-Public/CodonBERT/master/benchmarks/CodonBERT/data/fine-tune"373 374BENCHMARK_DATASETS = {375 'stability': {376 'url': f"{CODONBERT_BASE_URL}/mRNA_Stability.csv",377 'filename': 'mRNA_Stability.csv',378 'description': 'mRNA Stability (iCodon vertebrate mRNA half-life)',379 'source': 'Diez et al. 2022, Scientific Reports',380 'samples': 65356,381 'seq_length': '3-3066 nt (1-1022 codons)',382 'label': 'Half-life z-score (continuous, mean≈0, std≈1)',383 'metric': 'Spearman ρ',384 'columns': 'sequence, value, dataset, split',385 'species': 'Multi-vertebrate (human, mouse, frog, fish)',386 },387 'mrfp': {388 'url': f"{CODONBERT_BASE_URL}/mRFP_Expression.csv",389 'filename': 'mRFP_Expression.csv',390 'description': 'mRFP Protein Expression in E. coli',391 'source': 'Li et al. 2024, Genome Research (CodonBERT)',392 'samples': 1459,393 'seq_length': '678 nt (226 codons, fixed)',394 'label': 'Fluorescence intensity (log scale, range 7.4-11.4)',395 'metric': 'Spearman ρ',396 'columns': 'sequence, value, dataset, split',397 'species': 'E. coli (synthetic mRFP variants)',398 },399 'vaccine': {400 'url': f"{CODONBERT_BASE_URL}/CoV_Vaccine_Degradation.csv",401 'filename': 'CoV_Vaccine_Degradation.csv',402 'description': 'SARS-CoV-2 mRNA Vaccine Degradation',403 'source': 'CodonBERT benchmark (derived from Stanford OpenVaccine)',404 'samples': 2400,405 'seq_length': '81 nt (27 codons, fixed)',406 'label': 'Degradation score (z-normalized, range -7.2 to 6.5)',407 'metric': 'Spearman ρ',408 'columns': 'sequence, value, dataset, split',409 'species': 'Synthetic SARS-CoV-2 mRNA vaccine fragments',410 },411 'riboswitch': {412 'url': f"{CODONBERT_BASE_URL}/Tc-Riboswitches.csv",413 'filename': 'Tc-Riboswitches.csv',414 'description': 'Tetracycline Riboswitch Activity',415 'source': 'Groher et al. 2018 (via CodonBERT)',416 'samples': 355,417 'seq_length': '66-75 nt (22-25 codons)',418 'label': 'Switching factor (continuous, range -0.3 to 3.1)',419 'metric': 'Spearman ρ',420 'columns': 'sequence, value, dataset, split',421 'species': 'Synthetic tetracycline riboswitches',422 },423 'mlos': {424 'url': f"{CODONBERT_BASE_URL}/MLOS.csv",425 'filename': 'MLOS.csv',426 'description': 'MLOS Flu Vaccine Antigen Expression',427 'source': 'Ren et al. 2024 (HELM/MLOS)',428 'samples': 167,429 'seq_length': '~1700 nt (~567 codons)',430 'label': 'Expression level (continuous, range 0.3-2.2)',431 'metric': 'Spearman ρ',432 'columns': 'cds, value (no split column — uses random 70/15/15)',433 'species': 'Influenza haemagglutinin CDS variants',434 'notes': 'No pre-defined splits. Column name is "cds" not "sequence".',435 },436}437 438 439def download_benchmark_data(data_dir='./benchmark_data'):440 """Download all benchmark datasets."""441 os.makedirs(data_dir, exist_ok=True)442 443 for task_name, info in BENCHMARK_DATASETS.items():444 filepath = os.path.join(data_dir, info['filename'])445 if os.path.exists(filepath):446 size = os.path.getsize(filepath)447 print(f" ✓ {info['filename']} already exists ({size/1024:.1f} KB)")448 else:449 print(f" ↓ Downloading {info['filename']}...")450 try:451 urllib.request.urlretrieve(info['url'], filepath)452 size = os.path.getsize(filepath)453 print(f" ✓ Downloaded {info['filename']} ({size/1024:.1f} KB)")454 except Exception as e:455 print(f" ✗ Failed to download {info['filename']}: {e}")456 457 return data_dir458 459 460# ============================================================461# 4. AUDIT462# ============================================================463 464def audit_dataset(sequences, labels, name="dataset"):465 """Run a comprehensive audit on a list of sequences and labels."""466 print(f"\n{'='*60}")467 print(f"AUDIT: {name} ({len(sequences)} sequences)")468 print(f"{'='*60}")469 470 if len(sequences) == 0:471 print(" (empty)")472 return473 474 # ---- Sequence stats ----475 lengths_nt = [len(s) for s in sequences]476 lengths_codon = [len(s) // 3 for s in sequences]477 478 print(f"\n 📏 Sequence Lengths:")479 print(f" Nucleotides: min={min(lengths_nt)}, mean={np.mean(lengths_nt):.0f}, "480 f"median={np.median(lengths_nt):.0f}, max={max(lengths_nt)}")481 print(f" Codons: min={min(lengths_codon)}, mean={np.mean(lengths_codon):.0f}, "482 f"median={np.median(lengths_codon):.0f}, max={max(lengths_codon)}")483 484 # Length distribution buckets485 buckets = [0, 100, 300, 500, 1000, 2000, 3000, 10000]486 hist = np.histogram(lengths_codon, bins=buckets)[0]487 print(f" Codon length distribution:")488 for i, count in enumerate(hist):489 pct = 100 * count / len(sequences)490 bar = '█' * int(pct / 2)491 print(f" {buckets[i]:>5}-{buckets[i+1]:>5} codons: {count:>6} ({pct:>5.1f}%) {bar}")492 493 # Sequences > 2046 codons (CodonFM max)494 over_limit = sum(1 for c in lengths_codon if c > 2046)495 if over_limit > 0:496 print(f" ⚠️ {over_limit} sequences exceed CodonFM max (2046 codons) — will be truncated")497 498 # ---- Nucleotide composition ----499 all_chars = Counter()500 for s in sequences:501 all_chars.update(s.upper())502 total_bases = sum(all_chars.values())503 print(f"\n 🧬 Nucleotide Composition:")504 for base in ['A', 'U', 'G', 'C']:505 count = all_chars.get(base, 0)506 pct = 100 * count / total_bases507 print(f" {base}: {count:>12,} ({pct:.1f}%)")508 unexpected = {k: v for k, v in all_chars.items() if k not in 'AUGC'}509 if unexpected:510 print(f" ⚠️ Unexpected characters: {unexpected}")511 512 # Not divisible by 3513 not_div3 = sum(1 for s in sequences if len(s) % 3 != 0)514 if not_div3 > 0:515 print(f" ⚠️ {not_div3} sequences not divisible by 3")516 517 # ---- Codon usage ----518 codon_counts = Counter()519 for s in sequences[:5000]: # sample for speed520 codons = seq_to_codons(s)521 codon_counts.update(codons)522 523 print(f"\n 🔤 Codon Usage (top 10 / bottom 10 from {min(5000, len(sequences))} seqs):")524 sorted_codons = codon_counts.most_common()525 for codon, count in sorted_codons[:10]:526 aa = CODON_TABLE.get(codon, '?')527 print(f" {codon} ({aa:>9s}): {count:>8,}")528 print(f" ...")529 for codon, count in sorted_codons[-10:]:530 aa = CODON_TABLE.get(codon, '?')531 print(f" {codon} ({aa:>9s}): {count:>8,}")532 533 # Start/stop codon analysis534 starts_with_aug = sum(1 for s in sequences if s[:3].upper().replace('T', 'U') == 'AUG')535 stop_codons = {'UAA', 'UAG', 'UGA'}536 ends_with_stop = sum(1 for s in sequences537 if seq_to_codons(s)[-1] in stop_codons) if sequences else 0538 print(f"\n 🚦 Start/Stop Codons:")539 print(f" Starts with AUG: {starts_with_aug}/{len(sequences)} ({100*starts_with_aug/len(sequences):.1f}%)")540 print(f" Ends with stop: {ends_with_stop}/{len(sequences)} ({100*ends_with_stop/len(sequences):.1f}%)")541 542 # ---- Label stats ----543 labels_arr = np.array(labels, dtype=float)544 nan_count = np.isnan(labels_arr).sum()545 labels_clean = labels_arr[~np.isnan(labels_arr)]546 547 print(f"\n 📊 Label Distribution:")548 print(f" Count: {len(labels_arr)}, NaN: {nan_count}")549 if len(labels_clean) > 0:550 print(f" Min: {labels_clean.min():.4f}")551 print(f" Q1: {np.percentile(labels_clean, 25):.4f}")552 print(f" Median: {np.median(labels_clean):.4f}")553 print(f" Q3: {np.percentile(labels_clean, 75):.4f}")554 print(f" Max: {labels_clean.max():.4f}")555 print(f" Mean: {labels_clean.mean():.4f}")556 print(f" Std: {labels_clean.std():.4f}")557 print(f" Skew: {float(((labels_clean - labels_clean.mean()) ** 3).mean() / labels_clean.std() ** 3):.4f}")558 559 # ---- Duplicates ----560 unique_seqs = len(set(sequences))561 dup_count = len(sequences) - unique_seqs562 print(f"\n 🔍 Duplicates:")563 print(f" Unique sequences: {unique_seqs}")564 print(f" Duplicate sequences: {dup_count}")565 566 # ---- Outliers ----567 if len(labels_clean) > 0:568 q1, q3 = np.percentile(labels_clean, [25, 75])569 iqr = q3 - q1570 lower = q1 - 3 * iqr571 upper = q3 + 3 * iqr572 outliers = np.sum((labels_clean < lower) | (labels_clean > upper))573 print(f"\n ⚡ Outliers (>3 IQR):")574 print(f" Label outliers: {outliers}/{len(labels_clean)}")575 576 very_short = sum(1 for c in lengths_codon if c < 10)577 very_long = sum(1 for c in lengths_codon if c > 1000)578 print(f" Very short (<10 codons): {very_short}")579 print(f" Very long (>1000 codons): {very_long}")580 581 582def run_full_audit():583 """Run audit on all training and benchmark datasets."""584 print("=" * 70)585 print("FULL DATASET AUDIT")586 print("=" * 70)587 588 # Training datasets589 print("\n\n📚 TRAINING DATASETS")590 print("=" * 70)591 592 if load_dataset is not None:593 ds1 = load_dataset("mogam-ai/CDS-BART-mRNA-stability")594 for split in ['train', 'val', 'test']:595 audit_dataset(596 ds1[split]['seq'], ds1[split]['y'],597 f"mogam-ai/CDS-BART-mRNA-stability [{split}]"598 )599 600 ds2 = load_dataset("GleghornLab/mrna_stability_other")601 for split, hf_split in [('train', 'train'), ('val', 'valid'), ('test', 'test')]:602 audit_dataset(603 ds2[hf_split]['rna'], ds2[hf_split]['labels'],604 f"GleghornLab/mrna_stability_other [{split}]"605 )606 607 # Cross-dataset analysis608 print("\n\n📊 CROSS-DATASET ANALYSIS")609 print("=" * 60)610 ds1_all = set(ds1['train']['seq']) | set(ds1['val']['seq']) | set(ds1['test']['seq'])611 ds2_all = set(ds2['train']['rna']) | set(ds2['valid']['rna']) | set(ds2['test']['rna'])612 print(f" mogam-ai total unique: {len(ds1_all)}")613 print(f" GleghornLab total unique: {len(ds2_all)}")614 print(f" Overlap: {len(ds1_all & ds2_all)}")615 print(f" mogam-ai ⊂ GleghornLab: {ds1_all.issubset(ds2_all)}")616 print(f" GleghornLab-only: {len(ds2_all - ds1_all)}")617 else:618 print(" Skipping (datasets library not installed)")619 620 # Benchmark datasets621 print("\n\n📊 BENCHMARK DATASETS")622 print("=" * 70)623 624 if pd is not None:625 data_dir = download_benchmark_data()626 for task_name, info in BENCHMARK_DATASETS.items():627 filepath = os.path.join(data_dir, info['filename'])628 if not os.path.exists(filepath):629 continue630 df = pd.read_csv(filepath)631 df.columns = [c.lower().strip() for c in df.columns]632 633 seq_col = 'sequence' if 'sequence' in df.columns else 'cds'634 val_col = 'value'635 636 if seq_col in df.columns and val_col in df.columns:637 audit_dataset(638 df[seq_col].tolist(), df[val_col].tolist(),639 f"Benchmark: {task_name} ({info['filename']})"640 )641 else:642 print(" Skipping (pandas not installed)")643 644 645# ============================================================646# 5. VOCAB DISPLAY647# ============================================================648 649def show_vocab():650 """Display the full codon vocabulary with amino acid mapping."""651 print("=" * 70)652 print("CodonFM Tokenizer Vocabulary (vocab_size=69)")653 print("=" * 70)654 655 print("\n SPECIAL TOKENS:")656 print(f" {'ID':>4} {'Token':<10} {'Description'}")657 print(f" {'─'*4} {'─'*10} {'─'*30}")658 descriptions = {659 '[CLS]': 'Classification token (prepended)',660 '[SEP]': 'Separator token (appended)',661 '[MASK]': 'Mask token (for MLM pretraining)',662 '[PAD]': 'Padding token (pad_token_id=3)',663 '[UNK]': 'Unknown token (for invalid codons)',664 }665 for token, tid in sorted(SPECIAL_TOKENS.items(), key=lambda x: x[1]):666 print(f" {tid:>4} {token:<10} {descriptions.get(token, '')}")667 668 print(f"\n CODON TOKENS (64 codons → 20 amino acids + 3 stop):")669 print(f" {'ID':>4} {'Codon':<6} {'Amino Acid':<12} {'ID':>4} {'Codon':<6} {'Amino Acid':<12} "670 f"{'ID':>4} {'Codon':<6} {'Amino Acid':<12} {'ID':>4} {'Codon':<6} {'Amino Acid'}")671 print(f" {'─'*4} {'─'*6} {'─'*12} {'─'*4} {'─'*6} {'─'*12} "672 f"{'─'*4} {'─'*6} {'─'*12} {'─'*4} {'─'*6} {'─'*12}")673 674 items = sorted(CODON_TO_ID.items(), key=lambda x: x[1])675 for i in range(0, len(items), 4):676 row_parts = []677 for j in range(4):678 if i + j < len(items):679 codon, tid = items[i + j]680 aa = CODON_TABLE.get(codon, '?')681 row_parts.append(f" {tid:>4} {codon:<6} {aa:<12}")682 else:683 row_parts.append(f" {'':>4} {'':6} {'':12}")684 print("".join(row_parts))685 686 print(f"\n TOKENIZATION EXAMPLE:")687 example = "AUGGCAGCCGAGACUCGG"688 codons = seq_to_codons(example)689 tokens = tokenize_mRNA(example)690 print(f" Input: {example}")691 print(f" Codons: {' '.join(codons)}")692 print(f" Token IDs: {tokens['input_ids']}")693 decoded = [ID_TO_CODON.get(t, '?') for t in tokens['input_ids']]694 print(f" Decoded: {' '.join(decoded)}")695 696 print(f"\n CONFIG (matches nvidia/NV-CodonFM-Encodon-80M-v1/config.json):")697 print(f" vocab_size: 69")698 print(f" pad_token_id: 3 ([PAD])")699 print(f" max_position_embeddings: 2046 codons (~6138 nt)")700 print(f" position_embedding_type: rotary (RoPE, θ=10000)")701 702 703# ============================================================704# CLI705# ============================================================706 707def main():708 parser = argparse.ArgumentParser(709 description="Dataset setup for CodonFM-80M mRNA stability fine-tuning",710 formatter_class=argparse.RawDescriptionHelpFormatter,711 epilog="""712Examples:713 python data_setup.py --all # Download & audit everything714 python data_setup.py --training # Download training datasets715 python data_setup.py --training --export ./data # Download & export to CSV716 python data_setup.py --benchmark # Download benchmark datasets717 python data_setup.py --audit # Full audit of all datasets718 python data_setup.py --vocab # Show codon vocabulary719 python data_setup.py --preprocess # Preprocess & deduplicate720 """721 )722 723 parser.add_argument('--all', action='store_true', help='Download and audit everything')724 parser.add_argument('--training', action='store_true', help='Download training datasets from HF Hub')725 parser.add_argument('--benchmark', action='store_true', help='Download benchmark datasets from GitHub')726 parser.add_argument('--audit', action='store_true', help='Run full dataset audit')727 parser.add_argument('--preprocess', action='store_true', help='Preprocess training data (clean, deduplicate)')728 parser.add_argument('--vocab', action='store_true', help='Show codon vocabulary and tokenizer')729 parser.add_argument('--export', type=str, default=None, help='Export directory for preprocessed CSVs')730 parser.add_argument('--benchmark_dir', type=str, default='./benchmark_data',731 help='Directory for benchmark data')732 parser.add_argument('--use_both', action='store_true', default=True,733 help='Use both training datasets (default: True)')734 parser.add_argument('--mogam_only', action='store_true',735 help='Use only mogam-ai dataset (smaller, cleaner)')736 737 args = parser.parse_args()738 739 # Default: show help740 if not any([args.all, args.training, args.benchmark, args.audit, args.preprocess, args.vocab]):741 parser.print_help()742 return743 744 if args.vocab or args.all:745 show_vocab()746 747 if args.training or args.all:748 print("\n\n📦 DOWNLOADING TRAINING DATASETS")749 print("=" * 60)750 download_training_data(export_dir=args.export)751 752 if args.benchmark or args.all:753 print("\n\n📦 DOWNLOADING BENCHMARK DATASETS")754 print("=" * 60)755 download_benchmark_data(args.benchmark_dir)756 757 if args.preprocess or args.all:758 print("\n\n🔧 PREPROCESSING TRAINING DATA")759 print("=" * 60)760 use_both = not args.mogam_only761 clean_data = preprocess_training_data(use_both_datasets=use_both)762 if args.export:763 export_training_data(clean_data, args.export)764 765 if args.audit or args.all:766 print("\n\n🔍 RUNNING FULL AUDIT")767 run_full_audit()768 769 770if __name__ == '__main__':771 main()772