366 lines
15 KiB
Python
366 lines
15 KiB
Python
# 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 _chosen_target_quotas(
|
||
chosen_counts: Counter[str],
|
||
strength: float,
|
||
) -> dict[str, int]:
|
||
"""Return per-token occurrence caps for chosen-token regularisation."""
|
||
if not chosen_counts or strength <= 0:
|
||
return dict(chosen_counts)
|
||
|
||
# Prevent the largest outliers from dominating the distribution before
|
||
# applying the smoother median-threshold regularisation.
|
||
if len(chosen_counts) >= 10:
|
||
cap_value = sorted(chosen_counts.values(), reverse=True)[9]
|
||
capped = {
|
||
token: min(count, cap_value)
|
||
for token, count in chosen_counts.items()
|
||
}
|
||
else:
|
||
capped = dict(chosen_counts)
|
||
|
||
median = float(np.median(list(capped.values())))
|
||
return {
|
||
token: int(round(
|
||
count
|
||
if count <= median
|
||
else count * (median / count) ** strength
|
||
))
|
||
for token, count in capped.items()
|
||
}
|
||
|
||
|
||
def _trim_chosen_to_quotas(
|
||
rows: list[dict],
|
||
quotas: dict[str, int],
|
||
rng: np.random.Generator,
|
||
) -> list[dict]:
|
||
"""Uniformly retain chosen-token occurrences up to their global quotas."""
|
||
counts = Counter(
|
||
token
|
||
for row in rows
|
||
for token in (row.get("multi_chosen_decoded") or [])
|
||
)
|
||
if all(quotas.get(token, count) >= count for token, count in counts.items()):
|
||
return [dict(row) for row in rows]
|
||
|
||
# Randomise both row and slot traversal so quota allocation does not
|
||
# systematically favour early examples or a token's probability rank.
|
||
seen: Counter[str] = Counter()
|
||
trimmed: list[dict | None] = [None] * len(rows)
|
||
for row_idx_raw in rng.permutation(len(rows)):
|
||
row_idx = int(row_idx_raw)
|
||
row = rows[row_idx]
|
||
decoded = row.get("multi_chosen_decoded") or []
|
||
keep: set[int] = set()
|
||
for slot_idx_raw in rng.permutation(len(decoded)):
|
||
slot_idx = int(slot_idx_raw)
|
||
token = decoded[slot_idx]
|
||
if seen[token] < max(0, quotas.get(token, counts[token])):
|
||
keep.add(slot_idx)
|
||
seen[token] += 1
|
||
|
||
new_row = dict(row)
|
||
new_row["multi_chosen_decoded"] = [
|
||
token for slot_idx, token in enumerate(decoded) if slot_idx in keep
|
||
]
|
||
|
||
raw = row.get("multi_chosen_raw")
|
||
if isinstance(raw, list) and len(raw) == len(decoded):
|
||
new_row["multi_chosen_raw"] = [
|
||
token for slot_idx, token in enumerate(raw) if slot_idx in keep
|
||
]
|
||
|
||
trimmed[row_idx] = new_row
|
||
|
||
return [row for row in trimmed if row is not None]
|
||
|
||
|
||
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")
|
||
|
||
tgt_chosen = _chosen_target_quotas(
|
||
chosen_cts_orig,
|
||
chosen_reg_strength,
|
||
)
|
||
|
||
# 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),
|
||
)
|
||
|
||
rows = _trim_chosen_to_quotas(rows, tgt_chosen, rng)
|
||
chosen_cts_trimmed = Counter(
|
||
token
|
||
for row in rows
|
||
for token in (row["multi_chosen_decoded"] or [])
|
||
)
|
||
_log_top(chosen_cts_trimmed, "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
|
||
|
||
final_chosen_counts = Counter(
|
||
token
|
||
for row in rows
|
||
for token in (row["multi_chosen_decoded"] or [])
|
||
)
|
||
_log_top(final_chosen_counts, "FINAL CHOSEN TOKENS")
|
||
|
||
# ── 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)
|