Upload README.md with huggingface_hub

README.md

@@ -1,123 +1,115 @@
 ---
-license: mit
-base_model: zehui127/Omni-DNA-20M
+license: apache-2.0
 tags:
-  - dna
-  - genomics
-  - mutation
-  - horizontal-gene-transfer
-  - sad
-  - sequential-attenuation-denoising
+- dna
+- genomics
+- mutation-prediction
+- sad
+- omni-dna
+- hyenadna
+- causal-lm
+base_model: zehui127/Omni-DNA-20M
 pipeline_tag: text-generation
 ---
 
 # Omni-DNA SAD Mutation Model
 
-Fine-tuned **Omni-DNA-20M** for cross-domain horizontal gene transfer (HGT) mutation prediction using **SAD (Sequential Attenuation Denoising)**.
+Fine-tuned Omni-DNA-20M for cross-domain HGT (Horizontal Gene Transfer) mutation prediction using **SAD (Sequential Attenuation Denoising)**.
 
 ## Architecture
 
-- **Base model**: [zehui127/Omni-DNA-20M](https://huggingface.co/zehui127/Omni-DNA-20M) — OLMo-based causal LM, 20M parameters
-- **Pretraining scope**: Multi-task, cross-taxonomic-domain genomic data (broadest among candidates)
-- **Context**: 250 BPE tokens (~500-2000 bp)
+- **Base**: [Omni-DNA-20M](https://huggingface.co/zehui127/Omni-DNA-20M) (OLMo-based, 20M params, BPE tokenizer, 250-token context)
+- **Purpose**: Given a parent DNA sequence, predict the mutated child sequence
+- **Format**: `"mutate: {parent} -> {child}"` — instruction-tuned as a causal LM
 
 ## SAD Pipeline
 
 ```
-Stage 1 (ICI)                    Stage 2 (Attenuation)
-─────────────────────            ─────────────────────────
-Fine-tune on synthetic pairs     Fine-tune on real pairs
-Omni-DNA-20M (pretrained)  →    Stage 1 checkpoint   →    Final model
-8,112 pairs × 10 epochs         3,317 pairs × 5 epochs
-lr = 5e-5                        lr = 1e-5 (10x lower)
+Stage 1 (ICI — Interleaved Codon Interference):
+  Dual-model consensus generation (Omni-DNA-20M + HyenaDNA tiny-1k)
+  Every 3 codons, a 1-codon gap is introduced; both models predict the gap
+  Agreement = consensus (high quality), disagreement = contested (kept, weighted lower)
+  → Produces 8,112 synthetic mutation pairs
+
+Stage 2 (SAD — Sequential Attenuation Denoising):
+  Fine-tune on synthetic pairs first (builds broad mutation prior)
+  Then fine-tune on real data at 10x lower LR
+  Uncontradicted synthetic patterns persist, contradicted ones get attenuated
 ```
 
-**SAD Coefficient**: **4.89** (81,120 synthetic exposures / 16,585 real exposures)
-
-## ICI (Interleaved Codon Interleaving) — Synthetic Data Generation
+### Key Parameters
 
-Synthetic mutation pairs generated via dual-model consensus:
-- **FDI codon gap distance: 3** (gap every 3rd codon = 9 bp gap)
-- **Model A**: Omni-DNA-20M (BPE causal LM)
-- **Model B**: HyenaDNA tiny-1k (single-nucleotide causal LM, 451K params)
-- Models generate independently at each gap; agreement = consensus, disagreement = contested
-- 8 generation passes, gap intervals {3,4,5,6} × seeds {42,137}, deduplicated
-- Consensus rate: 2.0% (98% contested — models have fundamentally different representations)
+| Parameter | Value |
+|-----------|-------|
+| **ICI Codon Gap Distance** | 3 (every 3rd codon = 9 bp spacing) |
+| **SAD Coefficient** | 4.89 (81,120 synthetic exposures / 16,585 real exposures) |
+| Stage 1 LR | 5e-5, 10 epochs |
+| Stage 2 LR | 1e-5, 5 epochs |
+| Batch Size | 16 × 2 (effective 32) |
+| Precision | fp32 |
 
 ## Checkpoints
 
-| Path | Stage | Description |
-|------|-------|-------------|
-| `stage1_ici/` | Stage 1 | Weights after synthetic pre-training (ICI) |
-| `stage2_sad/` | Stage 2 | Weights after real-data attenuation (final) |
+| Path | Description |
+|------|-------------|
+| `stage1_ici/` | Weights after synthetic pre-training (ICI stage) |
+| `stage2_sad/` | Weights after real-data attenuation (final SAD model) |
 
-## Benchmark Results (200 test pairs)
+## Benchmarks (200 test pairs)
 
 | Model | Levenshtein ↓ | Similarity ↑ | Mutation Recall ↑ |
 |-------|:---:|:---:|:---:|
-| **Base Omni (no fine-tuning)** | 359.2 | 44.0% | **74.3%** |
+| Base Omni (no fine-tuning) | 359.2 | 44.0% | **74.3%** |
 | Stage 1 (synthetic only) | 260.1 | 45.1% | 67.3% |
-| **Stage 2 SAD (synthetic→real)** | 268.2 | 43.8% | 67.8% |
-
-### Analysis
-
-- Base Omni has strong zero-shot mutation-sensing (74.3% recall)
-- Stage 1 improved Levenshtein distance (359→260) but hurt recall (74.3→67.3%)
-- SAD attenuation barely recovered recall (67.3→67.8%)
-- **The SAD coefficient of 4.89 is too high** — synthetic patterns overwhelm real data correction
-
-### Data Distribution Mismatch
+| **Stage 2 SAD (final)** | **268.2** | **43.8%** | **67.8%** |
 
-| Dataset | Mean Mutation Rate |
-|---------|-------------------|
-| Synthetic (ICI) | 17.8% |
-| Real (train) | 3.9% |
-| Real (test) | 0.8% |
+### Known Issues
 
-The 4.5x mutation rate gap between synthetic and real data teaches the model to over-mutate. With a lower SAD coefficient (more real-data epochs), attenuation could correct this.
+1. **SAD coefficient too high**: At 4.89, synthetic exposures overwhelmed real data attenuation. Mutation recall dropped from base 74.3% → 67.8%. A lower coefficient (~1.5) is recommended for future runs.
+2. **Synthetic mutation rate mismatch**: Synthetic data has 17.8% mean mutation rate vs 3.9% in real data — model learns to over-mutate.
+3. **Omni `.generate()` broken after fine-tuning**: Model outputs PAD/EOS immediately. Use manual autoregressive loop with special token suppression instead (see Usage).
 
 ## Usage
 
 ```python
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Load Stage 2 (final) model
+# Load Stage 2 (final SAD model)
 tokenizer = AutoTokenizer.from_pretrained("Nhoodie/omni-dna-sad-mutation", subfolder="stage2_sad", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("Nhoodie/omni-dna-sad-mutation", subfolder="stage2_sad", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("Nhoodie/omni-dna-sad-mutation", subfolder="stage2_sad", trust_remote_code=True).to("cuda").eval()
 
-# IMPORTANT: Omni's .generate() outputs PAD immediately. Use manual autoregressive:
-def predict_child(model, tokenizer, parent, max_new_tokens=200, temperature=0.2):
+def predict_child(parent, max_new_tokens=200, temperature=0.2):
+    """Manual autoregressive generation (Omni .generate() outputs PAD after fine-tuning)"""
     prompt = f"mutate: {parent} -> "
-    input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"]
+    input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(model.device)
     generated = input_ids
     suppress = {0, 1, 2, 3}  # UNK, CLS, SEP, PAD
+
     with torch.inference_mode():
         for _ in range(max_new_tokens):
             logits = model(input_ids=generated).logits[:, -1, :] / max(temperature, 0.01)
-            for s in suppress: logits[0, s] = float("-inf")
-            next_token = torch.softmax(logits, dim=-1).multinomial(1)
+            for s in suppress:
+                logits[0, s] = float("-inf")
+            probs = torch.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
             generated = torch.cat([generated, next_token], dim=-1)
+
     text = tokenizer.decode(generated[0, input_ids.shape[1]:], skip_special_tokens=True)
     return "".join(c for c in text.upper() if c in "ACGT")
 
-child = predict_child(model, tokenizer, "ATGCGTACGTACGT...")
+parent = "ATGGCTAGCTGATCGATCGATCG..."
+child = predict_child(parent)
+print(child)
 ```
 
 ## Dataset
 
 Training data available at: [Nhoodie/omni-dna-sad-mutation-dataset](https://huggingface.co/datasets/Nhoodie/omni-dna-sad-mutation-dataset)
 
-## Training Details
-
-- **Hardware**: NVIDIA GTX 1080 8GB
-- **Framework**: PyTorch 2.6.0, transformers 4.49.0, ai2-olmo 0.6.0
-- **Stage 1**: 10 epochs, lr=5e-5, cosine schedule, batch=16×2, fp32
-- **Stage 2**: 5 epochs, lr=1e-5, cosine schedule, batch=16×2, fp32
-- **Total training time**: ~90 min (Stage 1) + ~30 min (Stage 2)
-
-## Known Issues
+## Training Hardware
 
-1. **Omni `.generate()` is broken** — model outputs PAD/EOS with very high confidence after fine-tuning. Must use manual autoregressive loop.
-2. **SAD coefficient too high** — 4.89 means synthetic patterns dominate. Needs more real-data epochs to attenuate effectively.
-3. **Synthetic mutation rate mismatch** — 17.8% synthetic vs 3.9% real creates distribution shift.
+- **GPU**: NVIDIA GTX 1080 8GB
+- **Stage 1**: ~8 min (8,112 synthetic pairs, 10 epochs)
+- **Stage 2**: ~9 min (3,317 real pairs, 5 epochs)
+- **Total**: <20 min training time