compute.py

"""
Real RecursiveMAS demo driver: runs Sequential-Light (Planner Qwen3-1.7B
+ Critic Llama-3.2-1B + Solver Qwen2.5-Math-1.5B + trained Outerlinks)
end-to-end on a single math problem, in two modes:

  - latent path: agents pass hidden states through trained adapters
  - text path:   agents pass English text by generating + re-tokenizing

Captures wall-clock timings, the actual generated texts (with per-token
pieces for streaming), and per-agent layer trajectories for visualization.
Writes everything to viz/trajectories.json.
"""

import json
import sys
import time
from argparse import Namespace
from pathlib import Path

import numpy as np
import torch
from sklearn.decomposition import PCA

# Make the cloned repo importable
REPO_ROOT = Path(__file__).parent / "recursivemas"
sys.path.insert(0, str(REPO_ROOT))

from system_loader import load_mas_system
from inference_utils import inference_mas as base
from prompts import (
    PLANNER_SLOT,
    REFINED_SLOT,
    build_math_planner_prompt,
    build_math_refiner_prompt,
    build_math_refiner_prompt_with_slot,
    build_math_solver_prompt,
    build_math_solver_prompt_with_slots,
)

DEVICE = "mps" if torch.backends.mps.is_available() else "cpu"
DEVICE_OBJ = torch.device(DEVICE)
DTYPE_STR = "float16"
LATENT_STEPS = 10

PROBLEM = (
    "Sara has $5.00. She buys 2 erasers that cost $0.50 each, and then "
    "spends the rest on pencils that cost $0.20 each. How many pencils "
    "does Sara buy?"
)

OUT = Path(__file__).parent / "viz" / "trajectories.json"


# -----------------------------------------------------------------------------
# Timing helpers
# -----------------------------------------------------------------------------

def sync():
    if DEVICE == "mps" and hasattr(torch.mps, "synchronize"):
        torch.mps.synchronize()
    elif DEVICE == "cuda":
        torch.cuda.synchronize()


def timed(label, fn):
    sync()
    t0 = time.perf_counter()
    result = fn()
    sync()
    elapsed = time.perf_counter() - t0
    print(f"  [{elapsed*1000:9.1f} ms]  {label}")
    return result, elapsed


# -----------------------------------------------------------------------------
# Trajectory capture: one forward pass per agent, gather per-layer last-token states
# -----------------------------------------------------------------------------

def capture_trajectory(agent, prompt_text):
    """One forward over the rendered chat prompt; returns list of (num_layers+1)
    last-token hidden states as np arrays."""
    rendered = base.render_chat_prompt(agent.tokenizer, prompt_text, enable_thinking=False)
    inputs = agent.tokenizer(rendered, return_tensors="pt").to(DEVICE_OBJ)
    with torch.no_grad():
        out = agent.model(
            input_ids=inputs.input_ids,
            attention_mask=inputs.attention_mask,
            output_hidden_states=True,
            return_dict=True,
        )
    return [hs[0, -1].float().cpu().numpy() for hs in out.hidden_states]


def capture_latent_trajectory(agent, slot_prompt_builder, slot_token, prev_latent, *args):
    """One forward over [prefix_embeds | prev_latent | suffix_embeds]; returns
    per-layer last-token hidden states. Mirrors what the agent processes in the
    real latent pipeline."""
    user_prompt = slot_prompt_builder(*args)
    seg_prefix, seg_suffix = base.split_prompt_ids_by_slots(
        agent.tokenizer, user_prompt, [slot_token], False
    )
    embed_layer = agent.model.get_input_embeddings()
    embed_dtype = embed_layer.weight.dtype
    prefix_embeds = base.token_ids_to_embeds(embed_layer, seg_prefix, device=DEVICE_OBJ, dtype=embed_dtype)
    suffix_embeds = base.token_ids_to_embeds(embed_layer, seg_suffix, device=DEVICE_OBJ, dtype=embed_dtype)
    prev_embed = prev_latent.to(device=DEVICE_OBJ, dtype=embed_dtype)
    seq = torch.cat([prefix_embeds, prev_embed, suffix_embeds], dim=0).unsqueeze(0)
    attn = torch.ones(1, seq.size(1), device=DEVICE_OBJ, dtype=torch.long)
    with torch.no_grad():
        out = agent.model(
            inputs_embeds=seq,
            attention_mask=attn,
            output_hidden_states=True,
            return_dict=True,
        )
    return [hs[0, -1].detach().float().cpu().numpy() for hs in out.hidden_states]


# -----------------------------------------------------------------------------
# Text path: each agent generates text, passed forward
# -----------------------------------------------------------------------------

def generate_text(agent, prompt_text, max_new=600):
    rendered = base.render_chat_prompt(agent.tokenizer, prompt_text, enable_thinking=False)
    inputs = agent.tokenizer(rendered, return_tensors="pt").to(DEVICE_OBJ)
    with torch.no_grad():
        out = agent.model.generate(
            input_ids=inputs.input_ids,
            attention_mask=inputs.attention_mask,
            max_new_tokens=max_new,
            do_sample=False,
            pad_token_id=agent.tokenizer.eos_token_id or agent.tokenizer.pad_token_id,
        )
    new_ids = out[0, inputs.input_ids.shape[1]:]
    return agent.tokenizer.decode(new_ids, skip_special_tokens=True).strip()


def run_text_path(mas, question):
    timings = {}
    planner = mas.agents["planner"]
    critic = mas.agents["critic"]
    solver = mas.agents["solver"]

    print("\n=== TEXT path ===")
    planner_text, t_p = timed(
        "planner.generate (writes plan)",
        lambda: generate_text(planner, build_math_planner_prompt(question), max_new=600),
    )
    timings["planner_text_ms"] = t_p * 1000

    critic_text, t_c = timed(
        "critic.generate (writes refined plan)",
        lambda: generate_text(critic, build_math_refiner_prompt(question, planner_text), max_new=600),
    )
    timings["critic_text_ms"] = t_c * 1000

    solver_text, t_s = timed(
        "solver.generate (writes answer)",
        lambda: generate_text(solver, build_math_solver_prompt(question, critic_text, args=None), max_new=600),
    )
    timings["solver_text_ms"] = t_s * 1000

    timings["total_ms"] = (t_p + t_c + t_s) * 1000
    return planner_text, critic_text, solver_text, timings


# -----------------------------------------------------------------------------
# Latent path: replicate inference_mas stage logic with pre-loaded models
# -----------------------------------------------------------------------------

def planner_latent_stage(planner, outer_12, question):
    """Run planner forward + latent rollout; return planner_to_critic latents."""
    user_prompt = build_math_planner_prompt(question)
    prompt_ids = base.render_chat_prompt_ids(planner.tokenizer, user_prompt, enable_thinking=False)
    pad_id = planner.tokenizer.pad_token_id
    if pad_id is None:
        pad_id = planner.tokenizer.eos_token_id
    input_ids, attn = base.pad_left_ids([prompt_ids], pad_id=pad_id, device=DEVICE_OBJ)
    embed_layer = planner.model.get_input_embeddings()
    embed_dtype = embed_layer.weight.dtype
    input_embeds = embed_layer(input_ids)
    hidden_rollout = base.autoregressive_latent_rollout(
        model=planner.model,
        rollout_inner_adapter=planner.inner_adapter,
        input_embeds=input_embeds,
        attention_mask=attn,
        latent_steps=LATENT_STEPS,
    )
    planner_self = base.run_inner_adapter(planner.inner_adapter, hidden_rollout, output_dtype=embed_dtype)
    planner_to_critic = base.run_outer_adapter(outer_12, planner_self, output_dtype=embed_dtype)
    return planner_to_critic[0].detach()  # [latent_steps, critic_hidden_dim]


def critic_latent_stage(critic, outer_23, question, planner_to_critic):
    user_prompt = build_math_refiner_prompt_with_slot(question)
    seg_prefix, seg_suffix = base.split_prompt_ids_by_slots(
        critic.tokenizer, user_prompt, [PLANNER_SLOT], False
    )
    embed_layer = critic.model.get_input_embeddings()
    embed_dtype = embed_layer.weight.dtype
    prefix_embeds = base.token_ids_to_embeds(embed_layer, seg_prefix, device=DEVICE_OBJ, dtype=embed_dtype)
    suffix_embeds = base.token_ids_to_embeds(embed_layer, seg_suffix, device=DEVICE_OBJ, dtype=embed_dtype)
    planner_embed = planner_to_critic.to(device=DEVICE_OBJ, dtype=embed_dtype)
    seq = torch.cat([prefix_embeds, planner_embed, suffix_embeds], dim=0).unsqueeze(0)
    attn = torch.ones(1, seq.size(1), device=DEVICE_OBJ, dtype=torch.long)
    hidden_rollout = base.autoregressive_latent_rollout(
        model=critic.model,
        rollout_inner_adapter=critic.inner_adapter,
        input_embeds=seq,
        attention_mask=attn,
        latent_steps=LATENT_STEPS,
    )
    critic_self = base.run_inner_adapter(critic.inner_adapter, hidden_rollout, output_dtype=embed_dtype)
    critic_to_solver = base.run_outer_adapter(outer_23, critic_self, output_dtype=embed_dtype)
    return critic_to_solver[0].detach()


def solver_latent_stage(solver, question, critic_to_solver, max_new=600):
    args = Namespace(mas_shape="chain", solver_pre_question=0, choice_old_prompt=0)
    user_prompt = build_math_solver_prompt_with_slots(question, args=args, mas_shape="chain")
    seg_prefix, seg_suffix = base.split_prompt_ids_by_slots(
        solver.tokenizer, user_prompt, [REFINED_SLOT], False
    )
    embed_layer = solver.model.get_input_embeddings()
    embed_dtype = embed_layer.weight.dtype
    prefix_embeds = base.token_ids_to_embeds(embed_layer, seg_prefix, device=DEVICE_OBJ, dtype=embed_dtype)
    suffix_embeds = base.token_ids_to_embeds(embed_layer, seg_suffix, device=DEVICE_OBJ, dtype=embed_dtype)
    critic_embed = critic_to_solver.to(device=DEVICE_OBJ, dtype=embed_dtype)
    seq = torch.cat([prefix_embeds, critic_embed, suffix_embeds], dim=0).unsqueeze(0)
    attn = torch.ones(1, seq.size(1), device=DEVICE_OBJ, dtype=torch.long)
    gen_kwargs = base.build_generation_kwargs(
        solver.tokenizer, max_new_tokens=max_new, do_sample=False, temperature=0.6, top_p=0.95,
    )
    with torch.no_grad():
        generated = solver.model.generate(inputs_embeds=seq, attention_mask=attn, **gen_kwargs)
    sequences = generated.sequences if hasattr(generated, "sequences") else generated
    prompt_len = attn.size(1)
    # Their pipeline accommodates two return formats:
    if sequences.size(1) > max_new:
        gen_ids = sequences[:, prompt_len:]
    else:
        gen_ids = sequences
    return solver.tokenizer.decode(gen_ids[0], skip_special_tokens=True).strip()


def run_latent_path(mas, question):
    timings = {}
    planner = mas.agents["planner"]
    critic = mas.agents["critic"]
    solver = mas.agents["solver"]

    print("\n=== LATENT path ===")
    planner_to_critic, t_p = timed(
        "planner latent rollout + adapter",
        lambda: planner_latent_stage(planner, mas.outer_adapters["outer_12"], question),
    )
    timings["planner_latent_ms"] = t_p * 1000

    critic_to_solver, t_c = timed(
        "critic latent rollout + adapter",
        lambda: critic_latent_stage(critic, mas.outer_adapters["outer_23"], question, planner_to_critic),
    )
    timings["critic_latent_ms"] = t_c * 1000

    solver_text, t_s = timed(
        "solver.generate from refiner_to_solver",
        lambda: solver_latent_stage(solver, question, critic_to_solver, max_new=600),
    )
    timings["solver_latent_ms"] = t_s * 1000

    timings["total_ms"] = (t_p + t_c + t_s) * 1000
    return solver_text, planner_to_critic, critic_to_solver, timings


# -----------------------------------------------------------------------------
# Main
# -----------------------------------------------------------------------------

def main():
    print(f"Loading Sequential-Light on {DEVICE} ({DTYPE_STR})...")
    mas, t_load = timed(
        "load_mas_system",
        lambda: load_mas_system(
            style="sequential_light",
            dataset="math500",
            device=DEVICE,
            dtype=DTYPE_STR,
            outer_dtype=DTYPE_STR,
        ),
    )
    for name, agent in mas.agents.items():
        print(f"  {name}: hidden={agent.hidden_size}, layers={len(agent.model.model.layers)}")

    # Warmup pass on each agent (avoids first-op compile cost polluting timings)
    print("\nWarmup...")
    for name, agent in mas.agents.items():
        warm_inputs = agent.tokenizer("warmup", return_tensors="pt").to(DEVICE_OBJ)
        with torch.no_grad():
            agent.model(input_ids=warm_inputs.input_ids, attention_mask=warm_inputs.attention_mask)
    sync()

    # Run TEXT path first so we have planner/critic/solver text outputs to use
    # for trajectory-capture inputs.
    planner_text, critic_text, solver_text_via_text, text_timings = run_text_path(mas, PROBLEM)
    print(f"\n[Planner text]\n{planner_text[:400]}{'...' if len(planner_text) > 400 else ''}\n")
    print(f"\n[Critic text]\n{critic_text[:400]}{'...' if len(critic_text) > 400 else ''}\n")
    print(f"\n[Solver text (via text path)]\n{solver_text_via_text[:600]}{'...' if len(solver_text_via_text) > 600 else ''}\n")

    # Run LATENT path
    solver_text_via_latent, planner_to_critic, critic_to_solver, latent_timings = run_latent_path(mas, PROBLEM)
    print(f"\n[Solver text (via LATENT — telepathy)]\n{solver_text_via_latent[:600]}{'...' if len(solver_text_via_latent) > 600 else ''}\n")

    # ---- Capture trajectories for BOTH paths separately ----
    # Alice's input is the same in both paths (just the question prompt), so
    # her trajectory is identical. Bob (critic) and Charlie (solver) receive
    # different inputs in each path, so their trajectories differ.
    print("\n=== Trajectories ===")
    planner_traj, _ = timed(
        "planner trajectory (shared)",
        lambda: capture_trajectory(mas.agents["planner"], build_math_planner_prompt(PROBLEM)),
    )
    critic_traj_text, _ = timed(
        "critic trajectory — TEXT path",
        lambda: capture_trajectory(mas.agents["critic"], build_math_refiner_prompt(PROBLEM, planner_text)),
    )
    critic_traj_latent, _ = timed(
        "critic trajectory — LATENT path",
        lambda: capture_latent_trajectory(
            mas.agents["critic"], build_math_refiner_prompt_with_slot, PLANNER_SLOT,
            planner_to_critic, PROBLEM,
        ),
    )
    solver_traj_text, _ = timed(
        "solver trajectory — TEXT path",
        lambda: capture_trajectory(mas.agents["solver"], build_math_solver_prompt(PROBLEM, critic_text, args=None)),
    )
    solver_args = Namespace(mas_shape="chain", solver_pre_question=0, choice_old_prompt=0)
    solver_traj_latent, _ = timed(
        "solver trajectory — LATENT path",
        lambda: capture_latent_trajectory(
            mas.agents["solver"], build_math_solver_prompt_with_slots, REFINED_SLOT,
            critic_to_solver, PROBLEM, solver_args, "chain",
        ),
    )

    # Per-agent PCA -> 3D fit jointly on BOTH paths so coords share an axis.
    def pca_3d_pair(traj_a, traj_b, offset):
        arr_a = np.array(traj_a); arr_b = np.array(traj_b)
        combined = np.vstack([arr_a, arr_b])
        comp = min(3, combined.shape[0], combined.shape[1])
        pca = PCA(n_components=comp)
        coords = pca.fit_transform(combined)
        if comp < 3:
            pad = np.zeros((coords.shape[0], 3 - comp))
            coords = np.concatenate([coords, pad], axis=1)
        max_r = np.max(np.linalg.norm(coords - coords.mean(axis=0), axis=1)) + 1e-8
        coords = (coords - coords.mean(axis=0)) / max_r * 0.7
        coords[:, 0] += offset
        n_a = arr_a.shape[0]
        return coords[:n_a].tolist(), coords[n_a:].tolist(), pca.explained_variance_ratio_.tolist()

    # Alice: same trajectory in both paths — fit on it once
    planner_3d_text,   planner_3d_latent,   planner_var = pca_3d_pair(planner_traj, planner_traj, offset=-2.2)
    critic_3d_text,    critic_3d_latent,    critic_var  = pca_3d_pair(critic_traj_text,  critic_traj_latent,  offset=0)
    solver_3d_text,    solver_3d_latent,    solver_var  = pca_3d_pair(solver_traj_text,  solver_traj_latent,  offset=2.2)

    # Token pieces for streaming display
    def to_pieces(tok, text):
        ids = tok(text, add_special_tokens=False).input_ids
        return [tok.decode([tid]) for tid in ids]

    planner_pieces = to_pieces(mas.agents["planner"].tokenizer, planner_text)
    critic_pieces  = to_pieces(mas.agents["critic"].tokenizer,  critic_text)
    solver_text_pieces = to_pieces(mas.agents["solver"].tokenizer, solver_text_via_text)
    solver_latent_pieces = to_pieces(mas.agents["solver"].tokenizer, solver_text_via_latent)

    # Timing summary
    print("\n=== timing summary ===")
    print(f"  latent path total: {latent_timings['total_ms']:9.1f} ms")
    print(f"  text   path total: {text_timings['total_ms']:9.1f} ms")
    speedup = text_timings["total_ms"] / max(latent_timings["total_ms"], 1e-9)
    print(f"  latent is {speedup:.2f}x faster")

    data = {
        "task": PROBLEM,
        "models": {
            "planner": mas.agents["planner"].repo_id,
            "critic":  mas.agents["critic"].repo_id,
            "solver":  mas.agents["solver"].repo_id,
        },
        "trajectories": {
            "latent": {
                "planner": planner_3d_latent,
                "critic":  critic_3d_latent,
                "solver":  solver_3d_latent,
            },
            "text": {
                "planner": planner_3d_text,
                "critic":  critic_3d_text,
                "solver":  solver_3d_text,
            },
        },
        "explained_variance": {
            "planner": planner_var,
            "critic":  critic_var,
            "solver":  solver_var,
        },
        "texts": {
            "planner":       planner_text,
            "critic":        critic_text,
            "solver_text":   solver_text_via_text,
            "solver_latent": solver_text_via_latent,
        },
        "token_pieces": {
            "planner":       planner_pieces,
            "critic":        critic_pieces,
            "solver_text":   solver_text_pieces,
            "solver_latent": solver_latent_pieces,
        },
        "timings": {
            "latent_path_total_ms": latent_timings["total_ms"],
            "text_path_total_ms":   text_timings["total_ms"],
            "per_stage": {
                "latent": latent_timings,
                "text":   text_timings,
            },
            "speedup":             speedup,
        },
        "comm_tokens": {
            # Latent path passes float tensors, no English in between.
            "latent_path": 0,
            # Text path: planner emits ~planner_pieces tokens, critic re-encodes those + emits ~critic_pieces, etc.
            "text_path":   len(planner_pieces) + len(critic_pieces),
        },
    }

    OUT.parent.mkdir(parents=True, exist_ok=True)
    with open(OUT, "w") as f:
        json.dump(data, f, indent=2)
    print(f"\nWrote {OUT}")


if __name__ == "__main__":
    main()