Upload plugins/mlintern/skills/ml-intern-harness/scripts/inspect_dataset.py with huggingface_hub

plugins/mlintern/skills/ml-intern-harness/scripts/inspect_dataset.py

@@ -0,0 +1,293 @@
+#!/usr/bin/env python3
+"""Inspect a Hugging Face dataset using the Dataset Viewer API.
+
+This mirrors the useful parts of upstream ml-intern's hf_inspect_dataset tool:
+status, configs/splits, schema, sample rows, and parquet file availability.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+import urllib.error
+import urllib.parse
+import urllib.request
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any
+
+
+BASE_URL = "https://datasets-server.huggingface.co"
+MAX_SAMPLE_VALUE_LEN = 150
+
+
+def fetch_json(path: str, params: dict[str, Any], token: str | None) -> dict[str, Any]:
+    query = urllib.parse.urlencode({k: v for k, v in params.items() if v is not None})
+    request = urllib.request.Request(f"{BASE_URL}{path}?{query}")
+    if token:
+        request.add_header("Authorization", f"Bearer {token}")
+    try:
+        with urllib.request.urlopen(request, timeout=30) as response:
+            return json.loads(response.read().decode("utf-8"))
+    except urllib.error.HTTPError as exc:
+        body = exc.read().decode("utf-8", errors="replace")
+        raise RuntimeError(f"{path} returned HTTP {exc.code}: {body[:500]}") from exc
+
+
+def type_name(feature: Any) -> str:
+    if isinstance(feature, str):
+        return feature
+    if not isinstance(feature, dict):
+        return type(feature).__name__
+    feature_type = feature.get("_type")
+    if feature_type == "ClassLabel":
+        names = feature.get("names") or []
+        if 0 < len(names) <= 5:
+            values = ", ".join(f"{name}={idx}" for idx, name in enumerate(names))
+            return f"ClassLabel ({values})"
+        return f"ClassLabel ({len(names)} classes)"
+    if feature_type:
+        return feature_type
+    if "dtype" in feature:
+        return str(feature["dtype"])
+    return json.dumps(feature, ensure_ascii=False)[:120]
+
+
+def extract_configs(splits_data: dict[str, Any]) -> list[dict[str, Any]]:
+    configs: dict[str, dict[str, Any]] = {}
+    for item in splits_data.get("splits", []):
+        config = item.get("config", "default")
+        split = item.get("split", "train")
+        row_count = item.get("num_rows") or item.get("num_examples")
+        configs.setdefault(config, {"name": config, "splits": []})
+        configs[config]["splits"].append({"name": split, "rows": row_count})
+    return list(configs.values())
+
+
+def format_status(data: dict[str, Any]) -> str:
+    available = [
+        key
+        for key in ("viewer", "preview", "search", "filter", "statistics")
+        if data.get(key)
+    ]
+    if available:
+        return f"## Status\nValid ({', '.join(available)})"
+    return "## Status\nDataset may have Dataset Viewer issues"
+
+
+def format_structure(configs: list[dict[str, Any]], max_rows: int = 20) -> str:
+    lines = ["## Structure (configs & splits)", "| Config | Split | Rows |", "|---|---|---:|"]
+    total = sum(len(config["splits"]) for config in configs)
+    shown = 0
+    for config in configs:
+        for split in config["splits"]:
+            if shown >= max_rows:
+                break
+            rows = split["rows"] if split["rows"] is not None else "?"
+            lines.append(f"| {config['name']} | {split['name']} | {rows} |")
+            shown += 1
+        if shown >= max_rows:
+            break
+    if total > shown:
+        lines.append("| ... | ... | ... |")
+        lines.append(f"_Showing {shown} of {total} config/split rows._")
+    return "\n".join(lines)
+
+
+def format_schema(info: dict[str, Any], config: str) -> str:
+    features = info.get("dataset_info", {}).get("features", {})
+    lines = [f"## Schema ({config})", "| Column | Type |", "|---|---|"]
+    if not features:
+        lines.append("| (none found) | unknown |")
+    for column, feature in features.items():
+        lines.append(f"| {column} | {type_name(feature)} |")
+    return "\n".join(lines)
+
+
+def maybe_json(value: Any) -> Any:
+    if isinstance(value, str):
+        try:
+            return json.loads(value)
+        except json.JSONDecodeError:
+            return value
+    return value
+
+
+def format_messages(messages: Any) -> str | None:
+    messages = maybe_json(messages)
+    if not isinstance(messages, list) or not messages:
+        return None
+    roles: set[str] = set()
+    keys: set[str] = set()
+    has_tool_calls = False
+    has_tool_results = False
+    example: dict[str, Any] | None = None
+    fallback: dict[str, Any] | None = None
+    for message in messages:
+        if not isinstance(message, dict):
+            continue
+        keys.update(message.keys())
+        if message.get("role"):
+            roles.add(str(message["role"]))
+        if message.get("tool_calls") or message.get("function_call"):
+            has_tool_calls = True
+            example = example or message
+        if message.get("role") in {"tool", "function"} or message.get("tool_call_id"):
+            has_tool_results = True
+        if message.get("role") == "assistant":
+            example = example or message
+        elif message.get("role") != "system":
+            fallback = fallback or message
+    example = example or fallback
+    lines = ["## Messages Column Format"]
+    lines.append(f"Roles: {', '.join(sorted(roles)) if roles else 'unknown'}")
+    common = ["role", "content", "tool_calls", "tool_call_id", "name", "function_call"]
+    lines.append("Message keys: " + ", ".join(f"{key} {'yes' if key in keys else 'no'}" for key in common))
+    if has_tool_calls:
+        lines.append("Tool calls: present")
+    if has_tool_results:
+        lines.append("Tool results: present")
+    if example:
+        cleaned = dict(example)
+        content = cleaned.get("content")
+        if isinstance(content, str) and len(content) > 100:
+            cleaned["content"] = content[:100] + "..."
+        lines.append("")
+        lines.append("Example message structure:")
+        lines.append("```json")
+        lines.append(json.dumps(cleaned, indent=2, ensure_ascii=False))
+        lines.append("```")
+    return "\n".join(lines)
+
+
+def format_samples(rows_data: dict[str, Any], config: str, split: str, limit: int) -> str:
+    rows = rows_data.get("rows", [])[:limit]
+    lines = [f"## Sample Rows ({config}/{split})"]
+    first_messages: Any = None
+    for idx, row_wrapper in enumerate(rows, 1):
+        row = row_wrapper.get("row", {})
+        lines.append(f"**Row {idx}:**")
+        for key, value in row.items():
+            if key.lower() == "messages" and first_messages is None:
+                first_messages = value
+            text = str(value)
+            if len(text) > MAX_SAMPLE_VALUE_LEN:
+                text = text[:MAX_SAMPLE_VALUE_LEN] + "..."
+            lines.append(f"- {key}: {text}")
+    if not rows:
+        lines.append("(no rows returned)")
+    message_section = format_messages(first_messages) if first_messages is not None else None
+    if message_section:
+        lines.append("")
+        lines.append(message_section)
+    return "\n".join(lines)
+
+
+def format_parquet(data: dict[str, Any], max_rows: int = 20) -> str | None:
+    files = data.get("parquet_files", [])
+    if not files:
+        return None
+    groups: dict[str, dict[str, int]] = {}
+    for item in files:
+        key = f"{item.get('config', 'default')}/{item.get('split', 'train')}"
+        groups.setdefault(key, {"count": 0, "size": 0})
+        groups[key]["count"] += 1
+        size = item.get("size") or 0
+        groups[key]["size"] += int(size) if isinstance(size, (int, float)) else 0
+    lines = ["## Files (Parquet)"]
+    for key, values in list(groups.items())[:max_rows]:
+        size_mb = values["size"] / (1024 * 1024)
+        lines.append(f"- {key}: {values['count']} file(s), {size_mb:.1f} MB")
+    if len(groups) > max_rows:
+        lines.append(f"- ... showing {max_rows} of {len(groups)} groups")
+    return "\n".join(lines)
+
+
+def compatibility_notes(features: dict[str, Any]) -> str:
+    columns = set(features)
+    lines = ["## Training Compatibility"]
+    checks = {
+        "SFT": bool({"messages", "text"} & columns or {"prompt", "completion"} <= columns),
+        "DPO": {"prompt", "chosen", "rejected"} <= columns,
+        "GRPO": "prompt" in columns,
+    }
+    for method, ok in checks.items():
+        lines.append(f"- {method}: {'looks compatible' if ok else 'columns not sufficient'}")
+    if "messages" in columns:
+        lines.append("- Chat data: inspect the sample message roles before choosing a trainer template.")
+    return "\n".join(lines)
+
+
+def inspect_dataset(dataset: str, config: str | None, split: str | None, sample_rows: int, token: str | None) -> str:
+    warnings: list[str] = []
+    with ThreadPoolExecutor(max_workers=3) as pool:
+        futures = {
+            pool.submit(fetch_json, "/is-valid", {"dataset": dataset}, token): "is-valid",
+            pool.submit(fetch_json, "/splits", {"dataset": dataset}, token): "splits",
+            pool.submit(fetch_json, "/parquet", {"dataset": dataset}, token): "parquet",
+        }
+        phase1: dict[str, Any] = {}
+        for future in as_completed(futures):
+            name = futures[future]
+            try:
+                phase1[name] = future.result()
+            except Exception as exc:
+                warnings.append(f"{name}: {exc}")
+
+    configs = extract_configs(phase1.get("splits", {}))
+    selected_config = config or (configs[0]["name"] if configs else "default")
+    selected_split = split or (configs[0]["splits"][0]["name"] if configs and configs[0]["splits"] else "train")
+
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        futures = {
+            pool.submit(fetch_json, "/info", {"dataset": dataset, "config": selected_config}, token): "info",
+            pool.submit(
+                fetch_json,
+                "/first-rows",
+                {"dataset": dataset, "config": selected_config, "split": selected_split},
+                token,
+            ): "first-rows",
+        }
+        phase2: dict[str, Any] = {}
+        for future in as_completed(futures):
+            name = futures[future]
+            try:
+                phase2[name] = future.result()
+            except Exception as exc:
+                warnings.append(f"{name}: {exc}")
+
+    features = phase2.get("info", {}).get("dataset_info", {}).get("features", {})
+    sections = [f"# {dataset}"]
+    if "is-valid" in phase1:
+        sections.append(format_status(phase1["is-valid"]))
+    if configs:
+        sections.append(format_structure(configs))
+    if "info" in phase2:
+        sections.append(format_schema(phase2["info"], selected_config))
+        sections.append(compatibility_notes(features))
+    if "first-rows" in phase2:
+        sections.append(format_samples(phase2["first-rows"], selected_config, selected_split, sample_rows))
+    parquet = format_parquet(phase1.get("parquet", {}))
+    if parquet:
+        sections.append(parquet)
+    if warnings:
+        sections.append("## Warnings\n" + "\n".join(f"- {warning}" for warning in warnings))
+    return "\n\n".join(sections)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("dataset", help="Dataset id, for example stanfordnlp/imdb")
+    parser.add_argument("--config", help="Config/subset name")
+    parser.add_argument("--split", help="Split for sample rows")
+    parser.add_argument("--sample-rows", type=int, default=3, help="Number of rows to show, max 10")
+    parser.add_argument("--token-env", default="HF_TOKEN", help="Environment variable containing an HF token")
+    args = parser.parse_args()
+    token = os.environ.get(args.token_env)
+    print(inspect_dataset(args.dataset, args.config, args.split, min(args.sample_rows, 10), token))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())