diff --git a/src/llm/specialise.py b/src/llm/specialise.py index 13a26c8..aade829 100644 --- a/src/llm/specialise.py +++ b/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)