init_knowledge.py

"""
One-time script: initialise knowledge/task_XXX/ for all 400 tasks.
Reads data/*.json + current scores + current ONNX cost breakdowns.
Run once:  /Users/yeyang/miniconda3/envs/neurogolf/bin/python init_knowledge.py
"""
import json, math, os, sys
from pathlib import Path

ROOT = Path(__file__).parent
DATA = ROOT / "data"
ONNX_DIR = ROOT / "submissions" / "s_6256_03" / "onnx"
KNOWLEDGE = ROOT / "knowledge"

# Color representations (compact single char)
COLOR_CHARS = ".123456789"  # 0=. 1-9=digit

def grid_to_str(grid):
    lines = []
    for row in grid:
        lines.append("".join(COLOR_CHARS[c] if c < 10 else "?" for c in row))
    return "\n".join(lines)

def grid_dims(grid):
    h = len(grid)
    w = len(grid[0]) if grid else 0
    return h, w

# Load scores
with open("/tmp/task_scores.json") as f:
    scores = json.load(f)

# Load ONNX sizes
onnx_sizes = {}
if ONNX_DIR.exists():
    for f in ONNX_DIR.glob("task*.onnx"):
        tid = int(f.stem.replace("task", ""))
        onnx_sizes[tid] = f.stat().st_size

def cost_from_score(sc):
    if sc <= 1.0: return None
    return math.exp(25.0 - sc)

def difficulty_tag(sc):
    if sc >= 22: return "EASY (top tier)"
    if sc >= 18: return "MEDIUM-EASY"
    if sc >= 15: return "MEDIUM"
    if sc >= 12: return "HARD"
    return "VERY HARD"

created, skipped = 0, 0
for tid in range(1, 401):
    task_file = DATA / f"task{tid:03d}.json"
    if not task_file.exists():
        print(f"  task{tid:03d}: no data file, skip")
        skipped += 1
        continue

    with open(task_file) as f:
        task = json.load(f)

    task_dir = KNOWLEDGE / f"task_{tid:03d}"
    task_dir.mkdir(parents=True, exist_ok=True)

    sc = scores.get(str(tid), 0.0)
    cost = cost_from_score(sc)
    onnx_size = onnx_sizes.get(tid, 0)
    n_train = len(task["train"])
    n_test = len(task["test"])
    n_arcgen = len(task.get("arc-gen", []))

    # ── description.md ────────────────────────────────────────────
    desc_path = task_dir / "description.md"
    if not desc_path.exists():
        lines = [f"# Task {tid:03d}"]
        lines.append(f"\n**Current score:** {sc:.4f}  |  **Difficulty:** {difficulty_tag(sc)}")
        if cost:
            lines.append(f"**Estimated cost:** {cost:,.0f} params+bytes  |  **ONNX size:** {onnx_size:,} B")
        lines.append(f"**Examples:** {n_train} train, {n_test} test, {n_arcgen} arc-gen")
        lines.append("")
        lines.append("## Transformation Rule")
        lines.append("")
        lines.append("*(Fill in after understanding the task)*")
        lines.append("")
        lines.append("## Train Examples")
        lines.append("")
        lines.append("Color key: `.`=black  `1`=blue  `2`=red  `3`=green  `4`=yellow  `5`=grey  `6`=magenta  `7`=orange  `8`=cyan  `9`=maroon")
        lines.append("")
        for i, ex in enumerate(task["train"]):
            ih, iw = grid_dims(ex["input"])
            oh, ow = grid_dims(ex["output"])
            lines.append(f"### Train {i+1}  ({ih}×{iw} → {oh}×{ow})")
            lines.append("```")
            lines.append("INPUT:")
            lines.append(grid_to_str(ex["input"]))
            lines.append("OUTPUT:")
            lines.append(grid_to_str(ex["output"]))
            lines.append("```")
            lines.append("")
        if task.get("test"):
            ex = task["test"][0]
            ih, iw = grid_dims(ex["input"])
            oh, ow = grid_dims(ex["output"])
            lines.append(f"## Test Example  ({ih}×{iw} → {oh}×{ow})")
            lines.append("```")
            lines.append("INPUT:")
            lines.append(grid_to_str(ex["input"]))
            lines.append("OUTPUT:")
            lines.append(grid_to_str(ex["output"]))
            lines.append("```")
        lines.append("")
        lines.append("## Category Tags")
        lines.append("")
        lines.append("*(e.g., color_replace, spatial_shift, pattern_tile, object_detect, border_fill)*")
        desc_path.write_text("\n".join(lines))

    # ── insights.md ────────────────────────────────────────────────
    ins_path = task_dir / "insights.md"
    if not ins_path.exists():
        ins_path.write_text(f"""# Task {tid:03d} — Iteration Log

## Current Best
- **Score:** {sc:.4f}
- **Cost:** {f"{cost:,.0f}" if cost else "0 (free!)"}
- **Source:** s_6256_03 (initial submission)
- **ONNX size:** {onnx_size:,} B

## Attempt History

| Version | Score | Delta | Description | Result |
|---------|-------|-------|-------------|--------|
| v0 (baseline) | {sc:.4f} | — | Initial submission | ✅ LB confirmed |

## Notes

*(Add insights as you iterate — what patterns work, what fails, hidden-test risks)*

## Score Target

To reach 18 pts → need cost < 5,958
To reach 20 pts → need cost < 799
To reach 22 pts → need cost < 107
""")

    # ── best_score.txt ─────────────────────────────────────────────
    score_path = task_dir / "best_score.txt"
    if not score_path.exists():
        score_path.write_text(f"{sc:.6f}\n")

    created += 1
    if created % 50 == 0:
        print(f"  {created}/400 created...")

print(f"\nDone. Created: {created}, Skipped: {skipped}")
print(f"Knowledge base at: {KNOWLEDGE}")