heavy training

.gitignore

@@ -2,10 +2,12 @@
 .env
 .env.*
 !.env.example
+.env.local
 
 # Python
 __pycache__/
 *.py[cod]
+*.pyc
 *.pyo
 *.pyd
 .Python
@@ -39,6 +41,7 @@ drive/
 wandb/
 lightning_logs/
 runs/
+checkpoints/
 
 # Genomics data and model artifacts
 data/raw/
@@ -67,6 +70,19 @@ outputs/
 artifacts/
 training/data/
 training/output/
+training/outputs/
+training/cache/
+training/data_cache/
+training/csv_files_large/
+training/csv_files_large_alt/
+training/csv_files_10k/
+training/csv_files_10k_alt/
+training/csv_files_20k/
+training/csv_files_20k_alt/
+training/local_dnabert2_patch/
+*.zip
+*.safetensors
+*.bin
 
 # OS/editor
 .DS_Store

training/COLAB_README.md

@@ -0,0 +1,33 @@
+# Colab Heavy Training Guide
+
+This workflow is for research and education only. The model output is not a
+medical diagnosis and should not be used for clinical decisions.
+
+## Steps
+
+1. Push this repository to GitHub.
+2. Open `training/colab_dnabert2_heavy_training.ipynb` in Google Colab.
+3. In the repo clone cell, replace:
+
+   ```python
+   REPO_URL = "PASTE_YOUR_GITHUB_REPO_URL_HERE"
+   ```
+
+   with your GitHub repository URL.
+
+4. Run the notebook cells one by one.
+5. Start with the 10k dataset:
+
+   ```bash
+   python training/prepare_larger_clinvar_dataset.py --target_total 10000
+   ```
+
+6. Only try the 20k dataset after the 10k dataset prepares and trains correctly.
+7. Save the final model and metrics to Google Drive from the notebook.
+
+## Notes
+
+- Use a Colab GPU runtime.
+- Do not add API keys or secrets to the notebook.
+- Dataset preparation uses caching and progress files, so reruns can resume.
+- The training script uses CUDA when available, MPS on Mac, and CPU as a slow fallback.

training/README.md

@@ -1,16 +1,19 @@
 # Training
 
-This folder contains the Google Colab training pipeline for fine-tuning DNABERT-2 on ClinVar-derived GRCh38 examples.
+This folder contains the training pipeline for fine-tuning DNABERT-2 on ClinVar-derived GRCh38 examples.
 
-Training should be run in Google Colab or another dedicated GPU notebook environment, not on the local development machine.
+The original path is Google Colab, but this project also includes Mac-friendly local scripts for small research runs and smoke tests.
 
 ## Contents
 
 - `01_prepare_clinvar_dataset.ipynb`: Colab notebook for downloading ClinVar GRCh38 VCF data and preparing binary SNV/small-indel CSV splits with sequence columns.
 - `colab_dnabert2_clinvar_finetune.ipynb`: Colab notebook for fine-tuning DNABERT-2 from `train_with_sequences.csv`, `val_with_sequences.csv`, and `test_with_sequences.csv`.
 - `requirements-colab.txt`: Python packages for the notebook.
+- `requirements-mac.txt`: Python packages for local Mac training.
 - `scripts/prepare_clinvar_dataset.py`: converts ClinVar GRCh38 VCF records into sequence classification examples.
 - `scripts/train_dnabert2_classifier.py`: fine-tunes DNABERT-2 with Hugging Face Transformers.
+- `train_smoke_test.py`: runs a tiny DNABERT-2 local smoke test.
+- `train_local_dnabert2.py`: runs Mac-friendly local DNABERT-2 fine-tuning.
 - `utils/clinvar_parser.py`: manual gzip VCF parsing and variant filtering helpers.
 - `utils/label_utils.py`: ClinVar clinical-significance label mapping helpers.
 - `utils/sequence_fetcher.py`: UCSC API and optional local FASTA sequence extraction helpers.
@@ -25,6 +28,70 @@ Use GRCh38 consistently:
 
 The fine-tuning notebook expects the sequence CSV files from the preparation notebook.
 
+## Dataset Check
+
+Before local training, verify the prepared CSV files:
+
+```bash
+python training/check_dataset.py
+```
+
+The checker expects `train_with_sequences.csv`, `val_with_sequences.csv`, and `test_with_sequences.csv` in `data/processed/`. It also falls back to `training/csv_files/` for sample files.
+
+## Mac Local Setup
+
+For local Mac training setup, see:
+
+```text
+training/setup_mac.md
+```
+
+Check the available PyTorch device with:
+
+```bash
+python training/check_device.py
+```
+
+Run a tiny smoke test first:
+
+```bash
+python training/train_smoke_test.py
+```
+
+Then run local DNABERT-2 fine-tuning:
+
+```bash
+python training/train_local_dnabert2.py
+```
+
+By default, local training uses `training/csv_files_large_alt/` when the 5,000-row alternate-sequence CSVs exist. It uses all available rows, trains for 5 epochs, freezes the DNABERT-2 encoder, applies class weights, crops long sequences around variant index 512, and tunes the final classification threshold on the validation split.
+
+Local Mac evaluation is memory-safe by default: epoch evaluation is disabled, validation/test prediction runs in small batches, and final metrics use an evaluation subset of 300 rows per split. To evaluate all validation/test rows after training, pass `--eval_subset_size 0`.
+
+To train the classifier plus the last encoder layer on Mac, use:
+
+```bash
+python training/train_local_dnabert2.py --unfreeze_last_n_layers 1
+```
+
+To build a larger balanced ClinVar dataset before training, use:
+
+```bash
+python training/prepare_larger_clinvar_dataset.py
+```
+
+Check the larger dataset before training:
+
+```bash
+python training/check_large_dataset.py
+```
+
+Then train from the larger alternate-sequence CSVs explicitly:
+
+```bash
+python training/train_local_dnabert2.py --train_csv training/csv_files_large_alt/train_with_alt_sequences.csv --val_csv training/csv_files_large_alt/val_with_alt_sequences.csv --test_csv training/csv_files_large_alt/test_with_alt_sequences.csv
+```
+
 ## Colab Flow
 
 1. Open `01_prepare_clinvar_dataset.ipynb` in Google Colab.

training/audit_variant_sequence_encoding.py

@@ -0,0 +1,416 @@
+#!/usr/bin/env python3
+"""Audit whether sequence strings contain REF or ALT alleles near the center.
+
+The prepared sequence files were fetched as roughly +/-512 bp around each
+variant, so the variant should begin around index 512 in the sequence string.
+This script checks that assumption for SNVs and small indels.
+"""
+
+from __future__ import annotations
+
+import argparse
+from dataclasses import dataclass
+from pathlib import Path
+
+import pandas as pd
+
+
+ORIGINAL_SPLIT_FILES = {
+    "train": "train_with_sequences.csv",
+    "val": "val_with_sequences.csv",
+    "test": "test_with_sequences.csv",
+}
+
+ALT_SPLIT_FILES = {
+    "train": "train_with_alt_sequences.csv",
+    "val": "val_with_alt_sequences.csv",
+    "test": "test_with_alt_sequences.csv",
+}
+
+REQUIRED_COLUMNS = {"sequence", "REF", "ALT"}
+OPTIONAL_EXAMPLE_COLUMNS = [
+    "variant_id",
+    "CHROM",
+    "POS",
+    "CLNHGVS",
+    "gene_symbol",
+    "GENEINFO",
+]
+
+CENTER_INDEX = 512
+SNIPPET_FLANK = 20
+MAX_INDEL_SIZE = 50
+VALID_BASES = set("ACGTN")
+
+
+@dataclass
+class AuditCounts:
+    total_rows_seen: int = 0
+    total_rows_checked: int = 0
+    snv_rows_checked: int = 0
+    indel_rows_checked: int = 0
+    reference_matches: int = 0
+    alternate_matches: int = 0
+    mismatches: int = 0
+    skipped_missing_values: int = 0
+    skipped_short_sequence: int = 0
+    skipped_multiple_alt: int = 0
+    skipped_symbolic_alt: int = 0
+    skipped_non_acgtn_allele: int = 0
+    skipped_not_snv_or_small_indel: int = 0
+
+    def add(self, other: "AuditCounts") -> None:
+        for field_name in self.__dataclass_fields__:
+            setattr(self, field_name, getattr(self, field_name) + getattr(other, field_name))
+
+
+@dataclass
+class DatasetChoice:
+    data_dir: Path
+    split_files: dict[str, str]
+    is_alternate_dataset: bool
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--data-dir",
+        type=Path,
+        default=None,
+        help=(
+            "Directory containing train/val/test CSV files. Defaults to "
+            "training/csv_files_20k_alt, training/csv_files_10k_alt, "
+            "training/csv_files_large_alt, training/csv_files_alt, "
+            "training/csv_files_20k, training/csv_files_10k, "
+            "training/csv_files_large, training/csv_files, then data/processed."
+        ),
+    )
+    parser.add_argument(
+        "--center-index",
+        type=int,
+        default=CENTER_INDEX,
+        help="0-based index where the variant is expected to start. Default: 512.",
+    )
+    return parser.parse_args()
+
+
+def project_root() -> Path:
+    return Path(__file__).resolve().parents[1]
+
+
+def has_all_files(directory: Path, split_files: dict[str, str]) -> bool:
+    return all((directory / filename).exists() for filename in split_files.values())
+
+
+def choose_dataset(root: Path, requested_dir: Path | None) -> DatasetChoice:
+    if requested_dir is not None:
+        data_dir = requested_dir.expanduser().resolve()
+        if has_all_files(data_dir, ALT_SPLIT_FILES):
+            return DatasetChoice(data_dir, ALT_SPLIT_FILES, True)
+        if has_all_files(data_dir, ORIGINAL_SPLIT_FILES):
+            return DatasetChoice(data_dir, ORIGINAL_SPLIT_FILES, False)
+        raise FileNotFoundError(
+            "The requested directory does not contain a complete alternate or original dataset:\n"
+            f"{data_dir}"
+        )
+
+    candidates: list[tuple[Path, dict[str, str], bool]] = [
+        (root / "training" / "csv_files_20k_alt", ALT_SPLIT_FILES, True),
+        (root / "training" / "csv_files_10k_alt", ALT_SPLIT_FILES, True),
+        (root / "training" / "csv_files_large_alt", ALT_SPLIT_FILES, True),
+        (root / "training" / "csv_files_alt", ALT_SPLIT_FILES, True),
+        (root / "training" / "csv_files_20k", ORIGINAL_SPLIT_FILES, False),
+        (root / "training" / "csv_files_10k", ORIGINAL_SPLIT_FILES, False),
+        (root / "training" / "csv_files_large", ORIGINAL_SPLIT_FILES, False),
+        (root / "training" / "csv_files", ORIGINAL_SPLIT_FILES, False),
+        (root / "data" / "processed", ORIGINAL_SPLIT_FILES, False),
+    ]
+
+    for directory, split_files, is_alternate_dataset in candidates:
+        if has_all_files(directory, split_files):
+            return DatasetChoice(directory, split_files, is_alternate_dataset)
+
+    searched = "\n".join(str(path) for path, _split_files, _is_alt in candidates)
+    raise FileNotFoundError(
+        "Could not find a complete alternate or original sequence dataset.\n"
+        f"Searched:\n{searched}"
+    )
+
+
+def normalize_sequence(value: object) -> str:
+    return str(value).strip().upper()
+
+
+def normalize_allele(value: object) -> str:
+    return str(value).strip().upper()
+
+
+def is_symbolic_alt(alt: str) -> bool:
+    return alt.startswith("<") or alt.endswith(">") or "[" in alt or "]" in alt
+
+
+def is_acgtn(value: str) -> bool:
+    return bool(value) and set(value).issubset(VALID_BASES)
+
+
+def is_snv(ref: str, alt: str) -> bool:
+    return len(ref) == 1 and len(alt) == 1
+
+
+def is_small_indel(ref: str, alt: str) -> bool:
+    return ref != alt and abs(len(ref) - len(alt)) <= MAX_INDEL_SIZE
+
+
+def center_snippet(sequence: str, center_index: int) -> str:
+    start = max(0, center_index - SNIPPET_FLANK)
+    end = min(len(sequence), center_index + SNIPPET_FLANK + 1)
+    snippet = sequence[start:end]
+    marker_position = center_index - start
+    if 0 <= marker_position < len(snippet):
+        return snippet[:marker_position] + "[" + snippet[marker_position] + "]" + snippet[marker_position + 1 :]
+    return snippet
+
+
+def classify_row(row: pd.Series, center_index: int, is_alternate_dataset: bool) -> tuple[str, str, str]:
+    sequence = normalize_sequence(row["sequence"])
+    ref = normalize_allele(row["REF"])
+    alt = normalize_allele(row["ALT"])
+
+    if not sequence or not ref or not alt or sequence == "NAN" or ref == "NAN" or alt == "NAN":
+        return "skipped_missing_values", "", ""
+
+    if "," in alt:
+        return "skipped_multiple_alt", "", ""
+
+    if is_symbolic_alt(alt):
+        return "skipped_symbolic_alt", "", ""
+
+    if not is_acgtn(ref) or not is_acgtn(alt):
+        return "skipped_non_acgtn_allele", "", ""
+
+    if len(sequence) <= center_index:
+        return "skipped_short_sequence", "", ""
+
+    if is_snv(ref, alt):
+        observed = sequence[center_index]
+        if is_alternate_dataset:
+            if observed == alt:
+                return "alternate_match_snv", observed, center_snippet(sequence, center_index)
+            if observed == ref:
+                return "reference_match_snv", observed, center_snippet(sequence, center_index)
+        else:
+            if observed == ref:
+                return "reference_match_snv", observed, center_snippet(sequence, center_index)
+            if observed == alt:
+                return "alternate_match_snv", observed, center_snippet(sequence, center_index)
+        return "mismatch_snv", observed, center_snippet(sequence, center_index)
+
+    if is_small_indel(ref, alt):
+        expected_allele = alt if is_alternate_dataset else ref
+        if len(sequence) < center_index + len(expected_allele):
+            return "skipped_short_sequence", "", ""
+
+        ref_window = sequence[center_index : center_index + len(ref)] if len(sequence) >= center_index + len(ref) else ""
+        alt_window = sequence[center_index : center_index + len(alt)] if len(sequence) >= center_index + len(alt) else ""
+
+        if is_alternate_dataset:
+            if alt_window == alt:
+                return "alternate_match_indel", alt_window, center_snippet(sequence, center_index)
+            if ref_window == ref:
+                return "reference_match_indel", ref_window, center_snippet(sequence, center_index)
+            return "mismatch_indel", alt_window, center_snippet(sequence, center_index)
+
+        if ref_window == ref:
+            return "reference_match_indel", ref_window, center_snippet(sequence, center_index)
+        if alt_window == alt:
+            return "alternate_match_indel", alt_window, center_snippet(sequence, center_index)
+        return "mismatch_indel", ref_window, center_snippet(sequence, center_index)
+
+    return "skipped_not_snv_or_small_indel", "", ""
+
+
+def update_counts(counts: AuditCounts, status: str) -> None:
+    if status.startswith("skipped_"):
+        current_value = getattr(counts, status)
+        setattr(counts, status, current_value + 1)
+        return
+
+    counts.total_rows_checked += 1
+
+    if status.endswith("_snv"):
+        counts.snv_rows_checked += 1
+    elif status.endswith("_indel"):
+        counts.indel_rows_checked += 1
+
+    if status.startswith("reference_match"):
+        counts.reference_matches += 1
+    elif status.startswith("alternate_match"):
+        counts.alternate_matches += 1
+    elif status.startswith("mismatch"):
+        counts.mismatches += 1
+
+
+def build_example(row: pd.Series, split_name: str, status: str, observed: str, snippet: str) -> dict[str, object]:
+    example = {
+        "split": split_name,
+        "status": status,
+        "REF": normalize_allele(row["REF"]),
+        "ALT": normalize_allele(row["ALT"]),
+        "observed_at_center": observed,
+        "center_sequence_snippet": snippet,
+    }
+
+    for column in OPTIONAL_EXAMPLE_COLUMNS:
+        if column in row.index:
+            example[column] = row[column]
+
+    return example
+
+
+def audit_split(
+    split_name: str,
+    csv_path: Path,
+    center_index: int,
+    is_alternate_dataset: bool,
+) -> tuple[AuditCounts, list[dict[str, object]]]:
+    print("=" * 80)
+    print(f"{split_name.upper()} SPLIT")
+    print("=" * 80)
+    print(f"Path: {csv_path}")
+
+    df = pd.read_csv(csv_path)
+    counts = AuditCounts(total_rows_seen=len(df))
+    examples: list[dict[str, object]] = []
+
+    missing_columns = sorted(REQUIRED_COLUMNS - set(df.columns))
+    if missing_columns:
+        print(f"ERROR: missing required columns: {missing_columns}")
+        print()
+        return counts, examples
+
+    for _, row in df.iterrows():
+        status, observed, snippet = classify_row(row, center_index, is_alternate_dataset)
+        update_counts(counts, status)
+
+        if len(examples) < 10 and not status.startswith("skipped_"):
+            examples.append(build_example(row, split_name, status, observed, snippet))
+
+    print_counts(counts)
+    print()
+    return counts, examples
+
+
+def print_counts(counts: AuditCounts) -> None:
+    print(f"Rows seen: {counts.total_rows_seen:,}")
+    print(f"Total rows checked: {counts.total_rows_checked:,}")
+    print(f"SNV rows checked: {counts.snv_rows_checked:,}")
+    print(f"Indel rows checked: {counts.indel_rows_checked:,}")
+    print(f"Reference allele matches: {counts.reference_matches:,}")
+    print(f"Alternate allele matches: {counts.alternate_matches:,}")
+    print(f"Mismatches: {counts.mismatches:,}")
+    print("Skipped rows:")
+    print(f"  missing sequence/REF/ALT: {counts.skipped_missing_values:,}")
+    print(f"  sequence too short for center check: {counts.skipped_short_sequence:,}")
+    print(f"  multiple ALT alleles: {counts.skipped_multiple_alt:,}")
+    print(f"  symbolic ALT allele: {counts.skipped_symbolic_alt:,}")
+    print(f"  non-ACGTN allele: {counts.skipped_non_acgtn_allele:,}")
+    print(f"  not SNV or small indel: {counts.skipped_not_snv_or_small_indel:,}")
+
+
+def print_examples(examples: list[dict[str, object]]) -> None:
+    print("=" * 80)
+    print("FIRST 10 CHECKED EXAMPLES")
+    print("=" * 80)
+
+    if not examples:
+        print("No checked examples available.")
+        print()
+        return
+
+    examples_df = pd.DataFrame(examples)
+    display_columns = [
+        column
+        for column in [
+            "split",
+            "variant_id",
+            "CHROM",
+            "POS",
+            "gene_symbol",
+            "REF",
+            "ALT",
+            "status",
+            "observed_at_center",
+            "center_sequence_snippet",
+            "CLNHGVS",
+            "GENEINFO",
+        ]
+        if column in examples_df.columns
+    ]
+
+    with pd.option_context("display.max_columns", None, "display.width", 220, "display.max_colwidth", 90):
+        print(examples_df[display_columns].to_string(index=False))
+    print()
+
+
+def print_conclusion(counts: AuditCounts) -> None:
+    print("=" * 80)
+    print("OVERALL SUMMARY")
+    print("=" * 80)
+    print_counts(counts)
+
+    checked = counts.total_rows_checked
+    if checked == 0:
+        conclusion = "Could not determine clearly."
+    else:
+        ref_fraction = counts.reference_matches / checked
+        alt_fraction = counts.alternate_matches / checked
+        print(f"Reference match fraction: {ref_fraction:.2%}")
+        print(f"Alternate match fraction: {alt_fraction:.2%}")
+
+        if ref_fraction >= 0.80 and counts.reference_matches > counts.alternate_matches:
+            conclusion = "The sequence appears to be reference sequence."
+        elif alt_fraction >= 0.80 and counts.alternate_matches > counts.reference_matches:
+            conclusion = "The sequence appears to contain alternate alleles."
+        else:
+            conclusion = "Could not determine clearly."
+
+    print()
+    print(f"Conclusion: {conclusion}")
+
+
+def main() -> None:
+    args = parse_args()
+    root = project_root()
+    dataset = choose_dataset(root, args.data_dir)
+
+    print("Variant sequence encoding audit")
+    print(f"Selected data directory: {dataset.data_dir}")
+    print(f"Auditing alternate dataset: {dataset.is_alternate_dataset}")
+    print(f"Expected variant start index: {args.center_index}")
+    print()
+
+    total_counts = AuditCounts()
+    all_examples: list[dict[str, object]] = []
+
+    for split_name, filename in dataset.split_files.items():
+        csv_path = dataset.data_dir / filename
+        if not csv_path.exists():
+            raise FileNotFoundError(f"Missing required CSV file: {csv_path}")
+
+        split_counts, split_examples = audit_split(
+            split_name,
+            csv_path,
+            args.center_index,
+            dataset.is_alternate_dataset,
+        )
+        total_counts.add(split_counts)
+        remaining_example_slots = 10 - len(all_examples)
+        if remaining_example_slots > 0:
+            all_examples.extend(split_examples[:remaining_example_slots])
+
+    print_examples(all_examples)
+    print_conclusion(total_counts)
+
+
+if __name__ == "__main__":
+    main()

training/build_alt_sequence_dataset.py

@@ -0,0 +1,316 @@
+#!/usr/bin/env python3
+"""Build alternate-allele sequence CSVs from reference sequence CSVs.
+
+The current prepared CSVs contain reference genome sequence around each
+variant. This script verifies that REF is present at index 512, replaces that
+REF allele with ALT, and writes new CSV files with mutated/alternate sequences.
+"""
+
+from __future__ import annotations
+
+import argparse
+from dataclasses import dataclass
+from pathlib import Path
+
+import pandas as pd
+
+
+FLANK_SIZE = 512
+VALID_BASES = set("ACGTN")
+
+SPLIT_FILES = {
+    "train": ("train_with_sequences.csv", "train_with_alt_sequences.csv"),
+    "val": ("val_with_sequences.csv", "val_with_alt_sequences.csv"),
+    "test": ("test_with_sequences.csv", "test_with_alt_sequences.csv"),
+}
+
+REQUIRED_COLUMNS = {"sequence", "REF", "ALT", "label"}
+USEFUL_METADATA_COLUMNS = [
+    "variant_id",
+    "CHROM",
+    "POS",
+    "ID",
+    "variant_type",
+    "gene_symbol",
+    "GENEINFO",
+    "CLNSIG",
+    "CLNHGVS",
+    "CLNVC",
+    "label_name",
+]
+
+
+@dataclass
+class BuildStats:
+    input_rows: int = 0
+    successful_rows: int = 0
+    failed_rows: int = 0
+    snv_count: int = 0
+    indel_count: int = 0
+    skipped_missing_values: int = 0
+    skipped_multiple_alt: int = 0
+    skipped_symbolic_alt: int = 0
+    skipped_non_acgtn: int = 0
+    skipped_short_sequence: int = 0
+    failed_ref_mismatch: int = 0
+
+    def add(self, other: "BuildStats") -> None:
+        for field_name in self.__dataclass_fields__:
+            setattr(self, field_name, getattr(self, field_name) + getattr(other, field_name))
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--data-dir",
+        type=Path,
+        default=None,
+        help="Directory containing the input CSV files. Defaults to data/processed, then training/csv_files.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default=None,
+        help="Directory for alternate-sequence CSV files. Default: training/csv_files_alt.",
+    )
+    return parser.parse_args()
+
+
+def project_root() -> Path:
+    return Path(__file__).resolve().parents[1]
+
+
+def resolve_path(root: Path, path: Path) -> Path:
+    return path.expanduser().resolve() if path.is_absolute() else (root / path).resolve()
+
+
+def choose_input_dir(root: Path, requested_dir: Path | None) -> Path:
+    if requested_dir is not None:
+        return resolve_path(root, requested_dir)
+
+    candidates = [
+        root / "data" / "processed",
+        root / "training" / "csv_files",
+    ]
+
+    for directory in candidates:
+        if all((directory / input_name).exists() for input_name, _output_name in SPLIT_FILES.values()):
+            return directory
+
+    searched = "\n".join(str(path) for path in candidates)
+    raise FileNotFoundError(
+        "Could not find all input sequence CSV files.\n"
+        f"Searched:\n{searched}"
+    )
+
+
+def choose_output_dir(root: Path, requested_dir: Path | None) -> Path:
+    if requested_dir is not None:
+        return resolve_path(root, requested_dir)
+    return root / "training" / "csv_files_alt"
+
+
+def normalize_text(value: object) -> str:
+    return str(value).strip().upper()
+
+
+def is_missing_text(value: str) -> bool:
+    return value == "" or value == "NAN" or value == "NONE"
+
+
+def is_symbolic_alt(alt: str) -> bool:
+    return alt.startswith("<") or alt.endswith(">") or "[" in alt or "]" in alt
+
+
+def contains_only_valid_bases(value: str) -> bool:
+    return bool(value) and set(value).issubset(VALID_BASES)
+
+
+def is_snv(ref: str, alt: str) -> bool:
+    return len(ref) == 1 and len(alt) == 1
+
+
+def is_small_indel_or_complex(ref: str, alt: str) -> bool:
+    return ref != alt and abs(len(ref) - len(alt)) <= 50
+
+
+def build_alt_sequence(row: pd.Series) -> tuple[dict[str, object] | None, str]:
+    ref_sequence = normalize_text(row["sequence"])
+    ref = normalize_text(row["REF"])
+    alt = normalize_text(row["ALT"])
+
+    if is_missing_text(ref_sequence) or is_missing_text(ref) or is_missing_text(alt):
+        return None, "skipped_missing_values"
+
+    if "," in alt:
+        return None, "skipped_multiple_alt"
+
+    if is_symbolic_alt(alt):
+        return None, "skipped_symbolic_alt"
+
+    if not (
+        contains_only_valid_bases(ref_sequence)
+        and contains_only_valid_bases(ref)
+        and contains_only_valid_bases(alt)
+    ):
+        return None, "skipped_non_acgtn"
+
+    if len(ref_sequence) < FLANK_SIZE + len(ref):
+        return None, "skipped_short_sequence"
+
+    upstream = ref_sequence[:FLANK_SIZE]
+    observed_ref = ref_sequence[FLANK_SIZE : FLANK_SIZE + len(ref)]
+    downstream = ref_sequence[FLANK_SIZE + len(ref) :]
+
+    if observed_ref != ref:
+        return None, "failed_ref_mismatch"
+
+    alt_sequence = upstream + alt + downstream
+    output_row = row.to_dict()
+    output_row["REF"] = ref
+    output_row["ALT"] = alt
+    output_row["ref_sequence"] = ref_sequence
+    output_row["alt_sequence"] = alt_sequence
+    output_row["sequence"] = alt_sequence
+    output_row["ref_center"] = observed_ref
+    output_row["alt_center"] = alt_sequence[FLANK_SIZE : FLANK_SIZE + len(alt)]
+    return output_row, "success"
+
+
+def ordered_columns(df: pd.DataFrame) -> list[str]:
+    important_columns = [
+        column
+        for column in USEFUL_METADATA_COLUMNS + ["REF", "ALT", "label", "ref_sequence", "alt_sequence", "sequence"]
+        if column in df.columns
+    ]
+    remaining_columns = [column for column in df.columns if column not in important_columns]
+    return important_columns + remaining_columns
+
+
+def update_failed_stat(stats: BuildStats, reason: str) -> None:
+    stats.failed_rows += 1
+    current_value = getattr(stats, reason)
+    setattr(stats, reason, current_value + 1)
+
+
+def process_split(split_name: str, input_path: Path, output_path: Path) -> tuple[BuildStats, pd.DataFrame]:
+    print("=" * 80)
+    print(f"{split_name.upper()} SPLIT")
+    print("=" * 80)
+    print(f"Input: {input_path}")
+    print(f"Output: {output_path}")
+
+    df = pd.read_csv(input_path)
+    missing_columns = sorted(REQUIRED_COLUMNS - set(df.columns))
+    if missing_columns:
+        raise ValueError(f"{input_path} is missing required columns: {missing_columns}")
+
+    stats = BuildStats(input_rows=len(df))
+    output_rows: list[dict[str, object]] = []
+
+    for _, row in df.iterrows():
+        output_row, status = build_alt_sequence(row)
+        if output_row is None:
+            update_failed_stat(stats, status)
+            continue
+
+        output_rows.append(output_row)
+        stats.successful_rows += 1
+
+        ref = output_row["REF"]
+        alt = output_row["ALT"]
+        if is_snv(ref, alt):
+            stats.snv_count += 1
+        elif is_small_indel_or_complex(ref, alt):
+            stats.indel_count += 1
+
+    output_df = pd.DataFrame(output_rows)
+    if not output_df.empty:
+        output_df = output_df[ordered_columns(output_df)]
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_df.to_csv(output_path, index=False)
+
+    print_split_summary(stats, output_df)
+    print_examples(output_df)
+    return stats, output_df
+
+
+def print_split_summary(stats: BuildStats, output_df: pd.DataFrame) -> None:
+    print(f"Input rows: {stats.input_rows:,}")
+    print(f"Successful rows: {stats.successful_rows:,}")
+    print(f"Failed rows: {stats.failed_rows:,}")
+    print(f"SNV count: {stats.snv_count:,}")
+    print(f"Indel count: {stats.indel_count:,}")
+
+    if not output_df.empty and "label" in output_df.columns:
+        print("Label distribution:")
+        print(output_df["label"].value_counts().sort_index().to_string())
+    else:
+        print("Label distribution: no successful rows")
+
+    if stats.failed_rows:
+        print("Failure reasons:")
+        print(f"  missing sequence/REF/ALT: {stats.skipped_missing_values:,}")
+        print(f"  multiple ALT alleles: {stats.skipped_multiple_alt:,}")
+        print(f"  symbolic ALT allele: {stats.skipped_symbolic_alt:,}")
+        print(f"  non-ACGTN sequence/REF/ALT: {stats.skipped_non_acgtn:,}")
+        print(f"  sequence too short: {stats.skipped_short_sequence:,}")
+        print(f"  REF mismatch at center: {stats.failed_ref_mismatch:,}")
+    print()
+
+
+def print_examples(output_df: pd.DataFrame) -> None:
+    print("First 5 examples:")
+    if output_df.empty:
+        print("No successful rows.")
+        print()
+        return
+
+    columns = [column for column in ["variant_id", "REF", "ALT", "label", "ref_center", "alt_center"] if column in output_df.columns]
+    examples = output_df[columns].head(5)
+    with pd.option_context("display.max_columns", None, "display.width", 160, "display.max_colwidth", 80):
+        print(examples.to_string(index=False))
+    print()
+
+
+def print_overall_summary(total_stats: BuildStats, all_outputs: list[pd.DataFrame]) -> None:
+    print("=" * 80)
+    print("OVERALL SUMMARY")
+    print("=" * 80)
+
+    combined = pd.concat(all_outputs, ignore_index=True) if all_outputs else pd.DataFrame()
+    print_split_summary(total_stats, combined)
+
+
+def main() -> None:
+    args = parse_args()
+    root = project_root()
+    input_dir = choose_input_dir(root, args.data_dir)
+    output_dir = choose_output_dir(root, args.output_dir)
+
+    print("Build alternate-allele sequence dataset")
+    print(f"Input directory: {input_dir}")
+    print(f"Output directory: {output_dir}")
+    print(f"Assumed flank size / variant center index: {FLANK_SIZE}")
+    print()
+
+    total_stats = BuildStats()
+    all_outputs: list[pd.DataFrame] = []
+
+    for split_name, (input_name, output_name) in SPLIT_FILES.items():
+        input_path = input_dir / input_name
+        output_path = output_dir / output_name
+        if not input_path.exists():
+            raise FileNotFoundError(f"Missing input CSV file: {input_path}")
+
+        split_stats, output_df = process_split(split_name, input_path, output_path)
+        total_stats.add(split_stats)
+        all_outputs.append(output_df)
+
+    print_overall_summary(total_stats, all_outputs)
+    print(f"Alternate-sequence CSV files saved to: {output_dir}")
+
+
+if __name__ == "__main__":
+    main()

training/check_dataset.py

@@ -0,0 +1,271 @@
+#!/usr/bin/env python3
+"""Check sequence CSV files before local model training.
+
+Expected files:
+- train_with_sequences.csv
+- val_with_sequences.csv
+- test_with_sequences.csv
+
+The script prefers data/processed/ and falls back to training/csv_files/ so it
+works with both the planned project layout and the current sample CSV location.
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from urllib.parse import unquote
+
+import pandas as pd
+
+
+SPLIT_FILES = {
+    "train": "train_with_sequences.csv",
+    "val": "val_with_sequences.csv",
+    "test": "test_with_sequences.csv",
+}
+
+REQUIRED_COLUMNS = {"sequence", "label"}
+VALID_LABELS = {0, 1}
+LABEL_MEANINGS = {
+    0: "Benign/Likely benign",
+    1: "Pathogenic/Likely pathogenic",
+}
+
+DROP_CLNSIG_TERMS = (
+    "conflicting",
+    "uncertain",
+    "risk",
+    "association",
+    "drug",
+    "protective",
+    "not provided",
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--data-dir",
+        type=Path,
+        default=None,
+        help="Directory containing train/val/test CSV files. Defaults to data/processed, then training/csv_files.",
+    )
+    parser.add_argument(
+        "--min-reasonable-length",
+        type=int,
+        default=200,
+        help="Warn when non-empty sequences are shorter than this length.",
+    )
+    return parser.parse_args()
+
+
+def project_root() -> Path:
+    return Path(__file__).resolve().parents[1]
+
+
+def choose_data_dir(root: Path, requested_dir: Path | None) -> Path:
+    if requested_dir is not None:
+        return requested_dir.expanduser().resolve()
+
+    preferred = root / "data" / "processed"
+    fallback = root / "training" / "csv_files"
+
+    if all((preferred / filename).exists() for filename in SPLIT_FILES.values()):
+        return preferred
+    if all((fallback / filename).exists() for filename in SPLIT_FILES.values()):
+        return fallback
+    return preferred
+
+
+def find_missing_files(data_dir: Path) -> list[Path]:
+    return [data_dir / filename for filename in SPLIT_FILES.values() if not (data_dir / filename).exists()]
+
+
+def sequence_lengths(sequence_series: pd.Series) -> pd.Series:
+    cleaned = sequence_series.fillna("").astype(str).str.strip()
+    return cleaned.str.len()
+
+
+def missing_sequence_count(sequence_series: pd.Series) -> int:
+    cleaned = sequence_series.fillna("").astype(str).str.strip()
+    return int((cleaned == "").sum())
+
+
+def normalize_clnsig(value: object) -> str:
+    decoded = unquote(str(value))
+    return (
+        decoded.replace("_", " ")
+        .replace("-", " ")
+        .replace("/", " ")
+        .replace("|", " ")
+        .replace(",", " ")
+        .strip()
+        .lower()
+    )
+
+
+def expected_label_from_clnsig(value: object) -> int | None:
+    normalized = normalize_clnsig(value)
+    if not normalized or normalized == "." or any(term in normalized for term in DROP_CLNSIG_TERMS):
+        return None
+
+    has_pathogenic = "pathogenic" in normalized
+    has_benign = "benign" in normalized
+    if has_pathogenic and has_benign:
+        return None
+    if has_pathogenic:
+        return 1
+    if has_benign:
+        return 0
+    return None
+
+
+def print_label_meaning() -> None:
+    print("Label encoding:")
+    for label, meaning in LABEL_MEANINGS.items():
+        print(f"  {label} = {meaning}")
+    print()
+
+
+def audit_split(split_name: str, csv_path: Path, min_reasonable_length: int) -> pd.DataFrame:
+    print("=" * 80)
+    print(f"{split_name.upper()} SPLIT")
+    print("=" * 80)
+    print(f"Path: {csv_path}")
+
+    df = pd.read_csv(csv_path)
+    print(f"Rows: {len(df):,}")
+    print(f"Columns: {list(df.columns)}")
+
+    missing_columns = sorted(REQUIRED_COLUMNS - set(df.columns))
+    if missing_columns:
+        print(f"ERROR: missing required columns: {missing_columns}")
+        print()
+        return df
+
+    labels = pd.to_numeric(df["label"], errors="coerce")
+    label_distribution = labels.value_counts(dropna=False).sort_index()
+    print("Label distribution:")
+    print(label_distribution.to_string())
+
+    invalid_labels = sorted(set(labels.dropna().astype(int).unique()) - VALID_LABELS)
+    if invalid_labels:
+        print(f"WARNING: unexpected label values found: {invalid_labels}")
+    else:
+        print("Label check: OK, labels are encoded as 0/1.")
+
+    if "label_name" in df.columns:
+        label_name = df["label_name"].fillna("").astype(str).str.lower()
+        label_name_mismatch = int(
+            (((labels == 0) & ~label_name.str.contains("benign")) | ((labels == 1) & ~label_name.str.contains("pathogenic"))).sum()
+        )
+        if label_name_mismatch:
+            print(f"WARNING: {label_name_mismatch:,} rows have label_name values that do not match label.")
+        else:
+            print("Label name check: OK.")
+
+    if "CLNSIG" in df.columns:
+        expected_labels = df["CLNSIG"].apply(expected_label_from_clnsig)
+        unsupported_clnsig = int(expected_labels.isna().sum())
+        comparable = expected_labels.notna() & labels.notna()
+        clnsig_mismatch = int((expected_labels[comparable].astype(int) != labels[comparable].astype(int)).sum())
+        print(f"CLNSIG rows that should be excluded before training: {unsupported_clnsig:,}")
+        if clnsig_mismatch:
+            print(f"WARNING: {clnsig_mismatch:,} rows have CLNSIG values that disagree with label.")
+        if unsupported_clnsig:
+            examples = df.loc[expected_labels.isna(), "CLNSIG"].dropna().astype(str).unique()[:5]
+            print(f"Excluded CLNSIG examples: {list(examples)}")
+
+    missing_sequences = missing_sequence_count(df["sequence"])
+    lengths = sequence_lengths(df["sequence"])
+    non_empty_lengths = lengths[lengths > 0]
+
+    print(f"Missing or empty sequence count: {missing_sequences:,}")
+    if non_empty_lengths.empty:
+        print("Sequence length min/mean/max: no non-empty sequences")
+    else:
+        print(
+            "Sequence length min/mean/max: "
+            f"{int(non_empty_lengths.min())} / "
+            f"{float(non_empty_lengths.mean()):.2f} / "
+            f"{int(non_empty_lengths.max())}"
+        )
+        short_count = int((non_empty_lengths < min_reasonable_length).sum())
+        if short_count:
+            print(f"WARNING: {short_count:,} sequences are shorter than {min_reasonable_length} bp.")
+
+    print("First 3 example rows:")
+    example_columns = [column for column in ["variant_id", "CHROM", "POS", "REF", "ALT", "label", "label_name"] if column in df.columns]
+    examples = df[example_columns].head(3).copy()
+    examples["sequence_length"] = lengths.head(3).to_list()
+    examples["sequence_preview"] = df["sequence"].fillna("").astype(str).str.slice(0, 80).head(3).to_list()
+    with pd.option_context("display.max_columns", None, "display.width", 160, "display.max_colwidth", 80):
+        print(examples.to_string(index=False))
+    print()
+    return df
+
+
+def print_overall_balance(frames: dict[str, pd.DataFrame]) -> None:
+    print("=" * 80)
+    print("OVERALL CLASS BALANCE")
+    print("=" * 80)
+
+    usable_frames = []
+    for split_name, df in frames.items():
+        if "label" not in df.columns:
+            continue
+        temp = df[["label"]].copy()
+        temp["split"] = split_name
+        usable_frames.append(temp)
+
+    if not usable_frames:
+        print("No label columns found, so balance cannot be checked.")
+        return
+
+    combined = pd.concat(usable_frames, ignore_index=True)
+    combined["label"] = pd.to_numeric(combined["label"], errors="coerce")
+    counts = combined["label"].value_counts().sort_index()
+    print(counts.to_string())
+
+    valid_counts = counts[counts.index.isin(list(VALID_LABELS))]
+    if len(valid_counts) == 2 and valid_counts.sum() > 0:
+        minority_fraction = float(valid_counts.min() / valid_counts.sum())
+        print(f"Minority class fraction: {minority_fraction:.2%}")
+        if minority_fraction < 0.10:
+            print("Balance note: very imbalanced. Consider class weights, sampling, or more data.")
+        elif minority_fraction < 0.25:
+            print("Balance note: moderately imbalanced, but usable for a demo with stratified metrics.")
+        else:
+            print("Balance note: reasonably balanced for a research demo.")
+    else:
+        print("Balance note: expected both labels 0 and 1.")
+
+
+def main() -> None:
+    args = parse_args()
+    root = project_root()
+    data_dir = choose_data_dir(root, args.data_dir)
+
+    print("Variant Risk Explainer dataset check")
+    print(f"Project root: {root}")
+    print(f"Data directory: {data_dir}")
+    print()
+    print_label_meaning()
+
+    missing_files = find_missing_files(data_dir)
+    if missing_files:
+        print("ERROR: missing expected CSV files:")
+        for path in missing_files:
+            print(f"  {path}")
+        raise SystemExit(1)
+
+    frames = {}
+    for split_name, filename in SPLIT_FILES.items():
+        frames[split_name] = audit_split(split_name, data_dir / filename, args.min_reasonable_length)
+
+    print_overall_balance(frames)
+
+
+if __name__ == "__main__":
+    main()

training/check_device.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+"""Check the best available PyTorch device for local Mac training."""
+
+from __future__ import annotations
+
+import torch
+
+
+def choose_device() -> str:
+    """Select CUDA first, then Apple MPS, then CPU."""
+    if torch.cuda.is_available():
+        return "cuda"
+
+    mps_backend = getattr(torch.backends, "mps", None)
+    if mps_backend is not None and mps_backend.is_available():
+        return "mps"
+
+    return "cpu"
+
+
+def main() -> None:
+    device = choose_device()
+
+    print("PyTorch device check")
+    print(f"Torch version: {torch.__version__}")
+    print(f"CUDA available: {torch.cuda.is_available()}")
+
+    mps_backend = getattr(torch.backends, "mps", None)
+    mps_available = bool(mps_backend is not None and mps_backend.is_available())
+    print(f"MPS available: {mps_available}")
+    print(f"Using device: {device}")
+
+    if device == "mps":
+        print("Apple Silicon acceleration is available through MPS.")
+    elif device == "cuda":
+        print("CUDA GPU acceleration is available.")
+    else:
+        print("WARNING: CPU training will be slow. Use small smoke tests locally, or use a GPU machine for full DNABERT-2 training.")
+
+
+if __name__ == "__main__":
+    main()

training/check_large_dataset.py

@@ -0,0 +1,265 @@
+#!/usr/bin/env python3
+"""Check the larger ClinVar datasets before DNABERT-2 training.
+
+This script checks generated larger datasets when present:
+- training/csv_files_20k_alt/ and training/csv_files_20k/
+- training/csv_files_10k_alt/ and training/csv_files_10k/
+- training/csv_files_large_alt/ and training/csv_files_large/
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from urllib.parse import unquote
+
+import pandas as pd
+
+
+DATASETS = [
+    (
+        "20k alternate-sequence dataset",
+        Path("training/csv_files_20k_alt"),
+        {
+            "train": "train_with_alt_sequences.csv",
+            "val": "val_with_alt_sequences.csv",
+            "test": "test_with_alt_sequences.csv",
+        },
+    ),
+    (
+        "20k reference-sequence dataset",
+        Path("training/csv_files_20k"),
+        {
+            "train": "train_with_sequences.csv",
+            "val": "val_with_sequences.csv",
+            "test": "test_with_sequences.csv",
+        },
+    ),
+    (
+        "10k alternate-sequence dataset",
+        Path("training/csv_files_10k_alt"),
+        {
+            "train": "train_with_alt_sequences.csv",
+            "val": "val_with_alt_sequences.csv",
+            "test": "test_with_alt_sequences.csv",
+        },
+    ),
+    (
+        "10k reference-sequence dataset",
+        Path("training/csv_files_10k"),
+        {
+            "train": "train_with_sequences.csv",
+            "val": "val_with_sequences.csv",
+            "test": "test_with_sequences.csv",
+        },
+    ),
+    (
+        "large alternate-sequence dataset",
+        Path("training/csv_files_large_alt"),
+        {
+            "train": "train_with_alt_sequences.csv",
+            "val": "val_with_alt_sequences.csv",
+            "test": "test_with_alt_sequences.csv",
+        },
+    ),
+    (
+        "large reference-sequence dataset",
+        Path("training/csv_files_large"),
+        {
+            "train": "train_with_sequences.csv",
+            "val": "val_with_sequences.csv",
+            "test": "test_with_sequences.csv",
+        },
+    ),
+]
+
+REQUIRED_COLUMNS = {"sequence", "label"}
+EXAMPLE_COLUMNS = ["variant_id", "REF", "ALT", "label"]
+SUSPICIOUS_CLNSIG_TERMS = (
+    "conflicting",
+    "uncertain",
+    "not provided",
+    "not_provided",
+    "risk_factor",
+    "risk factor",
+    "association",
+    "drug_response",
+    "drug response",
+    "protective",
+)
+
+
+def project_root() -> Path:
+    return Path(__file__).resolve().parents[1]
+
+
+def normalize_clnsig(value: object) -> str:
+    decoded = unquote(str(value))
+    return (
+        decoded.replace("_", " ")
+        .replace("-", " ")
+        .replace("/", " ")
+        .replace("|", " ")
+        .replace(",", " ")
+        .strip()
+        .lower()
+    )
+
+
+def suspicious_clnsig_count(df: pd.DataFrame) -> int:
+    if "CLNSIG" not in df.columns:
+        return 0
+    normalized = df["CLNSIG"].fillna("").apply(normalize_clnsig)
+    return int(normalized.apply(lambda value: any(term in value for term in SUSPICIOUS_CLNSIG_TERMS)).sum())
+
+
+def sequence_lengths(df: pd.DataFrame) -> pd.Series:
+    return df["sequence"].fillna("").astype(str).str.strip().str.len()
+
+
+def missing_sequence_count(df: pd.DataFrame) -> int:
+    cleaned = df["sequence"].fillna("").astype(str).str.strip()
+    return int((cleaned == "").sum())
+
+
+def print_sequence_length_summary(lengths: pd.Series) -> None:
+    non_empty = lengths[lengths > 0]
+    if non_empty.empty:
+        print("Sequence length min/mean/max: no non-empty sequences")
+        return
+
+    print(
+        "Sequence length min/mean/max: "
+        f"{int(non_empty.min())} / "
+        f"{float(non_empty.mean()):.2f} / "
+        f"{int(non_empty.max())}"
+    )
+
+
+def print_examples(df: pd.DataFrame, lengths: pd.Series) -> None:
+    print("First 3 examples:")
+    if df.empty:
+        print("  no rows")
+        return
+
+    columns = [column for column in EXAMPLE_COLUMNS if column in df.columns]
+    examples = df[columns].head(3).copy()
+    examples["sequence_length"] = lengths.head(3).to_list()
+    with pd.option_context("display.max_columns", None, "display.width", 160, "display.max_colwidth", 80):
+        print(examples.to_string(index=False))
+
+
+def check_split(split_name: str, csv_path: Path) -> tuple[pd.DataFrame | None, bool]:
+    print("-" * 80)
+    print(f"{split_name.upper()} SPLIT")
+    print("-" * 80)
+    print(f"File path: {csv_path}")
+
+    if not csv_path.exists():
+        print("ERROR: file is missing.")
+        print()
+        return None, False
+
+    df = pd.read_csv(csv_path)
+    print(f"Rows: {len(df):,}")
+    print(f"Columns: {list(df.columns)}")
+
+    missing_columns = sorted(REQUIRED_COLUMNS - set(df.columns))
+    if missing_columns:
+        print(f"ERROR: missing required columns: {missing_columns}")
+        print()
+        return df, False
+
+    labels = pd.to_numeric(df["label"], errors="coerce")
+    print("Label distribution:")
+    print(labels.value_counts(dropna=False).sort_index().to_string())
+
+    missing_sequences = missing_sequence_count(df)
+    lengths = sequence_lengths(df)
+    print(f"Missing sequence count: {missing_sequences:,}")
+    print_sequence_length_summary(lengths)
+
+    if "CLNSIG" in df.columns:
+        suspicious_count = suspicious_clnsig_count(df)
+        print(f"CLNSIG suspicious rows: {suspicious_count:,}")
+    else:
+        print("CLNSIG suspicious rows: CLNSIG column not present")
+
+    print_examples(df, lengths)
+    print()
+
+    valid_labels = labels.isin([0, 1]).all()
+    has_sequences = missing_sequences == 0 and int((lengths > 0).sum()) == len(df)
+    usable = len(df) > 0 and valid_labels and has_sequences
+    return df, usable
+
+
+def check_dataset(dataset_name: str, dataset_dir: Path, split_files: dict[str, str]) -> tuple[int, bool]:
+    print("=" * 80)
+    print(dataset_name.upper())
+    print("=" * 80)
+    print(f"Dataset directory: {dataset_dir}")
+
+    if not dataset_dir.exists():
+        print("Dataset directory is missing, skipping.")
+        print()
+        return 0, True
+
+    total_rows = 0
+    dataset_usable = True
+    frames = []
+
+    for split_name, filename in split_files.items():
+        df, split_usable = check_split(split_name, dataset_dir / filename)
+        dataset_usable = dataset_usable and split_usable
+        if df is not None:
+            total_rows += len(df)
+            if "label" in df.columns:
+                temp = df[["label"]].copy()
+                temp["split"] = split_name
+                frames.append(temp)
+
+    print(f"Total rows across train/val/test: {total_rows:,}")
+
+    if frames:
+        combined = pd.concat(frames, ignore_index=True)
+        combined["label"] = pd.to_numeric(combined["label"], errors="coerce")
+        print("Combined label distribution:")
+        print(combined["label"].value_counts(dropna=False).sort_index().to_string())
+
+    if dataset_usable:
+        print("Usability check: OK, this dataset has rows, labels, and non-empty sequences.")
+    else:
+        print("Usability check: FAILED, fix the issues above before training.")
+    print()
+
+    return total_rows, dataset_usable
+
+
+def main() -> None:
+    root = project_root()
+    print("Large ClinVar dataset check")
+    print(f"Project root: {root}")
+    print()
+
+    overall_usable = True
+    datasets_found = 0
+    for dataset_name, relative_dir, split_files in DATASETS:
+        dataset_dir = root / relative_dir
+        total_rows, usable = check_dataset(dataset_name, dataset_dir, split_files)
+        if dataset_dir.exists():
+            datasets_found += 1
+            overall_usable = overall_usable and usable and total_rows > 0
+
+    print("=" * 80)
+    print("FINAL RESULT")
+    print("=" * 80)
+    if datasets_found == 0:
+        print("No larger dataset folders were found yet. Run prepare_larger_clinvar_dataset.py first.")
+    elif overall_usable:
+        print("The larger dataset files found look usable for training.")
+    else:
+        print("At least one larger dataset found is not fully usable. Review the errors above.")
+
+
+if __name__ == "__main__":
+    main()

training/colab_dnabert2_clinvar_finetune.ipynb

@@ -204,7 +204,9 @@
       "source": [
         "## 6. Training Settings\n",
         "\n",
-        "Use conservative settings first. Increase epochs or batch size later if the run is stable. For small sample CSVs, expect overfitting; this is only a research demo."
+        "Use conservative settings first. Increase epochs or batch size later if the run is stable. For small sample CSVs, expect overfitting; this is only a research demo.\n",
+        "\n",
+        "If you only want to test the notebook on CPU, set `CPU_SMOKE_TEST = True`. That uses a tiny model and a small number of rows so you can see progress without waiting for DNABERT-2 CPU training."
       ]
     },
     {
@@ -213,17 +215,38 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "MODEL_NAME = 'zhihan1996/DNABERT-2-117M'\n",
-        "OUTPUT_DIR = PROJECT_ROOT / 'training' / 'output' / 'dnabert2-clinvar-grch38'\n",
+        "CPU_SMOKE_TEST = False\n",
+        "\n",
+        "if CPU_SMOKE_TEST:\n",
+        "    MODEL_NAME = 'hf-internal-testing/tiny-random-bert'\n",
+        "    OUTPUT_DIR = PROJECT_ROOT / 'training' / 'output' / 'cpu-smoke-test'\n",
+        "    MAX_LENGTH = 128\n",
+        "    EPOCHS = 1\n",
+        "    BATCH_SIZE = 8\n",
+        "    LEARNING_RATE = 5e-5\n",
+        "    MAX_TRAIN_SAMPLES = 64\n",
+        "    MAX_VAL_SAMPLES = 32\n",
+        "    MAX_TEST_SAMPLES = 32\n",
+        "    FORCE_CPU = True\n",
+        "else:\n",
+        "    MODEL_NAME = 'zhihan1996/DNABERT-2-117M'\n",
+        "    OUTPUT_DIR = PROJECT_ROOT / 'training' / 'output' / 'dnabert2-clinvar-grch38'\n",
+        "    MAX_LENGTH = 512\n",
+        "    EPOCHS = 1\n",
+        "    BATCH_SIZE = 4\n",
+        "    LEARNING_RATE = 2e-5\n",
+        "    MAX_TRAIN_SAMPLES = 0\n",
+        "    MAX_VAL_SAMPLES = 0\n",
+        "    MAX_TEST_SAMPLES = 0\n",
+        "    FORCE_CPU = False\n",
         "\n",
-        "MAX_LENGTH = 512\n",
-        "EPOCHS = 1\n",
-        "BATCH_SIZE = 4\n",
-        "LEARNING_RATE = 2e-5\n",
         "MIN_SEQUENCE_LENGTH = 200\n",
+        "LOGGING_STEPS = 5\n",
         "\n",
+        "print(f'CPU smoke test: {CPU_SMOKE_TEST}')\n",
         "print(f'Model: {MODEL_NAME}')\n",
-        "print(f'Output dir: {OUTPUT_DIR}')"
+        "print(f'Output dir: {OUTPUT_DIR}')\n",
+        "print(f'Max samples: train={MAX_TRAIN_SAMPLES}, val={MAX_VAL_SAMPLES}, test={MAX_TEST_SAMPLES}')"
       ]
     },
     {
@@ -232,7 +255,7 @@
       "source": [
         "## 7. Fine-tune DNABERT-2\n",
         "\n",
-        "This calls the repository training script using the CSV files. The script saves the best checkpoint and final model under `training/output/dnabert2-clinvar-grch38/final_model`."
+        "This calls the repository training script using the CSV files. Progress bars and step logs are enabled. In normal mode, the script saves the best checkpoint and final model under `training/output/dnabert2-clinvar-grch38/final_model`."
       ]
     },
     {
@@ -241,6 +264,8 @@
       "metadata": {},
       "outputs": [],
       "source": [
+        "force_cpu_arg = '--force-cpu' if FORCE_CPU else ''\n",
+        "\n",
         "!python training/scripts/train_dnabert2_classifier.py \\\n",
         "  --train-csv \"{TRAIN_CSV}\" \\\n",
         "  --val-csv \"{VAL_CSV}\" \\\n",
@@ -249,9 +274,14 @@
         "  --model-name \"{MODEL_NAME}\" \\\n",
         "  --max-length {MAX_LENGTH} \\\n",
         "  --min-sequence-length {MIN_SEQUENCE_LENGTH} \\\n",
+        "  --max-train-samples {MAX_TRAIN_SAMPLES} \\\n",
+        "  --max-val-samples {MAX_VAL_SAMPLES} \\\n",
+        "  --max-test-samples {MAX_TEST_SAMPLES} \\\n",
         "  --epochs {EPOCHS} \\\n",
         "  --batch-size {BATCH_SIZE} \\\n",
-        "  --learning-rate {LEARNING_RATE}"
+        "  --learning-rate {LEARNING_RATE} \\\n",
+        "  --logging-steps {LOGGING_STEPS} \\\n",
+        "  {force_cpu_arg}"
       ]
     },
     {

training/colab_dnabert2_heavy_training.ipynb

@@ -0,0 +1,326 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# DNABERT-2 Heavy Training on Google Colab\n",
+        "\n",
+        "This notebook prepares a larger ClinVar GRCh38 alternate-sequence dataset and fine-tunes DNABERT-2 on a Colab GPU.\n",
+        "\n",
+        "Research/education only. This is not a medical diagnosis tool."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 1. GPU Check"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import subprocess\n",
+        "import torch\n",
+        "\n",
+        "print(\"nvidia-smi output:\")\n",
+        "subprocess.run([\"nvidia-smi\"], check=False)\n",
+        "\n",
+        "cuda_available = torch.cuda.is_available()\n",
+        "print(\"torch CUDA available:\", cuda_available)\n",
+        "\n",
+        "if not cuda_available:\n",
+        "    raise RuntimeError(\"No GPU is available. In Colab, go to Runtime > Change runtime type > GPU, then rerun.\")\n",
+        "\n",
+        "print(\"GPU name:\", torch.cuda.get_device_name(0))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 2. Clone GitHub Repo"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "from pathlib import Path\n",
+        "import os\n",
+        "import subprocess\n",
+        "\n",
+        "REPO_URL = \"PASTE_YOUR_GITHUB_REPO_URL_HERE\"\n",
+        "\n",
+        "if REPO_URL == \"PASTE_YOUR_GITHUB_REPO_URL_HERE\":\n",
+        "    raise ValueError(\"Paste your GitHub repository URL into REPO_URL before running this cell.\")\n",
+        "\n",
+        "repo_name = REPO_URL.rstrip(\"/\").split(\"/\")[-1].replace(\".git\", \"\")\n",
+        "repo_path = Path(\"/content\") / repo_name\n",
+        "\n",
+        "if not repo_path.exists():\n",
+        "    subprocess.run([\"git\", \"clone\", REPO_URL, str(repo_path)], check=True)\n",
+        "else:\n",
+        "    print(f\"Repository already exists: {repo_path}\")\n",
+        "\n",
+        "os.chdir(repo_path)\n",
+        "print(\"Project root:\", Path.cwd())"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 3. Install Dependencies"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!python -m pip install --upgrade pip\n",
+        "!pip install -q torch transformers datasets accelerate scikit-learn pandas numpy safetensors tqdm biopython requests joblib"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 4. Optional Google Drive Mount"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "from pathlib import Path\n",
+        "\n",
+        "USE_GOOGLE_DRIVE = True\n",
+        "DRIVE_OUTPUT_DIR = Path(\"/content/drive/MyDrive/variant-risk-explainer\")\n",
+        "\n",
+        "if USE_GOOGLE_DRIVE:\n",
+        "    from google.colab import drive\n",
+        "    drive.mount(\"/content/drive\")\n",
+        "    DRIVE_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n",
+        "    print(\"Drive output folder:\", DRIVE_OUTPUT_DIR)\n",
+        "else:\n",
+        "    print(\"Google Drive mount skipped.\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 5. Prepare Larger ClinVar Dataset"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "TARGET_TOTAL = 10000\n",
+        "!python training/prepare_larger_clinvar_dataset.py --target_total {TARGET_TOTAL}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Optional: only run this 20k cell after 10k preparation and training work correctly."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "RUN_20K = False\n",
+        "\n",
+        "if RUN_20K:\n",
+        "    !python training/prepare_larger_clinvar_dataset.py --target_total 20000\n",
+        "else:\n",
+        "    print(\"20k preparation skipped. Set RUN_20K = True when ready.\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 6. Check Dataset"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!python training/check_large_dataset.py"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 7. Audit Alternate Sequence Encoding\n",
+        "\n",
+        "Expected conclusion: `The sequence appears to contain alternate alleles.`"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!python training/audit_variant_sequence_encoding.py"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 8. GPU Training"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "%%bash\n",
+        "python training/train_local_dnabert2.py \\\n",
+        "  --sample_size 0 \\\n",
+        "  --epochs 5 \\\n",
+        "  --batch_size 4 \\\n",
+        "  --grad_accum_steps 4 \\\n",
+        "  --learning_rate 2e-5 \\\n",
+        "  --freeze_encoder true \\\n",
+        "  --unfreeze_last_n_layers 4 \\\n",
+        "  --use_class_weights true \\\n",
+        "  --center_crop true \\\n",
+        "  --tune_threshold true \\\n",
+        "  --eval_accumulation_steps 8 \\\n",
+        "  --save_eval_each_epoch false \\\n",
+        "  --eval_subset_size 0"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 9. Full Evaluation"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!python training/evaluate_saved_model.py --tune_threshold true"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 10. Save Outputs to Google Drive"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "from pathlib import Path\n",
+        "import shutil\n",
+        "\n",
+        "source_dir = Path(\"training/outputs/dnabert2_clinvar\")\n",
+        "drive_output = Path(\"/content/drive/MyDrive/variant-risk-explainer\")\n",
+        "\n",
+        "if not drive_output.exists():\n",
+        "    print(\"Google Drive output folder is not available. Run the Drive mount cell first.\")\n",
+        "else:\n",
+        "    drive_output.mkdir(parents=True, exist_ok=True)\n",
+        "    items_to_copy = [\n",
+        "        source_dir / \"final_model\",\n",
+        "        source_dir / \"metrics.json\",\n",
+        "        source_dir / \"full_eval_metrics.json\",\n",
+        "        source_dir / \"final_dnabert2_clinvar_model.zip\",\n",
+        "    ]\n",
+        "\n",
+        "    for item in items_to_copy:\n",
+        "        if not item.exists():\n",
+        "            print(f\"Skipping missing item: {item}\")\n",
+        "            continue\n",
+        "\n",
+        "        destination = drive_output / item.name\n",
+        "        if item.is_dir():\n",
+        "            if destination.exists():\n",
+        "                shutil.rmtree(destination)\n",
+        "            shutil.copytree(item, destination)\n",
+        "        else:\n",
+        "            shutil.copy2(item, destination)\n",
+        "        print(f\"Copied {item} -> {destination}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 11. Download Model Artifacts"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "from pathlib import Path\n",
+        "from google.colab import files\n",
+        "\n",
+        "output_dir = Path(\"training/outputs/dnabert2_clinvar\")\n",
+        "zip_path = output_dir / \"final_dnabert2_clinvar_model.zip\"\n",
+        "\n",
+        "if not zip_path.exists():\n",
+        "    !cd training/outputs/dnabert2_clinvar && zip -r final_dnabert2_clinvar_model.zip final_model metrics.json full_eval_metrics.json\n",
+        "\n",
+        "files.download(str(zip_path))"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python",
+      "version": "3.10"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 5
+}

training/evaluate_saved_model.py

@@ -0,0 +1,598 @@
+#!/usr/bin/env python3
+"""Memory-safe full evaluation for a saved DNABERT-2 ClinVar model.
+
+This script does not train. It loads the saved model, evaluates the full
+validation and test CSV files in small batches, and writes metrics to JSON.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+from safetensors.torch import load_file as load_safetensors_file
+from sklearn.metrics import (
+    accuracy_score,
+    confusion_matrix,
+    f1_score,
+    matthews_corrcoef,
+    precision_score,
+    recall_score,
+    roc_auc_score,
+)
+from tqdm.auto import tqdm
+from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from training.train_smoke_test import (  # noqa: E402
+    LOCAL_DNABERT2_PATCH_DIR,
+    clear_local_patch_module_cache,
+    create_local_dnabert2_patch,
+    disable_flash_attention_on_config,
+    load_sequence_classification_model,
+)
+
+OUTPUT_DIR = PROJECT_ROOT / "training" / "outputs" / "dnabert2_clinvar"
+MODEL_DIR = OUTPUT_DIR / "final_model"
+TRAINING_METRICS_PATH = OUTPUT_DIR / "metrics.json"
+FULL_EVAL_METRICS_PATH = OUTPUT_DIR / "full_eval_metrics.json"
+
+ALT_SPLIT_FILES = {
+    "validation": "val_with_alt_sequences.csv",
+    "test": "test_with_alt_sequences.csv",
+}
+
+DATASET_CANDIDATES = [
+    PROJECT_ROOT / "training" / "csv_files_20k_alt",
+    PROJECT_ROOT / "training" / "csv_files_10k_alt",
+    PROJECT_ROOT / "training" / "csv_files_large_alt",
+    PROJECT_ROOT / "training" / "csv_files_alt",
+    PROJECT_ROOT / "data" / "processed",
+    PROJECT_ROOT / "training" / "csv_files",
+]
+
+SEQUENCE_COLUMN = "sequence"
+LABEL_COLUMN = "label"
+MAX_LENGTH = 512
+VARIANT_CENTER_INDEX = 512
+BATCH_SIZE = 1
+MPS_CACHE_EVERY = 25
+
+
+def parse_bool(value: str | bool) -> bool:
+    if isinstance(value, bool):
+        return value
+
+    normalized = value.strip().lower()
+    if normalized in {"true", "1", "yes", "y"}:
+        return True
+    if normalized in {"false", "0", "no", "n"}:
+        return False
+
+    raise argparse.ArgumentTypeError("Use true or false.")
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Memory-safe full validation/test evaluation for the saved DNABERT-2 ClinVar model.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--tune_threshold",
+        type=parse_bool,
+        nargs="?",
+        const=True,
+        default=True,
+        help="Tune the decision threshold on the full validation set.",
+    )
+    parser.add_argument(
+        "--threshold",
+        type=float,
+        default=None,
+        help="Use this fixed decision threshold instead of tuning.",
+    )
+    parser.add_argument(
+        "--threshold_min",
+        type=float,
+        default=0.1,
+        help="Minimum threshold to test when tuning.",
+    )
+    parser.add_argument(
+        "--threshold_max",
+        type=float,
+        default=0.9,
+        help="Maximum threshold to test when tuning.",
+    )
+    parser.add_argument(
+        "--threshold_step",
+        type=float,
+        default=0.01,
+        help="Threshold step size when tuning.",
+    )
+    return parser.parse_args()
+
+
+def validate_args(args: argparse.Namespace) -> None:
+    if args.threshold is not None and not 0.0 <= args.threshold <= 1.0:
+        raise ValueError("--threshold must be between 0 and 1.")
+
+    if args.threshold_step <= 0:
+        raise ValueError("--threshold_step must be greater than 0.")
+
+    if args.threshold_min > args.threshold_max:
+        raise ValueError("--threshold_min must be less than or equal to --threshold_max.")
+
+
+def resolve_project_path(path: Path) -> Path:
+    if path.is_absolute():
+        return path
+    return PROJECT_ROOT / path
+
+
+def choose_device() -> str:
+    if torch.cuda.is_available():
+        return "cuda"
+
+    mps_backend = getattr(torch.backends, "mps", None)
+    if mps_backend is not None and mps_backend.is_available():
+        return "mps"
+
+    return "cpu"
+
+
+def find_dataset_dir() -> Path:
+    for directory in DATASET_CANDIDATES:
+        if all((directory / filename).exists() for filename in ALT_SPLIT_FILES.values()):
+            return directory
+
+    searched = "\n".join(str(directory) for directory in DATASET_CANDIDATES)
+    raise FileNotFoundError(
+        "Could not find validation/test alternate-sequence CSV files.\n"
+        f"Searched:\n{searched}"
+    )
+
+
+def load_threshold() -> float:
+    if not TRAINING_METRICS_PATH.exists():
+        return 0.5
+
+    try:
+        metrics = json.loads(TRAINING_METRICS_PATH.read_text(encoding="utf-8"))
+    except json.JSONDecodeError:
+        return 0.5
+
+    threshold = metrics.get("selected_threshold")
+    if threshold is None:
+        return 0.5
+
+    try:
+        return float(threshold)
+    except (TypeError, ValueError):
+        return 0.5
+
+
+def clean_sequence(value: object) -> str:
+    return str(value).strip().upper()
+
+
+def crop_sequence_around_variant(sequence: str, max_length: int, variant_center_index: int) -> str:
+    if len(sequence) <= max_length:
+        return sequence
+
+    start = max(0, variant_center_index - max_length // 2)
+    end = start + max_length
+    if end > len(sequence):
+        end = len(sequence)
+        start = max(0, end - max_length)
+
+    return sequence[start:end]
+
+
+def load_eval_dataframe(csv_path: Path, split_name: str) -> pd.DataFrame:
+    df = pd.read_csv(csv_path)
+    required_columns = {SEQUENCE_COLUMN, LABEL_COLUMN}
+    missing_columns = sorted(required_columns - set(df.columns))
+    if missing_columns:
+        raise ValueError(f"{csv_path} is missing required columns: {missing_columns}")
+
+    df = df.copy()
+    df[LABEL_COLUMN] = pd.to_numeric(df[LABEL_COLUMN], errors="coerce")
+    df = df.loc[df[LABEL_COLUMN].isin([0, 1])].copy()
+    df[LABEL_COLUMN] = df[LABEL_COLUMN].astype(int)
+
+    df[SEQUENCE_COLUMN] = df[SEQUENCE_COLUMN].fillna("").apply(clean_sequence)
+    df = df.loc[df[SEQUENCE_COLUMN] != ""].copy()
+    df[SEQUENCE_COLUMN] = df[SEQUENCE_COLUMN].apply(
+        lambda sequence: crop_sequence_around_variant(sequence, MAX_LENGTH, VARIANT_CENTER_INDEX)
+    )
+
+    print(f"{split_name} CSV: {csv_path}")
+    print(f"{split_name} rows loaded for full evaluation: {len(df):,}")
+    print(f"{split_name} label distribution:")
+    print(df[LABEL_COLUMN].value_counts().sort_index().to_string())
+    print()
+
+    if df.empty:
+        raise ValueError(f"No usable rows remain for {split_name}.")
+
+    return df.reset_index(drop=True)
+
+
+def clear_mps_cache_if_needed(device: str) -> None:
+    if device != "mps":
+        return
+    mps_backend = getattr(torch, "mps", None)
+    if mps_backend is not None and hasattr(mps_backend, "empty_cache"):
+        mps_backend.empty_cache()
+
+
+def extract_logits(outputs) -> torch.Tensor:
+    if hasattr(outputs, "logits"):
+        return outputs.logits
+
+    if isinstance(outputs, (tuple, list)):
+        for item in outputs:
+            if torch.is_tensor(item) and item.ndim >= 2 and item.shape[-1] == 2:
+                return item
+        return outputs[0]
+
+    raise TypeError("Could not find logits in model outputs.")
+
+
+def local_patch_is_ready(patch_dir: Path) -> bool:
+    required_files = [
+        "config.json",
+        "configuration_bert.py",
+        "bert_layers.py",
+        "bert_padding.py",
+        "tokenizer.json",
+        "tokenizer_config.json",
+    ]
+    if not all((patch_dir / filename).exists() for filename in required_files):
+        return False
+
+    bert_layers_text = (patch_dir / "bert_layers.py").read_text(encoding="utf-8")
+    return "from .flash_attn_triton import" not in bert_layers_text
+
+
+def create_patch_from_project_root() -> Path:
+    previous_cwd = Path.cwd()
+    try:
+        os.chdir(PROJECT_ROOT)
+        patch_dir = create_local_dnabert2_patch()
+    finally:
+        os.chdir(previous_cwd)
+
+    return resolve_project_path(patch_dir)
+
+
+def get_local_patch_dir() -> Path:
+    patch_dir = resolve_project_path(LOCAL_DNABERT2_PATCH_DIR)
+    if patch_dir.exists() and local_patch_is_ready(patch_dir):
+        return patch_dir
+
+    print("Local Mac-safe DNABERT-2 patch was not found or is incomplete.")
+    return create_patch_from_project_root()
+
+
+def load_saved_state_dict(model_dir: Path) -> dict[str, torch.Tensor]:
+    safetensors_path = model_dir / "model.safetensors"
+    pytorch_path = model_dir / "pytorch_model.bin"
+
+    if safetensors_path.exists():
+        print(f"Loading fine-tuned weights: {safetensors_path}")
+        return load_safetensors_file(str(safetensors_path), device="cpu")
+
+    if pytorch_path.exists():
+        print(f"Loading fine-tuned weights: {pytorch_path}")
+        return torch.load(pytorch_path, map_location="cpu")
+
+    raise FileNotFoundError(
+        "Could not find saved model weights. Expected model.safetensors or pytorch_model.bin in "
+        f"{model_dir}"
+    )
+
+
+def load_saved_model_with_local_patch(model_dir: Path):
+    """Load Mac-safe DNABERT-2 code, then load the saved fine-tuned weights."""
+    patch_dir = get_local_patch_dir()
+    print(f"Using local Mac-safe DNABERT-2 code: {patch_dir}")
+    print("Triton/flash attention disabled for Mac.")
+
+    clear_local_patch_module_cache()
+    config = AutoConfig.from_pretrained(str(patch_dir), trust_remote_code=True)
+    config = disable_flash_attention_on_config(config)
+
+    saved_config_path = model_dir / "config.json"
+    if saved_config_path.exists():
+        saved_config = json.loads(saved_config_path.read_text(encoding="utf-8"))
+        if saved_config.get("id2label"):
+            config.id2label = {int(key): value for key, value in saved_config["id2label"].items()}
+        if saved_config.get("label2id"):
+            config.label2id = saved_config["label2id"]
+
+    model = load_sequence_classification_model(str(patch_dir), config)
+    state_dict = load_saved_state_dict(model_dir)
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+
+    if missing_keys:
+        print(f"Warning: missing keys while loading saved weights: {len(missing_keys)}")
+        print(missing_keys[:10])
+    if unexpected_keys:
+        print(f"Warning: unexpected keys while loading saved weights: {len(unexpected_keys)}")
+        print(unexpected_keys[:10])
+
+    print("Model loaded successfully.")
+    return model
+
+
+def load_saved_model(model_dir: Path, device: str):
+    if device == "cuda":
+        print("CUDA detected. Trying to load saved model directly first.")
+        try:
+            model = AutoModelForSequenceClassification.from_pretrained(
+                str(model_dir),
+                trust_remote_code=True,
+                low_cpu_mem_usage=False,
+            )
+            print("Saved model loaded successfully without the local Mac patch.")
+            return model
+        except Exception as error:
+            print("Direct saved-model load failed.")
+            print(f"Direct load error: {error}")
+            print("Falling back to local no-Triton DNABERT-2 code.")
+
+    return load_saved_model_with_local_patch(model_dir)
+
+
+def predict_in_small_batches(model, tokenizer, df: pd.DataFrame, device: str) -> tuple[np.ndarray, np.ndarray]:
+    device_object = torch.device(device)
+    model.to(device_object)
+    model.eval()
+
+    probabilities: list[float] = []
+    labels: list[int] = []
+
+    for row_number, row in enumerate(tqdm(df.itertuples(index=False), total=len(df), desc="Evaluating"), start=1):
+        sequence = getattr(row, SEQUENCE_COLUMN)
+        label = int(getattr(row, LABEL_COLUMN))
+
+        encoded = tokenizer(
+            sequence,
+            max_length=MAX_LENGTH,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+        )
+        encoded = {key: value.to(device_object) for key, value in encoded.items()}
+
+        with torch.no_grad():
+            outputs = model(**encoded)
+            logits = extract_logits(outputs)
+            probability = torch.softmax(logits.float(), dim=-1)[0, 1].detach().cpu().item()
+
+        probabilities.append(float(probability))
+        labels.append(label)
+
+        del encoded, outputs, logits
+        if row_number % MPS_CACHE_EVERY == 0:
+            clear_mps_cache_if_needed(device)
+
+    gc.collect()
+    clear_mps_cache_if_needed(device)
+
+    return np.asarray(probabilities, dtype=np.float64), np.asarray(labels, dtype=int)
+
+
+def metrics_at_threshold(probabilities: np.ndarray, labels: np.ndarray, threshold: float) -> dict:
+    predictions = (probabilities >= threshold).astype(int)
+    matrix = confusion_matrix(labels, predictions, labels=[0, 1]).astype(int)
+
+    metrics = {
+        "threshold": float(threshold),
+        "accuracy": float(accuracy_score(labels, predictions)),
+        "precision": float(precision_score(labels, predictions, zero_division=0)),
+        "recall": float(recall_score(labels, predictions, zero_division=0)),
+        "f1": float(f1_score(labels, predictions, zero_division=0)),
+        "mcc": float(matthews_corrcoef(labels, predictions)),
+        "auc_roc": None,
+        "confusion_matrix": matrix.tolist(),
+        "rows": int(len(labels)),
+    }
+
+    if len(np.unique(labels)) == 2:
+        try:
+            metrics["auc_roc"] = float(roc_auc_score(labels, probabilities))
+        except ValueError:
+            metrics["auc_roc"] = None
+
+    return metrics
+
+
+def build_threshold_grid(threshold_min: float, threshold_max: float, threshold_step: float) -> np.ndarray:
+    thresholds = np.arange(threshold_min, threshold_max + threshold_step / 2.0, threshold_step)
+    thresholds = thresholds[thresholds <= threshold_max + 1e-12]
+    return np.round(thresholds, 10)
+
+
+def tune_threshold_on_validation(
+    probabilities: np.ndarray,
+    labels: np.ndarray,
+    threshold_min: float,
+    threshold_max: float,
+    threshold_step: float,
+) -> tuple[float, dict]:
+    thresholds = build_threshold_grid(threshold_min, threshold_max, threshold_step)
+    if len(thresholds) == 0:
+        raise ValueError("No thresholds were generated. Check threshold_min, threshold_max, and threshold_step.")
+
+    best_threshold = float(thresholds[0])
+    best_metrics = metrics_at_threshold(probabilities, labels, best_threshold)
+
+    for threshold in thresholds[1:]:
+        candidate_metrics = metrics_at_threshold(probabilities, labels, float(threshold))
+        if candidate_metrics["mcc"] > best_metrics["mcc"]:
+            best_threshold = float(threshold)
+            best_metrics = candidate_metrics
+
+    tuning_summary = {
+        "threshold_min": float(threshold_min),
+        "threshold_max": float(threshold_max),
+        "threshold_step": float(threshold_step),
+        "thresholds_tested": int(len(thresholds)),
+        "best_threshold": best_threshold,
+        "best_validation_mcc": float(best_metrics["mcc"]),
+    }
+
+    print(f"Best full-validation threshold: {best_threshold:.4f}")
+    print(f"Best full-validation MCC: {best_metrics['mcc']:.4f}")
+    print()
+
+    return best_threshold, tuning_summary
+
+
+def choose_threshold(args: argparse.Namespace, probabilities: np.ndarray, labels: np.ndarray) -> tuple[float, dict]:
+    if args.threshold is not None:
+        print(f"Using threshold provided by --threshold: {args.threshold:.4f}")
+        print()
+        return float(args.threshold), {
+            "mode": "manual",
+            "selected_threshold": float(args.threshold),
+        }
+
+    if args.tune_threshold:
+        print("Tuning threshold on the full validation set.")
+        threshold, tuning_summary = tune_threshold_on_validation(
+            probabilities,
+            labels,
+            args.threshold_min,
+            args.threshold_max,
+            args.threshold_step,
+        )
+        tuning_summary["mode"] = "full_validation_mcc"
+        tuning_summary["selected_threshold"] = threshold
+        return threshold, tuning_summary
+
+    saved_threshold = load_threshold()
+    print("--tune_threshold is false and no --threshold was provided.")
+    print(f"Falling back to threshold from metrics.json/default: {saved_threshold:.4f}")
+    print()
+    return saved_threshold, {
+        "mode": "saved_or_default",
+        "selected_threshold": float(saved_threshold),
+    }
+
+
+def print_metrics(split_name: str, metrics: dict) -> None:
+    print("=" * 80)
+    print(f"{split_name.upper()} FULL EVALUATION")
+    print("=" * 80)
+    print(f"Rows: {metrics['rows']:,}")
+    print(f"Threshold used: {metrics['threshold']:.4f}")
+    for key in ["accuracy", "precision", "recall", "f1", "mcc", "auc_roc"]:
+        value = metrics[key]
+        if value is None:
+            print(f"{key}: n/a")
+        else:
+            print(f"{key}: {value:.4f}")
+
+    matrix = metrics["confusion_matrix"]
+    print("Confusion matrix:")
+    print("                 predicted_0  predicted_1")
+    print(f"actual_0         {matrix[0][0]:>11}  {matrix[0][1]:>11}")
+    print(f"actual_1         {matrix[1][0]:>11}  {matrix[1][1]:>11}")
+    print()
+
+
+def main() -> None:
+    args = parse_args()
+    validate_args(args)
+
+    if not MODEL_DIR.exists():
+        raise FileNotFoundError(f"Saved model directory not found: {MODEL_DIR}")
+
+    dataset_dir = find_dataset_dir()
+    device = choose_device()
+
+    print("Memory-safe saved model evaluation")
+    print(f"Saved model directory: {MODEL_DIR}")
+    print(f"Selected dataset directory: {dataset_dir}")
+    print(f"Selected device: {device}")
+    print(f"Tune threshold on full validation set: {args.tune_threshold}")
+    if args.threshold is not None:
+        print(f"Manual threshold requested: {args.threshold:.4f}")
+    else:
+        print(
+            "Threshold search range: "
+            f"{args.threshold_min:.4f} to {args.threshold_max:.4f} "
+            f"by {args.threshold_step:.4f}"
+        )
+    print("Using manual small-batch evaluation. No retraining. No HuggingFace Trainer evaluation.")
+    print()
+
+    print("Loading tokenizer and model.")
+    tokenizer = AutoTokenizer.from_pretrained(str(MODEL_DIR), trust_remote_code=True)
+    model = load_saved_model(MODEL_DIR, device)
+
+    all_metrics = {
+        "model_dir": str(MODEL_DIR),
+        "dataset_dir": str(dataset_dir),
+        "device": device,
+        "max_length": MAX_LENGTH,
+        "variant_center_index": VARIANT_CENTER_INDEX,
+        "batch_size": BATCH_SIZE,
+        "threshold_args": {
+            "tune_threshold": bool(args.tune_threshold),
+            "threshold": args.threshold,
+            "threshold_min": float(args.threshold_min),
+            "threshold_max": float(args.threshold_max),
+            "threshold_step": float(args.threshold_step),
+        },
+    }
+
+    predictions_by_split = {}
+    for split_name, filename in ALT_SPLIT_FILES.items():
+        csv_path = dataset_dir / filename
+        df = load_eval_dataframe(csv_path, split_name)
+        probabilities, labels = predict_in_small_batches(model, tokenizer, df, device)
+        predictions_by_split[split_name] = {
+            "probabilities": probabilities,
+            "labels": labels,
+        }
+
+    validation_predictions = predictions_by_split["validation"]
+    threshold, threshold_selection = choose_threshold(
+        args,
+        validation_predictions["probabilities"],
+        validation_predictions["labels"],
+    )
+
+    all_metrics["threshold"] = threshold
+    all_metrics["selected_threshold"] = threshold
+    all_metrics["threshold_selection"] = threshold_selection
+
+    for split_name, prediction_data in predictions_by_split.items():
+        probabilities = prediction_data["probabilities"]
+        labels = prediction_data["labels"]
+        split_metrics = metrics_at_threshold(probabilities, labels, threshold)
+        all_metrics[f"{split_name}_metrics"] = split_metrics
+        print_metrics(split_name, split_metrics)
+
+    FULL_EVAL_METRICS_PATH.parent.mkdir(parents=True, exist_ok=True)
+    FULL_EVAL_METRICS_PATH.write_text(json.dumps(all_metrics, indent=2), encoding="utf-8")
+    print(f"Saved full evaluation metrics to: {FULL_EVAL_METRICS_PATH}")
+
+
+if __name__ == "__main__":
+    main()

training/prepare_larger_clinvar_dataset.py

@@ -0,0 +1,742 @@
+#!/usr/bin/env python3
+"""Prepare a larger balanced ClinVar sequence dataset for DNABERT-2.
+
+The script prefers existing parsed/filtered CSV data when enough rows are
+available. If not, it downloads/parses the ClinVar GRCh38 VCF, samples a more
+balanced binary dataset, fetches GRCh38 reference sequence windows with resume
+support, and writes both reference-sequence and alternate-sequence CSV files.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gzip
+import shutil
+import sys
+from pathlib import Path
+from urllib.parse import unquote
+
+import pandas as pd
+import requests
+from sklearn.model_selection import train_test_split
+from tqdm.auto import tqdm
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from training.utils.clinvar_parser import (
+    add_variant_id,
+    classify_variant_type,
+    extract_gene_symbol,
+    has_multiple_alt,
+    is_sequence_allele,
+    is_symbolic_alt,
+    parse_info_field,
+)
+from training.utils.label_utils import assign_binary_label, label_name
+from training.utils.sequence_fetcher import build_sequence_fetcher, clean_sequence
+
+
+CLINVAR_GRCH38_VCF_URL = "https://ftp.ncbi.nlm.nih.gov/pub/clinvar/vcf_GRCh38/clinvar.vcf.gz"
+FLANK_SIZE_DEFAULT = 512
+MIN_SEQUENCE_LENGTH = 200
+PROGRESS_EVERY = 100
+
+ORIGINAL_SPLIT_FILES = {
+    "train": "train_with_sequences.csv",
+    "val": "val_with_sequences.csv",
+    "test": "test_with_sequences.csv",
+}
+
+ALT_SPLIT_FILES = {
+    "train": "train_with_alt_sequences.csv",
+    "val": "val_with_alt_sequences.csv",
+    "test": "test_with_alt_sequences.csv",
+}
+
+DROP_OUTPUT_COLUMNS = {
+    "_source_split",
+    "_source_kind",
+    "ref_sequence",
+    "alt_sequence",
+    "ref_center",
+    "alt_center",
+}
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--target_total", type=int, default=5000)
+    parser.add_argument("--max_pathogenic", type=int, default=None)
+    parser.add_argument("--max_benign", type=int, default=None)
+    parser.add_argument("--flank_size", type=int, default=FLANK_SIZE_DEFAULT)
+    parser.add_argument(
+        "--output_dir",
+        type=Path,
+        default=None,
+        help=(
+            "Reference-sequence output directory. If omitted, target_total 5000 writes "
+            "training/csv_files_large, 10000 writes training/csv_files_10k, and 20000 "
+            "writes training/csv_files_20k."
+        ),
+    )
+    parser.add_argument("--random_state", type=int, default=42)
+
+    parser.add_argument("--clinvar_vcf", type=Path, default=None, help="Optional local ClinVar GRCh38 VCF.GZ path.")
+    parser.add_argument("--fetch_mode", choices=["ucsc", "fasta"], default="ucsc")
+    parser.add_argument("--fasta_path", type=Path, default=None, help="Optional local GRCh38 FASTA for --fetch_mode fasta.")
+    parser.add_argument("--sequence_cache", type=Path, default=None, help="Optional JSON sequence cache path.")
+    parser.add_argument("--sleep_seconds", type=float, default=0.15)
+    parser.add_argument("--max_retries", type=int, default=3)
+    return parser.parse_args()
+
+
+def project_root() -> Path:
+    return PROJECT_ROOT
+
+
+def resolve_path(root: Path, path: Path | None) -> Path | None:
+    if path is None:
+        return None
+    return path.expanduser().resolve() if path.is_absolute() else (root / path).resolve()
+
+
+def alt_output_dir_for(output_dir: Path) -> Path:
+    return output_dir.with_name(f"{output_dir.name}_alt")
+
+
+def default_output_dir_for_target(target_total: int) -> Path:
+    if target_total == 5000:
+        return Path("training/csv_files_large")
+    if target_total == 10000:
+        return Path("training/csv_files_10k")
+    if target_total == 20000:
+        return Path("training/csv_files_20k")
+    return Path(f"training/csv_files_{target_total}")
+
+
+def has_complete_split(directory: Path, split_files: dict[str, str]) -> bool:
+    return all((directory / filename).exists() for filename in split_files.values())
+
+
+def read_split_directory(directory: Path, split_files: dict[str, str], source_kind: str) -> pd.DataFrame:
+    frames = []
+    for split_name, filename in split_files.items():
+        path = directory / filename
+        df = pd.read_csv(path)
+        df["_source_split"] = split_name
+        df["_source_kind"] = source_kind
+
+        # Alternate CSVs keep the true reference sequence in ref_sequence.
+        if source_kind == "alternate_split" and "ref_sequence" in df.columns:
+            df["sequence"] = df["ref_sequence"]
+
+        frames.append(df)
+    return pd.concat(frames, ignore_index=True)
+
+
+def existing_source_candidates(root: Path) -> list[tuple[str, pd.DataFrame]]:
+    candidates: list[tuple[str, pd.DataFrame]] = []
+
+    for path in [
+        root / "data" / "processed" / "clinvar_binary_variants.csv",
+        root / "training" / "csv_files" / "clinvar_binary_variants.csv",
+        root / "training" / "csv_files_20k" / "clinvar_binary_variants.csv",
+        root / "training" / "csv_files_10k" / "clinvar_binary_variants.csv",
+        root / "training" / "csv_files_large" / "clinvar_binary_variants.csv",
+    ]:
+        if path.exists():
+            candidates.append((str(path), pd.read_csv(path)))
+
+    for directory in [
+        root / "training" / "csv_files_20k",
+        root / "training" / "csv_files_10k",
+        root / "training" / "csv_files_large",
+        root / "training" / "csv_files",
+        root / "data" / "processed",
+    ]:
+        if has_complete_split(directory, ORIGINAL_SPLIT_FILES):
+            candidates.append((str(directory), read_split_directory(directory, ORIGINAL_SPLIT_FILES, "reference_split")))
+
+    for directory in [
+        root / "training" / "csv_files_20k_alt",
+        root / "training" / "csv_files_10k_alt",
+        root / "training" / "csv_files_large_alt",
+        root / "training" / "csv_files_alt",
+    ]:
+        if has_complete_split(directory, ALT_SPLIT_FILES):
+            candidates.append((str(directory), read_split_directory(directory, ALT_SPLIT_FILES, "alternate_split")))
+
+    return candidates
+
+
+def normalize_allele(value: object) -> str:
+    return str(value).strip().upper()
+
+
+def normalize_clnsig_display(value: object) -> str:
+    return unquote(str(value)).strip()
+
+
+def filter_clear_variants(df: pd.DataFrame) -> pd.DataFrame:
+    required = {"CHROM", "POS", "REF", "ALT"}
+    missing = sorted(required - set(df.columns))
+    if missing:
+        raise ValueError(f"Variant table is missing required columns: {missing}")
+
+    filtered = df.copy()
+    filtered["REF"] = filtered["REF"].apply(normalize_allele)
+    filtered["ALT"] = filtered["ALT"].apply(normalize_allele)
+    filtered["POS"] = pd.to_numeric(filtered["POS"], errors="coerce")
+    filtered = filtered.loc[filtered["POS"].notna()].copy()
+    filtered["POS"] = filtered["POS"].astype(int)
+
+    if "CLNSIG" in filtered.columns:
+        filtered["CLNSIG"] = filtered["CLNSIG"].apply(normalize_clnsig_display)
+        filtered["label"] = filtered["CLNSIG"].apply(assign_binary_label)
+    elif "label" in filtered.columns:
+        filtered["label"] = pd.to_numeric(filtered["label"], errors="coerce")
+        filtered = filtered.loc[filtered["label"].isin([0, 1])].copy()
+    else:
+        raise ValueError("Variant table needs either CLNSIG or label.")
+
+    filtered["variant_type"] = filtered.apply(lambda row: classify_variant_type(row["REF"], row["ALT"]), axis=1)
+
+    keep_mask = (
+        filtered["label"].notna()
+        & ~filtered["ALT"].apply(has_multiple_alt)
+        & ~filtered["ALT"].apply(is_symbolic_alt)
+        & filtered["REF"].apply(is_sequence_allele)
+        & filtered["ALT"].apply(is_sequence_allele)
+        & filtered["variant_type"].notna()
+    )
+    filtered = filtered.loc[keep_mask].copy()
+    filtered["label"] = filtered["label"].astype(int)
+
+    if "GENEINFO" in filtered.columns and "gene_symbol" not in filtered.columns:
+        filtered["gene_symbol"] = filtered["GENEINFO"].apply(extract_gene_symbol)
+    if "label_name" not in filtered.columns:
+        filtered["label_name"] = filtered["label"].apply(label_name)
+    if "variant_id" not in filtered.columns:
+        filtered["variant_id"] = filtered.apply(add_variant_id, axis=1)
+
+    filtered = filtered.drop_duplicates(subset=["variant_id"]).reset_index(drop=True)
+    return filtered
+
+
+def download_clinvar_vcf(vcf_path: Path) -> Path:
+    vcf_path.parent.mkdir(parents=True, exist_ok=True)
+    tmp_path = vcf_path.with_suffix(vcf_path.suffix + ".tmp")
+
+    print(f"Downloading ClinVar GRCh38 VCF to: {vcf_path}")
+    with requests.get(CLINVAR_GRCH38_VCF_URL, stream=True, timeout=60) as response:
+        response.raise_for_status()
+        total = int(response.headers.get("content-length", 0))
+        with tmp_path.open("wb") as handle, tqdm(total=total, unit="B", unit_scale=True, desc="Downloading VCF") as bar:
+            for chunk in response.iter_content(chunk_size=1024 * 1024):
+                if not chunk:
+                    continue
+                handle.write(chunk)
+                bar.update(len(chunk))
+
+    shutil.move(str(tmp_path), str(vcf_path))
+    return vcf_path
+
+
+def find_or_download_vcf(root: Path, requested_vcf: Path | None) -> Path:
+    if requested_vcf is not None:
+        vcf_path = resolve_path(root, requested_vcf)
+        if vcf_path is None or not vcf_path.exists():
+            raise FileNotFoundError(f"ClinVar VCF not found: {vcf_path}")
+        return vcf_path
+
+    candidates = [
+        root / "data" / "raw" / "clinvar.vcf.gz",
+        root / "training" / "data" / "raw" / "clinvar.vcf.gz",
+        root / "clinvar.vcf.gz",
+    ]
+    for path in candidates:
+        if path.exists():
+            return path
+
+    return download_clinvar_vcf(root / "data" / "raw" / "clinvar.vcf.gz")
+
+
+def parse_and_filter_vcf(vcf_path: Path) -> tuple[pd.DataFrame, int]:
+    records: list[dict[str, object]] = []
+    total_records = 0
+
+    print(f"Parsing ClinVar VCF: {vcf_path}")
+    with gzip.open(vcf_path, "rt", encoding="utf-8") as handle:
+        for line in tqdm(handle, desc="Parsing VCF records"):
+            if line.startswith("#"):
+                continue
+
+            total_records += 1
+            fields = line.rstrip("\n").split("\t")
+            if len(fields) < 8:
+                continue
+
+            chrom, pos_raw, clinvar_id, ref, alt, _qual, _filter, info_raw = fields[:8]
+            info = parse_info_field(info_raw)
+            clnsig = info.get("CLNSIG")
+            label = assign_binary_label(clnsig)
+            ref = ref.upper()
+            alt = alt.upper()
+            variant_type = classify_variant_type(ref, alt)
+
+            if (
+                label is None
+                or has_multiple_alt(alt)
+                or is_symbolic_alt(alt)
+                or not is_sequence_allele(ref)
+                or not is_sequence_allele(alt)
+                or variant_type is None
+            ):
+                continue
+
+            gene_info = info.get("GENEINFO")
+            record = {
+                "CHROM": chrom,
+                "POS": int(pos_raw),
+                "ID": clinvar_id,
+                "REF": ref,
+                "ALT": alt,
+                "variant_type": variant_type,
+                "gene_symbol": extract_gene_symbol(gene_info),
+                "GENEINFO": gene_info,
+                "CLNSIG": clnsig,
+                "CLNHGVS": info.get("CLNHGVS"),
+                "CLNVC": info.get("CLNVC"),
+                "label": int(label),
+                "label_name": label_name(label),
+            }
+            record["variant_id"] = add_variant_id(pd.Series(record))
+            records.append(record)
+
+    filtered = pd.DataFrame.from_records(records)
+    if not filtered.empty:
+        filtered = filtered.drop_duplicates(subset=["variant_id"]).reset_index(drop=True)
+    return filtered, total_records
+
+
+def choose_variant_source(root: Path, target_total: int, requested_vcf: Path | None) -> tuple[pd.DataFrame, str, int, int]:
+    best_source_name = ""
+    best_source = pd.DataFrame()
+    best_before = 0
+
+    for source_name, source_df in existing_source_candidates(root):
+        before = len(source_df)
+        filtered = filter_clear_variants(source_df)
+        print(f"Existing source candidate: {source_name}")
+        print(f"  rows before filtering: {before:,}")
+        print(f"  rows after filtering: {len(filtered):,}")
+
+        if len(filtered) > len(best_source):
+            best_source_name = source_name
+            best_source = filtered
+            best_before = before
+
+        if len(filtered) >= target_total:
+            return filtered, source_name, before, len(filtered)
+
+    if not best_source.empty and len(best_source) >= target_total:
+        return best_source, best_source_name, best_before, len(best_source)
+
+    if not best_source.empty:
+        print(
+            "Existing data is available but smaller than target_total. "
+            "Parsing ClinVar VCF to build a larger candidate set."
+        )
+
+    vcf_path = find_or_download_vcf(root, requested_vcf)
+    parsed_df, total_records = parse_and_filter_vcf(vcf_path)
+    return parsed_df, str(vcf_path), total_records, len(parsed_df)
+
+
+def validate_positive_int(name: str, value: int | None) -> None:
+    if value is not None and value <= 0:
+        raise ValueError(f"{name} must be positive when provided.")
+
+
+def sample_or_all(df: pd.DataFrame, n: int, random_state: int) -> pd.DataFrame:
+    if len(df) <= n:
+        return df.copy()
+    return df.sample(n=n, random_state=random_state)
+
+
+def balanced_sample(
+    df: pd.DataFrame,
+    target_total: int,
+    max_pathogenic: int | None,
+    max_benign: int | None,
+    random_state: int,
+) -> pd.DataFrame:
+    if target_total <= 0:
+        raise ValueError("--target_total must be positive.")
+    validate_positive_int("--max_pathogenic", max_pathogenic)
+    validate_positive_int("--max_benign", max_benign)
+
+    pathogenic = df.loc[df["label"] == 1].copy()
+    benign = df.loc[df["label"] == 0].copy()
+    target_per_class = max(1, target_total // 2)
+
+    pathogenic_limit = min(len(pathogenic), target_per_class)
+    if max_pathogenic is not None:
+        pathogenic_limit = min(pathogenic_limit, max_pathogenic)
+
+    pathogenic_sample = sample_or_all(pathogenic, pathogenic_limit, random_state)
+
+    benign_limit = min(len(benign), len(pathogenic_sample))
+    if max_benign is not None:
+        benign_limit = min(benign_limit, max_benign)
+
+    benign_sample = sample_or_all(benign, benign_limit, random_state)
+
+    sampled = pd.concat([pathogenic_sample, benign_sample], ignore_index=True)
+    sampled = sampled.sample(frac=1.0, random_state=random_state).reset_index(drop=True)
+    return sampled
+
+
+def sequence_matches_ref(sequence: str | None, ref: str, flank_size: int) -> bool:
+    sequence = clean_sequence(sequence)
+    if sequence is None or len(sequence) < flank_size + len(ref):
+        return False
+    return sequence[flank_size : flank_size + len(ref)] == ref
+
+
+def load_progress_sequences(progress_path: Path) -> dict[str, str]:
+    if not progress_path.exists():
+        return {}
+
+    progress_df = pd.read_csv(progress_path)
+    if "variant_id" not in progress_df.columns or "sequence" not in progress_df.columns:
+        return {}
+
+    sequences: dict[str, str] = {}
+    for row in progress_df.itertuples(index=False):
+        variant_id = str(getattr(row, "variant_id"))
+        sequence = clean_sequence(getattr(row, "sequence"))
+        if sequence is not None:
+            sequences[variant_id] = sequence
+    return sequences
+
+
+def save_progress(records: list[dict[str, object]], progress_path: Path) -> None:
+    progress_path.parent.mkdir(parents=True, exist_ok=True)
+    pd.DataFrame.from_records(records).to_csv(progress_path, index=False)
+
+
+def save_dataframe(df: pd.DataFrame, path: Path) -> Path:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    df.to_csv(path, index=False)
+    return path
+
+
+def save_fetcher_cache(fetcher) -> None:
+    if hasattr(fetcher, "save_cache"):
+        fetcher.save_cache()
+
+
+def fetch_reference_sequences(
+    df: pd.DataFrame,
+    output_dir: Path,
+    flank_size: int,
+    fetch_mode: str,
+    fasta_path: Path | None,
+    sequence_cache: Path | None,
+    sleep_seconds: float,
+    max_retries: int,
+) -> tuple[pd.DataFrame, dict[str, int]]:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    progress_path = output_dir / "sequence_progress.csv"
+    cache_path = sequence_cache or (output_dir / f"sequence_cache_flank{flank_size}.json")
+    progress_sequences = load_progress_sequences(progress_path)
+
+    fetcher = build_sequence_fetcher(
+        mode=fetch_mode,
+        fasta_path=fasta_path,
+        cache_path=cache_path,
+        sleep_seconds=sleep_seconds,
+        max_retries=max_retries,
+    )
+
+    records: list[dict[str, object]] = []
+    failed_fetches = 0
+    ref_mismatches = 0
+    too_short = 0
+    reused_source_sequences = 0
+    reused_progress_sequences = 0
+
+    for processed, row_dict in enumerate(tqdm(df.to_dict("records"), total=len(df), desc="Fetching reference sequences"), start=1):
+        variant_id = str(row_dict["variant_id"])
+        ref = str(row_dict["REF"]).upper()
+
+        sequence = clean_sequence(row_dict.get("sequence"))
+        if sequence_matches_ref(sequence, ref, flank_size):
+            reused_source_sequences += 1
+        else:
+            sequence = clean_sequence(progress_sequences.get(variant_id))
+            if sequence_matches_ref(sequence, ref, flank_size):
+                reused_progress_sequences += 1
+            else:
+                sequence = fetcher.fetch(row_dict["CHROM"], int(row_dict["POS"]), ref, flank_size)
+                sequence = clean_sequence(sequence)
+
+        if sequence is None:
+            failed_fetches += 1
+        elif len(sequence) < MIN_SEQUENCE_LENGTH:
+            too_short += 1
+        elif not sequence_matches_ref(sequence, ref, flank_size):
+            ref_mismatches += 1
+        else:
+            row_dict["sequence"] = sequence
+            records.append(row_dict)
+
+        if processed % PROGRESS_EVERY == 0:
+            save_progress(records, progress_path)
+            save_fetcher_cache(fetcher)
+            print(f"Saved sequence progress after {processed:,} variants: {progress_path}")
+
+    save_progress(records, progress_path)
+    save_fetcher_cache(fetcher)
+
+    stats = {
+        "input_rows": int(len(df)),
+        "output_rows": int(len(records)),
+        "failed_fetches": int(failed_fetches),
+        "too_short": int(too_short),
+        "ref_mismatches": int(ref_mismatches),
+        "reused_source_sequences": int(reused_source_sequences),
+        "reused_progress_sequences": int(reused_progress_sequences),
+    }
+    return pd.DataFrame.from_records(records), stats
+
+
+def stratified_train_val_test_split(df: pd.DataFrame, random_state: int) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
+    stratify = df["label"] if df["label"].nunique() == 2 and df["label"].value_counts().min() >= 3 else None
+    train_df, temp_df = train_test_split(
+        df,
+        test_size=0.30,
+        random_state=random_state,
+        stratify=stratify,
+    )
+
+    temp_stratify = temp_df["label"] if temp_df["label"].nunique() == 2 and temp_df["label"].value_counts().min() >= 2 else None
+    val_df, test_df = train_test_split(
+        temp_df,
+        test_size=0.50,
+        random_state=random_state,
+        stratify=temp_stratify,
+    )
+    return train_df.reset_index(drop=True), val_df.reset_index(drop=True), test_df.reset_index(drop=True)
+
+
+def clean_reference_output(df: pd.DataFrame) -> pd.DataFrame:
+    columns_to_drop = [column for column in DROP_OUTPUT_COLUMNS if column in df.columns]
+    return df.drop(columns=columns_to_drop, errors="ignore")
+
+
+def build_alt_sequences(df: pd.DataFrame, flank_size: int) -> tuple[pd.DataFrame, int]:
+    records: list[dict[str, object]] = []
+    failed = 0
+
+    for row_dict in df.to_dict("records"):
+        ref_sequence = clean_sequence(row_dict["sequence"])
+        ref = str(row_dict["REF"]).upper()
+        alt = str(row_dict["ALT"]).upper()
+
+        if ref_sequence is None or not sequence_matches_ref(ref_sequence, ref, flank_size):
+            failed += 1
+            continue
+
+        upstream = ref_sequence[:flank_size]
+        downstream = ref_sequence[flank_size + len(ref) :]
+        alt_sequence = upstream + alt + downstream
+
+        output_row = row_dict.copy()
+        output_row["ref_sequence"] = ref_sequence
+        output_row["alt_sequence"] = alt_sequence
+        output_row["sequence"] = alt_sequence
+        output_row["ref_center"] = ref_sequence[flank_size : flank_size + len(ref)]
+        output_row["alt_center"] = alt_sequence[flank_size : flank_size + len(alt)]
+        records.append(output_row)
+
+    return pd.DataFrame.from_records(records), failed
+
+
+def save_splits(
+    train_df: pd.DataFrame,
+    val_df: pd.DataFrame,
+    test_df: pd.DataFrame,
+    output_dir: Path,
+    alt_output_dir: Path,
+    flank_size: int,
+) -> tuple[dict[str, Path], dict[str, Path], int]:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    alt_output_dir.mkdir(parents=True, exist_ok=True)
+
+    reference_paths: dict[str, Path] = {}
+    alt_paths: dict[str, Path] = {}
+    total_alt_failures = 0
+
+    for split_name, split_df in [("train", train_df), ("val", val_df), ("test", test_df)]:
+        reference_path = output_dir / ORIGINAL_SPLIT_FILES[split_name]
+        reference_df = clean_reference_output(split_df)
+        reference_df.to_csv(reference_path, index=False)
+        reference_paths[split_name] = reference_path
+
+        alt_df, alt_failures = build_alt_sequences(reference_df, flank_size)
+        total_alt_failures += alt_failures
+        alt_path = alt_output_dir / ALT_SPLIT_FILES[split_name]
+        alt_df.to_csv(alt_path, index=False)
+        alt_paths[split_name] = alt_path
+
+    return reference_paths, alt_paths, total_alt_failures
+
+
+def print_distribution(title: str, df: pd.DataFrame) -> None:
+    print(title)
+    if df.empty or "label" not in df.columns:
+        print("  no rows")
+    else:
+        print(df["label"].value_counts().sort_index().to_string())
+    print()
+
+
+def warn_if_imbalanced(df: pd.DataFrame) -> None:
+    counts = df["label"].value_counts().sort_index()
+    if len(counts) < 2:
+        print("WARNING: selected sample contains only one class. Training will not be useful.")
+        return
+    if counts.iloc[0] != counts.iloc[1]:
+        print("WARNING: exact class balance was not possible with the available/capped variants.")
+        print(counts.to_string())
+        print()
+
+
+def main() -> None:
+    args = parse_args()
+    root = project_root()
+    output_dir_arg = args.output_dir or default_output_dir_for_target(args.target_total)
+    output_dir = resolve_path(root, output_dir_arg)
+    if output_dir is None:
+        raise ValueError("--output_dir is required")
+    alt_output_dir = alt_output_dir_for(output_dir)
+    requested_vcf = resolve_path(root, args.clinvar_vcf)
+    fasta_path = resolve_path(root, args.fasta_path)
+    sequence_cache = resolve_path(root, args.sequence_cache)
+
+    print("Prepare larger ClinVar dataset")
+    print(f"Target total rows: {args.target_total:,}")
+    print(f"Output reference directory: {output_dir}")
+    print(f"Output alternate directory: {alt_output_dir}")
+    print(f"Flank size: {args.flank_size}")
+    print()
+
+    variants_df, source_name, before_filter_count, after_filter_count = choose_variant_source(
+        root=root,
+        target_total=args.target_total,
+        requested_vcf=requested_vcf,
+    )
+
+    print("=" * 80)
+    print("FILTERING SUMMARY")
+    print("=" * 80)
+    print(f"Selected source: {source_name}")
+    print(f"Total variants before filtering: {before_filter_count:,}")
+    print(f"Total variants after filtering: {after_filter_count:,}")
+    print_distribution("Class distribution after filtering:", variants_df)
+    filtered_variants_path = save_dataframe(variants_df, output_dir / "clinvar_binary_variants.csv")
+    print(f"Saved filtered variant table: {filtered_variants_path}")
+    print()
+
+    if variants_df.empty:
+        raise RuntimeError("No variants remained after filtering. Check the input VCF/CSV and CLNSIG labels.")
+
+    sampled_df = balanced_sample(
+        variants_df,
+        target_total=args.target_total,
+        max_pathogenic=args.max_pathogenic,
+        max_benign=args.max_benign,
+        random_state=args.random_state,
+    )
+
+    print("=" * 80)
+    print("SAMPLING SUMMARY")
+    print("=" * 80)
+    print(f"Selected sample size: {len(sampled_df):,}")
+    print_distribution("Class distribution after balanced sampling:", sampled_df)
+    warn_if_imbalanced(sampled_df)
+    selected_variants_path = save_dataframe(sampled_df, output_dir / "selected_variants.csv")
+    print(f"Saved selected variant table: {selected_variants_path}")
+    print()
+
+    if sampled_df.empty or sampled_df["label"].nunique() < 2:
+        raise RuntimeError("Balanced sampling did not produce both classes. Adjust target_total/max caps.")
+
+    sequenced_df, fetch_stats = fetch_reference_sequences(
+        sampled_df,
+        output_dir=output_dir,
+        flank_size=args.flank_size,
+        fetch_mode=args.fetch_mode,
+        fasta_path=fasta_path,
+        sequence_cache=sequence_cache,
+        sleep_seconds=args.sleep_seconds,
+        max_retries=args.max_retries,
+    )
+
+    print("=" * 80)
+    print("SEQUENCE FETCH SUMMARY")
+    print("=" * 80)
+    print(f"Rows requested for sequence fetching: {fetch_stats['input_rows']:,}")
+    print(f"Rows with usable reference sequences: {fetch_stats['output_rows']:,}")
+    print(f"Failed sequence fetches: {fetch_stats['failed_fetches']:,}")
+    print(f"Too-short sequences: {fetch_stats['too_short']:,}")
+    print(f"Reference mismatches at variant center: {fetch_stats['ref_mismatches']:,}")
+    print(f"Reused source sequences: {fetch_stats['reused_source_sequences']:,}")
+    print(f"Reused progress sequences: {fetch_stats['reused_progress_sequences']:,}")
+    print_distribution("Class distribution after sequence filtering:", sequenced_df)
+
+    if sequenced_df.empty:
+        raise RuntimeError("No usable sequences were created. Check VCF/input data and sequence fetching settings.")
+
+    train_df, val_df, test_df = stratified_train_val_test_split(sequenced_df, args.random_state)
+
+    print("=" * 80)
+    print("SPLIT SUMMARY")
+    print("=" * 80)
+    print(f"Train rows: {len(train_df):,}")
+    print(f"Validation rows: {len(val_df):,}")
+    print(f"Test rows: {len(test_df):,}")
+    print_distribution("Train label distribution:", train_df)
+    print_distribution("Validation label distribution:", val_df)
+    print_distribution("Test label distribution:", test_df)
+
+    reference_paths, alt_paths, alt_failures = save_splits(
+        train_df=train_df,
+        val_df=val_df,
+        test_df=test_df,
+        output_dir=output_dir,
+        alt_output_dir=alt_output_dir,
+        flank_size=args.flank_size,
+    )
+
+    print("=" * 80)
+    print("OUTPUT SUMMARY")
+    print("=" * 80)
+    print(f"Alternate sequence build failures: {alt_failures:,}")
+    print("Reference-sequence output paths:")
+    for split_name, path in reference_paths.items():
+        print(f"  {split_name}: {path}")
+    print("Alternate-sequence output paths:")
+    for split_name, path in alt_paths.items():
+        print(f"  {split_name}: {path}")
+    print()
+    print("Larger ClinVar dataset preparation completed.")
+
+
+if __name__ == "__main__":
+    main()

training/requirements-colab.txt

@@ -8,7 +8,9 @@ pandas>=2.2.0
 pyfaidx>=0.8.1
 pysam>=0.22.0
 requests>=2.32.0
+safetensors>=0.4.3
 scikit-learn>=1.4.0
 torch>=2.2.0
 tqdm>=4.66.0
 transformers>=4.41.0
+joblib>=1.4.0

training/requirements-mac.txt

@@ -0,0 +1,11 @@
+torch
+transformers
+datasets
+accelerate
+scikit-learn
+pandas
+numpy
+evaluate
+einops
+tqdm
+safetensors

training/scripts/train_dnabert2_classifier.py

@@ -57,6 +57,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--batch-size", type=int, default=8)
     parser.add_argument("--learning-rate", type=float, default=2e-5)
     parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--logging-steps", type=int, default=10)
+    parser.add_argument("--max-train-samples", type=int, default=0, help="Optional cap for quick smoke tests. 0 means all rows.")
+    parser.add_argument("--max-val-samples", type=int, default=0, help="Optional cap for quick smoke tests. 0 means all rows.")
+    parser.add_argument("--max-test-samples", type=int, default=0, help="Optional cap for quick smoke tests. 0 means all rows.")
+    parser.add_argument("--max-steps", type=int, default=0, help="Optional Trainer max_steps for smoke tests. 0 means disabled.")
+    parser.add_argument("--force-cpu", action="store_true", help="Force CPU training even when a GPU is visible.")
     parser.add_argument("--no-clean-clnsig", action="store_true", help="Do not re-filter CLNSIG labels in CSV inputs.")
     return parser.parse_args()
 
@@ -196,6 +202,15 @@ def compute_metrics(eval_pred):
     }
 
 
+def limit_dataset(dataset: Dataset, max_samples: int, split_name: str, seed: int) -> Dataset:
+    """Optionally limit a dataset for fast CPU smoke tests."""
+    if max_samples <= 0 or len(dataset) <= max_samples:
+        return dataset
+    limited = dataset.shuffle(seed=seed).select(range(max_samples))
+    print(f"{split_name}: using {len(limited):,}/{len(dataset):,} rows for this run")
+    return limited
+
+
 def build_training_args(args: argparse.Namespace, output_dir: Path) -> TrainingArguments:
     kwargs = {
         "output_dir": str(output_dir / "checkpoints"),
@@ -212,6 +227,25 @@ def build_training_args(args: argparse.Namespace, output_dir: Path) -> TrainingA
     }
 
     signature = inspect.signature(TrainingArguments.__init__)
+    optional_kwargs = {
+        "disable_tqdm": False,
+        "logging_strategy": "steps",
+        "logging_steps": args.logging_steps,
+        "logging_first_step": True,
+        "save_total_limit": 2,
+    }
+    if args.max_steps > 0:
+        optional_kwargs["max_steps"] = args.max_steps
+    if args.force_cpu:
+        if "use_cpu" in signature.parameters:
+            optional_kwargs["use_cpu"] = True
+        elif "no_cuda" in signature.parameters:
+            optional_kwargs["no_cuda"] = True
+
+    for key, value in optional_kwargs.items():
+        if key in signature.parameters:
+            kwargs[key] = value
+
     if "eval_strategy" in signature.parameters:
         kwargs["eval_strategy"] = "epoch"
     else:
@@ -244,6 +278,10 @@ def main() -> None:
     else:
         raise ValueError("Provide either --dataset-jsonl or all of --train-csv, --val-csv, and --test-csv.")
 
+    train_dataset = limit_dataset(train_dataset, args.max_train_samples, "train", args.seed)
+    eval_dataset = limit_dataset(eval_dataset, args.max_val_samples, "val", args.seed)
+    test_dataset = limit_dataset(test_dataset, args.max_test_samples, "test", args.seed)
+
     tokenizer = AutoTokenizer.from_pretrained(args.model_name, trust_remote_code=True)
 
     def tokenize(batch):

training/setup_mac.md

@@ -0,0 +1,80 @@
+# Mac Local Training Setup
+
+This guide sets up a Mac-friendly Python environment for local DNABERT-2 experiments.
+
+Full DNABERT-2 training can be slow on a Mac, especially on CPU. Start with small smoke tests before running larger jobs.
+
+## CSV File Locations
+
+Training scripts should look for sequence CSVs in this order:
+
+1. `data/processed/`
+2. `training/csv_files/`
+
+Your current CSV files are in:
+
+```text
+training/csv_files/train_with_sequences.csv
+training/csv_files/val_with_sequences.csv
+training/csv_files/test_with_sequences.csv
+```
+
+That location is supported.
+
+## 1. Create Virtual Environment
+
+Run this from the project root:
+
+```bash
+python3 -m venv .venv
+```
+
+## 2. Activate It
+
+```bash
+source .venv/bin/activate
+```
+
+Your terminal prompt should now show `.venv`.
+
+## 3. Upgrade pip
+
+```bash
+python -m pip install --upgrade pip
+```
+
+## 4. Install Requirements
+
+```bash
+pip install -r training/requirements-mac.txt
+```
+
+## 5. Check PyTorch Device
+
+```bash
+python training/check_device.py
+```
+
+Expected output includes one of:
+
+```text
+Using device: mps
+```
+
+or:
+
+```text
+Using device: cpu
+```
+
+If you see `Using device: cpu`, the project will still work, but local training will be slow. Use very small smoke tests locally and use a GPU environment for full model training.
+
+## 6. Check Dataset
+
+Before training, verify the CSV files:
+
+```bash
+python training/check_dataset.py
+```
+
+This confirms that the `sequence` and `label` columns exist, labels are encoded as `0` and `1`, sequence values are present, and class balance is reasonable.

training/train_local_dnabert2.py

@@ -0,0 +1,1212 @@
+#!/usr/bin/env python3
+"""Fine-tune DNABERT-2 on prepared ClinVar sequence CSV files.
+
+This script works on CUDA, Mac MPS, and CPU. CUDA tries the standard
+Hugging Face DNABERT-2 model first; MPS and fallback paths use the local
+no-Triton patch from train_smoke_test.py.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import inspect
+import json
+import os
+import sys
+import zipfile
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+from datasets import Dataset
+from sklearn.metrics import (
+    accuracy_score,
+    confusion_matrix,
+    f1_score,
+    matthews_corrcoef,
+    precision_score,
+    recall_score,
+    roc_auc_score,
+)
+from sklearn.model_selection import train_test_split
+from tqdm.auto import tqdm
+from transformers import (
+    AutoConfig,
+    AutoTokenizer,
+    DataCollatorWithPadding,
+    Trainer,
+    TrainingArguments,
+    default_data_collator,
+)
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from training.train_smoke_test import (  # noqa: E402
+    LOCAL_DNABERT2_PATCH_DIR,
+    MODEL_NAME,
+    choose_device,
+    clear_local_patch_module_cache,
+    create_local_dnabert2_patch,
+    disable_flash_attention_on_config,
+    load_sequence_classification_model,
+    should_drop_clnsig,
+)
+
+
+ORIGINAL_SPLIT_FILES = {
+    "train": "train_with_sequences.csv",
+    "val": "val_with_sequences.csv",
+    "test": "test_with_sequences.csv",
+}
+
+ALT_SPLIT_FILES = {
+    "train": "train_with_alt_sequences.csv",
+    "val": "val_with_alt_sequences.csv",
+    "test": "test_with_alt_sequences.csv",
+}
+
+MIN_SEQUENCE_LENGTH = 50
+RANDOM_STATE = 42
+
+
+@dataclass
+class CsvSelection:
+    train_csv: Path
+    val_csv: Path
+    test_csv: Path
+    dataset_dir: Path
+    is_alt_dataset: bool
+    is_large_alt_dataset: bool
+    is_10k_alt_dataset: bool = False
+    is_20k_alt_dataset: bool = False
+    dataset_name: str = ""
+
+
+def parse_bool(value: str | bool) -> bool:
+    if isinstance(value, bool):
+        return value
+    normalized = value.strip().lower()
+    if normalized in {"true", "1", "yes", "y"}:
+        return True
+    if normalized in {"false", "0", "no", "n"}:
+        return False
+    raise argparse.ArgumentTypeError("Expected true or false.")
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--train_csv", type=Path, default=None)
+    parser.add_argument("--val_csv", type=Path, default=None)
+    parser.add_argument("--test_csv", type=Path, default=None)
+    parser.add_argument("--dataset_dir", type=Path, default=None)
+    parser.add_argument("--output_dir", type=Path, default=Path("training/outputs/dnabert2_clinvar"))
+    parser.add_argument("--sequence_column", type=str, default="sequence")
+    parser.add_argument("--sample_size", type=int, default=0)
+    parser.add_argument("--max_length", type=int, default=512)
+    parser.add_argument("--variant_center_index", type=int, default=512)
+    parser.add_argument("--epochs", type=float, default=5)
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--grad_accum_steps", type=int, default=8)
+    parser.add_argument("--learning_rate", type=float, default=2e-5)
+    parser.add_argument("--eval_accumulation_steps", type=int, default=1)
+    parser.add_argument("--eval_subset_size", type=int, default=0)
+    parser.add_argument("--save_eval_each_epoch", type=parse_bool, default=False)
+    parser.add_argument("--freeze_encoder", type=parse_bool, default=True)
+    parser.add_argument("--unfreeze_last_n_layers", type=int, default=0)
+    parser.add_argument("--freeze_embeddings", type=parse_bool, default=True)
+    parser.add_argument("--use_class_weights", type=parse_bool, default=True)
+    parser.add_argument("--center_crop", type=parse_bool, default=True)
+    parser.add_argument("--tune_threshold", type=parse_bool, default=True)
+    parser.add_argument("--resume_from_checkpoint", type=str, default=None)
+    return parser.parse_args()
+
+
+def resolve_path(path: Path) -> Path:
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def has_all_split_files(directory: Path, split_files: dict[str, str]) -> bool:
+    return all((directory / filename).exists() for filename in split_files.values())
+
+
+def infer_alt_dataset_from_paths(paths: list[Path]) -> bool:
+    return any("csv_files_alt" in path.parts or "with_alt_sequences" in path.name for path in paths)
+
+
+def infer_large_alt_dataset_from_paths(paths: list[Path]) -> bool:
+    return any(
+        directory_name in path.parts
+        for path in paths
+        for directory_name in ["csv_files_large_alt", "csv_files_10k_alt", "csv_files_20k_alt"]
+    )
+
+
+def infer_10k_alt_dataset_from_paths(paths: list[Path]) -> bool:
+    return any("csv_files_10k_alt" in path.parts for path in paths)
+
+
+def infer_20k_alt_dataset_from_paths(paths: list[Path]) -> bool:
+    return any("csv_files_20k_alt" in path.parts for path in paths)
+
+
+def selection_from_directory(directory: Path, dataset_name: str = "") -> CsvSelection:
+    if has_all_split_files(directory, ALT_SPLIT_FILES):
+        selected_paths = [
+            directory / ALT_SPLIT_FILES["train"],
+            directory / ALT_SPLIT_FILES["val"],
+            directory / ALT_SPLIT_FILES["test"],
+        ]
+        return CsvSelection(
+            train_csv=selected_paths[0],
+            val_csv=selected_paths[1],
+            test_csv=selected_paths[2],
+            dataset_dir=directory,
+            is_alt_dataset=True,
+            is_large_alt_dataset=infer_large_alt_dataset_from_paths(selected_paths),
+            is_10k_alt_dataset=infer_10k_alt_dataset_from_paths(selected_paths),
+            is_20k_alt_dataset=infer_20k_alt_dataset_from_paths(selected_paths),
+            dataset_name=dataset_name or directory.name,
+        )
+
+    if has_all_split_files(directory, ORIGINAL_SPLIT_FILES):
+        selected_paths = [
+            directory / ORIGINAL_SPLIT_FILES["train"],
+            directory / ORIGINAL_SPLIT_FILES["val"],
+            directory / ORIGINAL_SPLIT_FILES["test"],
+        ]
+        return CsvSelection(
+            train_csv=selected_paths[0],
+            val_csv=selected_paths[1],
+            test_csv=selected_paths[2],
+            dataset_dir=directory,
+            is_alt_dataset=False,
+            is_large_alt_dataset=False,
+            dataset_name=dataset_name or directory.name,
+        )
+
+    raise FileNotFoundError(
+        "Dataset directory does not contain a complete alternate or reference split:\n"
+        f"{directory}"
+    )
+
+
+def find_default_csv_selection() -> CsvSelection:
+    """Prefer larger alternate-sequence datasets, then smaller fallbacks."""
+    candidates = [
+        (PROJECT_ROOT / "training" / "csv_files_20k_alt", ALT_SPLIT_FILES, "20k alternate-sequence dataset"),
+        (PROJECT_ROOT / "training" / "csv_files_10k_alt", ALT_SPLIT_FILES, "10k alternate-sequence dataset"),
+        (PROJECT_ROOT / "training" / "csv_files_large_alt", ALT_SPLIT_FILES, "large alternate-sequence dataset"),
+        (PROJECT_ROOT / "training" / "csv_files_alt", ALT_SPLIT_FILES, "alternate-sequence dataset"),
+        (PROJECT_ROOT / "data" / "processed", ORIGINAL_SPLIT_FILES, "data/processed reference dataset"),
+        (PROJECT_ROOT / "training" / "csv_files", ORIGINAL_SPLIT_FILES, "reference-sequence dataset"),
+        (PROJECT_ROOT / "training" / "csv_files_large", ORIGINAL_SPLIT_FILES, "large reference-sequence dataset"),
+    ]
+
+    for directory, split_files, dataset_name in candidates:
+        if has_all_split_files(directory, split_files):
+            return selection_from_directory(directory, dataset_name)
+
+    searched = "\n".join(str(directory) for directory, _split_files, _dataset_name in candidates)
+    raise FileNotFoundError(
+        "Could not find train/val/test CSV files in the default locations.\n"
+        f"Searched:\n{searched}"
+    )
+
+
+def resolve_csv_paths(args: argparse.Namespace) -> CsvSelection:
+    default_selection = None
+    if args.dataset_dir is not None and not (args.train_csv or args.val_csv or args.test_csv):
+        return selection_from_directory(resolve_path(args.dataset_dir), args.dataset_dir.name)
+
+    if not (args.train_csv and args.val_csv and args.test_csv):
+        default_selection = find_default_csv_selection()
+
+    train_csv = resolve_path(args.train_csv) if args.train_csv else default_selection.train_csv
+    val_csv = resolve_path(args.val_csv) if args.val_csv else default_selection.val_csv
+    test_csv = resolve_path(args.test_csv) if args.test_csv else default_selection.test_csv
+
+    for csv_path in [train_csv, val_csv, test_csv]:
+        if not csv_path.exists():
+            raise FileNotFoundError(f"CSV file not found: {csv_path}")
+
+    selected_paths = [train_csv, val_csv, test_csv]
+    is_alt_dataset = (
+        default_selection.is_alt_dataset
+        if default_selection is not None
+        else infer_alt_dataset_from_paths(selected_paths)
+    )
+
+    if infer_alt_dataset_from_paths(selected_paths):
+        is_alt_dataset = True
+
+    is_large_alt_dataset = (
+        default_selection.is_large_alt_dataset
+        if default_selection is not None
+        else infer_large_alt_dataset_from_paths(selected_paths)
+    )
+
+    if infer_large_alt_dataset_from_paths(selected_paths):
+        is_large_alt_dataset = True
+
+    is_10k_alt_dataset = (
+        default_selection.is_10k_alt_dataset
+        if default_selection is not None
+        else infer_10k_alt_dataset_from_paths(selected_paths)
+    )
+    if infer_10k_alt_dataset_from_paths(selected_paths):
+        is_10k_alt_dataset = True
+
+    is_20k_alt_dataset = (
+        default_selection.is_20k_alt_dataset
+        if default_selection is not None
+        else infer_20k_alt_dataset_from_paths(selected_paths)
+    )
+    if infer_20k_alt_dataset_from_paths(selected_paths):
+        is_20k_alt_dataset = True
+
+    common_parent = selected_paths[0].parent
+    if any(path.parent != common_parent for path in selected_paths):
+        common_parent = PROJECT_ROOT
+
+    return CsvSelection(
+        train_csv=train_csv,
+        val_csv=val_csv,
+        test_csv=test_csv,
+        dataset_dir=common_parent,
+        is_alt_dataset=is_alt_dataset,
+        is_large_alt_dataset=is_large_alt_dataset,
+        is_10k_alt_dataset=is_10k_alt_dataset,
+        is_20k_alt_dataset=is_20k_alt_dataset,
+        dataset_name=default_selection.dataset_name if default_selection is not None else common_parent.name,
+    )
+
+
+def clean_sequence(value: object) -> str:
+    return str(value).strip().upper()
+
+
+def crop_sequence_around_variant(sequence: str, max_length: int, variant_center_index: int) -> str:
+    """Crop a long sequence while keeping the variant position in the window."""
+    if len(sequence) <= max_length:
+        return sequence
+
+    start = max(0, variant_center_index - max_length // 2)
+    end = start + max_length
+    if end > len(sequence):
+        end = len(sequence)
+        start = max(0, end - max_length)
+
+    return sequence[start:end]
+
+
+def apply_center_crop(
+    df: pd.DataFrame,
+    max_length: int,
+    variant_center_index: int,
+    split_name: str,
+    enabled: bool,
+) -> pd.DataFrame:
+    if not enabled:
+        print(f"{split_name} variant-centered crop: disabled")
+        return df.reset_index(drop=True)
+
+    df = df.copy()
+    original_lengths = df["sequence"].str.len()
+    cropped_count = int((original_lengths > max_length).sum())
+    df["sequence"] = df["sequence"].apply(
+        lambda sequence: crop_sequence_around_variant(sequence, max_length, variant_center_index)
+    )
+    cropped_lengths = df["sequence"].str.len()
+
+    print(
+        f"{split_name} variant-centered crop: "
+        f"{cropped_count:,}/{len(df):,} sequences cropped to {max_length} characters "
+        f"around input index {variant_center_index}"
+    )
+    print(
+        f"{split_name} sequence lengths after crop: "
+        f"min={int(cropped_lengths.min())}, "
+        f"mean={cropped_lengths.mean():.1f}, "
+        f"max={int(cropped_lengths.max())}"
+    )
+    print()
+    return df.reset_index(drop=True)
+
+
+def load_and_filter_dataframe(csv_path: Path, split_name: str, sequence_column: str) -> pd.DataFrame:
+    df = pd.read_csv(csv_path)
+    before_rows = len(df)
+
+    required_columns = {sequence_column, "label"}
+    missing_columns = sorted(required_columns - set(df.columns))
+    if missing_columns:
+        raise ValueError(f"{csv_path} is missing required columns: {missing_columns}")
+
+    df = df.copy()
+
+    if "CLNSIG" in df.columns:
+        df = df.loc[~df["CLNSIG"].fillna("").apply(should_drop_clnsig)].copy()
+
+    df["label"] = pd.to_numeric(df["label"], errors="coerce")
+    df = df.loc[df["label"].isin([0, 1])].copy()
+    df["label"] = df["label"].astype(int)
+
+    df["sequence"] = df[sequence_column].fillna("").apply(clean_sequence)
+    df = df.loc[df["sequence"] != ""].copy()
+    df = df.loc[df["sequence"].str.len() >= MIN_SEQUENCE_LENGTH].copy()
+
+    print(f"{split_name} CSV: {csv_path}")
+    print(f"{split_name} rows before filtering: {before_rows:,}")
+    print(f"{split_name} rows after filtering: {len(df):,}")
+    print(f"{split_name} label distribution after filtering:")
+    print(df["label"].value_counts().sort_index().to_string())
+    print()
+
+    if df.empty:
+        raise ValueError(f"No usable rows remain in {split_name}.")
+
+    return df.reset_index(drop=True)
+
+
+def stratified_sample(df: pd.DataFrame, sample_size: int, split_name: str) -> pd.DataFrame:
+    if sample_size <= 0 or len(df) <= sample_size:
+        return df.reset_index(drop=True)
+
+    label_counts = df["label"].value_counts()
+    if df["label"].nunique() == 2 and label_counts.min() >= 2 and sample_size >= 2:
+        try:
+            sampled, _unused = train_test_split(
+                df,
+                train_size=sample_size,
+                stratify=df["label"],
+                random_state=RANDOM_STATE,
+            )
+            sampled = sampled.sort_index()
+        except ValueError:
+            sampled = df.sample(n=sample_size, random_state=RANDOM_STATE).sort_index()
+    else:
+        sampled = df.sample(n=sample_size, random_state=RANDOM_STATE).sort_index()
+
+    print(f"{split_name} sampled rows: {len(sampled):,}/{len(df):,}")
+    print(f"{split_name} label distribution after sampling:")
+    print(sampled["label"].value_counts().sort_index().to_string())
+    print()
+    return sampled.reset_index(drop=True)
+
+
+def make_eval_subset(df: pd.DataFrame, eval_subset_size: int, split_name: str) -> pd.DataFrame:
+    if eval_subset_size <= 0 or len(df) <= eval_subset_size:
+        print(f"{split_name} evaluation rows: {len(df):,}")
+        print()
+        return df.reset_index(drop=True)
+
+    subset = stratified_sample(df, eval_subset_size, f"{split_name} evaluation")
+    print(f"{split_name} evaluation subset enabled: {len(subset):,}/{len(df):,}")
+    print()
+    return subset
+
+
+def dataframe_to_dataset(df: pd.DataFrame) -> Dataset:
+    dataset_df = df[["sequence", "label"]].rename(columns={"label": "labels"})
+    return Dataset.from_pandas(dataset_df, preserve_index=False)
+
+
+def tokenize_dataset(tokenizer, dataset: Dataset, max_length: int, split_name: str) -> Dataset:
+    def tokenize_batch(batch):
+        return tokenizer(
+            batch["sequence"],
+            max_length=max_length,
+            padding="max_length",
+            truncation=True,
+        )
+
+    tokenized = dataset.map(tokenize_batch, batched=True, remove_columns=["sequence"], desc=f"Tokenizing {split_name}")
+    return tokenized
+
+
+def extract_logits(eval_pred) -> tuple[np.ndarray, np.ndarray]:
+    if hasattr(eval_pred, "predictions"):
+        logits = eval_pred.predictions
+        labels = eval_pred.label_ids
+    else:
+        logits, labels = eval_pred
+
+    if isinstance(logits, (tuple, list)):
+        for item in logits:
+            candidate = np.asarray(item)
+            if candidate.ndim >= 2 and candidate.shape[-1] == 2:
+                logits = candidate
+                break
+        else:
+            logits = np.asarray(logits[0])
+    else:
+        logits = np.asarray(logits)
+
+    return logits, np.asarray(labels)
+
+
+def softmax(logits: np.ndarray) -> np.ndarray:
+    shifted = logits - np.max(logits, axis=-1, keepdims=True)
+    exp = np.exp(shifted)
+    return exp / np.sum(exp, axis=-1, keepdims=True)
+
+
+def extract_predictions_and_labels(eval_pred) -> tuple[np.ndarray, np.ndarray]:
+    if hasattr(eval_pred, "predictions"):
+        predictions = eval_pred.predictions
+        labels = eval_pred.label_ids
+    else:
+        predictions, labels = eval_pred
+
+    if isinstance(predictions, (tuple, list)):
+        for item in predictions:
+            candidate = np.asarray(item)
+            if candidate.ndim == 1 or (candidate.ndim >= 2 and candidate.shape[-1] in {1, 2}):
+                predictions = candidate
+                break
+        else:
+            predictions = np.asarray(predictions[0])
+    else:
+        predictions = np.asarray(predictions)
+
+    return predictions, np.asarray(labels)
+
+
+def probabilities_from_predictions(predictions: np.ndarray) -> np.ndarray:
+    predictions = np.asarray(predictions)
+    if predictions.ndim == 1:
+        return predictions.astype(np.float64)
+    if predictions.ndim >= 2 and predictions.shape[-1] == 1:
+        return predictions.reshape(-1).astype(np.float64)
+    if predictions.ndim >= 2 and predictions.shape[-1] == 2:
+        return softmax(predictions.astype(np.float64))[:, 1]
+    raise ValueError(f"Unsupported prediction shape for binary classification: {predictions.shape}")
+
+
+def class_predictions_from_scores(predictions: np.ndarray) -> np.ndarray:
+    predictions = np.asarray(predictions)
+    if predictions.ndim == 1 or (predictions.ndim >= 2 and predictions.shape[-1] == 1):
+        probabilities = probabilities_from_predictions(predictions)
+        return (probabilities >= 0.5).astype(int)
+    return np.argmax(predictions, axis=-1)
+
+
+def positive_class_probabilities(logits: np.ndarray) -> np.ndarray:
+    return probabilities_from_predictions(logits)
+
+
+def preprocess_logits_for_metrics(logits, labels):
+    """Keep only class-1 probabilities during Trainer evaluation to reduce memory."""
+    logits = extract_torch_logits(logits)
+    if logits.ndim >= 2 and logits.shape[-1] == 2:
+        return torch.softmax(logits.float(), dim=-1)[:, 1]
+    if logits.ndim >= 2 and logits.shape[-1] == 1:
+        return logits.squeeze(-1)
+    return logits
+
+
+def compute_metrics(eval_pred):
+    predictions, labels = extract_predictions_and_labels(eval_pred)
+    class_predictions = class_predictions_from_scores(predictions)
+
+    metrics = {
+        "accuracy": float(accuracy_score(labels, class_predictions)),
+        "precision": float(precision_score(labels, class_predictions, zero_division=0)),
+        "recall": float(recall_score(labels, class_predictions, zero_division=0)),
+        "f1": float(f1_score(labels, class_predictions, zero_division=0)),
+        "mcc": float(matthews_corrcoef(labels, class_predictions)),
+    }
+
+    if len(np.unique(labels)) == 2:
+        probabilities = probabilities_from_predictions(predictions)
+        try:
+            metrics["auc_roc"] = float(roc_auc_score(labels, probabilities))
+        except ValueError:
+            metrics["auc_roc"] = None
+
+    return metrics
+
+
+def tune_threshold_for_mcc(labels: np.ndarray, probabilities: np.ndarray) -> tuple[float, float]:
+    """Try thresholds from 0.10 to 0.90 and keep the best validation MCC."""
+    best_threshold = 0.5
+    best_mcc = -1.0
+    best_f1 = -1.0
+
+    for threshold in np.round(np.arange(0.10, 0.901, 0.01), 2):
+        predictions = (probabilities >= threshold).astype(int)
+        mcc = matthews_corrcoef(labels, predictions)
+        f1 = f1_score(labels, predictions, zero_division=0)
+
+        if mcc > best_mcc or (np.isclose(mcc, best_mcc) and f1 > best_f1):
+            best_threshold = float(threshold)
+            best_mcc = float(mcc)
+            best_f1 = float(f1)
+
+    return best_threshold, best_mcc
+
+
+def metrics_at_threshold(labels: np.ndarray, probabilities: np.ndarray, threshold: float) -> dict:
+    predictions = (probabilities >= threshold).astype(int)
+    matrix = confusion_matrix(labels, predictions, labels=[0, 1]).astype(int)
+
+    metrics = {
+        "selected_threshold": float(threshold),
+        "accuracy": float(accuracy_score(labels, predictions)),
+        "precision": float(precision_score(labels, predictions, zero_division=0)),
+        "recall": float(recall_score(labels, predictions, zero_division=0)),
+        "f1": float(f1_score(labels, predictions, zero_division=0)),
+        "mcc": float(matthews_corrcoef(labels, predictions)),
+        "auc_roc": None,
+        "confusion_matrix": matrix.tolist(),
+    }
+
+    if len(np.unique(labels)) == 2:
+        try:
+            metrics["auc_roc"] = float(roc_auc_score(labels, probabilities))
+        except ValueError:
+            metrics["auc_roc"] = None
+
+    return metrics
+
+
+def predict_probabilities(trainer: Trainer, dataset: Dataset) -> tuple[np.ndarray, np.ndarray]:
+    prediction_output = trainer.predict(dataset)
+    logits, labels = extract_logits(prediction_output)
+    probabilities = positive_class_probabilities(logits)
+    return labels.astype(int), probabilities
+
+
+def clear_mps_cache_if_needed(device: str) -> None:
+    if device != "mps":
+        return
+    mps_backend = getattr(torch, "mps", None)
+    if mps_backend is not None and hasattr(mps_backend, "empty_cache"):
+        mps_backend.empty_cache()
+
+
+def predict_in_small_batches(model, dataset: Dataset, batch_size: int = 1, device: str = "cpu") -> tuple[np.ndarray, np.ndarray]:
+    """Memory-safe prediction loop that immediately moves scores to CPU."""
+    device_object = torch.device(device)
+    model.to(device_object)
+    model.eval()
+
+    loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        collate_fn=default_data_collator,
+    )
+
+    all_probabilities: list[np.ndarray] = []
+    all_labels: list[np.ndarray] = []
+
+    for batch in tqdm(loader, desc="Memory-safe prediction"):
+        labels = batch.pop("labels")
+        batch = {key: value.to(device_object) for key, value in batch.items()}
+
+        with torch.no_grad():
+            outputs = model(**batch)
+            logits = extract_torch_logits(outputs)
+            probabilities = torch.softmax(logits.float(), dim=-1)[:, 1]
+
+        all_probabilities.append(probabilities.detach().cpu().numpy())
+        all_labels.append(labels.detach().cpu().numpy())
+
+        del outputs, logits, probabilities, labels, batch
+        clear_mps_cache_if_needed(device)
+
+    gc.collect()
+    clear_mps_cache_if_needed(device)
+
+    return np.concatenate(all_probabilities), np.concatenate(all_labels).astype(int)
+
+
+def print_confusion_matrix(split_name: str, matrix: list[list[int]]) -> None:
+    print(f"{split_name} confusion matrix")
+    print("                 predicted_0  predicted_1")
+    print(f"actual_0         {matrix[0][0]:>11}  {matrix[0][1]:>11}")
+    print(f"actual_1         {matrix[1][0]:>11}  {matrix[1][1]:>11}")
+
+
+def print_threshold_metrics(split_name: str, metrics: dict) -> None:
+    print(f"{split_name} metrics at threshold {metrics['selected_threshold']:.2f}")
+    for key in ["accuracy", "precision", "recall", "f1", "mcc", "auc_roc"]:
+        value = metrics[key]
+        if value is None:
+            print(f"{key}: n/a")
+        else:
+            print(f"{key}: {value:.4f}")
+    print_confusion_matrix(split_name, metrics["confusion_matrix"])
+    print()
+
+
+def compute_class_weights(train_df: pd.DataFrame) -> torch.Tensor:
+    """Balanced class weights: total_rows / (num_classes * class_count)."""
+    counts = train_df["label"].value_counts().sort_index()
+    if not {0, 1}.issubset(set(counts.index)):
+        print("Class weights requested, but both classes are not present. Using equal weights.")
+        return torch.tensor([1.0, 1.0], dtype=torch.float32)
+
+    total = float(counts.sum())
+    weights = [
+        total / (2.0 * float(counts.loc[0])),
+        total / (2.0 * float(counts.loc[1])),
+    ]
+    return torch.tensor(weights, dtype=torch.float32)
+
+
+def maybe_compute_class_weights(train_df: pd.DataFrame, enabled: bool) -> torch.Tensor | None:
+    if not enabled:
+        print("Class weights: disabled")
+        print()
+        return None
+
+    class_weights = compute_class_weights(train_df)
+    print("Class weights: enabled")
+    print(f"label 0 weight: {class_weights[0].item():.4f}")
+    print(f"label 1 weight: {class_weights[1].item():.4f}")
+    print()
+    return class_weights
+
+
+def extract_torch_logits(outputs) -> torch.Tensor:
+    if hasattr(outputs, "logits"):
+        return outputs.logits
+
+    if isinstance(outputs, (tuple, list)):
+        for item in outputs:
+            if torch.is_tensor(item) and item.ndim >= 2 and item.shape[-1] == 2:
+                return item
+        return outputs[0]
+
+    raise TypeError("Could not find logits in model outputs.")
+
+
+class WeightedLossTrainer(Trainer):
+    """Trainer that uses weighted cross entropy for imbalanced labels."""
+
+    def __init__(self, *args, class_weights: torch.Tensor | None = None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.class_weights = class_weights
+
+    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
+        labels = inputs.pop("labels")
+        outputs = model(**inputs)
+        logits = extract_torch_logits(outputs)
+
+        weight = self.class_weights.to(logits.device) if self.class_weights is not None else None
+        loss_fn = torch.nn.CrossEntropyLoss(weight=weight)
+        loss = loss_fn(logits.view(-1, logits.shape[-1]), labels.view(-1).long())
+
+        if return_outputs:
+            return loss, outputs
+        return loss
+
+
+def local_patch_is_ready(patch_dir: Path) -> bool:
+    required_files = [
+        "config.json",
+        "configuration_bert.py",
+        "bert_layers.py",
+        "bert_padding.py",
+        "tokenizer.json",
+        "tokenizer_config.json",
+    ]
+    if not all((patch_dir / filename).exists() for filename in required_files):
+        return False
+
+    bert_layers = (patch_dir / "bert_layers.py").read_text(encoding="utf-8")
+    return "from .flash_attn_triton import" not in bert_layers and "getattr(self.alibi, 'is_meta', False)" in bert_layers
+
+
+def create_patch_from_project_root() -> Path:
+    """Run the smoke-test patch creator from the project root."""
+    previous_cwd = Path.cwd()
+    try:
+        os.chdir(PROJECT_ROOT)
+        patch_dir = create_local_dnabert2_patch()
+    finally:
+        os.chdir(previous_cwd)
+    return resolve_path(patch_dir)
+
+
+def load_model_from_source(model_source: str | Path):
+    config = AutoConfig.from_pretrained(str(model_source), trust_remote_code=True)
+    config = disable_flash_attention_on_config(config)
+    return load_sequence_classification_model(str(model_source), config)
+
+
+def load_mac_safe_model() -> tuple[object, Path | str]:
+    patch_dir = resolve_path(LOCAL_DNABERT2_PATCH_DIR)
+
+    if patch_dir.exists() and local_patch_is_ready(patch_dir):
+        print(f"Using local patched DNABERT-2: {patch_dir}")
+        clear_local_patch_module_cache()
+        model = load_model_from_source(patch_dir)
+        print("Triton/flash attention disabled for Mac.")
+        print("Model loaded successfully.")
+        return model, patch_dir
+
+    print("No usable local patched DNABERT-2 found. Trying DNABERT-2 with eager attention...")
+    try:
+        model = load_model_from_source(MODEL_NAME)
+        print("Triton/flash attention disabled for Mac.")
+        print("Model loaded successfully.")
+        return model, MODEL_NAME
+    except Exception as eager_error:
+        print("Direct eager DNABERT-2 load failed. Creating local Mac-safe patch.")
+        print(f"Direct load error: {eager_error}")
+
+    patch_dir = create_patch_from_project_root()
+    print(f"Using local patched DNABERT-2: {patch_dir}")
+    model = load_model_from_source(patch_dir)
+    print("Triton/flash attention disabled for Mac.")
+    print("Model loaded successfully.")
+    return model, patch_dir
+
+
+def load_dnabert2_model(device: str) -> tuple[object, Path | str]:
+    """Load DNABERT-2 using the best strategy for the selected device."""
+    if device == "cuda":
+        print("CUDA detected. Trying standard Hugging Face DNABERT-2 first.")
+        try:
+            model = load_model_from_source(MODEL_NAME)
+            print("Model loaded successfully from Hugging Face.")
+            return model, MODEL_NAME
+        except Exception as error:
+            print("Standard Hugging Face DNABERT-2 load failed.")
+            print(f"Direct load error: {error}")
+            print("Falling back to the local no-Triton DNABERT-2 patch.")
+
+    return load_mac_safe_model()
+
+
+def get_module_by_path(model, path: str):
+    current = model
+    for part in path.split("."):
+        if not hasattr(current, part):
+            return None
+        current = getattr(current, part)
+    return current
+
+
+def find_encoder_layers(model) -> tuple[str | None, object | None]:
+    layer_paths = [
+        "bert.encoder.layer",
+        "encoder.layer",
+        "base_model.encoder.layer",
+        "bert.encoder.layers",
+    ]
+
+    for path in layer_paths:
+        layers = get_module_by_path(model, path)
+        if layers is not None and hasattr(layers, "__len__") and hasattr(layers, "__getitem__"):
+            return f"model.{path}", layers
+
+    return None, None
+
+
+def find_embedding_module(model):
+    embedding_paths = [
+        "bert.embeddings",
+        "embeddings",
+        "base_model.embeddings",
+        "base_model.bert.embeddings",
+    ]
+
+    for path in embedding_paths:
+        embeddings = get_module_by_path(model, path)
+        if embeddings is not None and hasattr(embeddings, "parameters"):
+            return f"model.{path}", embeddings
+
+    return None, None
+
+
+def unfreeze_classifier_parameters(model) -> list[str]:
+    trainable_names = []
+    for name, parameter in model.named_parameters():
+        if name.startswith("classifier") or ".classifier" in name:
+            parameter.requires_grad = True
+            trainable_names.append(name)
+    return trainable_names
+
+
+def unfreeze_pooler_parameters(model) -> list[str]:
+    trainable_names = []
+    for name, parameter in model.named_parameters():
+        if name.startswith("pooler") or ".pooler" in name:
+            parameter.requires_grad = True
+            trainable_names.append(name)
+    return trainable_names
+
+
+def print_trainable_parameter_summary(model) -> None:
+    total_params = sum(parameter.numel() for parameter in model.parameters())
+    trainable_params = sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)
+    trainable_percentage = (trainable_params / total_params * 100.0) if total_params else 0.0
+    trainable_names = [name for name, parameter in model.named_parameters() if parameter.requires_grad]
+
+    print(f"Total parameters: {total_params:,}")
+    print(f"Trainable parameters: {trainable_params:,}")
+    print(f"Trainable percentage: {trainable_percentage:.4f}%")
+    print("First 20 trainable parameter names:")
+    for name in trainable_names[:20]:
+        print(f"  {name}")
+    if len(trainable_names) > 20:
+        print(f"  ... {len(trainable_names) - 20:,} more")
+    print()
+
+
+def freeze_encoder_if_requested(
+    model,
+    freeze_encoder: bool,
+    unfreeze_last_n_layers: int,
+    freeze_embeddings: bool,
+    device: str,
+) -> None:
+    if unfreeze_last_n_layers < 0:
+        raise ValueError("--unfreeze_last_n_layers must be 0 or greater.")
+
+    if not freeze_encoder:
+        for parameter in model.parameters():
+            parameter.requires_grad = True
+        print("Encoder frozen: false")
+        print("All model parameters are trainable.")
+        print_trainable_parameter_summary(model)
+        return
+
+    for parameter in model.parameters():
+        parameter.requires_grad = False
+
+    classifier_names = unfreeze_classifier_parameters(model)
+    print("Encoder frozen: true")
+
+    if unfreeze_last_n_layers == 0:
+        print("Partial unfreezing: disabled")
+        print("Trainable modules: classifier head only")
+        if not classifier_names:
+            print("WARNING: no classifier parameters were found to unfreeze.")
+        print_trainable_parameter_summary(model)
+        return
+
+    if device == "mps":
+        print("WARNING: unfreezing encoder layers on MPS may be slower and may use more memory.")
+
+    print(f"Partial unfreezing: last {unfreeze_last_n_layers} encoder layer(s)")
+    pooler_names = unfreeze_pooler_parameters(model)
+    if pooler_names:
+        print("Pooler unfrozen: true")
+    else:
+        print("Pooler unfrozen: false, no pooler parameters found")
+
+    layer_path, layers = find_encoder_layers(model)
+    if layers is None:
+        print("Encoder layer path found: none")
+        print("WARNING: no supported encoder layer path was found. Only classifier/pooler parameters are trainable.")
+    else:
+        layer_count = len(layers)
+        layers_to_unfreeze = min(unfreeze_last_n_layers, layer_count)
+        start_index = layer_count - layers_to_unfreeze
+        print(f"Encoder layer path found: {layer_path}")
+        print(f"Encoder layers found: {layer_count}")
+        print(f"Encoder layer indexes unfrozen: {start_index} to {layer_count - 1}")
+
+        for layer in list(layers)[start_index:]:
+            for parameter in layer.parameters():
+                parameter.requires_grad = True
+
+    if freeze_embeddings:
+        print("Embeddings frozen: true")
+    else:
+        embedding_path, embeddings = find_embedding_module(model)
+        if embeddings is None:
+            print("Embeddings frozen: false requested, but no embedding module was found")
+        else:
+            for parameter in embeddings.parameters():
+                parameter.requires_grad = True
+            print(f"Embeddings frozen: false")
+            print(f"Embedding path found: {embedding_path}")
+
+    if not classifier_names:
+        print("WARNING: no classifier parameters were found to unfreeze.")
+
+    print_trainable_parameter_summary(model)
+
+
+def make_training_arguments(args: argparse.Namespace, output_dir: Path, device: str) -> TrainingArguments:
+    eval_batch_size = args.batch_size if device == "cuda" else 1
+    base_kwargs = {
+        "output_dir": str(output_dir),
+        "num_train_epochs": args.epochs,
+        "learning_rate": args.learning_rate,
+        "per_device_train_batch_size": args.batch_size,
+        "per_device_eval_batch_size": eval_batch_size,
+        "gradient_accumulation_steps": args.grad_accum_steps,
+        "eval_accumulation_steps": args.eval_accumulation_steps,
+        "save_strategy": "epoch",
+        "load_best_model_at_end": args.save_eval_each_epoch,
+        "save_total_limit": 2,
+        "logging_steps": 20,
+        "report_to": "none",
+        "dataloader_num_workers": 0,
+        "dataloader_pin_memory": device == "cuda",
+        "remove_unused_columns": False,
+        "fp16": device == "cuda",
+        "bf16": False,
+    }
+
+    if args.save_eval_each_epoch:
+        base_kwargs["metric_for_best_model"] = "eval_mcc"
+        base_kwargs["greater_is_better"] = True
+        eval_strategy = "epoch"
+    else:
+        eval_strategy = "no"
+
+    try:
+        return TrainingArguments(**base_kwargs, eval_strategy=eval_strategy)
+    except TypeError:
+        return TrainingArguments(**base_kwargs, evaluation_strategy=eval_strategy)
+
+
+def make_trainer(
+    model,
+    tokenizer,
+    training_args,
+    train_dataset: Dataset,
+    val_dataset: Dataset,
+    class_weights: torch.Tensor | None,
+) -> Trainer:
+    trainer_kwargs = {
+        "model": model,
+        "args": training_args,
+        "train_dataset": train_dataset,
+        "eval_dataset": val_dataset,
+        "data_collator": DataCollatorWithPadding(tokenizer=tokenizer),
+        "compute_metrics": compute_metrics,
+    }
+
+    trainer_class = Trainer
+    if class_weights is not None:
+        trainer_class = WeightedLossTrainer
+        trainer_kwargs["class_weights"] = class_weights
+
+    trainer_signature = inspect.signature(Trainer.__init__)
+    if "preprocess_logits_for_metrics" in trainer_signature.parameters:
+        trainer_kwargs["preprocess_logits_for_metrics"] = preprocess_logits_for_metrics
+    if "processing_class" in trainer_signature.parameters:
+        trainer_kwargs["processing_class"] = tokenizer
+    else:
+        trainer_kwargs["tokenizer"] = tokenizer
+
+    return trainer_class(**trainer_kwargs)
+
+
+def save_metrics(output_dir: Path, metrics: dict) -> Path:
+    metrics_path = output_dir / "metrics.json"
+    metrics_path.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
+    return metrics_path
+
+
+def zip_final_model(output_dir: Path, final_model_dir: Path) -> Path:
+    zip_path = output_dir / "final_dnabert2_clinvar_model.zip"
+    if zip_path.exists():
+        zip_path.unlink()
+
+    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as archive:
+        for file_path in final_model_dir.rglob("*"):
+            if file_path.is_file():
+                archive.write(file_path, arcname=file_path.relative_to(final_model_dir.parent))
+
+    return zip_path
+
+
+def main() -> None:
+    args = parse_args()
+    output_dir = resolve_path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    final_model_dir = output_dir / "final_model"
+
+    csv_selection = resolve_csv_paths(args)
+    device = choose_device()
+
+    print("DNABERT-2 ClinVar training")
+    print(f"Selected dataset directory: {csv_selection.dataset_dir}")
+    print(f"Selected dataset name: {csv_selection.dataset_name}")
+    print(f"Using 20k alt-sequence dataset: {csv_selection.is_20k_alt_dataset}")
+    print(f"Using 10k alt-sequence dataset: {csv_selection.is_10k_alt_dataset}")
+    print(f"Using large alt-sequence dataset: {csv_selection.is_large_alt_dataset}")
+    print(f"Selected train CSV path: {csv_selection.train_csv}")
+    print(f"Selected validation CSV path: {csv_selection.val_csv}")
+    print(f"Selected test CSV path: {csv_selection.test_csv}")
+    print(f"Selected sequence column: {args.sequence_column}")
+    print(f"Using alt-sequence dataset: {csv_selection.is_alt_dataset}")
+    print(f"Selected device: {device}")
+    if device == "cuda":
+        print("CUDA training enabled. fp16 will be enabled in TrainingArguments.")
+    if device == "mps":
+        print("MPS training enabled. fp16/bf16 stay disabled for Mac compatibility.")
+    if device == "cpu":
+        print("WARNING: CPU training will be slow. Consider using sample_size > 0.")
+    print(f"Output path: {output_dir}")
+    print()
+
+    train_df = load_and_filter_dataframe(csv_selection.train_csv, "train", args.sequence_column)
+    val_df = load_and_filter_dataframe(csv_selection.val_csv, "validation", args.sequence_column)
+    test_df = load_and_filter_dataframe(csv_selection.test_csv, "test", args.sequence_column)
+
+    if args.sample_size > 0:
+        train_df = stratified_sample(train_df, args.sample_size, "train")
+        eval_sample_size = min(200, args.sample_size)
+        val_df = stratified_sample(val_df, min(eval_sample_size, len(val_df)), "validation")
+        test_df = stratified_sample(test_df, min(eval_sample_size, len(test_df)), "test")
+
+    train_df = apply_center_crop(
+        train_df,
+        args.max_length,
+        args.variant_center_index,
+        "train",
+        args.center_crop,
+    )
+    val_df = apply_center_crop(
+        val_df,
+        args.max_length,
+        args.variant_center_index,
+        "validation",
+        args.center_crop,
+    )
+    test_df = apply_center_crop(
+        test_df,
+        args.max_length,
+        args.variant_center_index,
+        "test",
+        args.center_crop,
+    )
+    class_weights = maybe_compute_class_weights(train_df, args.use_class_weights)
+    eval_val_df = make_eval_subset(val_df, args.eval_subset_size, "validation")
+    eval_test_df = make_eval_subset(test_df, args.eval_subset_size, "test")
+
+    print(f"Final train rows: {len(train_df):,}")
+    print(f"Final validation rows: {len(val_df):,}")
+    print(f"Final test rows: {len(test_df):,}")
+    print(f"Validation rows used for evaluation: {len(eval_val_df):,}")
+    print(f"Test rows used for evaluation: {len(eval_test_df):,}")
+    print(f"Save/evaluate each epoch: {args.save_eval_each_epoch}")
+    print(f"Eval accumulation steps: {args.eval_accumulation_steps}")
+    print("Using memory-safe evaluation mode.")
+    print()
+
+    model, model_path = load_dnabert2_model(device)
+    freeze_encoder_if_requested(
+        model,
+        args.freeze_encoder,
+        args.unfreeze_last_n_layers,
+        args.freeze_embeddings,
+        device,
+    )
+
+    print(f"Loading tokenizer from: {model_path}")
+    tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True)
+
+    train_dataset = tokenize_dataset(tokenizer, dataframe_to_dataset(train_df), args.max_length, "train")
+    val_dataset = tokenize_dataset(tokenizer, dataframe_to_dataset(eval_val_df), args.max_length, "validation")
+    test_dataset = tokenize_dataset(tokenizer, dataframe_to_dataset(eval_test_df), args.max_length, "test")
+
+    training_args = make_training_arguments(args, output_dir, device)
+    trainer = make_trainer(model, tokenizer, training_args, train_dataset, val_dataset, class_weights)
+
+    print("Starting training.")
+    train_output = trainer.train(resume_from_checkpoint=args.resume_from_checkpoint)
+
+    print("Predicting validation split for threshold tuning.")
+    validation_probabilities, validation_labels = predict_in_small_batches(
+        trainer.model,
+        val_dataset,
+        batch_size=1,
+        device=device,
+    )
+    if args.tune_threshold:
+        selected_threshold, best_validation_mcc = tune_threshold_for_mcc(validation_labels, validation_probabilities)
+        print(f"Selected threshold from validation MCC: {selected_threshold:.2f}")
+        print(f"Best validation MCC during threshold tuning: {best_validation_mcc:.4f}")
+    else:
+        selected_threshold = 0.5
+        best_validation_mcc = None
+        print("Threshold tuning disabled. Using threshold: 0.50")
+    print()
+
+    validation_metrics = metrics_at_threshold(validation_labels, validation_probabilities, selected_threshold)
+    print_threshold_metrics("validation", validation_metrics)
+
+    print("Predicting test split with selected threshold.")
+    test_probabilities, test_labels = predict_in_small_batches(
+        trainer.model,
+        test_dataset,
+        batch_size=1,
+        device=device,
+    )
+    test_metrics = metrics_at_threshold(test_labels, test_probabilities, selected_threshold)
+    print_threshold_metrics("test", test_metrics)
+
+    final_model_dir.mkdir(parents=True, exist_ok=True)
+    trainer.save_model(str(final_model_dir))
+    tokenizer.save_pretrained(str(final_model_dir))
+
+    metrics = {
+        "model_name": MODEL_NAME,
+        "model_path": str(model_path),
+        "train_csv": str(csv_selection.train_csv),
+        "validation_csv": str(csv_selection.val_csv),
+        "test_csv": str(csv_selection.test_csv),
+        "dataset_dir": str(csv_selection.dataset_dir),
+        "dataset_name": csv_selection.dataset_name,
+        "sequence_column": args.sequence_column,
+        "is_alt_sequence_dataset": csv_selection.is_alt_dataset,
+        "is_large_alt_sequence_dataset": csv_selection.is_large_alt_dataset,
+        "is_10k_alt_sequence_dataset": csv_selection.is_10k_alt_dataset,
+        "is_20k_alt_sequence_dataset": csv_selection.is_20k_alt_dataset,
+        "device": device,
+        "freeze_encoder": args.freeze_encoder,
+        "unfreeze_last_n_layers": args.unfreeze_last_n_layers,
+        "freeze_embeddings": args.freeze_embeddings,
+        "use_class_weights": args.use_class_weights,
+        "class_weights": class_weights.tolist() if class_weights is not None else None,
+        "center_crop": args.center_crop,
+        "tune_threshold": args.tune_threshold,
+        "save_eval_each_epoch": args.save_eval_each_epoch,
+        "eval_accumulation_steps": args.eval_accumulation_steps,
+        "eval_subset_size": args.eval_subset_size,
+        "selected_threshold": selected_threshold,
+        "best_validation_mcc_for_threshold": best_validation_mcc,
+        "sample_size": args.sample_size,
+        "epochs": args.epochs,
+        "max_length": args.max_length,
+        "variant_center_index": args.variant_center_index,
+        "train_metrics": train_output.metrics,
+        "train_rows": len(train_df),
+        "validation_rows": len(val_df),
+        "test_rows": len(test_df),
+        "validation_eval_rows": len(eval_val_df),
+        "test_eval_rows": len(eval_test_df),
+        "validation_metrics": validation_metrics,
+        "test_metrics": test_metrics,
+    }
+    metrics_path = save_metrics(output_dir, metrics)
+    zip_path = zip_final_model(output_dir, final_model_dir)
+
+    print("Final metrics:")
+    print(json.dumps(metrics, indent=2))
+    print(f"Saved final model to: {final_model_dir}")
+    print(f"Saved metrics to: {metrics_path}")
+    print(f"Created zip file: {zip_path}")
+    print("Local DNABERT-2 training completed successfully.")
+
+
+if __name__ == "__main__":
+    main()

training/train_smoke_test.py

@@ -0,0 +1,462 @@
+#!/usr/bin/env python3
+"""Run a tiny local DNABERT-2 smoke test on Mac.
+
+This script checks that the CSV files, tokenizer, model loading, Trainer setup,
+device selection, and model saving all work before attempting a larger run.
+"""
+
+from __future__ import annotations
+
+import inspect
+import json
+import shutil
+from pathlib import Path
+from urllib.parse import unquote
+
+import numpy as np
+import pandas as pd
+import torch
+from datasets import Dataset
+from huggingface_hub import snapshot_download
+from sklearn.metrics import accuracy_score, f1_score, matthews_corrcoef
+from transformers import (
+    AutoConfig,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    DataCollatorWithPadding,
+    Trainer,
+    TrainingArguments,
+)
+
+
+MODEL_NAME = "zhihan1996/DNABERT-2-117M"
+OUTPUT_DIR = Path("training/outputs/smoke_test")
+FINAL_MODEL_DIR = OUTPUT_DIR / "final_model"
+LOCAL_DNABERT2_PATCH_DIR = Path("training/local_dnabert2_patch")
+
+TRAIN_FILENAME = "train_with_sequences.csv"
+VAL_FILENAME = "val_with_sequences.csv"
+
+MAX_TRAIN_ROWS = 64
+MAX_VAL_ROWS = 32
+MIN_SEQUENCE_LENGTH = 50
+MAX_LENGTH = 512
+
+DROP_CLNSIG_VALUES = (
+    "Conflicting_classifications_of_pathogenicity",
+    "Conflicting_interpretations",
+    "Uncertain_significance",
+    "not_provided",
+    "risk_factor",
+    "association",
+    "drug_response",
+    "protective",
+)
+
+FLASH_DISABLE_FLAGS = {
+    "use_flash_attn": False,
+    "use_flash_attention": False,
+    "flash_attn": False,
+    "attn_implementation": "eager",
+    "attention_implementation": "eager",
+}
+
+
+def project_root() -> Path:
+    return Path(__file__).resolve().parents[1]
+
+
+def find_dataset_dir(root: Path) -> Path:
+    """Prefer data/processed, then fall back to training/csv_files."""
+    candidates = [
+        root / "data" / "processed",
+        root / "training" / "csv_files",
+    ]
+
+    for directory in candidates:
+        train_path = directory / TRAIN_FILENAME
+        val_path = directory / VAL_FILENAME
+        if train_path.exists() and val_path.exists():
+            return directory
+
+    searched = "\n".join(str(path) for path in candidates)
+    raise FileNotFoundError(
+        "Could not find train/validation CSV files. Searched:\n"
+        f"{searched}\n"
+        f"Required files: {TRAIN_FILENAME}, {VAL_FILENAME}"
+    )
+
+
+def choose_device() -> str:
+    if torch.cuda.is_available():
+        return "cuda"
+
+    mps_backend = getattr(torch.backends, "mps", None)
+    if mps_backend is not None and mps_backend.is_available():
+        return "mps"
+
+    return "cpu"
+
+
+def disable_flash_attention_on_config(config):
+    """Set common config flags that disable flash/triton attention."""
+    for field_name, value in FLASH_DISABLE_FLAGS.items():
+        setattr(config, field_name, value)
+    config.pad_token_id = getattr(config, "pad_token_id", None) or 3
+    config.num_labels = 2
+    config.id2label = {0: "benign_or_likely_benign", 1: "pathogenic"}
+    config.label2id = {"benign_or_likely_benign": 0, "pathogenic": 1}
+    return config
+
+
+def normalize_text(value: object) -> str:
+    return unquote(str(value)).strip().lower()
+
+
+def should_drop_clnsig(value: object) -> bool:
+    text = normalize_text(value)
+    return any(drop_value.lower() in text for drop_value in DROP_CLNSIG_VALUES)
+
+
+def clean_sequence(value: object) -> str:
+    return str(value).strip().upper()
+
+
+def load_and_filter_csv(path: Path, split_name: str, max_rows: int) -> Dataset:
+    df = pd.read_csv(path)
+    before_rows = len(df)
+
+    required_columns = {"sequence", "label"}
+    missing_columns = sorted(required_columns - set(df.columns))
+    if missing_columns:
+        raise ValueError(f"{path} is missing required columns: {missing_columns}")
+
+    df = df.copy()
+
+    if "CLNSIG" in df.columns:
+        df = df.loc[~df["CLNSIG"].fillna("").apply(should_drop_clnsig)].copy()
+
+    df["label"] = pd.to_numeric(df["label"], errors="coerce")
+    df = df.loc[df["label"].isin([0, 1])].copy()
+    df["label"] = df["label"].astype(int)
+
+    df["sequence"] = df["sequence"].fillna("").apply(clean_sequence)
+    df = df.loc[df["sequence"] != ""].copy()
+    df = df.loc[df["sequence"].str.len() >= MIN_SEQUENCE_LENGTH].copy()
+
+    after_filter_rows = len(df)
+    df = df.head(max_rows).copy()
+
+    print(f"{split_name} CSV: {path}")
+    print(f"{split_name} rows before filtering: {before_rows:,}")
+    print(f"{split_name} rows after filtering: {after_filter_rows:,}")
+    print(f"{split_name} rows used for smoke test: {len(df):,}")
+    print(f"{split_name} label distribution:")
+    print(df["label"].value_counts().sort_index().to_string())
+    print()
+
+    if df.empty:
+        raise ValueError(f"No usable rows remain for {split_name}.")
+
+    dataset_df = df[["sequence", "label"]].rename(columns={"label": "labels"})
+    return Dataset.from_pandas(dataset_df, preserve_index=False)
+
+
+def tokenize_datasets(tokenizer, train_dataset: Dataset, val_dataset: Dataset) -> tuple[Dataset, Dataset]:
+    def tokenize_batch(batch):
+        return tokenizer(
+            batch["sequence"],
+            max_length=MAX_LENGTH,
+            padding="max_length",
+            truncation=True,
+        )
+
+    train_dataset = train_dataset.map(tokenize_batch, batched=True, remove_columns=["sequence"])
+    val_dataset = val_dataset.map(tokenize_batch, batched=True, remove_columns=["sequence"])
+    return train_dataset, val_dataset
+
+
+def compute_metrics(eval_pred):
+    if hasattr(eval_pred, "predictions"):
+        logits = eval_pred.predictions
+        labels = eval_pred.label_ids
+    else:
+        logits, labels = eval_pred
+
+    if isinstance(logits, (tuple, list)):
+        for item in logits:
+            candidate = np.asarray(item)
+            if candidate.ndim >= 2 and candidate.shape[-1] == 2:
+                logits = candidate
+                break
+        else:
+            logits = np.asarray(logits[0])
+    else:
+        logits = np.asarray(logits)
+
+    labels = np.asarray(labels)
+    predictions = np.argmax(logits, axis=-1)
+    return {
+        "accuracy": accuracy_score(labels, predictions),
+        "f1": f1_score(labels, predictions, zero_division=0),
+        "matthews_corrcoef": matthews_corrcoef(labels, predictions),
+    }
+
+
+def make_training_arguments() -> TrainingArguments:
+    base_kwargs = {
+        "output_dir": str(OUTPUT_DIR),
+        "num_train_epochs": 1,
+        "learning_rate": 2e-5,
+        "per_device_train_batch_size": 1,
+        "per_device_eval_batch_size": 1,
+        "gradient_accumulation_steps": 4,
+        "save_strategy": "epoch",
+        "logging_steps": 5,
+        "save_total_limit": 1,
+        "report_to": "none",
+        "dataloader_num_workers": 0,
+        "dataloader_pin_memory": False,
+        "remove_unused_columns": False,
+        "fp16": False,
+        "bf16": False,
+    }
+
+    try:
+        return TrainingArguments(**base_kwargs, eval_strategy="epoch")
+    except TypeError:
+        return TrainingArguments(**base_kwargs, evaluation_strategy="epoch")
+
+
+def make_trainer(model, tokenizer, training_args, train_dataset: Dataset, val_dataset: Dataset) -> Trainer:
+    trainer_kwargs = {
+        "model": model,
+        "args": training_args,
+        "train_dataset": train_dataset,
+        "eval_dataset": val_dataset,
+        "data_collator": DataCollatorWithPadding(tokenizer=tokenizer),
+        "compute_metrics": compute_metrics,
+    }
+
+    trainer_signature = inspect.signature(Trainer.__init__)
+    if "processing_class" in trainer_signature.parameters:
+        trainer_kwargs["processing_class"] = tokenizer
+    else:
+        trainer_kwargs["tokenizer"] = tokenizer
+
+    return Trainer(**trainer_kwargs)
+
+
+def load_dnabert2_with_eager_attention():
+    """Approach A: ask DNABERT-2 remote code to use eager attention."""
+    print("Trying DNABERT-2 with eager attention...")
+    config = AutoConfig.from_pretrained(MODEL_NAME, trust_remote_code=True)
+    config = disable_flash_attention_on_config(config)
+    print("Triton/flash attention disabled for Mac.")
+
+    model = load_sequence_classification_model(MODEL_NAME, config)
+    print("Model loaded successfully.")
+    return model
+
+
+def load_sequence_classification_model(model_source: str, config):
+    """Load a classifier, retrying without num_labels for older custom classes."""
+    common_kwargs = {
+        "config": config,
+        "trust_remote_code": True,
+        "low_cpu_mem_usage": False,
+    }
+    try:
+        return AutoModelForSequenceClassification.from_pretrained(
+            model_source,
+            num_labels=2,
+            **common_kwargs,
+        )
+    except TypeError as error:
+        if "num_labels" not in str(error):
+            raise
+        return AutoModelForSequenceClassification.from_pretrained(
+            model_source,
+            **common_kwargs,
+        )
+
+
+def patch_bert_layers_for_mac(source_path: Path, destination_path: Path) -> None:
+    """Remove the flash_attn_triton import so Transformers does not require Triton."""
+    source_text = source_path.read_text(encoding="utf-8")
+    flash_import_block = """try:
+    from .flash_attn_triton import flash_attn_qkvpacked_func
+except ImportError as e:
+    flash_attn_qkvpacked_func = None
+"""
+    patched_text = source_text.replace(
+        flash_import_block,
+        "# Mac-safe local patch: always use the PyTorch attention fallback.\n"
+        "flash_attn_qkvpacked_func = None\n",
+    )
+
+    if "from .flash_attn_triton import" in patched_text:
+        raise RuntimeError("Could not patch flash_attn_triton import from bert_layers.py.")
+
+    # Newer Transformers versions may instantiate custom models under a meta
+    # device context. The original ALiBi code can then multiply a CPU tensor by
+    # a meta tensor. Keep both tensors on the same device.
+    patched_text = patched_text.replace(
+        "        slopes = torch.Tensor(_get_alibi_head_slopes(n_heads)).to(device)\n"
+        "        alibi = slopes.unsqueeze(1).unsqueeze(1) * -relative_position\n",
+        "        slope_device = device if device is not None else relative_position.device\n"
+        "        slopes = torch.tensor(_get_alibi_head_slopes(n_heads), device=slope_device)\n"
+        "        alibi = slopes.unsqueeze(1).unsqueeze(1) * -relative_position\n",
+    )
+    patched_text = patched_text.replace(
+        "        elif self.alibi.device != hidden_states.device:\n"
+        "            # Device catch-up\n"
+        "            self.alibi = self.alibi.to(hidden_states.device)\n",
+        "        elif getattr(self.alibi, 'is_meta', False) or self.alibi.device != hidden_states.device:\n"
+        "            # Device catch-up. Under newer Transformers, the buffer may be created\n"
+        "            # on the meta device during low-level loading, so rebuild it instead of\n"
+        "            # copying it.\n"
+        "            self.rebuild_alibi_tensor(size=self._current_alibi_size, device=hidden_states.device)\n",
+    )
+
+    destination_path.write_text(patched_text, encoding="utf-8")
+
+
+def clear_local_patch_module_cache() -> None:
+    """Clear Transformers' cached dynamic module for the local patch."""
+    cache_dir = Path.home() / ".cache" / "huggingface" / "modules" / "transformers_modules" / "local_dnabert2_patch"
+    if cache_dir.exists():
+        shutil.rmtree(cache_dir)
+
+
+def create_local_dnabert2_patch() -> Path:
+    """Approach B: create a local DNABERT-2 copy without the Triton import."""
+    print("Creating local Mac-safe DNABERT-2 patch...")
+    snapshot_path = Path(
+        snapshot_download(
+            MODEL_NAME,
+            allow_patterns=[
+                "config.json",
+                "configuration_bert.py",
+                "bert_layers.py",
+                "bert_padding.py",
+                "tokenizer.json",
+                "tokenizer_config.json",
+                "pytorch_model.bin",
+                "model.safetensors",
+                "generation_config.json",
+            ],
+        )
+    )
+
+    if LOCAL_DNABERT2_PATCH_DIR.exists():
+        shutil.rmtree(LOCAL_DNABERT2_PATCH_DIR)
+    LOCAL_DNABERT2_PATCH_DIR.mkdir(parents=True, exist_ok=True)
+
+    files_to_copy = [
+        "config.json",
+        "configuration_bert.py",
+        "bert_padding.py",
+        "tokenizer.json",
+        "tokenizer_config.json",
+        "generation_config.json",
+        "pytorch_model.bin",
+        "model.safetensors",
+    ]
+    for filename in files_to_copy:
+        source = snapshot_path / filename
+        if source.exists():
+            shutil.copy2(source, LOCAL_DNABERT2_PATCH_DIR / filename)
+
+    patch_bert_layers_for_mac(snapshot_path / "bert_layers.py", LOCAL_DNABERT2_PATCH_DIR / "bert_layers.py")
+    clear_local_patch_module_cache()
+
+    config_path = LOCAL_DNABERT2_PATCH_DIR / "config.json"
+    config_json = json.loads(config_path.read_text(encoding="utf-8"))
+    config_json.update(FLASH_DISABLE_FLAGS)
+    config_json["pad_token_id"] = config_json.get("pad_token_id") or 3
+    config_json["num_labels"] = 2
+    config_json["id2label"] = {"0": "benign_or_likely_benign", "1": "pathogenic"}
+    config_json["label2id"] = {"benign_or_likely_benign": 0, "pathogenic": 1}
+    config_path.write_text(json.dumps(config_json, indent=2), encoding="utf-8")
+
+    print(f"Local Mac-safe DNABERT-2 patch ready: {LOCAL_DNABERT2_PATCH_DIR}")
+    return LOCAL_DNABERT2_PATCH_DIR
+
+
+def load_dnabert2_from_local_patch():
+    """Load the local patched copy that avoids flash_attn_triton."""
+    patch_dir = create_local_dnabert2_patch()
+    config = AutoConfig.from_pretrained(str(patch_dir), trust_remote_code=True)
+    config = disable_flash_attention_on_config(config)
+
+    model = load_sequence_classification_model(str(patch_dir), config)
+    print("Triton/flash attention disabled for Mac.")
+    print("Model loaded successfully.")
+    return model
+
+
+def load_mac_safe_dnabert2_model():
+    """Load DNABERT-2 without requiring Triton on Mac."""
+    try:
+        return load_dnabert2_with_eager_attention()
+    except Exception as eager_error:
+        print("Approach A failed. Falling back to local Mac-safe DNABERT-2 patch.")
+        print(f"Approach A error: {eager_error}")
+
+    try:
+        return load_dnabert2_from_local_patch()
+    except Exception as patch_error:
+        raise RuntimeError(
+            "DNABERT-2 could not be loaded without Triton. "
+            "The Mac-safe eager and local patch strategies both failed."
+        ) from patch_error
+
+
+def main() -> None:
+    root = project_root()
+    dataset_dir = find_dataset_dir(root)
+    train_csv = dataset_dir / TRAIN_FILENAME
+    val_csv = dataset_dir / VAL_FILENAME
+    device = choose_device()
+
+    print("DNABERT-2 local smoke test")
+    print(f"Selected CSV directory: {dataset_dir}")
+    print(f"Selected train CSV path: {train_csv}")
+    print(f"Selected validation CSV path: {val_csv}")
+    print(f"Selected device: {device}")
+    if device == "cpu":
+        print("WARNING: CPU training will be slow. This script uses a tiny subset only.")
+    print()
+
+    train_dataset = load_and_filter_csv(train_csv, "train", MAX_TRAIN_ROWS)
+    val_dataset = load_and_filter_csv(val_csv, "validation", MAX_VAL_ROWS)
+
+    print(f"Loading tokenizer: {MODEL_NAME}")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+
+    print(f"Loading model: {MODEL_NAME}")
+    model = load_mac_safe_dnabert2_model()
+
+    train_dataset, val_dataset = tokenize_datasets(tokenizer, train_dataset, val_dataset)
+
+    training_args = make_training_arguments()
+    trainer = make_trainer(model, tokenizer, training_args, train_dataset, val_dataset)
+
+    print("Starting 1-epoch smoke-test training.")
+    trainer.train()
+
+    print("Running validation evaluation.")
+    metrics = trainer.evaluate()
+    print(metrics)
+
+    FINAL_MODEL_DIR.mkdir(parents=True, exist_ok=True)
+    trainer.save_model(str(FINAL_MODEL_DIR))
+    tokenizer.save_pretrained(str(FINAL_MODEL_DIR))
+
+    print(f"Saved final model to: {FINAL_MODEL_DIR}")
+    print("Smoke test completed successfully.")
+
+
+if __name__ == "__main__":
+    main()