seeker
227 lines
8.5 KiB
Python
227 lines
8.5 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from dataclasses import asdict
|
|
from functools import partial
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
from safetensors.torch import load_file
|
|
from transformers import LogitsProcessorList, StoppingCriteriaList
|
|
|
|
from alpamayo_lite.action_space import RelativeDeltaActionSpace
|
|
from alpamayo_lite.config import load_config
|
|
from alpamayo_lite.data import load_manifest
|
|
from alpamayo_lite.diffusion import FlowMatching
|
|
from alpamayo_lite.evaluation import (
|
|
save_prediction_arrays,
|
|
save_trajectory_plot,
|
|
trajectory_metrics,
|
|
)
|
|
from alpamayo_lite.expert import ExpertArchitecture, KVCachedDiffusionExpert, PrefixKV
|
|
from alpamayo_lite.prompts import build_vlm_prefix
|
|
from alpamayo_lite.runtime import (
|
|
empty_device_cache,
|
|
parameter_counts,
|
|
resolve_dtype,
|
|
seed_everything,
|
|
select_device,
|
|
write_json,
|
|
)
|
|
from alpamayo_lite.trajectory_tokens import (
|
|
DiscreteTrajectoryMask,
|
|
StopOneTokenAfterBoundary,
|
|
mask_after_boundary,
|
|
)
|
|
from alpamayo_lite.vlm import load_vlm, process_image_text, vlm_text_config
|
|
|
|
|
|
def _resolve(root: Path, value: str) -> Path:
|
|
path = Path(value).expanduser()
|
|
return path.resolve() if path.is_absolute() else (root / path).resolve()
|
|
|
|
|
|
def _read_json(path: Path) -> dict[str, Any]:
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", default="configs/local_overfit.yaml")
|
|
parser.add_argument("--vlm-dir", default="outputs/stage1_vlm_trajectory")
|
|
parser.add_argument("--stage2-dir", default="outputs/stage2_expert")
|
|
parser.add_argument("--output-dir", default="outputs/structured_vlm_handoff")
|
|
parser.add_argument("--max-new-tokens", type=int, default=96)
|
|
args = parser.parse_args()
|
|
|
|
cfg = load_config(args.config)
|
|
seed_everything(cfg.seed)
|
|
device = select_device()
|
|
dtype = resolve_dtype(cfg.vlm.dtype, device)
|
|
vlm_dir = _resolve(cfg.root, args.vlm_dir)
|
|
stage2_dir = _resolve(cfg.root, args.stage2_dir)
|
|
output_dir = _resolve(cfg.root, args.output_dir)
|
|
sample = load_manifest(cfg.resolve(cfg.paths.manifest))[0]
|
|
history, future = sample.load_trajectory()
|
|
|
|
vocabulary = _read_json(vlm_dir / "trajectory_vocabulary.json")
|
|
trajectory_start = int(vocabulary["start_id"])
|
|
trajectory_size = int(vocabulary["size"])
|
|
boundary_id = int(vocabulary["special_token_ids"]["<|traj_future_start|>"])
|
|
|
|
processor, vlm = load_vlm(vlm_dir, cfg.vlm, device, dtype)
|
|
vlm.requires_grad_(False)
|
|
vlm.eval()
|
|
inputs = process_image_text(
|
|
processor,
|
|
sample.load_image(),
|
|
build_vlm_prefix(history),
|
|
device,
|
|
dtype,
|
|
)
|
|
prompt_length = inputs["input_ids"].shape[1]
|
|
generated = vlm.generate(
|
|
**inputs,
|
|
max_new_tokens=args.max_new_tokens,
|
|
do_sample=False,
|
|
use_cache=True,
|
|
return_dict_in_generate=True,
|
|
output_logits=True,
|
|
logits_processor=LogitsProcessorList(
|
|
[DiscreteTrajectoryMask(trajectory_start, trajectory_size)]
|
|
),
|
|
stopping_criteria=StoppingCriteriaList([StopOneTokenAfterBoundary(boundary_id)]),
|
|
pad_token_id=processor.tokenizer.pad_token_id,
|
|
eos_token_id=processor.tokenizer.eos_token_id,
|
|
)
|
|
new_ids = generated.sequences[:, prompt_length:]
|
|
visible_sequences = mask_after_boundary(
|
|
generated.sequences,
|
|
boundary_id,
|
|
processor.tokenizer.pad_token_id,
|
|
)
|
|
visible_new_ids = visible_sequences[:, prompt_length:]
|
|
prefix = PrefixKV.from_transformers_cache(generated.past_key_values)
|
|
boundary_positions = torch.nonzero(new_ids[0] == boundary_id, as_tuple=False).flatten()
|
|
if boundary_positions.numel() == 0:
|
|
raise RuntimeError("Structured VLM did not generate the diffusion handoff boundary")
|
|
boundary_index = int(boundary_positions[0])
|
|
expected_prefix_length = prompt_length + boundary_index + 1
|
|
if prefix.sequence_length != expected_prefix_length:
|
|
raise RuntimeError(
|
|
f"Generation cache does not end at the handoff boundary: "
|
|
f"cache={prefix.sequence_length}, expected={expected_prefix_length}"
|
|
)
|
|
|
|
expert_config = _read_json(stage2_dir / "expert_config.json")
|
|
architecture = ExpertArchitecture(**expert_config["architecture"])
|
|
expert = KVCachedDiffusionExpert(vlm_text_config(vlm), architecture).to(
|
|
device=device, dtype=dtype
|
|
)
|
|
expert.load_state_dict(load_file(stage2_dir / "expert.safetensors"), strict=True)
|
|
expert.requires_grad_(False)
|
|
expert.eval()
|
|
flow = FlowMatching(
|
|
action_shape=tuple(expert_config["action_shape"]),
|
|
inference_steps=int(expert_config["inference_steps"]),
|
|
)
|
|
action_space = RelativeDeltaActionSpace(
|
|
torch.tensor(expert_config["action_mean"], device=device),
|
|
torch.tensor(expert_config["action_scale"], device=device),
|
|
).to(device)
|
|
seed_everything(cfg.seed + 1)
|
|
predicted_action = flow.sample(
|
|
batch_size=1,
|
|
step_fn=partial(expert, prefix=prefix),
|
|
device=device,
|
|
dtype=dtype,
|
|
).float()
|
|
history_xy = history[None, :, :2].to(device)
|
|
target_xy = future[None, :, :2].to(device)
|
|
target_action = action_space.trajectory_to_action(history_xy, target_xy)
|
|
predicted_xy = action_space.action_to_trajectory(predicted_action, history_xy)
|
|
|
|
post_boundary_index = boundary_index + 1
|
|
returned_masked_count = None
|
|
if 0 <= post_boundary_index < len(generated.logits or ()):
|
|
post_boundary_logits = generated.logits[post_boundary_index][0]
|
|
returned_masked_count = int(
|
|
torch.isneginf(
|
|
post_boundary_logits[trajectory_start : trajectory_start + trajectory_size]
|
|
).sum().cpu()
|
|
)
|
|
|
|
result = {
|
|
"stage": "structured_vlm_to_diffusion_handoff",
|
|
"device": str(device),
|
|
"dtype": str(dtype),
|
|
"vlm_dir": str(vlm_dir),
|
|
"stage2_dir": str(stage2_dir),
|
|
"prompt_length": prompt_length,
|
|
"generated_token_count_including_one_cache_update_token": int(new_ids.shape[1]),
|
|
"boundary_index_in_generated_tokens": boundary_index,
|
|
"expected_cache_length_through_boundary": expected_prefix_length,
|
|
"visible_vlm_output": processor.tokenizer.decode(
|
|
visible_new_ids[0], skip_special_tokens=False
|
|
),
|
|
"boundary_generated": bool((new_ids == boundary_id).any()),
|
|
"visible_discrete_trajectory_tokens": int(
|
|
(
|
|
(visible_new_ids >= trajectory_start)
|
|
& (visible_new_ids < trajectory_start + trajectory_size)
|
|
).sum().cpu()
|
|
),
|
|
"returned_post_boundary_masked_logits": returned_masked_count,
|
|
"handoff_prefix": {
|
|
"layers": prefix.num_layers,
|
|
"sequence_length": prefix.sequence_length,
|
|
"first_key_shape": list(prefix.keys[0].shape),
|
|
"all_finite": bool(
|
|
all(
|
|
torch.isfinite(key).all() and torch.isfinite(value).all()
|
|
for key, value in zip(prefix.keys, prefix.values)
|
|
)
|
|
),
|
|
},
|
|
"expert_architecture": asdict(architecture),
|
|
"expert_parameters": parameter_counts(expert),
|
|
"diffusion_output_shape": list(predicted_action.shape),
|
|
"diffusion_output_finite": bool(torch.isfinite(predicted_action).all()),
|
|
"action_mse": float(
|
|
torch.nn.functional.mse_loss(predicted_action, target_action).cpu()
|
|
),
|
|
**trajectory_metrics(predicted_xy, target_xy),
|
|
"quality_note": (
|
|
"This verifies exact structured-generation KV compatibility. The expert was trained "
|
|
"with the earlier plain-text Stage-1 prefix, so trajectory quality is diagnostic only."
|
|
),
|
|
}
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
write_json(output_dir / "handoff_metrics.json", result)
|
|
save_prediction_arrays(
|
|
output_dir / "handoff_prediction.npz",
|
|
history[:, :2],
|
|
target_action[0],
|
|
predicted_action[0],
|
|
target_xy[0],
|
|
predicted_xy[0],
|
|
)
|
|
save_trajectory_plot(
|
|
output_dir / "handoff_trajectory.png",
|
|
history[:, :2],
|
|
target_xy[0],
|
|
predicted_xy[0],
|
|
prediction_label="diffusion from structured VLM KV",
|
|
title="Structured VLM handoff to diffusion",
|
|
)
|
|
print(json.dumps(result, indent=2, ensure_ascii=False))
|
|
del expert, vlm
|
|
empty_device_cache(device)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|