done with preprocessing

training/README.md

@@ -6,8 +6,8 @@ Training should be run in Google Colab or another dedicated GPU notebook environ
 
 ## Contents
 
-- `01_prepare_clinvar_dataset.ipynb`: Colab notebook for downloading ClinVar GRCh38 VCF data and preparing binary SNV/small-indel CSV splits.
-- `colab_dnabert2_clinvar_finetune.ipynb`: notebook entry point for Colab.
+- `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.
 - `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.
@@ -20,17 +20,18 @@ Training should be run in Google Colab or another dedicated GPU notebook environ
 Use GRCh38 consistently:
 
 - ClinVar GRCh38 VCF: `https://ftp.ncbi.nlm.nih.gov/pub/clinvar/vcf_GRCh38/clinvar.vcf.gz`
-- GRCh38 reference FASTA: provide a local or Google Drive path in Colab.
+- Sequence extraction uses the UCSC hg38 API by default.
+- Optional GRCh38 reference FASTA: provide a local or Google Drive path in Colab if using local FASTA sequence extraction.
 
-The preparation script validates reference alleles against the FASTA and skips records that do not match.
+The fine-tuning notebook expects the sequence CSV files from the preparation notebook.
 
 ## Colab Flow
 
-1. Open `colab_dnabert2_clinvar_finetune.ipynb` in Google Colab.
-2. Install `requirements-colab.txt`.
-3. Download ClinVar GRCh38 VCF.
-4. Mount Google Drive or upload a GRCh38 FASTA.
-5. Run dataset preparation.
+1. Open `01_prepare_clinvar_dataset.ipynb` in Google Colab.
+2. Run dataset preparation and sequence extraction.
+3. Save or download `train_with_sequences.csv`, `val_with_sequences.csv`, and `test_with_sequences.csv`.
+4. Upload those CSV files into `training/csv_files/` or `data/processed/`.
+5. Open `colab_dnabert2_clinvar_finetune.ipynb`.
 6. Run DNABERT-2 fine-tuning.
 7. Export the saved model directory.
 

training/colab_dnabert2_clinvar_finetune.ipynb

@@ -4,11 +4,11 @@
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "# Variant Risk Explainer: ClinVar + DNABERT-2 Fine-tuning\n",
+        "# Variant Risk Explainer: DNABERT-2 Fine-tuning\n",
         "\n",
-        "Research and education only. This notebook does not produce a clinically validated diagnostic model.\n",
+        "This notebook trains from the sequence CSV files produced by `01_prepare_clinvar_dataset.ipynb`.\n",
         "\n",
-        "Use GRCh38 for both ClinVar coordinates and reference FASTA sequence extraction."
+        "Research and education only. This does not produce a clinically validated diagnostic model."
       ]
     },
     {
@@ -17,7 +17,7 @@
       "source": [
         "## 1. Runtime\n",
         "\n",
-        "In Colab, choose `Runtime -> Change runtime type -> GPU` before running the training cells."
+        "In Colab, choose `Runtime -> Change runtime type -> GPU` before running training."
       ]
     },
     {
@@ -34,9 +34,9 @@
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## 2. Install dependencies\n",
+        "## 2. Find Repository Root\n",
         "\n",
-        "Upload or clone this repository into `/content/variant-risk-explainer`, then install the Colab requirements."
+        "Clone or upload the full repository first. The notebook will search `/content` for the training script, so the GitHub repo folder name can be different."
       ]
     },
     {
@@ -45,15 +45,43 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "%cd /content/variant-risk-explainer\n",
-        "!pip install -q -r training/requirements-colab.txt"
+        "from pathlib import Path\n",
+        "import sys\n",
+        "\n",
+        "\n",
+        "def find_project_root() -> Path:\n",
+        "    candidates = [Path.cwd(), Path('/content/variant-risk-explainer')]\n",
+        "    for candidate in candidates:\n",
+        "        script = candidate / 'training' / 'scripts' / 'train_dnabert2_classifier.py'\n",
+        "        if script.exists():\n",
+        "            return candidate\n",
+        "\n",
+        "    content_root = Path('/content')\n",
+        "    if content_root.exists():\n",
+        "        matches = list(content_root.glob('*/training/scripts/train_dnabert2_classifier.py'))\n",
+        "        if matches:\n",
+        "            return matches[0].parents[2]\n",
+        "\n",
+        "    raise FileNotFoundError(\n",
+        "        'Could not find training/scripts/train_dnabert2_classifier.py. Clone or upload the full repository first.'\n",
+        "    )\n",
+        "\n",
+        "\n",
+        "PROJECT_ROOT = find_project_root()\n",
+        "if str(PROJECT_ROOT) not in sys.path:\n",
+        "    sys.path.insert(0, str(PROJECT_ROOT))\n",
+        "\n",
+        "%cd {PROJECT_ROOT}\n",
+        "print(f'Project root: {PROJECT_ROOT}')"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## 3. Download ClinVar GRCh38 VCF"
+        "## 3. Install Dependencies\n",
+        "\n",
+        "This installs the Colab training dependencies. If Colab asks you to restart the runtime after installation, restart and run the notebook again from the top."
       ]
     },
     {
@@ -62,18 +90,26 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "!mkdir -p training/data\n",
-        "!wget -q -O training/data/clinvar_grch38.vcf.gz https://ftp.ncbi.nlm.nih.gov/pub/clinvar/vcf_GRCh38/clinvar.vcf.gz\n",
-        "!ls -lh training/data/clinvar_grch38.vcf.gz"
+        "!pip install -q -r training/requirements-colab.txt"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## 4. Provide GRCh38 FASTA\n",
+        "## 4. Locate Sequence CSV Files\n",
         "\n",
-        "Mount Google Drive or upload a GRCh38 FASTA. The FASTA must match GRCh38 and should be indexed. If the `.fai` index is missing, the next cell creates it."
+        "Upload your CSVs into one of these locations before running this cell:\n",
+        "\n",
+        "- `training/csv_files/`\n",
+        "- `data/processed/`\n",
+        "- `training/data/processed/`\n",
+        "\n",
+        "Required file names:\n",
+        "\n",
+        "- `train_with_sequences.csv`\n",
+        "- `val_with_sequences.csv`\n",
+        "- `test_with_sequences.csv`"
       ]
     },
     {
@@ -82,11 +118,40 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from google.colab import drive\n",
-        "drive.mount('/content/drive')\n",
+        "import pandas as pd\n",
+        "\n",
+        "CSV_SEARCH_DIRS = [\n",
+        "    PROJECT_ROOT / 'training' / 'csv_files',\n",
+        "    PROJECT_ROOT / 'data' / 'processed',\n",
+        "    PROJECT_ROOT / 'training' / 'data' / 'processed',\n",
+        "]\n",
+        "\n",
+        "\n",
+        "def find_csv(filename: str) -> Path:\n",
+        "    for directory in CSV_SEARCH_DIRS:\n",
+        "        candidate = directory / filename\n",
+        "        if candidate.exists():\n",
+        "            return candidate\n",
+        "    searched = '\\n'.join(str(directory / filename) for directory in CSV_SEARCH_DIRS)\n",
+        "    raise FileNotFoundError(f'Could not find {filename}. Searched:\\n{searched}')\n",
+        "\n",
+        "\n",
+        "TRAIN_CSV = find_csv('train_with_sequences.csv')\n",
+        "VAL_CSV = find_csv('val_with_sequences.csv')\n",
+        "TEST_CSV = find_csv('test_with_sequences.csv')\n",
+        "\n",
+        "print(f'Train CSV: {TRAIN_CSV}')\n",
+        "print(f'Val CSV:   {VAL_CSV}')\n",
+        "print(f'Test CSV:  {TEST_CSV}')"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 5. Check CSV Quality\n",
         "\n",
-        "REFERENCE_FASTA = '/content/drive/MyDrive/genomics/GRCh38.primary_assembly.genome.fa.gz'\n",
-        "print(REFERENCE_FASTA)"
+        "This checks required columns, sequence lengths, bad sequence characters, and labels. The training script will also defensively drop uncertain, conflicting, risk, association, drug-response, protective, and not-provided CLNSIG rows."
       ]
     },
     {
@@ -95,18 +160,51 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "import os\n",
-        "if not os.path.exists(REFERENCE_FASTA + '.fai'):\n",
-        "    import pysam\n",
-        "    pysam.faidx(REFERENCE_FASTA)\n",
-        "print('FASTA index ready')"
+        "from training.utils.label_utils import assign_binary_label\n",
+        "\n",
+        "REQUIRED_COLUMNS = {'sequence', 'label'}\n",
+        "\n",
+        "\n",
+        "def inspect_split(name: str, path: Path) -> pd.DataFrame:\n",
+        "    df = pd.read_csv(path)\n",
+        "    missing = REQUIRED_COLUMNS - set(df.columns)\n",
+        "    if missing:\n",
+        "        raise ValueError(f'{name} is missing columns: {sorted(missing)}')\n",
+        "\n",
+        "    bad_sequence_chars = int((~df['sequence'].astype(str).str.upper().str.match(r'^[ACGTN]+$')).sum())\n",
+        "    cleanable_rows = None\n",
+        "    if 'CLNSIG' in df.columns:\n",
+        "        clean_labels = df['CLNSIG'].apply(assign_binary_label)\n",
+        "        cleanable_rows = int(clean_labels.notna().sum())\n",
+        "\n",
+        "    print(f'\\n{name}')\n",
+        "    print(f'rows: {len(df):,}')\n",
+        "    print(f'label counts: {df[\"label\"].value_counts(dropna=False).sort_index().to_dict()}')\n",
+        "    print(\n",
+        "        'sequence length min/mean/max:',\n",
+        "        int(df['sequence'].astype(str).str.len().min()),\n",
+        "        round(float(df['sequence'].astype(str).str.len().mean()), 2),\n",
+        "        int(df['sequence'].astype(str).str.len().max()),\n",
+        "    )\n",
+        "    print(f'bad sequence rows: {bad_sequence_chars:,}')\n",
+        "    if cleanable_rows is not None:\n",
+        "        print(f'rows after defensive CLNSIG cleaning: {cleanable_rows:,}')\n",
+        "    display(df.head(3))\n",
+        "    return df\n",
+        "\n",
+        "\n",
+        "train_preview = inspect_split('train', TRAIN_CSV)\n",
+        "val_preview = inspect_split('val', VAL_CSV)\n",
+        "test_preview = inspect_split('test', TEST_CSV)"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "Optional: store the FASTA path in an environment variable for scripts or later notebook cells."
+        "## 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."
       ]
     },
     {
@@ -115,17 +213,26 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "import os\n",
-        "os.environ['REFERENCE_FASTA_PATH'] = REFERENCE_FASTA"
+        "MODEL_NAME = 'zhihan1996/DNABERT-2-117M'\n",
+        "OUTPUT_DIR = PROJECT_ROOT / 'training' / 'output' / 'dnabert2-clinvar-grch38'\n",
+        "\n",
+        "MAX_LENGTH = 512\n",
+        "EPOCHS = 1\n",
+        "BATCH_SIZE = 4\n",
+        "LEARNING_RATE = 2e-5\n",
+        "MIN_SEQUENCE_LENGTH = 200\n",
+        "\n",
+        "print(f'Model: {MODEL_NAME}')\n",
+        "print(f'Output dir: {OUTPUT_DIR}')"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## 5. Prepare ClinVar sequence examples\n",
+        "## 7. Fine-tune DNABERT-2\n",
         "\n",
-        "For a quick smoke test, keep `--max-records 2000`. For a full run, remove it."
+        "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`."
       ]
     },
     {
@@ -134,21 +241,24 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "!python training/scripts/prepare_clinvar_dataset.py \\\n",
-        "  --clinvar-vcf training/data/clinvar_grch38.vcf.gz \\\n",
-        "  --reference-fasta \"$REFERENCE_FASTA\" \\\n",
-        "  --output-jsonl training/data/clinvar_grch38_dnabert2.jsonl \\\n",
-        "  --window-size 251 \\\n",
-        "  --max-records 2000"
+        "!python training/scripts/train_dnabert2_classifier.py \\\n",
+        "  --train-csv \"{TRAIN_CSV}\" \\\n",
+        "  --val-csv \"{VAL_CSV}\" \\\n",
+        "  --test-csv \"{TEST_CSV}\" \\\n",
+        "  --output-dir \"{OUTPUT_DIR}\" \\\n",
+        "  --model-name \"{MODEL_NAME}\" \\\n",
+        "  --max-length {MAX_LENGTH} \\\n",
+        "  --min-sequence-length {MIN_SEQUENCE_LENGTH} \\\n",
+        "  --epochs {EPOCHS} \\\n",
+        "  --batch-size {BATCH_SIZE} \\\n",
+        "  --learning-rate {LEARNING_RATE}"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## 6. Fine-tune DNABERT-2\n",
-        "\n",
-        "Use a small smoke-test run first. Increase epochs, remove `--max-records` during preparation, and tune batch size based on GPU memory for real experiments."
+        "## 8. Inspect Saved Model Metadata"
       ]
     },
     {
@@ -157,21 +267,23 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "!python training/scripts/train_dnabert2_classifier.py \\\n",
-        "  --dataset-jsonl training/data/clinvar_grch38_dnabert2.jsonl \\\n",
-        "  --output-dir training/output/dnabert2-clinvar-grch38 \\\n",
-        "  --epochs 1 \\\n",
-        "  --batch-size 4 \\\n",
-        "  --max-length 256"
+        "import json\n",
+        "\n",
+        "metadata_path = OUTPUT_DIR / 'final_model' / 'variant_risk_metadata.json'\n",
+        "if metadata_path.exists():\n",
+        "    metadata = json.loads(metadata_path.read_text())\n",
+        "    print(json.dumps(metadata, indent=2))\n",
+        "else:\n",
+        "    print(f'Metadata not found yet: {metadata_path}')"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## 7. Export model artifact\n",
+        "## 9. Export Model Artifact\n",
         "\n",
-        "The backend expects a Hugging Face model directory containing tokenizer and model files. Copy or zip `training/output/dnabert2-clinvar-grch38/final_model`."
+        "Zip the final model directory so you can download it from Colab and use it with the FastAPI backend later."
       ]
     },
     {
@@ -180,8 +292,15 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "!zip -r dnabert2-clinvar-grch38-final-model.zip training/output/dnabert2-clinvar-grch38/final_model\n",
-        "!ls -lh dnabert2-clinvar-grch38-final-model.zip"
+        "ZIP_PATH = PROJECT_ROOT / 'dnabert2-clinvar-grch38-final-model.zip'\n",
+        "!cd \"{PROJECT_ROOT}\" && zip -qr \"{ZIP_PATH}\" training/output/dnabert2-clinvar-grch38/final_model\n",
+        "print(f'Wrote {ZIP_PATH}')\n",
+        "\n",
+        "try:\n",
+        "    from google.colab import files\n",
+        "    files.download(str(ZIP_PATH))\n",
+        "except Exception:\n",
+        "    print('Download helper is only available inside Google Colab.')"
       ]
     }
   ],
@@ -195,7 +314,8 @@
       "name": "python3"
     },
     "language_info": {
-      "name": "python"
+      "name": "python",
+      "version": "3.x"
     }
   },
   "nbformat": 4,

training/csv_files/README.md

@@ -0,0 +1,11 @@
+# Sequence CSV Files
+
+Place the sequence-enriched ClinVar split files here for Colab training:
+
+- `train_with_sequences.csv`
+- `val_with_sequences.csv`
+- `test_with_sequences.csv`
+
+These files are generated by `training/01_prepare_clinvar_dataset.ipynb`.
+
+Large generated CSVs should usually be stored outside Git unless you intentionally want to version the sample dataset.

training/requirements-colab.txt

@@ -1,6 +1,7 @@
 accelerate>=0.30.0
 biopython>=1.83
 datasets>=2.19.0
+einops>=0.8.0
 evaluate>=0.4.2
 numpy>=1.26.0
 pandas>=2.2.0

training/scripts/train_dnabert2_classifier.py

@@ -7,10 +7,17 @@ Run this script in Google Colab or another GPU notebook environment.
 from __future__ import annotations
 
 import argparse
+import inspect
 import json
+import sys
 from pathlib import Path
 
+PROJECT_ROOT = Path(__file__).resolve().parents[2]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
 import numpy as np
+import pandas as pd
 from datasets import Dataset
 from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
 from transformers import (
@@ -21,26 +28,36 @@ from transformers import (
     TrainingArguments,
 )
 
+from training.utils.label_utils import assign_binary_label
+
 
 ID_TO_LABEL = {
-    0: "likely_benign",
-    1: "likely_pathogenic",
+    0: "benign_or_likely_benign",
+    1: "pathogenic",
 }
 LABEL_TO_ID = {label: idx for idx, label in ID_TO_LABEL.items()}
+VALID_BASES = set("ACGTN")
 
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--dataset-jsonl", required=True, help="Prepared JSONL from prepare_clinvar_dataset.py.")
+    parser.add_argument("--dataset-jsonl", help="Legacy prepared JSONL from prepare_clinvar_dataset.py.")
+    parser.add_argument("--train-csv", help="CSV containing sequence and label columns for training.")
+    parser.add_argument("--val-csv", help="CSV containing sequence and label columns for validation.")
+    parser.add_argument("--test-csv", help="CSV containing sequence and label columns for final evaluation.")
     parser.add_argument("--output-dir", required=True, help="Directory for saved model artifacts.")
     parser.add_argument("--model-name", default="zhihan1996/DNABERT-2-117M", help="Hugging Face base model.")
-    parser.add_argument("--max-length", type=int, default=256)
+    parser.add_argument("--sequence-column", default="sequence")
+    parser.add_argument("--label-column", default="label")
+    parser.add_argument("--min-sequence-length", type=int, default=200)
+    parser.add_argument("--max-length", type=int, default=512)
     parser.add_argument("--test-size", type=float, default=0.1)
     parser.add_argument("--eval-size", type=float, default=0.1)
     parser.add_argument("--epochs", type=float, default=2.0)
     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("--no-clean-clnsig", action="store_true", help="Do not re-filter CLNSIG labels in CSV inputs.")
     return parser.parse_args()
 
 
@@ -64,6 +81,95 @@ def load_examples(path: str) -> Dataset:
     return Dataset.from_list(records)
 
 
+def clean_sequence(value: object) -> str:
+    sequence = "".join(base for base in str(value).upper() if base in VALID_BASES)
+    return sequence
+
+
+def load_csv_split(
+    path: str,
+    sequence_column: str,
+    label_column: str,
+    min_sequence_length: int,
+    clean_clnsig: bool,
+) -> Dataset:
+    df = pd.read_csv(path)
+    missing_columns = [column for column in [sequence_column, label_column] if column not in df.columns]
+    if missing_columns:
+        raise ValueError(f"{path} is missing required columns: {missing_columns}")
+
+    start_rows = len(df)
+    df = df.copy()
+    df[sequence_column] = df[sequence_column].map(clean_sequence)
+    df = df[df[sequence_column].str.len() >= min_sequence_length].copy()
+
+    if clean_clnsig and "CLNSIG" in df.columns:
+        df["_clean_label"] = df["CLNSIG"].apply(assign_binary_label)
+        df = df[df["_clean_label"].notna()].copy()
+        df[label_column] = df["_clean_label"].astype(int)
+
+    df[label_column] = pd.to_numeric(df[label_column], errors="coerce")
+    df = df[df[label_column].isin([0, 1])].copy()
+    df[label_column] = df[label_column].astype(int)
+    if df.empty:
+        raise ValueError(f"No usable rows remain after cleaning {path}.")
+
+    records = []
+    for _, row in df.iterrows():
+        records.append(
+            {
+                "sequence": row[sequence_column],
+                "labels": int(row[label_column]),
+                "variant_id": row.get("variant_id"),
+                "chromosome": row.get("CHROM"),
+                "position": row.get("POS"),
+                "gene": row.get("gene_symbol"),
+                "clnsig": row.get("CLNSIG"),
+            }
+        )
+
+    print(
+        f"{path}: kept {len(records):,}/{start_rows:,} rows "
+        f"with label counts {df[label_column].value_counts().sort_index().to_dict()}"
+    )
+    return Dataset.from_list(records)
+
+
+def load_csv_splits(args: argparse.Namespace) -> tuple[Dataset, Dataset, Dataset]:
+    required = {
+        "--train-csv": args.train_csv,
+        "--val-csv": args.val_csv,
+        "--test-csv": args.test_csv,
+    }
+    missing = [name for name, value in required.items() if not value]
+    if missing:
+        raise ValueError(f"CSV training requires: {', '.join(missing)}")
+
+    clean_clnsig = not args.no_clean_clnsig
+    train_dataset = load_csv_split(
+        args.train_csv,
+        sequence_column=args.sequence_column,
+        label_column=args.label_column,
+        min_sequence_length=args.min_sequence_length,
+        clean_clnsig=clean_clnsig,
+    )
+    eval_dataset = load_csv_split(
+        args.val_csv,
+        sequence_column=args.sequence_column,
+        label_column=args.label_column,
+        min_sequence_length=args.min_sequence_length,
+        clean_clnsig=clean_clnsig,
+    )
+    test_dataset = load_csv_split(
+        args.test_csv,
+        sequence_column=args.sequence_column,
+        label_column=args.label_column,
+        min_sequence_length=args.min_sequence_length,
+        clean_clnsig=clean_clnsig,
+    )
+    return train_dataset, eval_dataset, test_dataset
+
+
 def split_dataset(dataset: Dataset, test_size: float, eval_size: float, seed: int):
     train_eval = dataset.train_test_split(test_size=test_size, seed=seed, stratify_by_column="labels")
     eval_fraction = eval_size / (1.0 - test_size)
@@ -76,7 +182,11 @@ def split_dataset(dataset: Dataset, test_size: float, eval_size: float, seed: in
 
 
 def compute_metrics(eval_pred):
-    logits, labels = eval_pred
+    if hasattr(eval_pred, "predictions"):
+        logits = eval_pred.predictions
+        labels = eval_pred.label_ids
+    else:
+        logits, labels = eval_pred
     predictions = np.argmax(logits, axis=-1)
     return {
         "accuracy": accuracy_score(labels, predictions),
@@ -86,18 +196,53 @@ def compute_metrics(eval_pred):
     }
 
 
+def build_training_args(args: argparse.Namespace, output_dir: Path) -> TrainingArguments:
+    kwargs = {
+        "output_dir": str(output_dir / "checkpoints"),
+        "learning_rate": args.learning_rate,
+        "per_device_train_batch_size": args.batch_size,
+        "per_device_eval_batch_size": args.batch_size,
+        "num_train_epochs": args.epochs,
+        "save_strategy": "epoch",
+        "load_best_model_at_end": True,
+        "metric_for_best_model": "f1",
+        "greater_is_better": True,
+        "report_to": "none",
+        "seed": args.seed,
+    }
+
+    signature = inspect.signature(TrainingArguments.__init__)
+    if "eval_strategy" in signature.parameters:
+        kwargs["eval_strategy"] = "epoch"
+    else:
+        kwargs["evaluation_strategy"] = "epoch"
+
+    return TrainingArguments(**kwargs)
+
+
 def main() -> None:
     args = parse_args()
     output_dir = Path(args.output_dir)
     output_dir.mkdir(parents=True, exist_ok=True)
 
-    dataset = load_examples(args.dataset_jsonl)
-    train_dataset, eval_dataset, test_dataset = split_dataset(
-        dataset,
-        test_size=args.test_size,
-        eval_size=args.eval_size,
-        seed=args.seed,
-    )
+    if args.train_csv or args.val_csv or args.test_csv:
+        train_dataset, eval_dataset, test_dataset = load_csv_splits(args)
+        dataset_source = {
+            "train_csv": args.train_csv,
+            "val_csv": args.val_csv,
+            "test_csv": args.test_csv,
+        }
+    elif args.dataset_jsonl:
+        dataset = load_examples(args.dataset_jsonl)
+        train_dataset, eval_dataset, test_dataset = split_dataset(
+            dataset,
+            test_size=args.test_size,
+            eval_size=args.eval_size,
+            seed=args.seed,
+        )
+        dataset_source = {"dataset_jsonl": args.dataset_jsonl}
+    else:
+        raise ValueError("Provide either --dataset-jsonl or all of --train-csv, --val-csv, and --test-csv.")
 
     tokenizer = AutoTokenizer.from_pretrained(args.model_name, trust_remote_code=True)
 
@@ -116,30 +261,23 @@ def main() -> None:
         trust_remote_code=True,
     )
 
-    training_args = TrainingArguments(
-        output_dir=str(output_dir / "checkpoints"),
-        learning_rate=args.learning_rate,
-        per_device_train_batch_size=args.batch_size,
-        per_device_eval_batch_size=args.batch_size,
-        num_train_epochs=args.epochs,
-        evaluation_strategy="epoch",
-        save_strategy="epoch",
-        load_best_model_at_end=True,
-        metric_for_best_model="f1",
-        greater_is_better=True,
-        report_to="none",
-        seed=args.seed,
-    )
+    training_args = build_training_args(args, output_dir)
 
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=train_dataset,
-        eval_dataset=eval_dataset,
-        tokenizer=tokenizer,
-        data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
-        compute_metrics=compute_metrics,
-    )
+    trainer_kwargs = {
+        "model": model,
+        "args": training_args,
+        "train_dataset": train_dataset,
+        "eval_dataset": eval_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
+
+    trainer = Trainer(**trainer_kwargs)
 
     trainer.train()
     test_metrics = trainer.evaluate(test_dataset, metric_key_prefix="test")
@@ -152,6 +290,11 @@ def main() -> None:
         "base_model": args.model_name,
         "genome_build": "GRCh38",
         "labels": ID_TO_LABEL,
+        "dataset_source": dataset_source,
+        "train_rows": len(train_dataset),
+        "eval_rows": len(eval_dataset),
+        "test_rows": len(test_dataset),
+        "max_length": args.max_length,
         "test_metrics": test_metrics,
         "research_only": True,
         "disclaimer": "For research and education only. Not for medical diagnosis.",

training/utils/label_utils.py

@@ -6,10 +6,15 @@ from urllib.parse import unquote
 
 
 DROP_LABEL_TERMS = (
+    "conflicting",
     "conflicting interpretations",
+    "conflicting classifications",
+    "uncertain",
     "uncertain significance",
+    "risk",
     "risk factor",
     "association",
+    "drug",
     "drug response",
     "protective",
     "not provided",