specialise: extract train_lora_on_tasks — reusable answer-only SFT primitive

Factor the LoRA SFT loop out of train_specialist into train_lora_on_tasks(base,
tasks, out_dir, ...), which fine-tunes a fresh adapter on an arbitrary list of
(prompt, answer) pairs. train_specialist now generates its family's tasks and
delegates. Backward-compatible (same behaviour for the merge/moe/directed
experiments); the primitive is available for future loops that train on
arbitrary or model-generated pairs.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

src/llm/specialise.py

@@ -27,15 +27,25 @@ def _encode(tok, task, device):
     return ids, labels
 
 
-def train_specialist(base_name: str, family: str, out_dir: str, *, n_train: int = 600,
-                     epochs: int = 3, lr: float = 2e-4, batch_size: int = 8, r: int = 16,
-                     alpha: int = 32, seed: int = 0, device: str = "cuda", hard: bool = False) -> str:
-    """Fine-tune a LoRA specialist on ``family`` and save the adapter to ``out_dir``.
+def train_lora_on_tasks(base_name: str, tasks: list, out_dir: str, *, epochs: int = 3,
+                        lr: float = 2e-4, batch_size: int = 8, r: int = 16, alpha: int = 32,
+                        seed: int = 0, device: str = "cuda") -> str:
+    """Fine-tune a fresh LoRA adapter on an arbitrary list of ``tasks`` and save it to ``out_dir``.
+
+    The reusable answer-only SFT primitive: each task supplies a ``.prompt`` and a ``.answer`` (which
+    may be a *ground-truth* answer, for specialisation, or a *model-generated* one, for the
+    self-consumption loop). The prompt tokens are masked out of the loss. A new adapter is trained from
+    the frozen base each call (replace dynamics), so a lineage of generations is a pure function of the
+    data each generation supplies.
 
     Args:
-        hard (bool): train on the harder task variant (must match the eval difficulty).
+        base_name (str): HF id of the frozen base model.
+        tasks (list): objects with ``.prompt`` and ``.answer`` string attributes.
+        out_dir (str): where to save the adapter.
+        epochs, lr, batch_size, r, alpha, seed, device: standard SFT knobs.
 
-    Returns the adapter directory path.
+    Returns:
+        str: ``out_dir``.
     """
     import torch
     from peft import LoraConfig, get_peft_model
@@ -52,7 +62,6 @@ def train_specialist(base_name: str, family: str, out_dir: str, *, n_train: int
     model = get_peft_model(model, lora)
     model.train()
 
-    tasks = make_tasks(family, n_train, seed=seed, hard=hard)
     encoded = [_encode(tok, t, device) for t in tasks]
     opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=lr)
     rng = np.random.default_rng(seed)
@@ -77,3 +86,18 @@ def train_specialist(base_name: str, family: str, out_dir: str, *, n_train: int
     Path(out_dir).mkdir(parents=True, exist_ok=True)
     model.save_pretrained(out_dir)
     return out_dir
+
+
+def train_specialist(base_name: str, family: str, out_dir: str, *, n_train: int = 600,
+                     epochs: int = 3, lr: float = 2e-4, batch_size: int = 8, r: int = 16,
+                     alpha: int = 32, seed: int = 0, device: str = "cuda", hard: bool = False) -> str:
+    """Fine-tune a LoRA specialist on ``family`` and save the adapter to ``out_dir``.
+
+    Args:
+        hard (bool): train on the harder task variant (must match the eval difficulty).
+
+    Returns the adapter directory path.
+    """
+    tasks = make_tasks(family, n_train, seed=seed, hard=hard)
+    return train_lora_on_tasks(base_name, tasks, out_dir, epochs=epochs, lr=lr,
+                               batch_size=batch_size, r=r, alpha=alpha, seed=seed, device=device)