Add fine-tuning module

.gitignore

@@ -61,4 +61,17 @@ notebooks/**
 data/**
 
 # Notes
-notes.md
+notes.md
+# Sweep results
+sweep/results.parquet
+sweep/sweep.log
+sweep/state.json
+
+# Fine-tuning runs and local configs
+runs/
+!runs/.gitkeep
+configs/single/**
+configs/sweep/**
+!configs/single/example.yaml
+!configs/single/test.yaml
+!configs/sweep/example.yaml

.pre-commit-config.yaml

@@ -74,16 +74,3 @@ repos:
       - id: pydoclint
         args: ["--config=pyproject.toml"]
         exclude: ^tests/
-
-  # Cognitive complexity limit via flake8-cognitive-complexity (keep for src; exclude tests)
-  - repo: https://github.com/pycqa/flake8
-    rev: 7.3.0
-    hooks:
-      - id: flake8
-        additional_dependencies:
-          - flake8-cognitive-complexity
-        args:
-          - --max-line-length=79
-          - --extend-ignore=E203
-          - --max-cognitive-complexity=10
-        exclude: ^tests/

README.md

@@ -1,3 +1,14 @@
-# VOICE - Fine-tuning for Stylistic Fidelity
+# 🗣️ VOICE
+
+...
 
 ---
+
+## Documentation
+
+```
+docs/
+├── 00_data.md         — Included datasets
+├── 01_evals.md        — Evaluation suite: scoring and interpretation
+└── 02_stylometry.md   — Stylometric metrics: definition and catalogue
+```

configs/single/example.yaml

@@ -0,0 +1,87 @@
+# =============================================================================
+# Example single-run configuration
+#
+# Run with:
+#   voice finetune single configs/single/example.yaml
+#
+# All axolotl settings live under the `axolotl:` key. The CLI writes a
+# frozen copy to runs/{run_name}/config.yaml before launching training,
+# injecting output_dir automatically.
+# =============================================================================
+
+axolotl:
+
+  # Model
+  base_model: meta-llama/Llama-3.1-8B-Instruct
+
+  # Train split used for fine-tuning; validation and test splits used by
+  # EvalCompletionsPlugin and post-training alignment scoring.
+  datasets:
+    - path: AccelerateScience/bo-press-conference-qa
+      type: chat_template
+      field_messages: messages
+      roles_to_train:
+        - assistant
+      train_on_eos: last
+      eot_tokens: ["<|eot_id|>"]
+
+  special_tokens:
+    pad_token: "<|eot_id|>"
+
+  hf_use_auth_token: true
+
+  chat_template: tokenizer_default
+  val_set_size: 0.0
+
+  # Validation dataset for per-epoch eval_loss and completions generation.
+  test_datasets:
+    - path: AccelerateScience/bo-press-conference-qa
+      split: validation
+      type: chat_template
+      field_messages: messages
+
+  # Sequence
+  sequence_len: 2048
+  sample_packing: true
+
+  # Adapter
+  adapter: lora
+  lora_r: 16
+  lora_alpha: 16
+  lora_dropout: 0.0
+  lora_target_modules: [q_proj, k_proj, v_proj, o_proj]
+
+  # Training
+  num_epochs: 3
+  learning_rate: 1.0e-4
+  micro_batch_size: 4
+  gradient_accumulation_steps: 4
+
+  # Precision & performance
+  bf16: true
+  tf32: true
+  flash_attention: true
+  gradient_checkpointing: true
+  lora_mlp_kernel: false
+  lora_qkv_kernel: false
+  lora_o_kernel: false
+
+  # Optimiser
+  optimizer: adamw_torch_fused
+  lr_scheduler: cosine
+  warmup_ratio: 0.03
+  weight_decay: 0.0
+
+  # Evaluation and checkpointing
+  eval_strategy: "epoch"
+  save_strategy: epoch
+  save_total_limit: 3
+
+  # Plugin that generates completions at each epoch end (recommended).
+  plugins:
+    - voice.finetune.callbacks.EvalCompletionsPlugin
+
+  # W&B
+  use_wandb: true
+  wandb_project: VOICE
+  wandb_entity: accelerate-science

configs/sweep/example.yaml

@@ -0,0 +1,127 @@
+# =============================================================================
+# Example sweep configuration
+#
+# Two top-level sections:
+#   sweep:   hyperparameter grid axes (read by `voice finetune sweep`)
+#   axolotl: shared axolotl settings merged into every per-run config
+#
+# The CLI expands the grid into per-run configs, injecting per-run values
+# (learning_rate, lora_r, lora_alpha, lora_target_modules,
+# micro_batch_size, gradient_accumulation_steps, output_dir) on top of
+# the axolotl section below.
+#
+# Grid: 3 LRs × 4 ranks × 1 micro_batch_size × 2 accum steps
+#       × 2 target layer groups = 48 training runs.
+#
+# Run with:
+#   voice finetune sweep configs/sweep/example.yaml
+#   voice finetune sweep configs/sweep/example.yaml --resume
+# =============================================================================
+
+# --- Sweep grid -------------------------------------------------------------
+sweep:
+  learning_rate: [1e-3, 5.0e-4, 1e-4]
+  lora_r: [4, 8, 16, 32]
+
+  micro_batch_size: [2]
+  gradient_accumulation_steps: [2, 4]
+
+  target_layers:
+    - name: mlp
+      modules: [gate_proj, up_proj, down_proj]
+    - name: attention
+      modules: [q_proj, k_proj, v_proj, o_proj]
+
+# lora_alpha is set equal to lora_r per run (not swept separately).
+
+# --- Shared axolotl settings ------------------------------------------------
+axolotl:
+
+  # Model
+  base_model: meta-llama/Llama-3.1-8B-Instruct
+
+  # Train split used for fine-tuning; validation split used by the
+  # EvalCompletionsPlugin for per-epoch completions; test split scored after
+  # each epoch for alignment evaluation.
+  datasets:
+    - path: AccelerateScience/gwb-press-conference-qa
+      type: chat_template
+      field_messages: messages
+      roles_to_train:
+        - assistant
+      train_on_eos: last
+      # Llama 3.1's chat template ends turns with <|eot_id|> (128009), not
+      # <|end_of_text|> (128001). Without this, axolotl defaults eot_tokens
+      # to eos_token which never appears, so train_on_eos: last silently
+      # does nothing.
+      eot_tokens: ["<|eot_id|>"]
+
+  special_tokens:
+    pad_token: "<|eot_id|>"
+
+  hf_use_auth_token: true
+
+  # Llama 3.1 Instruct ships with its own Jinja chat template; axolotl
+  # picks it up automatically.
+  chat_template: tokenizer_default
+  val_set_size: 0.0
+
+  # Axolotl computes eval_loss after each epoch and the EvalCompletionsPlugin
+  # generates completions for scoring.
+  test_datasets:
+    - path: AccelerateScience/gwb-press-conference-qa
+      split: validation
+      type: chat_template
+      field_messages: messages
+
+  # Sequence
+  # Bush answers are long (~700-900 tokens including system + question).
+  # 2048 gives comfortable headroom without wasting VRAM.
+  sequence_len: 2048
+  sample_packing: true
+
+  # lora_r, lora_alpha, lora_target_modules injected per run.
+  adapter: lora
+  lora_dropout: 0.05
+
+  # 3 epochs × 48 runs = 144 data points.
+  num_epochs: 3
+
+  # Precision & performance
+  bf16: true
+  tf32: true
+  flash_attention: true
+  gradient_checkpointing: true
+
+  # Disable axolotl's auto-enabled LoRA kernel optimizations. The MLP
+  # SwiGLU Triton kernel fails to compile on our host (gcc cannot link
+  # libcuda), which crashes all attention_mlp runs.
+  lora_mlp_kernel: false
+  lora_qkv_kernel: false
+  lora_o_kernel: false
+
+  # Optimiser
+  optimizer: adamw_torch_fused
+  lr_scheduler: cosine
+  warmup_ratio: 0.03
+  weight_decay: 0.0
+
+  # Evaluate on the validation split after each epoch.
+  eval_strategy: "epoch"
+
+  # Save a checkpoint after each epoch so adapters are available for
+  # per-epoch alignment scoring (intermediate epochs treated as separate
+  # model candidates).
+  save_strategy: epoch
+  save_total_limit: 3
+
+  # Generates completions at each epoch end using the live in-memory model
+  # (no train/inference distribution mismatch).
+  # Dataset is read from datasets[0].path above.
+  plugins:
+    - voice.finetune.callbacks.EvalCompletionsPlugin
+
+  # Run name for W&B is set by the CLI; project and entity configured here.
+  use_wandb: true
+  wandb_project: VOICE
+  wandb_entity: accelerate-science

cspell/library-words.txt

@@ -9,3 +9,18 @@ ylabel
 frameon
 allclose
 linalg
+mathcal
+mathbb
+mathbf
+dotenv
+wandb
+jinja
+adamw
+checkpointing
+proc
+peft
+pyyaml
+libcuda
+PYTHONPATH
+venv
+cuda

cspell/project-words.txt

@@ -12,3 +12,14 @@ bonf
 bonferroni
 aeiou
 stylometrically
+lora
+resumability
+resumable
+hyperparams
+redef
+evals
+submatrix
+vocab
+finetune
+accum
+nonpositive

docs/00_data.md

@@ -0,0 +1,22 @@
+# Datasets
+
+---
+
+Two example datasets are included, each containing press conference Q&A transcripts split into train, validation and test sets:
+
+| President | HuggingFace |
+|---|---|
+| Barack Obama | [`AccelerateScience/bo-press-conference-qa`](https://huggingface.co/datasets/AccelerateScience/bo-press-conference-qa) |
+| George W. Bush | [`AccelerateScience/gwb-press-conference-qa`](https://huggingface.co/datasets/AccelerateScience/gwb-press-conference-qa) |
+
+Each example is a single JSONL record in chat format:
+
+```json
+{
+  "messages": [
+    {"role": "system",    "content": "You are president Barack Obama ..."},
+    {"role": "user",      "content": "<press conference question>"},
+    {"role": "assistant", "content": "<president's answer>"}
+  ]
+}
+```

docs/01_stylometry.md

@@ -0,0 +1,52 @@
+# Stylometric Metrics
+
+---
+
+A **stylometric metric** is any function
+
+$$f : \mathcal{T} \rightarrow \mathbb{R}$$
+
+that maps a text string to a real scalar capturing some surface property of writing style. VOICE treats each metric as a distribution over a corpus: given a set of texts, $f$ is applied to each one to produce a sample from the author's stylometric distribution for that feature.
+
+## Implemented Metrics
+
+### Word Length Distribution
+Moments of the per-word character length distribution.
+
+| Metric | Description |
+|---|---|
+| `avg_word_length` | Mean word length |
+| `std_word_length` | Standard deviation of word length |
+| `skew_word_length` | Skewness of word length |
+| `kurtosis_word_length` | Kurtosis of word length |
+
+### Vocabulary Richness
+Type–token and word statistics measuring lexical diversity.
+
+| Metric | Description |
+|---|---|
+| `type_token_ratio` | Unique tokens / total tokens |
+| `moving_avg_type_token_ratio` | TTR averaged over a sliding window (MATTR) |
+| `hapax_legomena_ratio` | Fraction of words appearing exactly once |
+| `dis_legomena_ratio` | Fraction of words appearing exactly twice |
+| `tri_legomena_ratio` | Fraction of words appearing exactly three times |
+
+### Function Words
+| Metric | Description |
+|---|---|
+| `function_word_ratio` | Proportion of tokens drawn from a closed function-word list |
+
+### Character N-gram Diversity
+Type–token ratio and MATTR computed over character $n$-grams for $n \in \{3, 4, 5\}$.
+
+| Metric | Description |
+|---|---|
+| `char_{n}gram_type_token_ratio` | Character $n$-gram TTR |
+| `char_{n}gram_moving_avg_type_token_ratio` | Character $n$-gram MATTR |
+
+### Text Length
+| Metric | Description |
+|---|---|
+| `num_words` | Total word count |
+
+> **Note:** `num_words` may be deprecated in a future release. It does not tend to be used as a signature for authorship attribution in the broader stylometry literature.