Spaces:
Sleeping
Sleeping
| """ | |
| Hugging Face Spaces entrypoint — Gradio UI for Toy-IR compiler RL demo. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import sys | |
| import textwrap | |
| from pathlib import Path | |
| import gradio as gr | |
| _ROOT = Path(__file__).resolve().parent | |
| if str(_ROOT) not in sys.path: | |
| sys.path.insert(0, str(_ROOT)) | |
| from runtime_core import ( | |
| CompilerOptimizationEnv, | |
| Deliverable2_Formatter, | |
| MOCK_PASSES, | |
| MockEngine, | |
| SAMPLE_PROGRAM, | |
| ) | |
| from train import run_toy_training | |
| DEFAULT_IR = json.dumps(SAMPLE_PROGRAM, indent=2) | |
| def translate_ir(ir_json: str) -> tuple[str, str]: | |
| try: | |
| data = json.loads(ir_json.strip() or "[]") | |
| except json.JSONDecodeError as e: | |
| return "", f"Invalid JSON: {e}" | |
| if not isinstance(data, list): | |
| return "", "JSON root must be a list of instruction objects." | |
| d2 = Deliverable2_Formatter.translate_state(data) | |
| env = CompilerOptimizationEnv(MockEngine(), MOCK_PASSES, max_steps=10) | |
| env.reset(data) | |
| internal = env.state() | |
| return d2, internal | |
| def parse_llm_output(llm_text: str) -> str: | |
| try: | |
| arr = Deliverable2_Formatter.extract_action_array(llm_text) | |
| return json.dumps(arr, indent=2) | |
| except ValueError as e: | |
| return str(e) | |
| def parse_action_line(line: str) -> list[str]: | |
| line = line.strip() | |
| if not line: | |
| return [] | |
| try: | |
| got = Deliverable2_Formatter.extract_action_array(line) | |
| return [str(x).strip() for x in got] | |
| except ValueError: | |
| parts = [p.strip().strip("\"'") for p in line.split(",") if p.strip()] | |
| return [p.lower() for p in parts] | |
| def run_episode(ir_json: str, actions_multiline: str, max_steps: int) -> str: | |
| try: | |
| program = json.loads(ir_json.strip() or "[]") | |
| except json.JSONDecodeError as e: | |
| return f"Invalid program JSON: {e}" | |
| if not isinstance(program, list): | |
| return "Program must be a JSON list." | |
| lines = [ln for ln in actions_multiline.splitlines() if ln.strip()] | |
| actions: list[str] = [] | |
| for ln in lines: | |
| actions.extend(parse_action_line(ln)) | |
| if not actions: | |
| return "No actions parsed. Enter JSON arrays or comma-separated pass names." | |
| engine = MockEngine() | |
| env = CompilerOptimizationEnv(engine, MOCK_PASSES, max_steps=int(max_steps)) | |
| obs0 = env.reset(program) | |
| log = [ | |
| f"Baseline cycles: {env.previous_cycles}", | |
| f"Initial observation (env):\n{obs0}", | |
| "", | |
| f"Deliverable 2 pseudo-assembly:\n{Deliverable2_Formatter.translate_state(program)}", | |
| "", | |
| "--- steps ---", | |
| ] | |
| for i, act in enumerate(actions, start=1): | |
| r = env.step(act) | |
| log.append( | |
| f"{i}. {act!r} → reward={r.reward:+.3f} done={r.done} cycles={env.previous_cycles}" | |
| ) | |
| if r.info: | |
| slim = {k: v for k, v in r.info.items() if k in ("delta_pct", "error", "no_op", "reason", "terminal_bonus")} | |
| if slim: | |
| log.append(f" info: {slim}") | |
| if r.done: | |
| break | |
| log.append("") | |
| log.append("Episode summary:") | |
| log.append(json.dumps(env._episode_summary(), indent=2)) | |
| return "\n".join(log) | |
| def build_demo() -> gr.Blocks: | |
| with gr.Blocks(title="Compiler Tetris — Toy-IR RL Demo") as demo: | |
| gr.Markdown( | |
| textwrap.dedent( | |
| """ | |
| # Compiler optimization (Toy-IR) — interactive demo | |
| This Space runs **CPU-only** demos: **Deliverable 2** state translation and action parsing, | |
| plus the **CompilerOptimizationEnv** loop. A **toy REINFORCE** tab trains a tiny | |
| stateless policy over the mock passes (see `train.py`). Full **GRPO + LLM + Unsloth** | |
| still belongs on Colab or a GPU machine. | |
| """ | |
| ).strip() | |
| ) | |
| with gr.Tabs(): | |
| with gr.Tab("Translate IR (D2)"): | |
| gr.Markdown("Paste Toy-IR as a JSON **list** of instruction dicts (`op`, `args`, `dest`, …).") | |
| ir_in = gr.Textbox(label="IR JSON", value=DEFAULT_IR, lines=12, max_lines=24) | |
| btn_t = gr.Button("Translate", variant="primary") | |
| out_d2 = gr.Textbox(label="Deliverable2 pseudo-assembly (LLM-facing)", lines=10) | |
| out_env = gr.Textbox(label="Env internal pseudo-asm (with type hints)", lines=10) | |
| btn_t.click(translate_ir, [ir_in], [out_d2, out_env]) | |
| with gr.Tab("Parse LLM output (D2)"): | |
| gr.Markdown( | |
| "Paste model output. A JSON array is preferred; bracket extraction handles light noise." | |
| ) | |
| llm_in = gr.Textbox( | |
| label="LLM output", | |
| value='Here is the plan: ["constant_folding", "dead_code_elimination"]', | |
| lines=4, | |
| ) | |
| out_parse = gr.Textbox(label="Parsed array or error", lines=6) | |
| gr.Button("Parse", variant="primary").click(parse_llm_output, [llm_in], [out_parse]) | |
| with gr.Tab("Env episode (mock engine)"): | |
| gr.Markdown( | |
| "One JSON program + actions: each line can be a JSON array or comma-separated names." | |
| ) | |
| ir_ep = gr.Textbox(label="Program JSON", value=DEFAULT_IR, lines=10) | |
| acts = gr.Textbox( | |
| label="Actions (one JSON array or comma-list per line)", | |
| value='["constant_folding", "dead_code_elimination"]\nloop_unrolling', | |
| lines=5, | |
| ) | |
| ms = gr.Slider(1, 20, value=10, step=1, label="max_steps") | |
| out_ep = gr.Textbox(label="Log", lines=20) | |
| gr.Button("Run episode", variant="primary").click(run_episode, [ir_ep, acts, ms], [out_ep]) | |
| with gr.Tab("Toy training (REINFORCE)"): | |
| gr.Markdown( | |
| textwrap.dedent( | |
| """ | |
| Trains a **stateless** categorical policy over the three mock passes | |
| (`constant_folding`, `dead_code_elimination`, `loop_unrolling`) using | |
| **REINFORCE** in pure Python. Same logic as: `python train.py --episodes 50`. | |
| This is a CPU smoke run, not an LLM. For real GRPO, use your project notebooks | |
| on a GPU. | |
| """ | |
| ).strip() | |
| ) | |
| tr_ep = gr.Slider(5, 200, value=50, step=1, label="episodes") | |
| tr_ms = gr.Slider(2, 20, value=8, step=1, label="max_steps per episode") | |
| tr_seed = gr.Number(value=0, label="random seed", precision=0) | |
| tr_lr = gr.Slider(0.01, 0.5, value=0.15, step=0.01, label="learning rate") | |
| out_tr = gr.Textbox(label="Training log", lines=18) | |
| gr.Button("Run training", variant="primary").click( | |
| run_toy_training, [tr_ep, tr_ms, tr_seed, tr_lr], [out_tr] | |
| ) | |
| gr.Markdown( | |
| "Source notebooks in the parent repo: `compiler_optimization_grpo.ipynb`, " | |
| "`role2_deliverable3_training_loop (2) (1).ipynb`, `compiler_tetris (1).ipynb`, `metahack1 (1).ipynb`." | |
| ) | |
| return demo | |
| if __name__ == "__main__": | |
| port = int(os.environ.get("PORT", "7860")) | |
| build_demo().launch(server_name="0.0.0.0", server_port=port) | |