File size: 5,637 Bytes
e58f648 | 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 | """RAID loader (liamdugan/raid) — multi-domain, multi-generator AI-text data.
RAID is sorted by (domain, model, attack). Schema: text=`generation`, label=`model`
('human' = human, else = AI), `domain`, `attack` ('none' = clean).
IMPORTANT: HF only PARTIALLY converts RAID to parquet (it's huge), and because RAID is
domain-sorted the partial parquet contains only the *early* domains. So we never hardcode which
domains exist — `load_raid_records` returns domain-tagged records for whatever is present, and the
caller splits train/OOD from the domains actually observed. Shards are cached + column-pruned.
"""
from __future__ import annotations
import os
import random
import tempfile
from collections import Counter, defaultdict
import requests
_PARQUET_API = "https://datasets-server.huggingface.co/parquet?dataset=liamdugan/raid"
_CACHE_DIR = os.path.join(tempfile.gettempdir(), "raid_parquet")
_NEEDED = ["generation", "model", "domain", "attack"]
def _shard_files(split: str):
meta = requests.get(_PARQUET_API, timeout=60).json()
files = [f for f in meta.get("parquet_files", [])
if f.get("split") == split and f.get("config") == "raid"]
if not files:
raise RuntimeError(f"No RAID parquet shards for split={split}; api keys={list(meta)}")
files.sort(key=lambda f: f["url"])
print(f"[raid] split='{split}': {len(files)} shard(s), partial-conversion="
f"{requests.get(_PARQUET_API, timeout=60).json().get('partial')}")
return files
def _download_shard(f) -> str:
os.makedirs(_CACHE_DIR, exist_ok=True)
path = os.path.join(_CACHE_DIR, f["split"] + "_" + f["filename"])
if os.path.exists(path) and os.path.getsize(path) > 1000:
return path
print(f"[raid] downloading {f['filename']} ({f.get('size',0)/1e6:.0f} MB)...")
b = requests.get(f["url"], timeout=900).content
with open(path, "wb") as fh:
fh.write(b)
return path
def _read_shard(path):
import pandas as pd
try:
import pyarrow.parquet as pq
names = pq.read_schema(path).names
cols = [c for c in _NEEDED if c in names]
if len(cols) == len(_NEEDED):
return pd.read_parquet(path, columns=cols)
except Exception:
pass
return pd.read_parquet(path)
def load_raid_records(n_per_domain_class=300, split="train", attack="none",
min_words=30, max_words=600, max_shards=10, seed=42, verbose=True):
"""Return a list of {text, label(0=human,1=ai), domain, model} tagged records, balanced up to
`n_per_domain_class` per (domain, label) across whatever domains the partial parquet contains."""
counts = defaultdict(int)
records = []
raw_domains, raw_attacks = set(), set()
for si, f in enumerate(_shard_files(split)[:max_shards]):
df = _read_shard(_download_shard(f))
for col in _NEEDED:
if col not in df.columns:
raise RuntimeError(f"RAID shard missing '{col}'; columns={list(df.columns)}")
dom = df["domain"].astype(str).str.strip().str.lower()
atk = df["attack"].astype(str).str.strip().str.lower()
mdl = df["model"].astype(str).str.strip().str.lower()
raw_domains.update(dom.unique().tolist()); raw_attacks.update(atk.unique().tolist())
if si == 0 and verbose:
print(f"[raid] shard0 shape={df.shape} domains={sorted(set(dom))[:12]} "
f"attacks={sorted(set(atk))[:6]}")
mask = df["generation"].notna()
if attack is not None:
mask &= (atk == str(attack).lower())
wc = df["generation"].str.split().str.len()
mask &= (wc >= min_words) & (wc <= max_words)
idx = df.index[mask].tolist()
random.Random(seed + si).shuffle(idx)
gen = df["generation"]
for i in idx:
d = dom[i]; label = 0 if mdl[i] == "human" else 1
if counts[(d, label)] >= n_per_domain_class:
continue
counts[(d, label)] += 1
records.append({"text": gen[i].strip(), "label": label, "domain": d, "model": mdl[i]})
if verbose:
print(f"[raid] after {f['filename']}: domains={sorted({k[0] for k in counts})}")
if not records:
raise RuntimeError(f"RAID collected 0 records. pre-filter domains={sorted(raw_domains)} "
f"attacks={sorted(raw_attacks)} attack-requested='{attack}'.")
if verbose:
print(f"[raid] records/domain: {dict(Counter(r['domain'] for r in records))}")
print(f"[raid] label balance: {dict(Counter(r['label'] for r in records))}")
return records
def load_raid(n_per_class=1500, split="train", domains=None, holdout_domains=None,
attack="none", min_words=30, max_words=600, max_shards=10, seed=42, verbose=True):
"""(human, ai) convenience wrapper over load_raid_records, with optional domain include/holdout."""
recs = load_raid_records(n_per_domain_class=max(50, n_per_class // 3), split=split, attack=attack,
min_words=min_words, max_words=max_words, max_shards=max_shards,
seed=seed, verbose=verbose)
if domains:
ds = {d.lower() for d in domains}; recs = [r for r in recs if r["domain"] in ds]
if holdout_domains:
hs = {d.lower() for d in holdout_domains}; recs = [r for r in recs if r["domain"] not in hs]
human = [r["text"] for r in recs if r["label"] == 0]
ai = [r["text"] for r in recs if r["label"] == 1]
random.Random(seed + 2).shuffle(human); random.Random(seed + 3).shuffle(ai)
return human[:n_per_class], ai[:n_per_class]
|