initial upload

utils/config_loader.py

@@ -0,0 +1,256 @@
+# utils/config_loader.py
+
+import argparse
+import copy
+import logging
+from pathlib import Path
+from typing import Dict, List, Sequence, Any
+
+import yaml
+
+logger = logging.getLogger(__name__)
+
+
+_ALWAYS: Sequence[str] = (
+    # minimal required keys for the pipeline to run at all
+    "experiment_base_dir",
+    "human_profile_path",
+    "num_iterations",
+    "min_word_len_for_analysis",
+    "log_level",
+    "model_id",
+)
+
+_VLLM: Sequence[str] = (
+    "vllm_model_id",
+    "vllm_port",
+    "vllm_hf_token",
+    "vllm_cuda_visible_devices",
+    "vllm_gpu_memory_utilization",
+    "vllm_max_model_len",
+    "vllm_dtype",
+    "vllm_extra_args",
+    "vllm_env", 
+)
+
+_GENERATION: Sequence[str] = (
+    "generation_api_key",
+    "generation_api_base_url",       # needed if you do local or remote calls
+    "generation_model_id",
+    "generation_max_new_tokens",
+    "generation_threads",
+    "generation_max_prompts",
+    "generation_hf_dataset_name",
+    "generation_hf_dataset_split",
+    "generation_logging_level",
+    "generation_chat_template_model_id",
+    "generation_param_chunk_size",
+    "generation_param_top_logprobs_count",
+    "generation_param_temperature",
+    "generation_param_top_p",
+    "generation_param_top_k",
+    "generation_param_min_p",
+    "generation_param_timeout",
+    "generation_param_stop_sequences",
+    "generation_ngram_remove_stopwords",
+    "generation_ngram_language",
+    "generation_force_backtrack",
+    "generation_prompt_template",
+    "generation_system_prompt"
+
+)
+
+_NGRAM: Sequence[str] = (
+    "top_k_bigrams",
+    "top_k_trigrams",
+    "dict_bigrams_initial",
+    "dict_bigrams_subsequent",
+    "nodict_bigrams_initial",
+    "nodict_bigrams_subsequent",
+    "dict_trigrams_initial",
+    "dict_trigrams_subsequent",
+    "nodict_trigrams_initial",
+    "nodict_trigrams_subsequent",
+    "extra_ngrams_to_ban",
+)
+
+_SLOP: Sequence[str] = (
+    "min_phrase_freq_to_keep",
+    "top_n_initial_slop_ban",
+    "top_n_subsequent_slop_ban",
+    "extra_slop_phrases_to_ban",
+    "banned_slop_phrases_filename",
+)
+
+_OVERREP: Sequence[str] = (
+    "top_k_words_for_overrep_analysis",
+    "dict_overrep_initial",
+    "dict_overrep_subsequent",
+    "nodict_overrep_initial",
+    "nodict_overrep_subsequent",
+)
+
+_FINETUNE: Sequence[str] = (
+    "finetune_mode",
+    "finetune_ftpo_dataset",
+    "finetune_base_model_id",
+    "finetune_max_seq_length",
+    "finetune_load_in_4bit",
+    "finetune_lora_r",
+    "finetune_lora_alpha",
+    "finetune_lora_dropout",
+    "finetune_weight_decay",
+    "finetune_target_modules",
+    "finetune_gradient_checkpointing",
+    "finetune_chat_template",
+    "finetune_batch_size",
+    "finetune_gradient_accumulation_steps",
+    "finetune_warmup_ratio",
+    "finetune_num_epochs",
+    "finetune_learning_rate",
+    "finetune_auto_learning_rate",
+    "finetune_beta",
+    "finetune_output_dir_suffix",
+    "finetune_save_merged_16bit",
+    "finetune_save_gguf_q8_0",
+    "finetune_max_train_examples",    
+    "finetune_cuda_visible_devices",
+    "ftpo_sample_rejected_regularisation_strength",
+    "ftpo_sample_chosen_regularisation_strength",
+    "ftpo_sample_min_chosen_tokens",
+)
+
+
+def _deep_update(dst: Dict, src: Dict) -> Dict:
+    """Recursively merge src into dst (src wins)."""
+    for k, v in src.items():
+        if k in dst and isinstance(v, dict) and isinstance(dst[k], dict):
+            _deep_update(dst[k], v)
+        else:
+            dst[k] = copy.deepcopy(v)
+    return dst
+
+def load_pipeline_config(config_path: Path) -> Dict[str, Any]:
+    """Load config from a YAML file, or return empty dict if missing/invalid."""
+    if config_path and config_path.exists():
+        try:
+            with config_path.open('r', encoding='utf-8') as f:
+                data = yaml.safe_load(f) or {}
+            logger.info("Loaded configuration from %s", config_path)
+            return data
+        except Exception as e:
+            logger.error("Could not load %s: %s – using empty config", config_path, e)
+    else:
+        logger.info("Config file %s not found – using empty config", config_path)
+    return {}
+
+def merge_config_with_cli_args(config: Dict[str, Any], cli_args: argparse.Namespace) -> Dict[str, Any]:
+    """
+    Merges every possible CLI parameter from your old DEFAULT_CONFIG
+    into 'config' if the user actually provided it (i.e. it's not None).
+    Also merges housekeeping flags (config_file, resume_from_dir, log_level).
+    """
+    merged = copy.deepcopy(config)
+
+    # 1. Housekeeping arguments (not originally in DEFAULT_CONFIG, but we keep them if set)
+    if getattr(cli_args, 'config_file', None) is not None:
+        merged['config_file'] = cli_args.config_file
+    if getattr(cli_args, 'resume_from_dir', None) is not None:
+        merged['resume_from_dir'] = cli_args.resume_from_dir
+    if getattr(cli_args, 'log_level', None) is not None:
+        merged['log_level'] = cli_args.log_level
+
+    # 2. Booleans that map from CLI flags to known keys in config
+    if getattr(cli_args, 'run_finetune', None) is not None:
+        merged['finetune_enabled'] = cli_args.run_finetune
+    if getattr(cli_args, 'manage_vllm', None) is not None:
+        merged['manage_vllm'] = cli_args.manage_vllm
+    if getattr(cli_args, 'generation_step_enabled', None) is not None:
+        merged['generation_step_enabled'] = cli_args.generation_step_enabled
+    if getattr(cli_args, "finetune_cuda_visible_devices", None) is not None:
+        merged["finetune_cuda_visible_devices"] = cli_args.finetune_cuda_visible_devices
+
+
+    # 3. All remaining keys from the old DEFAULT_CONFIG
+    _all_groups: Sequence[Sequence[str]] = (
+        _ALWAYS,
+        _VLLM,
+        _GENERATION,
+        _NGRAM,
+        _SLOP,
+        _OVERREP,
+        _FINETUNE,
+    )
+    all_config_keys: List[str] = [k for group in _all_groups for k in group]
+
+     # 4. fallback for per-stage model IDs  <-- add this block
+    for key in (
+        "vllm_model_id",
+        "generation_model_id",
+        "generation_chat_template_model_id",
+        "finetune_base_model_id",
+    ):
+        if not merged.get(key):                 # None, "", or missing
+            merged[key] = merged.get("model_id")
+
+    # Overwrite config if user specified a value
+    for key in all_config_keys:
+        cli_val = getattr(cli_args, key, None)
+        if cli_val is not None:
+            merged[key] = cli_val
+
+    return merged
+
+# ---------------------------------------------------------------------------
+# Validate with partial requirements depending on which features are enabled
+# ---------------------------------------------------------------------------
+
+def _missing(cfg: Dict[str, Any], keys: Sequence[str]) -> List[str]:
+    return [k for k in keys if k not in cfg or cfg[k] is None]
+
+def validate_config(cfg: Dict[str, Any]) -> None:
+    """Raise ValueError if any required config is missing based on pipeline flags."""
+    missing = []
+    # always
+    missing.extend(_missing(cfg, _ALWAYS))
+
+    # vllm
+    if cfg.get("manage_vllm", False):
+        missing.extend(_missing(cfg, _VLLM))
+
+    # generation
+    if cfg.get("generation_step_enabled", True):
+        missing.extend(_missing(cfg, _GENERATION))
+
+    # n-gram ban
+    if cfg.get("enable_ngram_ban", False):
+        missing.extend(_missing(cfg, _NGRAM))
+
+    # slop phrase ban
+    if cfg.get("enable_slop_phrase_ban", False):
+        missing.extend(_missing(cfg, _SLOP))
+
+    # over-rep analysis
+    if cfg.get("compute_overrep_words", False):
+        missing.extend(_missing(cfg, _OVERREP))
+
+    # finetuning
+    if cfg.get("finetune_enabled", False):
+        missing.extend(_missing(cfg, _FINETUNE))
+
+    if missing:
+        raise ValueError(
+            f"Configuration is incomplete; missing these keys: {', '.join(sorted(set(missing)))}"
+        )
+    logger.info("Configuration validated – all required keys present (for enabled features).")
+
+def load_merge_validate(config_path: Path, cli_args: argparse.Namespace) -> Dict[str, Any]:
+    """
+    1) Load YAML from config_path,
+    2) Merge in any CLI flags user typed,
+    3) Validate that all needed keys for enabled features are present.
+    """
+    cfg = load_pipeline_config(config_path)
+    cfg = merge_config_with_cli_args(cfg, cli_args)
+    validate_config(cfg)
+    return cfg

utils/dataset_helpers.py

@@ -0,0 +1,290 @@
+# utils/dataset_helpers.py
+from __future__ import annotations
+import logging, os
+from pathlib import Path
+from collections import Counter, defaultdict
+from typing import Collection, Optional
+from datetime import datetime, timezone
+import json
+
+import numpy as np
+from datasets import load_dataset
+
+logger = logging.getLogger(__name__)
+
+# tokens we want to watch closely
+_WATCH = [" nodded", " leaned"]
+
+
+def load_ftpo_multi_dataset(
+    path: Path,
+    tokenizer,
+    *,
+    experiment_run_dir: Path | None = None,
+    max_seq_len: int = 4096,
+    rejected_reg_strength: float = 0.0,
+    chosen_reg_strength: float = 0.0,
+    min_chosen_tokens: int = 1,
+    max_train_examples: int | None = None,
+    stop_words: Optional[Collection[str]] = None,
+    num_proc: int | None = None,
+    batch_size: int = 512,
+):
+    """
+    Parallel loader for “multi-chosen” FTPO JSONL with dual regularisation.
+    Logs the counts of `_WATCH` tokens at every major stage.
+    """
+
+    if min_chosen_tokens < 1:
+        min_chosen_tokens = 1
+
+    # ------------------------------------------------------------------
+    # helpers
+    # ------------------------------------------------------------------
+    rng = np.random.default_rng(3407)
+
+    def _median_threshold(cts: Counter[str], strength: float) -> dict[str, float]:
+        if not cts or strength <= 0:
+            return {}
+        med = float(np.median(list(cts.values())))
+        return {t: 1.0 if c <= med else (med / c) ** strength for t, c in cts.items()}
+
+    def _log_top(cts: Counter[str], what: str) -> None:
+        head = ", ".join(f"{tok!r}:{cnt}" for tok, cnt in cts.most_common(20))
+        logger.info(f"[ftpo-loader] {what} top-20 → {head}")
+        logger.info(
+            "          ↳ watch «%s»: %s   «%s»: %s",
+            _WATCH[0], cts[_WATCH[0]],
+            _WATCH[1], cts[_WATCH[1]],
+        )
+
+    # ------------------------------------------------------------------
+    # stop-word list (unchanged)
+    # ------------------------------------------------------------------
+    if stop_words is None:
+        stop_words = {
+            "the","a","an","in","on","at","by","for","to","of","and","or","but",
+            "if","then","else","when","where","how","why","what","who","whom",
+            "this","that","these","those","is","are","was","were","be","being",
+            "been","have","has","had","do","does","did","will","would","shall",
+            "should","can","could","may","might","must"
+        }
+    stop_words = {w.lower() for w in stop_words}
+
+    # ------------------------------------------------------------------
+    # 0️⃣  raw load + shuffle
+    # ------------------------------------------------------------------
+    raw = load_dataset("json", data_files=str(path), split="train").shuffle(seed=3407)
+    rows = list(raw)
+    if not rows:
+        raise ValueError(f"{path} contained no rows")
+
+    rej_counts = Counter(r["rejected_decoded"] for r in rows)
+    _log_top(rej_counts, "BEFORE")
+
+    # ────────────────────────────────────────────────────────────────
+    # 1️⃣  Capture ORIGINAL rejected-token distribution & ratios
+    #     (no rows removed, no chosen trimming yet)
+    # ────────────────────────────────────────────────────────────────
+    rej_cts_orig = Counter(r["rejected_decoded"] for r in rows)
+    _log_top(rej_cts_orig, "PRE-NORMALISATION")
+
+    # convert to fractional “weights” via median-threshold regularisation
+    med = float(np.median(list(rej_cts_orig.values())))
+    w_rej = {tok: 1.0 if c <= med else (med / c) ** rejected_reg_strength
+            for tok, c in rej_cts_orig.items()}
+    
+    # normalised ratios we *want* to keep in the final dataset
+    total_weighted = sum(w_rej[t] * c for t, c in rej_cts_orig.items())
+    ratio_rej = {tok: (w_rej[tok] * cnt) / total_weighted
+                for tok, cnt in rej_cts_orig.items()}
+
+    # ────────────────────────────────────────────────────────────────
+    # 2️⃣  Chosen-token trimming  (build quotas *before* we cut)
+    # ────────────────────────────────────────────────────────────────
+    chosen_cts_orig = Counter(tok
+                        for r in rows
+                        for tok in (r["multi_chosen_decoded"] or []))
+
+    _log_top(chosen_cts_orig, "ORIGINAL CHOSEN TOKENS")
+
+    # Trim the peak: cap top tokens to match the 10th highest count
+    if len(chosen_cts_orig) >= 10:
+        top_counts = sorted(chosen_cts_orig.values(), reverse=True)
+        cap_value = top_counts[9]  # 10th highest count
+        chosen_cts_capped = Counter()
+        for tok, cnt in chosen_cts_orig.items():
+            chosen_cts_capped[tok] = min(cnt, cap_value)
+    else:
+        chosen_cts_capped = chosen_cts_orig.copy()
+
+    # Now calculate regularization on the capped distribution
+    med_chosen = float(np.median(list(chosen_cts_capped.values())))
+    w_chosen = {tok: 1.0 if c <= med_chosen
+            else (med_chosen / c) ** chosen_reg_strength
+            for tok, c in chosen_cts_capped.items()}
+
+    tgt_chosen = {tok: int(round(c * w_chosen.get(tok, 1.0)))
+                for tok, c in chosen_cts_capped.items()}
+
+    # Log the target quotas
+    quota_items = sorted(tgt_chosen.items(), key=lambda x: x[1], reverse=True)[:20]
+    quota_str = ", ".join(f"{tok!r}:{quota}" for tok, quota in quota_items)
+    logger.info(f"[ftpo-loader] CHOSEN TARGET QUOTAS top-20 → {quota_str}")
+    logger.info(
+        "          ↳ watch quotas «%s»: %s (was %s)   «%s»: %s (was %s)",
+        _WATCH[0], tgt_chosen.get(_WATCH[0], 0), chosen_cts_orig.get(_WATCH[0], 0),
+        _WATCH[1], tgt_chosen.get(_WATCH[1], 0), chosen_cts_orig.get(_WATCH[1], 0),
+    )
+
+    _log_top(Counter(r["rejected_decoded"] for r in rows), "POST-CHOSEN")
+
+    # ────────────────────────────────────────────────────────────────
+    # 3️⃣  Apply min_chosen_tokens row filter
+    # ────────────────────────────────────────────────────────────────
+    rows = [r for r in rows if len(r["multi_chosen_decoded"]) >= min_chosen_tokens]
+
+    _log_top(Counter(r["rejected_decoded"] for r in rows), "POST-MIN-FILTER")
+
+    # ────────────────────────────────────────────────────────────────
+    # 4️⃣  Row-level quota sampling **now** that trimming & filtering
+    #     are done.  Scale the original ratios to the remaining size.
+    # ────────────────────────────────────────────────────────────────
+    N_final = max_train_examples or len(rows)
+    target_rows = {tok: int(round(ratio_rej[tok] * N_final))
+                for tok in ratio_rej}
+
+    rng.shuffle(rows)
+    selected, seen = [], defaultdict(int)
+    selected_indices = set()  # Track indices instead of row objects
+    
+    # First pass: try to fill quotas
+    for i, r in enumerate(rows):
+        tok = r["rejected_decoded"]
+        if seen[tok] < target_rows.get(tok, 0):
+            selected.append(r)
+            selected_indices.add(i)
+            seen[tok] += 1
+        if len(selected) >= N_final:
+            break
+    
+    # Second pass: if still short, keep adding while maintaining proportions
+    if len(selected) < N_final:
+        # Build index of remaining rows by token
+        remaining_by_token = defaultdict(list)
+        for i, r in enumerate(rows):
+            if i not in selected_indices:
+                remaining_by_token[r["rejected_decoded"]].append((i, r))
+        
+        # Keep adding until we reach N_final
+        while len(selected) < N_final:
+            # Find token that's furthest below its target ratio AND has rows available
+            best_tok = None
+            best_ratio_diff = -1
+            
+            for tok, available_rows in remaining_by_token.items():
+                if not available_rows:  # Skip tokens with no remaining rows
+                    continue
+                    
+                current_ratio = seen[tok] / len(selected) if len(selected) > 0 else 0
+                target_ratio = ratio_rej.get(tok, 0)
+                ratio_diff = target_ratio - current_ratio
+                
+                if ratio_diff > best_ratio_diff:
+                    best_ratio_diff = ratio_diff
+                    best_tok = tok
+            
+            # If no tokens have available rows, we're done
+            if best_tok is None:
+                break
+            
+            # Add one row for the most underrepresented token
+            idx, r = remaining_by_token[best_tok].pop()
+            selected.append(r)
+            selected_indices.add(idx)
+            seen[best_tok] += 1
+    
+    rows = selected
+
+    # ── Dump the final row subset exactly as it was read (no tokenisation) ──
+    if experiment_run_dir is not None:
+        ts = datetime.now(timezone.utc).astimezone()\
+                 .strftime("%Y-%m-%d_%H-%M-%S")
+        dump_file = experiment_run_dir / f"ftpo_training_set_used_{ts}.jsonl"
+        try:
+            with open(dump_file, "w", encoding="utf-8") as fh:
+                for r in rows:
+                    fh.write(json.dumps(r, ensure_ascii=False) + "\n")
+            logger.info("[ftpo-loader] dumped %d rows → %s", len(rows), dump_file)
+        except Exception as e:
+            logger.warning("[ftpo-loader] failed to dump training rows: %s", e)
+
+
+
+    _log_top(Counter(r["rejected_decoded"] for r in rows), "AFTER-SAMPLING")
+    logger.info("[ftpo-loader] kept %d rows after quota sampling", len(rows))
+
+
+    # ------------------------------------------------------------------
+    # 5️⃣  tokenisation (unchanged section)
+    # ------------------------------------------------------------------
+    from datasets import Dataset
+    ds = Dataset.from_list(rows)
+
+    tokenizer.truncation_side = "left"
+    num_proc = num_proc or max(1, int(os.cpu_count() / 4))
+
+    def _tok(batch):
+        out_prompt, out_chosen, out_rej, out_valid = [], [], [], []
+
+        prompt_tok = tokenizer(
+            batch["context_with_chat_template"],
+            add_special_tokens=False,
+            truncation=False,
+            return_attention_mask=False,
+        ).input_ids
+
+        for p_ids, chosen_surf, rej_surf in zip(
+            prompt_tok, batch["multi_chosen_decoded"], batch["rejected_decoded"]
+        ):
+            chosen_surf = chosen_surf or []
+            chosen_tok_ids = [tokenizer(t, add_special_tokens=False).input_ids
+                              for t in chosen_surf]
+            rej_tok_ids = tokenizer(rej_surf, add_special_tokens=False).input_ids
+
+            valid = (
+                chosen_tok_ids
+                and all(len(t) == 1 for t in chosen_tok_ids)
+                and len(rej_tok_ids) == 1
+                and rej_surf.strip().lower() not in stop_words
+                and len(p_ids) + 1 <= max_seq_len
+            )
+            if valid and rej_tok_ids[0] in [t[0] for t in chosen_tok_ids]:
+                valid = False
+
+            out_valid.append(valid)
+            if valid:
+                out_prompt.append(p_ids)
+                out_chosen.append([t[0] for t in chosen_tok_ids])
+                out_rej.append(rej_tok_ids[0])
+            else:
+                out_prompt.append([0]); out_chosen.append([0]); out_rej.append(0)
+
+        return {
+            "prompt_ids":        out_prompt,
+            "chosen_ids":        out_chosen,
+            "rejected_token_id": out_rej,
+            "__valid":           out_valid,
+        }
+
+    ds = ds.map(
+        _tok, batched=True, batch_size=batch_size,
+        remove_columns=ds.column_names,
+        num_proc=num_proc, desc="tokenising",
+    )
+    ds = ds.filter(lambda ex: ex["__valid"], num_proc=num_proc, desc="filter")
+    ds = ds.remove_columns("__valid")
+    if len(ds) == 0:
+        raise ValueError("no ftpo samples survived length / sanity checks")
+
+    return ds.shuffle(seed=3407)

utils/fix_gemma.py

@@ -0,0 +1,116 @@
+#!/usr/bin/env python3
+"""
+Gemma-3 key repair utility
+==========================
+
+Repairs checkpoints whose weight names are in either of the two incorrect
+forms:
+
+1.  model.language_model.embed_tokens.weight          (# leading "model.")
+2.  language_model.embed_tokens.weight                (# missing ".model.")
+
+to the correct form:
+
+    language_model.model.embed_tokens.weight
+
+Usage:
+    python repair_gemma3_keys.py /path/to/checkpoint_dir
+"""
+
+import sys, json, shutil
+from pathlib import Path
+from safetensors.torch import safe_open, save_file
+
+BAD_LEADING  = "model."                    # variant 1
+GOOD_PREFIX  = "language_model."
+GOOD_FULL    = "language_model.model."
+OUT_SUFFIX   = "_repaired"
+
+# ----------------------------------------------------------------------
+# key transformation ----------------------------------------------------
+# ----------------------------------------------------------------------
+
+def fix_key(key: str) -> str:
+    """
+    1) strip a leading "model." if present
+    2) ensure "language_model." is followed by "model."
+    """
+    # step 1 – drop wrapper prefix once
+    if key.startswith(BAD_LEADING):
+        key = key[len(BAD_LEADING):]
+
+    # step 2 – insert ".model." if missing
+    if key.startswith(GOOD_PREFIX) and not key.startswith(GOOD_FULL):
+        key = GOOD_FULL + key[len(GOOD_PREFIX):]
+
+    return key
+
+# ----------------------------------------------------------------------
+# shard processing ------------------------------------------------------
+# ----------------------------------------------------------------------
+
+def repair_shard(src: Path, dst: Path) -> None:
+    """
+    Re-write a .safetensors shard with corrected keys.
+    """
+    corrected = {}
+
+    with safe_open(src, framework="pt", device="cpu") as f:
+        for old_key in f.keys():
+            corrected[fix_key(old_key)] = f.get_tensor(old_key)
+
+    save_file(corrected, dst, metadata={"format": "pt"})
+
+# ----------------------------------------------------------------------
+# driver ----------------------------------------------------------------
+# ----------------------------------------------------------------------
+
+def main():
+    if len(sys.argv) != 2:
+        sys.exit("Usage: repair_gemma3_keys.py <checkpoint_dir>")
+
+    src_dir = Path(sys.argv[1]).expanduser().resolve()
+    if not src_dir.is_dir():
+        sys.exit(f"Directory not found: {src_dir}")
+
+    out_dir = src_dir.with_name(src_dir.name + OUT_SUFFIX)
+    out_dir.mkdir(exist_ok=True)
+
+    index_path = src_dir / "model.safetensors.index.json"
+    if not index_path.is_file():
+        sys.exit("model.safetensors.index.json not found in checkpoint dir.")
+
+    # ---- load index ----------------------------------------------------
+    with open(index_path, "r") as f:
+        index = json.load(f)
+
+    new_weight_map   = {}
+    processed_shards = set()
+
+    # ---- process every tensor key -------------------------------------
+    for old_key, shard_name in index["weight_map"].items():
+        new_key = fix_key(old_key)
+        new_weight_map[new_key] = shard_name
+
+        if shard_name in processed_shards:
+            continue
+        processed_shards.add(shard_name)
+        repair_shard(src_dir / shard_name, out_dir / shard_name)
+
+    index["weight_map"] = new_weight_map
+
+    # ---- write new index ----------------------------------------------
+    with open(out_dir / "model.safetensors.index.json", "w") as f:
+        json.dump(index, f, indent=2)
+
+    # ---- copy auxiliary files -----------------------------------------
+    for fp in src_dir.iterdir():
+        if fp.name == "model.safetensors.index.json" or fp.suffix == ".safetensors":
+            continue
+        shutil.copy2(fp, out_dir / fp.name)
+
+    print(f"✓ Repaired checkpoint written to {out_dir}")
+
+if __name__ == "__main__":
+    main()
+

utils/fs_helpers.py

@@ -0,0 +1,150 @@
+import nltk
+import datetime
+from pathlib import Path
+import logging
+import sys
+import json
+import shutil
+from typing import List
+
+logger = logging.getLogger(__name__)
+
+def set_from_json(path: Path) -> set[str]:
+    """Return a hashable set of strings from a ban-list file that may be
+       either  ["foo bar", …]  or  [["foo bar", 1], …]."""
+    if not path or not path.is_file():
+        return set()
+    try:
+        raw = json.loads(path.read_text("utf-8"))
+    except Exception:
+        return set()
+
+    if not isinstance(raw, list):
+        return set()
+
+    out = set()
+    for item in raw:
+        if isinstance(item, list) and item:        # slop format
+            out.add(str(item[0]))
+        else:                                      # plain string
+            out.add(str(item))
+    return out
+
+
+def merge_custom_bans_into_file(path: Path, extra_items: List[str]) -> None:
+    """Merge extra_items into `path`, preserving original on-disk format."""
+
+    # 1) read whatever is already there
+    try:
+        current_raw = json.loads(path.read_text("utf-8")) if path.exists() else []
+    except Exception:
+        current_raw = []
+
+    if not isinstance(current_raw, list):
+        current_raw = []
+
+    # 2) normalise existing items → plain strings
+    existing: set[str] = set()
+    slop_format = False               # do we need to write back [[phrase,1]] ?
+
+    for entry in current_raw:
+        if isinstance(entry, list):   # slop-phrase style [phrase, freq]
+            slop_format = True
+            if entry:                 # non-empty list
+                existing.add(str(entry[0]))
+        else:                         # plain string
+            existing.add(str(entry))
+
+    # 3) merge & sort
+    merged = sorted(existing | set(map(str, extra_items)))
+
+    # 4) write back in the same shape we found
+    if slop_format:
+        payload = [[p, 1] for p in merged]
+    else:
+        payload = merged
+
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, indent=2, ensure_ascii=False), "utf-8")
+
+
+def download_nltk_resource(resource_id: str, resource_name: str):
+    """Downloads NLTK resource if not found."""
+    try:
+        nltk.data.find(resource_id)
+        logger.debug(f"NLTK '{resource_name}' resource found.")
+    except LookupError:
+        logger.info(f"NLTK '{resource_name}' resource not found. Downloading...")
+        try:
+            nltk.download(resource_name, quiet=True)
+            logger.info(f"NLTK '{resource_name}' resource downloaded.")
+        except Exception as e:
+            logger.warning(f"Could not automatically download NLTK '{resource_name}' resource: {e}. "
+                           "Manual download might be required (e.g., python -m nltk.downloader punkt stopwords)")
+    except Exception as e:
+        logger.warning(f"Error checking NLTK '{resource_name}' resource: {e}.")
+
+def create_experiment_dir(base_dir_path: Path, resume_dir: Path | None = None) -> Path:
+    """
+    Determine the directory that the pipeline should work in.
+
+    • When --resume-from-dir is supplied we MUST use that exact path.
+      If the directory is missing or not a directory, raise immediately.
+
+    • When no resume dir is given, create a new timestamped directory under
+      *base_dir_path* (parents created as needed) and return it.
+    """
+    if resume_dir is not None:
+        if resume_dir.is_dir():
+            logger.info(f"Resuming experiment in existing directory: {resume_dir.resolve()}")
+            return resume_dir
+        # hard-fail: the user explicitly asked to resume here
+        raise FileNotFoundError(
+            f"--resume-from-dir was set to '{resume_dir}', "
+            "but that path does not exist or is not a directory."
+        )
+
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    experiment_dir = base_dir_path / f"run_{timestamp}"
+    experiment_dir.mkdir(parents=True, exist_ok=False)
+    logger.info(f"Created experiment directory: {experiment_dir.resolve()}")
+    return experiment_dir
+
+
+def ensure_antislop_vllm_config_exists(antislop_vllm_dir: Path):
+    """
+    Copies antislop-vllm/config-example.yaml to config.yaml if config.yaml is absent.
+    This is a helper for users, but the main pipeline will pass params via CLI.
+    """
+    cfg_path = antislop_vllm_dir / "config.yaml"
+    example_path = antislop_vllm_dir / "config-example.yaml"
+
+    if not cfg_path.exists():
+        if example_path.exists():
+            try:
+                shutil.copy(example_path, cfg_path)
+                logger.info(f"Copied {example_path} to {cfg_path} for antislop-vllm (user convenience).")
+            except Exception as e:
+                logger.warning(f"Could not copy antislop-vllm config example: {e}")
+        else:
+            logger.debug("antislop-vllm/config-example.yaml not found. No default config.yaml created for it.")
+    else:
+        logger.debug("antislop-vllm/config.yaml already exists.")
+
+
+###############################################################################
+# NLTK helpers
+###############################################################################
+CORE_NLTK_RESOURCES = [
+    ("tokenizers/punkt",       "punkt"),        # sentence + word tokeniser data
+    ("tokenizers/punkt_tab",   "punkt_tab"),    # new in NLTK 3.9+, used by PunktTokenizer
+    ("corpora/stopwords",      "stopwords"),    # obvious
+]
+
+def ensure_core_nltk_resources() -> None:
+    """
+    Download the three NLTK resources our pipeline needs *once* at start-up.
+    Safe to call multiple times – it‘s a no-op if they’re already present.
+    """
+    for resource_id, resource_name in CORE_NLTK_RESOURCES:
+        download_nltk_resource(resource_id, resource_name)

utils/merge_from_lora.py

@@ -0,0 +1,67 @@
+#!/usr/bin/env python3
+"""
+Merge a LoRA adapter (saved by your finetune run) into the full-precision
+base model and write the merged fp16 weights to disk.
+
+Requires:
+  pip install unsloth peft transformers accelerate
+"""
+
+from pathlib import Path
+import torch
+from unsloth import FastLanguageModel
+from peft import PeftModel
+from transformers import AutoTokenizer
+
+# ---------------------------------------------------------------------
+# Adjust these three paths if your directory layout is different.
+# ---------------------------------------------------------------------
+BASE_MODEL   = "unsloth/gemma-3-4b-it"
+ADAPTER_DIR  = (
+    "results/auto_antislop_runs/run_20250608_102159/"
+    "finetuned_model_ftpo_exp01/lora_adapters"
+)
+OUT_DIR      = (
+    "results/auto_antislop_runs/run_20250608_102159/"
+    "finetuned_model_ftpo_exp01/merged_manual_fp16"
+)
+
+# ---------------------------------------------------------------------
+def main() -> None:
+    print("→ loading base model …")
+    base_model, _ = FastLanguageModel.from_pretrained(
+        model_name      = BASE_MODEL,
+        max_seq_length  = 4096,          # keep consistent with training
+        load_in_4bit    = False,         # full-precision
+        dtype           = torch.float16,
+        device_map      = {"": "cpu"},   # CPU merge; change to {"": 0} for GPU
+    )
+
+    print("→ plugging in LoRA adapter …")
+    peft_model = PeftModel.from_pretrained(
+        base_model,
+        ADAPTER_DIR,
+        device_map = {"": "cpu"},
+    )
+
+    print("→ merging and unloading …")
+    merged_model = peft_model.merge_and_unload()  # returns a plain nn.Module
+
+    Path(OUT_DIR).mkdir(parents=True, exist_ok=True)
+
+    print(f"→ saving merged model to {OUT_DIR} …")
+    merged_model.save_pretrained(
+        OUT_DIR,
+        safe_serialization = True,       # *.safetensors shards
+        max_shard_size     = "5GB",
+    )
+
+    # save the tokenizer so the directory is immediately usable
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+    tokenizer.save_pretrained(OUT_DIR)
+
+    print("✓ done")
+
+# ---------------------------------------------------------------------
+if __name__ == "__main__":
+    main()

utils/model_helpers.py

@@ -0,0 +1,228 @@
+# ---------------------------------------------------------------------
+#  helper: ensure Gemma-3 checkpoints use  language_model.model.…  keys
+# ---------------------------------------------------------------------
+import os, json, logging
+from pathlib import Path
+from safetensors.torch import safe_open, save_file
+
+log = logging.getLogger(__name__)
+
+def fix_gemma3_checkpoint(ckpt_dir: str | Path) -> None:
+    """
+    If `ckpt_dir` is a Gemma-3 checkpoint whose tensor keys look like
+        language_model.embed_tokens.weight
+    instead of
+        language_model.model.embed_tokens.weight
+    rewrite the shards and index file in-place.
+
+    No-op when:
+      • model_type ≠ 'gemma3'
+      • keys are already correct
+      • required files are missing
+    """
+    ckpt_dir = Path(ckpt_dir)
+    index_file = ckpt_dir / "model.safetensors.index.json"
+    config_file = ckpt_dir / "config.json"
+    if not index_file.is_file() or not config_file.is_file():
+        return  # nothing to do
+
+    # ── guard: only patch Gemma-3 checkpoints ───────────────────────────
+    try:
+        with open(config_file) as f:
+            cfg = json.load(f)
+        if (cfg.get("model_type") or "").lower() != "gemma3":
+            return
+    except Exception as e:
+        log.warning("Could not read %s (%s); skipping fix.", config_file, e)
+        return
+
+    # ── scan weight map ─────────────────────────────────────────────────
+    with open(index_file) as f:
+        idx = json.load(f)
+
+    wm = idx["weight_map"]
+    broken = [
+        k for k in wm
+        if k.startswith("language_model.") and not k.startswith("language_model.model.")
+    ]
+    if not broken:
+        return                                      # already fine
+
+    log.info("Repairing Gemma-3 key prefixes in %s", ckpt_dir)
+
+    def _fixed(k: str) -> str:
+        if k.startswith("language_model.") and not k.startswith("language_model.model."):
+            return "language_model.model." + k[len("language_model."):]
+        return k
+
+    # ── rewrite every shard exactly once ────────────────────────────────
+    repaired_shards = set()
+    for old_key, shard_name in wm.items():
+        wm[_fixed(old_key)] = wm.pop(old_key)       # update key in dict
+        if shard_name in repaired_shards:
+            continue
+        repaired_shards.add(shard_name)
+
+        src = ckpt_dir / shard_name
+        tmp = ckpt_dir / (shard_name + ".tmp")
+
+        fixed_tensors = {}
+        with safe_open(src, framework="pt", device="cpu") as f:
+            for k in f.keys():
+                fixed_tensors[_fixed(k)] = f.get_tensor(k)
+
+        save_file(fixed_tensors, tmp, metadata={"format": "pt"})
+        tmp.replace(src)                            # atomic overwrite
+
+    # ── write new index ────────────────────────────────────────────────
+    with open(index_file, "w") as f:
+        json.dump(idx, f, indent=2)
+
+    log.info("✓ Gemma-3 checkpoint repaired.")
+
+
+# fully detie lm_head from embeddings so safetensors can flatten
+def detie_lm_head(model):
+    """
+    Untie the logits projection from the input embeddings and register it
+    *exactly* where the model (and loaders like vLLM) expect it.
+
+    Works with HF models whose output head is either `lm_head` or some
+    nested attribute (e.g. `language_model.output_projection` in Gemma-3).
+    """
+    import torch
+    from types import SimpleNamespace
+
+    emb = model.get_input_embeddings()          # nn.Embedding
+    old_head = model.get_output_embeddings()    # whatever Linear HF exposes
+
+    # nothing to do if they are already separate tensors
+    if old_head.weight.data_ptr() != emb.weight.data_ptr():
+        return
+
+    vocab_size, hidden_size = emb.weight.shape
+    new_head = torch.nn.Linear(hidden_size, vocab_size, bias=False)
+    new_head.weight = torch.nn.Parameter(emb.weight.detach().clone())
+    new_head.to(next(model.parameters()).dtype)
+
+    # ------------------------------------------------------------------
+    # find the *attribute path* of the existing output head
+    # ------------------------------------------------------------------
+    path = None
+    for name, module in model.named_modules():
+        if module is old_head:
+            path = name            # e.g. "lm_head" or "language_model.output_projection"
+            break
+    if path is None:               # very unusual, but fall back to "lm_head"
+        path = "lm_head"
+
+    print('!!', name)
+
+    # ------------------------------------------------------------------
+    # install the new head at that path
+    # ------------------------------------------------------------------
+    def set_by_path(root, dotted_name, value):
+        parts = dotted_name.split(".")
+        parent = root
+        for p in parts[:-1]:
+            parent = getattr(parent, p)
+        setattr(parent, parts[-1], value)
+
+    set_by_path(model, path, new_head)
+
+    # HF convenience: if the public attribute `lm_head` *is not* the main path,
+    # mirror it so code expecting `model.lm_head` still works. This does *not*
+    # duplicate weights – both names reference the same nn.Linear instance.
+    #if path != "lm_head":
+    #    model.lm_head = new_head
+
+    model.config.tie_word_embeddings = False
+
+
+# --------------------------------------------------------------------
+#  restore Gemma-3’s weight-tying and ensure only the embed_tokens
+#  key lands in the safetensors index (no lm_head, no duplication)
+# --------------------------------------------------------------------
+# ---------------------------------------------------------------
+#  Gemma-3: keep weight-tying *and* give vLLM the path it wants
+# ---------------------------------------------------------------
+def retie_gemma3_and_prune_alias(model):
+    """
+    Re-establish tying between embeddings and logits projection and ensure
+    the projection is reachable at `language_model.output_projection`.
+    Removes the top-level `lm_head` alias so the serializer never emits
+    an `lm_head.*` key.
+
+    Call just before `save_pretrained(...)`.
+    """
+    import torch.nn as nn
+
+    if (getattr(model.config, "model_type", "") or "").lower() != "gemma3":
+        return  # skip for anything that isn't Gemma-3
+
+    emb  = model.get_input_embeddings()          # nn.Embedding
+    proj = getattr(model, "lm_head", None)       # HF always defines this
+
+    if proj is None or not isinstance(proj, nn.Linear):
+        raise RuntimeError("Could not find lm_head on Gemma-3 model")
+
+    # ── tie weights if they were detied earlier ─────────────────────────
+    if proj.weight.data_ptr() != emb.weight.data_ptr():
+        proj.weight = emb.weight      # share storage again
+    model.config.tie_word_embeddings = True
+
+    # ── ensure wrapper + attribute for vLLM ────────────────────────────
+    # 1. make / fetch `model.language_model`
+    if not hasattr(model, "language_model"):
+        wrapper = nn.Module()
+        model.add_module("language_model", wrapper)
+    else:
+        wrapper = model.language_model
+
+    # 2. register projection inside wrapper
+    wrapper.add_module("output_projection", proj)
+
+    # ── drop the top-level alias so it won't be serialised ─────────────
+    if hasattr(model, "lm_head"):
+        delattr(model, "lm_head")
+
+
+# ------------------------------------------------------------------
+#  Gemma-3 helper: detie + relabel head for vLLM + safetensors
+# ------------------------------------------------------------------
+def prepare_gemma3_for_save(model):
+    """
+    • Makes the output projection an independent tensor if it still shares
+      storage with the embeddings.
+    • Registers it at `language_model.lm_head` (the path vLLM uses).
+    • Deletes the top-level `lm_head` alias so no `lm_head.*` key is saved.
+    • Sets `tie_word_embeddings=False` so Transformers knows they’re untied.
+    """
+    import torch.nn as nn, torch
+
+    if (getattr(model.config, "model_type", "") or "").lower() != "gemma3":
+        return
+
+    emb  = model.get_input_embeddings()
+    head = model.get_output_embeddings()           # this is model.lm_head
+
+    # 1. Detie if they still share storage
+    if head.weight.data_ptr() == emb.weight.data_ptr():
+        vocab, hidden = emb.weight.shape
+        new_head = nn.Linear(hidden, vocab, bias=False)
+        new_head.weight = nn.Parameter(emb.weight.detach().clone())
+        new_head.to(next(model.parameters()).dtype)
+        head = new_head
+
+    # 2. Ensure `language_model` wrapper exists
+    if not hasattr(model, "language_model"):
+        model.add_module("language_model", nn.Module())
+
+    # 3. Register under vLLM path
+    #model.language_model.add_module("lm_head", head)
+
+    # 4. Drop the alias so no `lm_head.*` key lands in the state-dict
+    if hasattr(model, "lm_head"):
+        delattr(model, "lm_head")
+
+    model.config.tie_word_embeddings = False

utils/trainer_dataloaders.py

@@ -0,0 +1,348 @@
+# utils/trainer_dataloaders.py
+
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from datasets import load_dataset, Dataset
+from utils.dataset_helpers import load_ftpo_multi_dataset
+
+if TYPE_CHECKING:
+    from transformers import AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+def load_and_prepare_dataset(config: dict, experiment_run_dir: Path, tokenizer: "AutoTokenizer") -> Dataset | None:
+    """
+    Loads and prepares the dataset based on the finetuning mode specified in the config.
+
+    Args:
+        config (dict): The experiment configuration dictionary.
+        experiment_run_dir (Path): The directory for the current experiment run.
+        tokenizer (AutoTokenizer): The tokenizer to use for processing.
+
+    Returns:
+        Dataset or None: The prepared Hugging Face dataset, or None if loading fails.
+    """
+    mode = config.get("finetune_mode", "ftpo").lower()
+    max_seq_length = config['finetune_max_seq_length']
+    dpo_dataset_hf = None
+
+    if mode == "dpo":
+        # full-sequence preference pairs: rejected is baseline; chosen is the generation made with antislop
+        dataset_path = experiment_run_dir / "dpo_pairs_dataset.jsonl"
+        if not dataset_path.is_file():
+            logger.error(f"DPO dataset not found at {dataset_path}")
+            return None
+
+        dpo_dataset_hf = load_dataset(
+            "json",
+            data_files=str(dataset_path),
+            split="train"
+        )
+
+        # ----------------------------------------------------------
+        #   discard rows whose prompt+continuation would overflow
+        # ----------------------------------------------------------
+        def _within_len(example):
+            prompt_ids = tokenizer(example["prompt"],
+                                add_special_tokens=False).input_ids
+            chosen_ids = tokenizer(example["chosen"],
+                                add_special_tokens=False).input_ids
+            rejected_ids = tokenizer(example["rejected"],
+                                    add_special_tokens=False).input_ids
+            max_len = config['finetune_max_seq_length']
+            return (
+                len(prompt_ids) + len(chosen_ids) <= max_len
+                and
+                len(prompt_ids) + len(rejected_ids) <= max_len
+            )
+
+        before = len(dpo_dataset_hf)
+        dpo_dataset_hf = dpo_dataset_hf.filter(_within_len)
+        after  = len(dpo_dataset_hf)
+        logger.info(f"DPO length filter: kept {after}/{before} examples "
+                    f"(max_seq_len = {config['finetune_max_seq_length']})")
+
+        if after == 0:
+            raise ValueError("every DPO sample exceeded finetune_max_seq_length")
+
+
+        dpo_dataset_hf = dpo_dataset_hf.shuffle(seed=config.get("finetune_shuffle_seed", 3407))
+        max_train = config.get("finetune_max_train_examples")
+        if isinstance(max_train, int) and max_train > 0 and len(dpo_dataset_hf) > max_train:
+            dpo_dataset_hf = dpo_dataset_hf.select(range(max_train))
+            logger.info(f"Capped training dataset to {max_train} examples.")
+
+        # ── filter malformed rows (prompt / chosen / rejected missing) ──
+        req_cols = {"prompt", "chosen", "rejected"}
+        before_len = len(dpo_dataset_hf)
+        dpo_dataset_hf = dpo_dataset_hf.filter(
+            lambda x: all(col in x and x[col] for col in req_cols)
+        )
+        after_len = len(dpo_dataset_hf)
+        if after_len == 0:
+            logger.error("All rows in DPO dataset were filtered out. Check contents.")
+            return None
+        if after_len < before_len:
+            logger.info(f"Filtered out {before_len - after_len} malformed rows; "
+                        f"{after_len} remain.")
+        logger.info(f"DPO dataset ready with {after_len} samples.")
+
+    elif mode == "ftpo":
+        if config.get("finetune_ftpo_dataset"):
+            dataset_path = Path(config["finetune_ftpo_dataset"])
+        else:
+            ftpo_files = sorted(experiment_run_dir.glob("iter_*_ftpo_pairs.jsonl"))
+            if not ftpo_files:
+                logger.error("No ftpo files found for ftpo.")
+                return None
+            dataset_path = ftpo_files[-1]
+
+        # ------------------------------------------------------------------
+        # FTPO dataset with dual regularisation + built-in size cap
+        # ------------------------------------------------------------------
+        dpo_dataset_hf = load_ftpo_multi_dataset(
+            dataset_path,
+            tokenizer,
+            experiment_run_dir = experiment_run_dir,
+            max_seq_len         = max_seq_length,
+            # balance *rejected* tokens
+            rejected_reg_strength = config.get("ftpo_sample_rejected_regularisation_strength", 0.8),
+            # balance *chosen* tokens
+            chosen_reg_strength   = config.get("ftpo_sample_chosen_regularisation_strength", 0.2),
+            # hard floor on |chosen|
+            min_chosen_tokens     = config.get("ftpo_sample_min_chosen_tokens", 3),
+            # overall training-set cap (used for per-token quotas too)
+            max_train_examples    = config.get("finetune_max_train_examples"),
+        )
+
+        # loader already returns a shuffled dataset; an extra shuffle is fine but optional
+        dpo_dataset_hf = dpo_dataset_hf.shuffle(seed=config.get("finetune_shuffle_seed", 3407))
+
+
+        
+        # ──────────────────────────────────────────────────────────────
+        # [DEBUG] Inspect last-5 prompt tokens + chosen / rejected token
+        #         –– prints up to 50 ftpo examples for a quick sanity check.
+        #         –– gated by new config flag `finetune_debug_ftpo_tokens`.
+        # ──────────────────────────────────────────────────────────────
+        if False:
+            sample_n = min(50, len(dpo_dataset_hf))
+            print(f"\n🔎 ftpo debug: showing {sample_n} examples "
+                "(last-5 prompt tokens, chosen ▸ rejected)\n")
+            for i, ex in enumerate(dpo_dataset_hf.select(range(sample_n))):
+                tail_prompt = tokenizer.convert_ids_to_tokens(ex["prompt_ids"][-5:])
+                chosen_tok  = tokenizer.convert_ids_to_tokens([ex["chosen_ids"][0]])[0]
+                rejected_tok = tokenizer.convert_ids_to_tokens([ex["rejected_token_id"]])[0]
+                tail_str = " ".join(tail_prompt)
+                print(f"{i:03d}: … {tail_str}  →  {chosen_tok} ▸ {rejected_tok}")
+            print("\n—— end ftpo debug ——\n")
+
+    elif mode == "dpo_final_token":
+        # ------------------------------------------------------------
+        # 1. Build the raw dataset **exactly** the same way FTPO does
+        # ------------------------------------------------------------
+        if config.get("finetune_ftpo_dataset"):
+            dataset_path = Path(config["finetune_ftpo_dataset"])
+        else:
+            ftpo_files = sorted(experiment_run_dir.glob("iter_*_ftpo_pairs.jsonl"))
+            if not ftpo_files:
+                logger.error("No ftpo files found for dpo_final_token.")
+                return None
+            dataset_path = ftpo_files[-1]
+
+        ftpo_ds = load_ftpo_multi_dataset(
+            dataset_path,
+            tokenizer,
+            experiment_run_dir = experiment_run_dir,
+            max_seq_len                  = max_seq_length,
+            rejected_reg_strength        = config.get("ftpo_sample_rejected_regularisation_strength", 0.8),
+            chosen_reg_strength          = config.get("ftpo_sample_chosen_regularisation_strength", 0.2),
+            min_chosen_tokens            = config.get("ftpo_sample_min_chosen_tokens", 3),
+            max_train_examples           = config.get("finetune_max_train_examples"),
+        )
+
+        # ------------------------------------------------------------
+        # 2. Convert each row into a *single-token* DPO pair
+        # ------------------------------------------------------------
+        pairs = []
+        pad_id = tokenizer.pad_token_id
+
+        for ex in ftpo_ds:
+            # ––– recover the left-padded prompt as text –––
+            prompt_ids = [tid for tid in ex["prompt_ids"] if tid != pad_id]
+            prompt_txt = tokenizer.decode(prompt_ids, skip_special_tokens=False)
+
+            # ––– single-token continuations –––
+            chosen_txt   = tokenizer.decode(
+                [ex["chosen_ids"][0]], skip_special_tokens=False
+            )
+            rejected_txt = tokenizer.decode(
+                [ex["rejected_token_id"]], skip_special_tokens=False
+            )
+
+            pairs.append(
+                {
+                    "prompt":   prompt_txt,
+                    "chosen":   chosen_txt,     # continuation only!
+                    "rejected": rejected_txt,   # continuation only!
+                }
+            )
+
+        dpo_dataset_hf = Dataset.from_list(pairs)
+
+        # ── DEBUG: inspect a few prompt / chosen / rejected triples ──────────────
+        def _show_examples(ds, n=3):
+            for i, ex in enumerate(ds.select(range(n))):
+                print(f"\n── example {i} ──")
+                print("PROMPT:\n",   ex["prompt"])
+                print("CHOSEN:\n",   ex["chosen"])
+                print("REJECTED:\n", ex["rejected"])
+                print("-" * 40)
+
+        _show_examples(dpo_dataset_hf, n=3)
+
+        # ------------------------------------------------------------
+        # 3. Apply the *same* length filter & book-keeping as vanilla DPO
+        # ------------------------------------------------------------
+        def _within_len(example):
+            p = tokenizer(example["prompt"],  add_special_tokens=False).input_ids
+            c = tokenizer(example["chosen"],  add_special_tokens=False).input_ids
+            r = tokenizer(example["rejected"],add_special_tokens=False).input_ids
+            return len(p) + len(c) <= max_seq_length and len(p) + len(r) <= max_seq_length
+
+        before = len(dpo_dataset_hf)
+        dpo_dataset_hf = dpo_dataset_hf.filter(_within_len)
+        after  = len(dpo_dataset_hf)
+        logger.info(f"dpo_final_token length filter: kept {after}/{before} examples "
+                    f"(max_seq_len = {max_seq_length})")
+
+        if after == 0:
+            raise ValueError("every sample exceeded finetune_max_seq_length")
+        
+        max_train = config.get("finetune_max_train_examples")
+        if isinstance(max_train, int) and max_train > 0 and len(dpo_dataset_hf) > max_train:
+            dpo_dataset_hf = dpo_dataset_hf.select(range(max_train))
+            logger.info(f"Capped training dataset to {max_train} examples.")
+
+    # ─────────────────────────────────────────────────────────────────────
+    #  ORPO — single-token pairs (prompt, chosen, rejected)
+    #      Mode value:  "orpo_final_token"
+    # ─────────────────────────────────────────────────────────────────────
+    elif mode == "orpo_final_token":
+        # 1) Construct the FTPO dataset exactly as in the ftpo branch
+        if config.get("finetune_ftpo_dataset"):
+            dataset_path = Path(config["finetune_ftpo_dataset"])
+        else:
+            ftpo_files = sorted(experiment_run_dir.glob("iter_*_ftpo_pairs.jsonl"))
+            if not ftpo_files:
+                logger.error("No ftpo files found for orpo_final_token.")
+                return None
+            dataset_path = ftpo_files[-1]
+
+        ftpo_ds = load_ftpo_multi_dataset(
+            dataset_path,
+            tokenizer,
+            experiment_run_dir = experiment_run_dir,
+            max_seq_len               = max_seq_length,
+            rejected_reg_strength     = config.get("ftpo_sample_rejected_regularisation_strength", 0.8),
+            chosen_reg_strength       = config.get("ftpo_sample_chosen_regularisation_strength", 0.2),
+            min_chosen_tokens         = config.get("ftpo_sample_min_chosen_tokens", 3),
+            max_train_examples        = config.get("finetune_max_train_examples"),
+        )
+
+        # 2) Convert to (prompt, chosen, rejected) triples — one per row
+        pairs   = []
+        pad_id  = tokenizer.pad_token_id
+
+        for ex in ftpo_ds:
+            prompt_ids = [tid for tid in ex["prompt_ids"] if tid != pad_id]
+            prompt_txt = tokenizer.decode(prompt_ids, skip_special_tokens=False)
+
+            chosen_txt   = tokenizer.decode([ex["chosen_ids"][0]], skip_special_tokens=False)
+            rejected_txt = tokenizer.decode([ex["rejected_token_id"]], skip_special_tokens=False)
+
+            pairs.append({"prompt": prompt_txt,
+                        "chosen": chosen_txt,
+                        "rejected": rejected_txt})
+
+        dpo_dataset_hf = Dataset.from_list(pairs)
+
+        # 3) Length filter / shuffle / cap  (reuse helper)
+        def _within_len(ex):
+            p = tokenizer(ex["prompt"],   add_special_tokens=False).input_ids
+            c = tokenizer(ex["chosen"],   add_special_tokens=False).input_ids
+            r = tokenizer(ex["rejected"], add_special_tokens=False).input_ids
+            return len(p) + len(c) <= max_seq_length and len(p) + len(r) <= max_seq_length
+
+        before = len(dpo_dataset_hf)
+        dpo_dataset_hf = dpo_dataset_hf.filter(_within_len)
+        logger.info(f"orpo_final_token length filter: kept {len(dpo_dataset_hf)}/{before} samples")
+
+        dpo_dataset_hf = dpo_dataset_hf.shuffle(seed=config.get("finetune_shuffle_seed", 3407))
+        max_train = config.get("finetune_max_train_examples")
+        if isinstance(max_train, int) and max_train > 0 and len(dpo_dataset_hf) > max_train:
+            dpo_dataset_hf = dpo_dataset_hf.select(range(max_train))
+            logger.info(f"Capped training dataset to {max_train} examples.")
+
+
+    # ─────────────────────────────────────────────────────────────────────
+    #  KTO — single-token, unpaired (prompt, completion, label)
+    #      Mode value:  "kto_final_token"
+    # ─────────────────────────────────────────────────────────────────────
+    elif mode == "kto_final_token":
+        # 1) build the FTPO dataset exactly as before …  (unchanged)
+        ftpo_ds = load_ftpo_multi_dataset(
+            dataset_path,
+            tokenizer,
+            experiment_run_dir = experiment_run_dir,
+            max_seq_len               = max_seq_length,
+            rejected_reg_strength     = config.get("ftpo_sample_rejected_regularisation_strength", 0.8),
+            chosen_reg_strength       = config.get("ftpo_sample_chosen_regularisation_strength", 0.2),
+            min_chosen_tokens         = config.get("ftpo_sample_min_chosen_tokens", 3),
+            max_train_examples        = config.get("finetune_max_train_examples"),
+        )
+
+        # 2) ONE positive + ONE negative row per prompt  ──────────────────────
+        rows, pad_id = [], tokenizer.pad_token_id
+        for ex in ftpo_ds:
+            prompt_ids = [tid for tid in ex["prompt_ids"] if tid != pad_id]
+            prompt_txt = tokenizer.decode(prompt_ids, skip_special_tokens=False)
+
+            if not ex["chosen_ids"]:
+                continue                                 # skip degenerate prompt
+
+            # positive (first chosen id)
+            pos_txt = tokenizer.decode([ex["chosen_ids"][0]], skip_special_tokens=False)
+            rows.append({"prompt": prompt_txt, "completion": pos_txt, "label": True})
+
+            # negative
+            neg_txt = tokenizer.decode([ex["rejected_token_id"]], skip_special_tokens=False)
+            rows.append({"prompt": prompt_txt, "completion": neg_txt, "label": False})
+
+        dpo_dataset_hf = Dataset.from_list(rows)
+
+        # 3) length filter  ───────────────────────────────────────────────────
+        def _within_len(ex):
+            p = tokenizer(ex["prompt"],     add_special_tokens=False).input_ids
+            c = tokenizer(ex["completion"], add_special_tokens=False).input_ids
+            return len(p) + len(c) <= max_seq_length
+
+        before = len(dpo_dataset_hf)
+        dpo_dataset_hf = dpo_dataset_hf.filter(_within_len)
+        logger.info(f"kto_final_token length filter: kept {len(dpo_dataset_hf)}/{before} samples")
+
+        # 4) cap first, then shuffle  ─────────────────────────────────────────
+        max_train = config.get("finetune_max_train_examples")
+        if isinstance(max_train, int) and max_train > 0 and len(dpo_dataset_hf) > max_train:
+            dpo_dataset_hf = dpo_dataset_hf.select(range(max_train))
+            logger.info(f"Capped training dataset to {max_train} examples.")
+
+        #dpo_dataset_hf = dpo_dataset_hf.shuffle(seed=config.get("finetune_shuffle_seed", 3407))
+
+    else:
+        logger.error(f"Unknown finetune_mode '{mode}'. Use 'dpo' or 'ftpo'.")
+        return None
+
+    return dpo_dataset_hf

utils/vllm_manager.py

@@ -0,0 +1,180 @@
+import os
+import sys
+import subprocess
+import time
+import requests
+import logging
+from pathlib import Path, PurePath
+import tempfile, textwrap
+from typing import Optional, List
+
+logger = logging.getLogger(__name__)
+
+def _show_tail(log_path: Path, *, n_lines: int = 300) -> None:
+    """Dump the last *n_lines* of *log_path* to the logger."""
+    try:
+        if log_path.is_file():
+            tail = log_path.read_text(encoding="utf-8").splitlines()[-n_lines:]
+            logger.error(
+                "──── vLLM stdout/stderr (last %d lines) ────\n%s\n────────────────────────────────────────",
+                n_lines, "\n".join(tail),
+            )
+    except Exception as exc:                           # pragma: no cover
+        logger.error("Could not read vLLM log: %s", exc)
+        
+def is_vllm_server_alive(port: int, api_base_path: str = "/v1") -> bool:
+    """Checks if the vLLM server is responsive."""
+    health_url = f"http://127.0.0.1:{port}/health" # Standard vLLM health endpoint
+    # Fallback for older vLLM or if /health is not available, try listing models
+    models_url = f"http://127.0.0.1:{port}{api_base_path.rstrip('/')}/models"
+    
+    try:
+        response = requests.get(health_url, timeout=2)
+        if response.status_code == 200:
+            logger.debug(f"vLLM server on port {port} is healthy (via /health).")
+            return True
+    except (requests.exceptions.ConnectionError, requests.exceptions.ReadTimeout):
+        logger.debug(f"vLLM /health endpoint on port {port} not responding. Trying /models.")
+    
+    try:
+        response = requests.get(models_url, timeout=2)
+        # Expect 200 and a JSON response, typically with a 'data' list
+        if response.status_code == 200 and isinstance(response.json(), dict):
+            logger.debug(f"vLLM server on port {port} is alive (via /models).")
+            return True
+    except (requests.exceptions.ConnectionError, requests.exceptions.ReadTimeout, requests.exceptions.JSONDecodeError):
+        logger.debug(f"vLLM /models endpoint on port {port} not responding or invalid response.")
+        
+    return False
+
+
+def start_vllm_server(
+    model_id: str,
+    port: int,
+    hf_token: Optional[str],
+    cuda_visible_devices: str,
+    gpu_memory_utilization: float,
+    max_model_len: int,
+    dtype: str,
+    vllm_extra_args: Optional[List[str]] = None,
+    extra_env: Optional[dict[str, str]] = None,
+    wait_timeout: int = 720,
+    uvicorn_log_level: str = "error",
+    quiet_stdout: bool = True,
+    log_to_file: bool | Path = True,
+) -> Optional[subprocess.Popen]:
+    """Starts the vLLM API server."""
+    if is_vllm_server_alive(port):
+        logger.info(f"vLLM server already running on port {port}.")
+        return None # Indicate it was already running
+
+    cmd = [
+        #sys.executable, "-m", "vllm.entrypoints.openai.api_server", # Corrected entrypoint
+        #"--model", model_id,
+        "vllm", "serve", model_id,
+        "--port", str(port),
+        "--host", "127.0.0.1",
+        "--gpu-memory-utilization", str(gpu_memory_utilization),
+        "--max-model-len", str(max_model_len),
+        "--dtype", dtype,
+        "--disable-log-requests", # Cleaner logs during generation
+        "--uvicorn-log-level", uvicorn_log_level.lower(),
+    ]
+    if hf_token:
+        cmd.extend(["--hf-token", hf_token])
+    if vllm_extra_args:
+        cmd.extend(vllm_extra_args)
+
+    env = os.environ.copy()
+    env["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices
+    env["HIP_VISIBLE_DEVICES"] = cuda_visible_devices
+    if extra_env:
+        # stringify to avoid type issues
+        env.update({k: str(v) for k, v in extra_env.items()})
+        logger.debug(f"vLLM extra env → {extra_env}")
+
+    logger.info("Starting vLLM server...")
+    logger.info(f"Command: {' '.join(cmd)}")
+    
+    # ------------- stdout / stderr routing -----------------
+    if quiet_stdout:
+        if log_to_file is True:
+            tmp = Path(tempfile.gettempdir()) / f"vllm_{port}_{int(time.time())}.log"
+        elif log_to_file:
+            tmp = Path(PurePath(log_to_file)).expanduser().resolve()
+            tmp.parent.mkdir(parents=True, exist_ok=True)
+        else:
+            tmp = None                                        # swallow completely
+
+        stdout_target = tmp.open("w") if tmp else subprocess.DEVNULL
+        stderr_target = stdout_target if tmp else subprocess.DEVNULL
+        if tmp:
+            logger.info("vLLM stdout/stderr → %s", tmp)
+    else:
+        stdout_target = None          # inherit terminal
+        stderr_target = None
+    # --------------------------------------------------------
+
+    try:
+        server_proc = subprocess.Popen(
+            cmd,
+            env=env,
+            stdout=stdout_target,
+            stderr=stderr_target,
+            text=True,
+        )
+    except FileNotFoundError:
+        logger.error("vLLM not found. Is it installed (pip install vllm)?")
+        return None
+    except Exception as e:
+        logger.error("Failed to start vLLM: %s", e)
+        return None
+
+
+
+    logger.info(f"Waiting for vLLM server to become ready on port {port} (timeout: {wait_timeout}s)...")
+    start_time = time.time()
+    while time.time() - start_time < wait_timeout:
+        if server_proc.poll() is not None: # Process terminated
+            logger.error(f"vLLM server process terminated prematurely with code {server_proc.returncode}.")
+            if quiet_stdout and tmp:
+                _show_tail(tmp)
+            # Try to get some output if possible (might not work well without pipes)
+            # stdout, stderr = server_proc.communicate()
+            # if stdout: logger.error(f"vLLM stdout: {stdout.decode(errors='ignore')}")
+            # if stderr: logger.error(f"vLLM stderr: {stderr.decode(errors='ignore')}")
+            return None
+        if is_vllm_server_alive(port):
+            logger.info(f"🚀 vLLM server ready at http://127.0.0.1:{port}")
+            return server_proc
+        time.sleep(5) # Check every 5 seconds
+
+    logger.error(f"vLLM server failed to start or become healthy within {wait_timeout} seconds.")
+    if server_proc.poll() is None: # If still running, terminate it
+        logger.info("Terminating unresponsive vLLM server process...")
+        server_proc.terminate()
+        try:
+            server_proc.wait(timeout=10)
+        except subprocess.TimeoutExpired:
+            logger.warning("vLLM server did not terminate gracefully, killing.")
+            server_proc.kill()
+    return None
+
+def stop_vllm_server(server_proc: Optional[subprocess.Popen]):
+    """Stops the vLLM server process if it was started by this script."""
+    if server_proc and server_proc.poll() is None: # Check if process exists and is running
+        logger.info("Stopping managed vLLM server...")
+        server_proc.terminate()
+        try:
+            server_proc.wait(timeout=30) # Wait for graceful shutdown
+            logger.info("vLLM server stopped.")
+        except subprocess.TimeoutExpired:
+            logger.warning("vLLM server did not terminate gracefully after 30s, killing.")
+            server_proc.kill()
+            logger.info("vLLM server killed.")
+        except Exception as e:
+            logger.error(f"Error while stopping vLLM server: {e}")
+    elif server_proc and server_proc.poll() is not None:
+        logger.debug("Managed vLLM server was already stopped.")
+    else:
+        logger.debug("No managed vLLM server process to stop.")

utils/whitelist.py

@@ -0,0 +1,163 @@
+# utils/whitelist.py
+"""
+Constructs a global whitelist of strings that must never be placed in
+any ban list.
+
+Sources
+-------
+1. All special-token texts exposed by the model’s tokenizer.
+2. Every phrase (entire line) – and every word inside those phrases –
+   that appears *after* the assistant-message placeholder in a single
+   user→assistant chat-template example.
+3. Optional user-supplied strings from the YAML / CLI configuration.
+"""
+
+from __future__ import annotations
+import json
+import threading
+from pathlib import Path
+from typing import Iterable, Set
+
+from transformers import AutoTokenizer
+from slop_forensics.utils import normalize_text as normalise_keep_marks
+from slop_forensics.utils import extract_words
+
+# ------------------------------------------------------------------------------
+# Helper class
+# ------------------------------------------------------------------------------
+
+class WhitelistBuilder:
+    """
+    Static helper for creating and persisting the whitelist.
+
+    All strings are:
+
+    * converted to lowercase
+    * normalised via `normalise_keep_marks`
+    * deduplicated
+    """
+
+    _tokenizer_cache: dict[str, "AutoTokenizer"] = {}
+    _cache_lock = threading.Lock()
+
+    # ------------------------------------------------------------------ #
+    # Public API                                                         #
+    # ------------------------------------------------------------------ #
+
+    @classmethod
+    def build(
+        cls,
+        model_id: str,
+        *,
+        extra_user_items: Iterable[str] | None = None,
+    ) -> Set[str]:
+        """
+        Parameters
+        ----------
+        model_id
+            Hugging Face model ID or local checkpoint directory.
+        extra_user_items
+            Arbitrary strings provided by the user that must also be whitelisted.
+
+        Returns
+        -------
+        Set[str]
+            Normalised whitelist entries (lower-cased, no duplicates, no blanks).
+        """
+        tokenizer = cls._get_tokenizer(model_id)
+        whitelist: set[str] = set()
+
+        # 1. Special token texts -----------------------------------------
+        special_token_texts = [
+            tokenizer.bos_token,
+            tokenizer.eos_token,
+            tokenizer.unk_token,
+            tokenizer.pad_token,
+            tokenizer.cls_token,
+            tokenizer.sep_token,
+            *(tokenizer.additional_special_tokens or []),
+        ]
+        for raw_text in special_token_texts:
+            if not raw_text:
+                continue
+            cls._add_phrase_and_words(whitelist, raw_text)
+
+        # 2. Tail of the chat template -----------------------------------
+        template_tail_text = cls._get_chat_template_tail(tokenizer)
+        for line in template_tail_text.splitlines():
+            cls._add_phrase_and_words(whitelist, line)
+
+        # 3. User-supplied extras ----------------------------------------
+        if extra_user_items:
+            for item in extra_user_items:
+                cls._add_phrase_and_words(whitelist, str(item))
+
+        # Final clean-up: remove any empty strings that might have slipped in
+        whitelist.discard("")
+        return whitelist
+
+    @classmethod
+    def write(cls, file_path: Path, whitelist: Iterable[str]) -> None:
+        """Write the whitelist to *file_path* as pretty-printed JSON."""
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        file_path.write_text(
+            json.dumps(sorted(whitelist), indent=2, ensure_ascii=False),
+            encoding="utf-8",
+        )
+
+    # ------------------------------------------------------------------ #
+    # Internal helpers                                                   #
+    # ------------------------------------------------------------------ #
+
+    @classmethod
+    def _get_tokenizer(cls, model_id: str):
+        """Thread-safe one-time load of the tokenizer."""
+        with cls._cache_lock:
+            tokenizer = cls._tokenizer_cache.get(model_id)
+            if tokenizer is None:
+                tokenizer = AutoTokenizer.from_pretrained(
+                    model_id, trust_remote_code=True
+                )
+                cls._tokenizer_cache[model_id] = tokenizer
+            return tokenizer
+
+    @staticmethod
+    def _add_phrase_and_words(target_set: set[str], raw_text: str) -> None:
+        """
+        Normalise *raw_text*, add the whole phrase, then add each individual
+        word extracted from the phrase.
+        """
+        normalised = normalise_keep_marks(raw_text)
+        if not normalised:
+            return
+        target_set.add(normalised)
+        target_set.update(extract_words(normalised))
+
+    @staticmethod
+    def _get_chat_template_tail(tokenizer) -> str:
+        """
+        Build one user→assistant chat-template instance and return only the
+        text *after* the assistant placeholder.  That is the scaffold the
+        model tends to emit, so its words must be whitelisted.
+        """
+        placeholder_user = "__USER__"
+        placeholder_assistant = "__ASSISTANT__"
+
+        messages = [
+            {"role": "user",      "content": placeholder_user},
+            {"role": "assistant", "content": placeholder_assistant},
+        ]
+
+        full_template: str = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+
+        assistant_pos = full_template.find(placeholder_assistant)
+        if assistant_pos == -1:
+            # Fallback: return the whole template if the placeholder wasn't found
+            return full_template.strip()
+
+        tail = full_template[assistant_pos + len(placeholder_assistant):].strip()
+        return tail