File size: 10,480 Bytes
a4019dd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 | from __future__ import annotations
import json
import random
from collections import defaultdict
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any
from datasets import Dataset, load_dataset
from .config import DataConfig
@dataclass(slots=True)
class SourcePair:
pair_id: str
text_id: str
source_id: str
dataset_name: str
source: str
model: str
text_type: str
cosine_score: float | None
ai_text: str
human_text: str
@dataclass(slots=True)
class BinaryEvalRow:
row_id: str
text: str
label: int
text_type: str
model: str
source_id: str
def _valid_text(text: str, *, min_text_chars: int) -> bool:
return isinstance(text, str) and len(text.strip()) >= min_text_chars
def _normalize_text(text: Any) -> str:
return str(text or "").strip()
def _find_local_arrow_file(root: Path, split: str) -> Path:
direct_path = root / f"editlens_iclr-{split}.arrow"
if direct_path.exists():
return direct_path
matches = sorted(root.rglob(f"editlens_iclr-{split}.arrow"))
if not matches:
raise FileNotFoundError(f"Missing local dataset file for split={split!r} under {root}")
return matches[0]
def _pick_best_row(rows: list[dict[str, Any]], *, text_key: str) -> dict[str, Any]:
if not rows:
raise ValueError("Cannot pick from an empty row list.")
rows = sorted(
rows,
key=lambda row: (
_normalize_text(row.get("prompt")) == "",
_normalize_text(row.get("title")) == "",
_normalize_text(row.get(text_key)) == "",
),
)
return rows[0]
def load_pangram_rows(config: DataConfig, *, split: str) -> Dataset:
if config.pangram.local_dataset_path is not None:
arrow_path = _find_local_arrow_file(config.pangram.local_dataset_path, split)
return Dataset.from_file(str(arrow_path))
return load_dataset(config.pangram.dataset_name, split=split)
def load_raid_rows(config: DataConfig, *, split: str) -> Dataset:
return load_dataset(config.raid.dataset_name, split=split)
def build_pangram_binary_pairs(config: DataConfig) -> list[SourcePair]:
rows = load_pangram_rows(config, split=config.pangram.dataset_split)
ai_rows = []
human_by_text_id: dict[str, list[dict[str, Any]]] = defaultdict(list)
human_by_source_id: dict[str, list[dict[str, Any]]] = defaultdict(list)
for row in rows:
text = _normalize_text(row.get("text"))
if not _valid_text(text, min_text_chars=config.min_text_chars):
continue
text_type = _normalize_text(row.get("text_type"))
if text_type in config.pangram.human_text_types:
text_id = _normalize_text(row.get("text_id"))
source_id = _normalize_text(row.get("source_id"))
if text_id:
human_by_text_id[text_id].append(dict(row))
if source_id:
human_by_source_id[source_id].append(dict(row))
elif text_type in config.pangram.ai_text_types:
ai_rows.append(dict(row))
pairs: list[SourcePair] = []
for row in ai_rows:
pair_source_id = _normalize_text(row.get("source_id"))
if not pair_source_id:
continue
candidates = human_by_text_id.get(pair_source_id)
if not candidates:
candidates = human_by_source_id.get(pair_source_id)
if not candidates:
continue
human_row = _pick_best_row(candidates, text_key="text")
pairs.append(
SourcePair(
pair_id=f"pangram::{pair_source_id}::{_normalize_text(row.get('text_id'))}",
text_id=_normalize_text(row.get("text_id")),
source_id=pair_source_id,
dataset_name="pangram",
source=_normalize_text(row.get("source")),
model=_normalize_text(row.get("model")),
text_type=_normalize_text(row.get("text_type")),
cosine_score=float(row["cosine_score"]) if row.get("cosine_score") is not None else None,
ai_text=_normalize_text(row.get("text")),
human_text=_normalize_text(human_row.get("text")),
)
)
return pairs
def build_raid_binary_pairs(config: DataConfig) -> list[SourcePair]:
rows = load_raid_rows(config, split=config.raid.dataset_split)
human_rows: list[dict[str, Any]] = []
ai_rows: list[dict[str, Any]] = []
for row in rows:
if config.raid.require_attack_none and _normalize_text(row.get("attack")) not in {"", "none"}:
continue
text = _normalize_text(row.get("generation"))
if not _valid_text(text, min_text_chars=config.min_text_chars):
continue
model = _normalize_text(row.get("model"))
if model == config.raid.human_model_name:
human_rows.append(dict(row))
else:
ai_rows.append(dict(row))
human_by_id: dict[str, list[dict[str, Any]]] = defaultdict(list)
human_by_source_id: dict[str, list[dict[str, Any]]] = defaultdict(list)
for row in human_rows:
row_id = _normalize_text(row.get("id"))
source_id = _normalize_text(row.get("source_id"))
if row_id:
human_by_id[row_id].append(row)
if source_id:
human_by_source_id[source_id].append(row)
pairs: list[SourcePair] = []
for row in ai_rows:
pair_source_id = _normalize_text(row.get("source_id"))
if not pair_source_id:
continue
candidates = human_by_id.get(pair_source_id)
if not candidates:
candidates = human_by_source_id.get(pair_source_id)
if not candidates:
continue
human_row = _pick_best_row(candidates, text_key="generation")
pairs.append(
SourcePair(
pair_id=f"raid::{pair_source_id}::{_normalize_text(row.get('model'))}::{_normalize_text(row.get('id'))}",
text_id=_normalize_text(row.get("id")),
source_id=pair_source_id,
dataset_name="raid",
source=_normalize_text(row.get("domain")),
model=_normalize_text(row.get("model")),
text_type="ai_generated",
cosine_score=None,
ai_text=_normalize_text(row.get("generation")),
human_text=_normalize_text(human_row.get("generation")),
)
)
return pairs
def _take_pairs(
pairs: list[SourcePair],
*,
take: int,
seed: int,
) -> tuple[list[SourcePair], list[SourcePair]]:
rng = random.Random(seed)
shuffled = list(pairs)
rng.shuffle(shuffled)
if len(shuffled) < take:
raise ValueError(f"Need at least {take} pairs, found {len(shuffled)}.")
return shuffled[:take], shuffled[take:]
def build_training_and_eval_splits(
config: DataConfig,
*,
seed: int,
) -> tuple[list[SourcePair], list[SourcePair], list[SourcePair], dict[str, int]]:
source_pools: dict[str, list[SourcePair]] = {}
if config.pangram.enabled:
source_pools["pangram"] = build_pangram_binary_pairs(config)
if config.raid.enabled:
source_pools["raid"] = build_raid_binary_pairs(config)
train_pairs: list[SourcePair] = []
holdout_candidates: list[SourcePair] = []
raid_eval_pairs: list[SourcePair] = []
metadata = {f"{name}_pairs_available": len(pairs) for name, pairs in source_pools.items()}
if config.raid.enabled:
raid_eval_pairs, remaining_raid = _take_pairs(
source_pools["raid"],
take=config.raid.eval_holdout_pairs,
seed=seed + 100,
)
source_pools["raid"] = remaining_raid
if config.pangram.enabled:
selected, remaining = _take_pairs(
source_pools["pangram"],
take=config.pangram.train_pairs,
seed=seed + 1,
)
train_pairs.extend(selected)
holdout_candidates.extend(remaining)
if config.raid.enabled:
selected, remaining = _take_pairs(
source_pools["raid"],
take=config.raid.train_pairs,
seed=seed + 2,
)
train_pairs.extend(selected)
holdout_candidates.extend(remaining)
holdout_pairs, _ = _take_pairs(
holdout_candidates,
take=config.training_holdout_pairs,
seed=seed + 3,
)
random.Random(seed + 4).shuffle(train_pairs)
random.Random(seed + 5).shuffle(holdout_pairs)
metadata.update(
{
"train_pairs_from_pangram": sum(pair.dataset_name == "pangram" for pair in train_pairs),
"train_pairs_from_raid": sum(pair.dataset_name == "raid" for pair in train_pairs),
"holdout_pairs_from_pangram": sum(pair.dataset_name == "pangram" for pair in holdout_pairs),
"holdout_pairs_from_raid": sum(pair.dataset_name == "raid" for pair in holdout_pairs),
"raid_eval_pairs_from_raid": len(raid_eval_pairs),
}
)
return train_pairs, holdout_pairs, raid_eval_pairs, metadata
def save_pairs(path: Path, pairs: list[SourcePair]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps([asdict(pair) for pair in pairs], indent=2), encoding="utf-8")
def load_pairs(path: Path) -> list[SourcePair]:
rows = json.loads(path.read_text(encoding="utf-8"))
return [SourcePair(**row) for row in rows]
def load_binary_eval_rows(
config: DataConfig,
*,
split: str,
positive_text_types: set[str],
negative_text_types: set[str],
) -> list[BinaryEvalRow]:
rows = load_pangram_rows(config, split=split)
payload: list[BinaryEvalRow] = []
for index, row in enumerate(rows):
text_type = str(row.get("text_type", "")).strip()
if text_type in positive_text_types:
label = 1
elif text_type in negative_text_types:
label = 0
else:
continue
text = str(row.get("text", "")).strip()
if not _valid_text(text, min_text_chars=config.min_text_chars):
continue
payload.append(
BinaryEvalRow(
row_id=str(row.get("text_id", index)),
text=text,
label=label,
text_type=text_type,
model=str(row.get("model", "")),
source_id=str(row.get("source_id", "")),
)
)
return payload
|