RL-Explorer / app /envs /processors.py
AdithyaSK's picture
AdithyaSK HF Staff
Deploy HF RL Explorer
da5cba1 verified
Raw History Blame Contribute Delete
20.9 kB
"""The processors that ship with the explorer, most specific first. Add yours here (see CUSTOM_ENVS.md):
class MyFormat(Processor):
id, name, framework = "my-format", "My format", "My format"
about = "One sentence: what a row is and how it's graded."
def match(self, ds): return 0.9 if ds.has("my_column") else 0
def view(self, row, i, ds, roles): ... return {"title", "id", "chips", "sections", "withheld", "glance", "run"?}
and list it in PROCESSORS. `card` and `roles` have sensible defaults (base.Processor), so most only need `view`.
"""
from __future__ import annotations
import base64
import binascii
import gzip
import io
import posixpath
import re
import tarfile
import tomllib
from typing import Any
from .base import Dataset, Processor, first_line, is_messages, parse_json, section, text_of, withhold
from .nemogym import HIDE as NEMO_HIDE
from .nemogym import NemoGym
from .verifiers_rows import SWE_HIDE
from .verifiers_rows import Verifiers
from .verl_rows import HIDE as VERL_HIDE
from .verl_rows import Verl
# ── Harbor tasks packed into rows (a tar.gz per row), e.g. open-thoughts/TaskTrove ──────────────────────────
BLOB_COLS = ("task_binary", "task_tar", "tar", "archive", "task_archive", "files_tar", "bundle")
MAX_TAR = 25 * 2**20
def unpack(blob: Any) -> dict[str, bytes]:
"""A row's packed task as {path: bytes}: base64 text, a list of byte values, or bytes; tar, tar.gz or zip."""
if isinstance(blob, str):
try:
data = base64.b64decode(blob, validate=False)
except (binascii.Error, ValueError):
return {}
elif isinstance(blob, list):
data = bytes(x & 0xFF for x in blob[:MAX_TAR])
elif isinstance(blob, (bytes, bytearray)):
data = bytes(blob)
else:
return {}
if len(data) > MAX_TAR:
return {}
files: dict[str, bytes] = {}
try:
if data[:2] == b"PK":
import zipfile
with zipfile.ZipFile(io.BytesIO(data)) as z:
for n in z.namelist()[:3000]:
if not n.endswith("/") and z.getinfo(n).file_size < 4 * 2**20:
files[n] = z.read(n)
else:
with tarfile.open(fileobj=io.BytesIO(data), mode="r:*") as t:
total = 0
for m in t.getmembers()[:3000]:
if m.isfile() and m.size < 4 * 2**20 and total < 60 * 2**20:
f = t.extractfile(m)
if f:
files[m.name] = f.read()
total += m.size
except (tarfile.TarError, OSError, EOFError, gzip.BadGzipFile, ValueError):
return {}
# a folder around everything (task_0/task.toml): strip it, so paths read like a Harbor task's own
names = [n.lstrip("./") for n in files]
files = dict(zip(names, files.values()))
tomls = [n for n in names if posixpath.basename(n) == "task.toml"]
root = posixpath.dirname(min(tomls, key=lambda n: n.count("/"))) if tomls else ""
if root:
files = {n[len(root) + 1:]: b for n, b in files.items() if n.startswith(root + "/")}
return {n: b for n, b in files.items() if n and not any(p.startswith("..") for p in n.split("/"))}
class PackedHarbor(Processor):
id, name, framework = "harbor-packed", "Harbor tasks in rows", "Harbor (packed)"
about = "Each row packs one whole Harbor task (task.toml, instruction.md, tests/, environment/) into an archive; it's unpacked when opened."
def blob(self, ds: Dataset) -> str | None:
return next((c for c in ds.columns if c.lower() in BLOB_COLS), None)
def match(self, ds: Dataset) -> float:
col = self.blob(ds)
if not col or not any(r.get(col) for r in ds.sample[:20]):
return 0
if Traces().conv(ds): # recorded rollouts that also pack the task they ran on: traces, with the task inside
return 0.5
return 0.97 if any(t in ds.tags for t in ("library:harbor", "harbor")) else 0.8
def roles(self, ds: Dataset) -> dict[str, Any]:
"""The preamble the first tasks' instructions share ("Environment Setup … # Terminal Automation Request"),
so a title can be the task's own first line."""
from .base import _common_prefix
col = self.blob(ds)
texts = []
for r in ds.sample[:20]:
b = r.get(col)
if isinstance(b, (str, list, bytes)) and len(b) < 2 * 2**20:
texts.append(unpack(b).get("instruction.md", b"").decode("utf-8", "replace"))
return {"task": None, "id": "path", "title": None, "answer": [], "grading": [], "environment": [],
"_prefix": _common_prefix(texts) if len([t for t in texts if t]) >= 3 else ""}
def _title(self, instruction: str, meta: dict[str, Any], roles: dict[str, Any], fallback: str) -> tuple[str, str]:
p = roles.get("_prefix") or ""
own = instruction[len(p):].lstrip() if p and instruction.startswith(p) else instruction
title = str(meta.get("title") or "") or first_line(own) or fallback
return title, re.sub(r"\s+", " ", own).strip()
def card(self, row: dict[str, Any], i: int, roles: dict[str, Any]) -> dict[str, Any]:
ident = next((str(row[c]) for c in ("path", "task_id", "id", "name") if row.get(c) is not None), None)
blob = next((row[c] for c in row if c.lower() in BLOB_COLS), None)
files = unpack(blob) if isinstance(blob, (str, list, bytes)) and len(blob) < 2 * 2**20 else {}
instruction = files.get("instruction.md", b"").decode("utf-8", "replace")
meta: dict[str, Any] = {}
try:
meta = (tomllib.loads(files["task.toml"].decode("utf-8", "replace")).get("metadata") or {}) if "task.toml" in files else {}
except tomllib.TOMLDecodeError:
pass
title, rest = self._title(instruction, meta, roles, ident or f"Task {i}")
chips = [str(meta[k]) for k in ("difficulty", "category") if isinstance(meta.get(k), str)][:2]
snippet = rest[len(title.rstrip("…")):].strip() if rest.startswith(title.rstrip("…")[:60]) else rest
return {"i": i, "id": ident, "title": title[:200], "snippet": snippet[:260], "chips": chips,
"_text": "" if meta.get("title") or roles.get("_prefix") else instruction}
def view(self, row: dict[str, Any], i: int, ds: Dataset, roles: dict[str, Any]) -> dict[str, Any]:
from .. import catalog
files = unpack(row.get(self.blob(ds)))
ident = next((str(row[c]) for c in ("path", "task_id", "id", "name") if row.get(c) is not None), None)
if not files:
return {"title": ident or f"Task {i}", "id": ident, "chips": [], "withheld": [], "glance": [["Row", f"{i:,}"]],
"sections": [section("task", "The task", "This row's archive couldn't be opened (too large, or not a tar/zip).", kind="markdown")]}
texts = {n: b.decode("utf-8", "replace") for n, b in files.items() if b"\x00" not in b[:4096]}
toml_text = texts.get("task.toml", "")
try:
doc = tomllib.loads(toml_text) if toml_text else {}
except tomllib.TOMLDecodeError:
doc = {}
meta = doc.get("metadata") or {}
instruction = texts.get("instruction.md", "")
title, _ = self._title(instruction, meta, roles, ident or f"Task {i}")
sections = [section("task", "The task", instruction or "This task has no instruction.md.", kind="markdown", note="instruction.md, word for word")]
test_sh = texts.get("tests/test.sh")
if test_sh and not catalog._withheld("tests/test.sh", test_sh):
sections.append(section("grading", "How it's graded", {"path": "tests/test.sh", "text": catalog._mask_answers("tests/test.sh", test_sh)},
kind="code", note="tests/test.sh, run after the agent"))
docker = texts.get("environment/Dockerfile")
if docker:
sections.append(section("environment", "What it runs in", {"path": "environment/Dockerfile", "text": docker}, kind="code"))
clean_meta, gone = withhold({k: v for k, v in meta.items() if not catalog.CONTACT_KEY.search(k)})
if clean_meta:
sections.append(section("metadata", "Metadata", clean_meta, kind="value"))
tree = [{"path": n, "size": len(b), "withheld": catalog._withheld(n, texts.get(n))} for n, b in sorted(files.items())]
sections.append(section("files", "Files", {"tree": tree}, kind="files", note=f"{len(tree)} files, unpacked from the row"))
env = doc.get("environment") or {}
from .. import catalog as _catalog
look = {"env": {"image": env.get("docker_image"), "network": _catalog._network(env),
"compose": any(re.match(r"environment/(docker-)?compose\.ya?ml$", n) for n in files)}}
runnable = _catalog.runnable(look, texts)
return {"packed_run": {"runnable": runnable, "verifier_env": sorted(((doc.get("verifier") or {}).get("env") or {}).keys()),
"toml": toml_text or ""},
"title": title[:200], "id": ident, "chips": [str(meta[k]) for k in ("difficulty", "category") if isinstance(meta.get(k), str)][:3],
"sections": sections, "withheld": gone + [t["path"] for t in tree if t["withheld"]],
"glance": [["Row", f"{i:,}"], ["Task", ident or "–"], ["Runs in", env.get("docker_image") or ("a Dockerfile" if docker else "–")],
["Graded by", "tests/test.sh" if test_sh else "–"], ["Files", f"{len(tree):,}"]]}
def file(self, row: dict[str, Any], ds: Dataset, rel: str) -> dict[str, Any]:
"""One file of the unpacked task, for the file viewer (answers withheld, as for any Harbor task)."""
from .. import catalog
files = unpack(row.get(self.blob(ds)))
if rel not in files:
return {"path": rel, "error": "not a file in this task"}
b = files[rel]
if b"\x00" in b[:4096]:
return {"path": rel, "size": len(b), "binary": True}
text = b.decode("utf-8", "replace")
if catalog._withheld(rel, text):
return {"path": rel, "size": len(b), "withheld": True}
if rel == "task.toml":
text = catalog._masked_toml(text)
return {"path": rel, "size": len(b), "text": catalog._mask_answers(rel, text)[:512 * 1024], "truncated": len(text) > 512 * 1024}
# ── XiaomiMiMo/MiMo-V2.6-RL-oss, the raw form the MiMo explorer reads ────────────────────────────────────
MIMO_HIDE = re.compile(r"^(patch|gold_patch|fix_patch|solution_patch|poc|poc_\w+|reference_\w+)$", re.I)
def _slug(task_id: str) -> str: # mimo_harbor.adapter.slug: how the Harbor conversion names its task folders
return "".join(c if c.isalnum() or c in "-_." else "-" for c in task_id.lower()).strip("-.")
class MiMo(Processor):
id, name, framework = "mimo", "MiMo-V2.6 RL", "MiMo raw"
about = ("Xiaomi's MiMo-V2.6 RL environments as released: one verl-style row per environment (a chat prompt, a "
"reward style, and an instance with its image, tests or rubric). Each has a Harbor twin that runs here.")
def match(self, ds: Dataset) -> float:
if ds.spec.lower() == "xiaomimimo/mimo-v2.6-rl-oss":
return 1.0
# MiMo's rows are verl-shaped with an agent_name; what sets them apart is the instance packed in extra_info
packed = any(isinstance(parse_json(r.get("extra_info")), dict) and "instance_json" in parse_json(r.get("extra_info")) for r in ds.sample[:10])
return 0.92 if ds.has("data_source", "ability", "agent_name", "extra_info") and packed else 0
def instance(self, row: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
extra = parse_json(row.get("extra_info")) or {}
inst = parse_json(extra.get("instance_json")) if isinstance(extra, dict) else {}
return (extra if isinstance(extra, dict) else {}), (inst if isinstance(inst, dict) else {})
def task_id(self, row: dict[str, Any], config: str) -> str | None:
extra, inst = self.instance(row)
tid = inst.get("instance_id") or extra.get("instance_id")
if not tid and config == "music" and extra.get("src_id") is not None:
tid = f"music-{extra['src_id']}"
return str(tid) if tid else None
def card(self, row: dict[str, Any], i: int, roles: dict[str, Any]) -> dict[str, Any]:
extra, inst = self.instance(row)
prompt = text_of(row.get("prompt"))
title = inst.get("display_title") or first_line(inst.get("description") or "") or first_line(prompt) or f"Row {i}"
chips = [x for x in (row.get("ability"), inst.get("category"), inst.get("subcategory"), extra.get("lang")) if isinstance(x, str) and x][:4]
tid = inst.get("instance_id") or extra.get("instance_id") or (f"music-{extra['src_id']}" if extra.get("src_id") is not None else None)
return {"i": i, "id": tid, "title": str(title)[:200], "snippet": re.sub(r"\s+", " ", prompt)[:260] if first_line(prompt) != title else "", "chips": chips}
def view(self, row: dict[str, Any], i: int, ds: Dataset, roles: dict[str, Any]) -> dict[str, Any]:
config = roles.get("_config", "")
extra, inst = self.instance(row)
inst = {k: v for k, v in inst.items() if not MIMO_HIDE.match(k)}
clean_inst, gone = withhold({k: parse_json(v) for k, v in inst.items()})
prompt = parse_json(row.get("prompt"))
rm = parse_json(row.get("reward_model")) or {}
sections = [section("task", "The task", prompt, note="the prompt, as the model got it")]
grading: dict[str, Any] = {"Reward style": rm.get("style") if isinstance(rm, dict) else None}
for k in ("test_command", "verifier_timeout_sec", "description"):
if clean_inst.get(k) is not None:
grading[k] = clean_inst.pop(k)
tests = clean_inst.pop("test_patch", None)
rubric = clean_inst.pop("tests_files", None)
sections.append(section("grading", "How it's graded", {k: v for k, v in grading.items() if v is not None}, kind="value",
note="from reward_model and the instance"))
if tests:
sections.append(section("tests", "Hidden tests", {"path": "test.patch", "text": tests}, kind="code", note="applied after the agent, then run"))
if rubric:
sections.append(section("rubric", "Verifier files", rubric, kind="value", note="the rubric's questions; its anchors are withheld"))
env = {k: clean_inst.pop(k) for k in ("docker_image", "cwd", "cpus", "memory_mb", "storage_mb", "allow_internet", "agent_timeout_sec") if k in clean_inst}
if env:
sections.append(section("environment", "What it runs in", env, kind="value"))
if config == "music":
sections.append(section("music", "The brief", {k: v for k, v in extra.items() if k not in ("index",)}, kind="value"))
rest = {k: v for k, v in clean_inst.items() if k not in ("problem_statement", "instance_id", "dataset_type")}
if rest:
sections.append(section("data", "The rest of the instance", rest, kind="value"))
tid = self.task_id(row, config)
run = twin(config, tid)
return {"title": self.card(row, i, roles)["title"], "id": tid, "chips": self.card(row, i, roles)["chips"], "sections": sections,
"withheld": sorted(set(gone + [f"instance.{k}" for k in self.instance(row)[1] if MIMO_HIDE.match(k)])),
"glance": [["Domain", config or "–"], ["Ability", row.get("ability") or "–"], ["Agent", row.get("agent_name") or "–"],
["Reward", (rm.get("style") if isinstance(rm, dict) else None) or "–"], ["Image", env.get("docker_image") or "–"]],
"run": run}
def twin(config: str, tid: str | None) -> dict[str, Any] | None:
"""A MiMo row's Harbor task in the FineEnvs conversion, when this explorer has that dataset's index."""
from .. import catalog
if not tid:
return None
path = f"tasks/{_slug(tid)}"
for name in ([config, "terminal"] if config == "general" else [config]):
spec = f"FineEnvs/MiMo-V2.6-RL-harbor-{name}"
idx = catalog._read_index(spec)
if idx and any(r["path"] == path for r in idx["tasks"]):
return {"dataset": spec, "path": path, "label": f"{spec.split('/')[1]} Β· {path}"}
return None
# ── verl / SkyRL RL rows (data_source, prompt, ability, reward_model | env_class, reward_spec): verl_rows.py ─────
# ── NeMo Gym (nvidia/Nemotron-RL-*): responses_create_params, agent_ref and the verifier's fields ──────────────
# read in nemogym.py, with NeMo Gym's own configs: which resources server grades a row, and how to run it
# ── agent traces (conversations / messages with the agent's turns), e.g. open-thoughts/AgentTrove ──────────
class Traces(Processor):
id, name, framework = "traces", "Agent traces", "Traces"
about = "Each row is a recorded rollout: the conversation between an agent and its environment, with its result."
def conv(self, ds: Dataset) -> str | None:
for c in ("conversations", "messages", "trajectory", "conversation"):
if c in ds.columns and any(is_messages(parse_json(r.get(c))) and any(m.get("role") == "assistant" for m in parse_json(r.get(c))) for r in ds.sample[:5]):
return c
return None
def match(self, ds: Dataset) -> float:
c = self.conv(ds)
if not c:
return 0
return 0.85 if any(x in ds.columns for x in ("agent", "model", "trial_name", "result", "reward")) else 0.5
def view(self, row, i, ds, roles):
c = self.conv(ds)
clean, gone = withhold({k: parse_json(v) for k, v in row.items() if k != c})
conv = parse_json(row.get(c)) or []
turns = sum(1 for m in conv if isinstance(m, dict) and m.get("role") == "assistant")
result = {k: clean.pop(k) for k in ("result", "reward", "score", "success", "passed", "judgment") if k in clean}
sections = [section("trace", "The trajectory", conv, kind="messages", note=f"{len(conv)} messages, {turns} from the agent")]
packed = PackedHarbor()
blob = packed.blob(ds)
if blob and row.get(blob): # the task it ran on, packed alongside
for k in (blob,):
clean.pop(k, None)
task = packed.view(row, i, ds, packed.roles(ds))
files = next((x for x in task["sections"] if x["kind"] == "files"), None)
if files:
sections.append({**files, "title": "The task's files"})
gone = gone + [w for w in task["withheld"] if w not in gone]
if result:
sections.append(section("result", "Result", result, kind="value"))
sections.append(section("data", "About this run", clean, kind="value"))
first_user = next((m.get("content") for m in conv if isinstance(m, dict) and m.get("role") == "user"), "")
title = str(clean.get("trial_name") or clean.get("task") or first_line(first_user if isinstance(first_user, str) else "") or f"Trace {i}")
return {"title": title[:200], "id": str(clean.get("run_id") or clean.get("id") or i), "chips": [str(clean[k]) for k in ("agent", "model") if clean.get(k)][:3],
"sections": sections, "withheld": gone, "glance": [["Row", f"{i:,}"], ["Agent turns", f"{turns:,}"], *[[k.title(), str(v)[:60]] for k, v in result.items()],
*[[k.title(), str(clean[k])[:60]] for k in ("agent", "model") if clean.get(k)]]}
def file(self, row: dict[str, Any], ds: Dataset, rel: str) -> dict[str, Any]:
return PackedHarbor().file(row, ds, rel)
def card(self, row, i, roles):
title = str(row.get("trial_name") or row.get("task") or f"Trace {i}")
chips = [str(row[k]) for k in ("agent", "model", "result") if isinstance(row.get(k), str) and row.get(k)][:3]
return {"i": i, "id": str(row.get("run_id") or i), "title": title[:200], "snippet": "", "chips": chips}
# ── data for Verifiers environments (Prime's Environments Hub): verifiers_rows.py ─────────────────────────────
class Generic(Processor):
pass
PROCESSORS: list[Processor] = [PackedHarbor(), MiMo(), NemoGym(), Verl(), Traces(), Verifiers(), Generic()]
# fields each format's reader knows hold answers though their names don't say so: search never matches them either
HIDDEN: frozenset[str] = NEMO_HIDE | VERL_HIDE | SWE_HIDE