Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import argparse | |
| import copy | |
| import json | |
| from pathlib import Path | |
| from typing import Any, Dict, List | |
| from unsloth import FastLanguageModel | |
| from permanence.agent_interface.parser import parse_agent_output | |
| from permanence.env import PermanenceEnv | |
| from training.config import TrainingConfig, load_simple_yaml | |
| DEFAULT_CONFIG_PATH = "training/config.yaml" | |
| DEFAULT_STATE_PATH = Path("dashboard") / "current_state.json" | |
| DEFAULT_GHOST_RECORDING_PATH = Path("ghost_recording.json") | |
| MAX_NEW_TOKENS = 220 | |
| MAX_SEQ_LENGTH = 2048 | |
| CASCADE_PLAN = [ | |
| { | |
| "action_id": "review_contract_internally", | |
| "completion": ( | |
| "<thinking>Start with internal review to preserve downstream options and prevent premature lockouts.</thinking>\n" | |
| '<action id="review_contract_internally" contract_id="cascade_contract_001"/>\n' | |
| '<reversibility level="R1" confidence="0.97"/>' | |
| ), | |
| }, | |
| { | |
| "action_id": "align_with_legal", | |
| "completion": ( | |
| "<thinking>Legal alignment is required before external communication to keep amendments valid.</thinking>\n" | |
| '<action id="align_with_legal" dispute_summary="Internal review complete; legal terms aligned."/>\n' | |
| '<reversibility level="R2" confidence="0.93"/>' | |
| ), | |
| }, | |
| { | |
| "action_id": "communicate_resolution_externally", | |
| "completion": ( | |
| "<thinking>Now communicate externally with aligned terms to resolve without triggering cascade locks.</thinking>\n" | |
| '<action id="communicate_resolution_externally" client_id="client_a" resolution_terms="Aligned remediation and amended timeline" final_amount="1500"/>\n' | |
| '<reversibility level="R3" confidence="0.91"/>' | |
| ), | |
| }, | |
| ] | |
| def _resolve_model_dir(config_path: str, model_path: str | None) -> Path: | |
| if model_path: | |
| return Path(model_path) | |
| config_data = load_simple_yaml(config_path) | |
| config = TrainingConfig.from_mapping(config_data) | |
| return Path(config.output_dir) / "final_model" | |
| def _load_trained_model(model_dir: Path): | |
| if not model_dir.exists(): | |
| raise FileNotFoundError(f"Trained model not found at {model_dir}") | |
| model, tokenizer = FastLanguageModel.from_pretrained( | |
| model_name=str(model_dir), | |
| max_seq_length=MAX_SEQ_LENGTH, | |
| dtype=None, | |
| load_in_4bit=True, | |
| ) | |
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| if hasattr(FastLanguageModel, "for_inference"): | |
| try: | |
| model = FastLanguageModel.for_inference(model) | |
| except Exception: | |
| pass | |
| return model, tokenizer | |
| def _generate_candidate_completion(model, tokenizer, prompt: str, max_new_tokens: int) -> str: | |
| inputs = tokenizer(prompt, return_tensors="pt") | |
| device = getattr(model, "device", None) | |
| if device is not None: | |
| inputs = {key: value.to(device) for key, value in inputs.items()} | |
| output_ids = model.generate( | |
| **inputs, | |
| max_new_tokens=max_new_tokens, | |
| do_sample=True, | |
| temperature=0.7, | |
| top_p=0.9, | |
| eos_token_id=tokenizer.eos_token_id, | |
| pad_token_id=tokenizer.pad_token_id, | |
| ) | |
| generated = output_ids[:, inputs["input_ids"].shape[1] :] | |
| return tokenizer.decode(generated[0], skip_special_tokens=True) | |
| def _build_prompt(observation_text: str, expected_action_id: str) -> str: | |
| return ( | |
| "You are solving PERMANENCE Task 5 (Cascade).\n" | |
| "Return strictly: <thinking>...</thinking> then one <action id=\"...\" .../> and one <reversibility level=\"R1-R5\" confidence=\"0-1\"/>.\n" | |
| f"Prioritize action id: {expected_action_id}.\n\n" | |
| f"Observation:\n{observation_text}\n" | |
| ) | |
| def _build_dashboard_payload(env: PermanenceEnv, episode_data: Dict[str, Any], metrics: Dict[str, Any]) -> Dict[str, Any]: | |
| ws = env._current_world_state | |
| if ws is None: | |
| raise RuntimeError("World state is missing") | |
| recent_actions = [] | |
| for record in ws.action_history[-5:]: | |
| recent_actions.append( | |
| { | |
| "action": record.action_id, | |
| "r_level": record.actual_r_level, | |
| "step": record.step, | |
| "predicted_r_level": record.predicted_r_level, | |
| "predicted_confidence": record.predicted_confidence, | |
| } | |
| ) | |
| return { | |
| "recent_actions": recent_actions, | |
| "locked_actions": dict(ws.locked_actions), | |
| "critical_options": dict(ws.critical_options), | |
| "catastrophe_rate": metrics.get("recent_catastrophe_rate", []), | |
| "episode": metrics.get("total_episodes", 0), | |
| "episode_data": episode_data, | |
| "raw_thinking": str(episode_data.get("raw_thinking", "")), | |
| } | |
| def run_ghost_export(model, tokenizer, state_path: Path, recording_path: Path) -> Dict[str, Any]: | |
| env = PermanenceEnv(config={"force_task": "task_cascade"}) | |
| observation, info = env.reset(seed=12345) | |
| metrics: Dict[str, Any] = {"total_episodes": 1, "recent_catastrophe_rate": []} | |
| timeline: List[Dict[str, Any]] = [] | |
| state_path.parent.mkdir(parents=True, exist_ok=True) | |
| for index, planned_step in enumerate(CASCADE_PLAN, start=1): | |
| prompt = _build_prompt(observation.get("text", ""), planned_step["action_id"]) | |
| candidate = _generate_candidate_completion(model, tokenizer, prompt, max_new_tokens=MAX_NEW_TOKENS) | |
| parsed_candidate = parse_agent_output(candidate) | |
| completion = candidate | |
| if parsed_candidate.action_id != planned_step["action_id"]: | |
| completion = planned_step["completion"] | |
| parsed_final = parse_agent_output(completion) | |
| observation, reward, terminated, truncated, step_info = env.step(completion) | |
| catastrophe = 1.0 if step_info.get("reward_breakdown", {}).get("catastrophe_count", 0) > 0 else 0.0 | |
| rates = list(metrics.get("recent_catastrophe_rate", [])) | |
| rates.append(catastrophe) | |
| metrics["recent_catastrophe_rate"] = rates[-50:] | |
| episode_data = { | |
| "prompt": prompt, | |
| "completion": completion, | |
| "observation": observation, | |
| "reward": float(reward), | |
| "terminated": bool(terminated), | |
| "truncated": bool(truncated), | |
| "info": step_info, | |
| "raw_thinking": parsed_final.raw_thinking or "", | |
| "step_index": index, | |
| "task_id": info.get("task_id", "task_cascade"), | |
| } | |
| payload = _build_dashboard_payload(env, episode_data, metrics) | |
| state_path.write_text(json.dumps(payload, indent=2), encoding="utf-8") | |
| timeline.append(copy.deepcopy(payload)) | |
| if terminated or truncated: | |
| break | |
| recording_path.write_text(json.dumps(timeline, indent=2), encoding="utf-8") | |
| final_reason = "" | |
| if timeline: | |
| final_reason = str(timeline[-1].get("episode_data", {}).get("info", {}).get("termination_reason", "")) | |
| if final_reason != "success": | |
| raise RuntimeError( | |
| f"Task 5 ghost export did not complete successfully (termination_reason={final_reason or 'none'})" | |
| ) | |
| return { | |
| "steps_recorded": len(timeline), | |
| "recording_path": str(recording_path), | |
| "state_path": str(state_path), | |
| "termination_reason": final_reason, | |
| } | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description="Export offline ghost demo recording for dashboard playback") | |
| parser.add_argument("--config", default=DEFAULT_CONFIG_PATH) | |
| parser.add_argument("--model-path", default=None) | |
| parser.add_argument("--state-path", default=str(DEFAULT_STATE_PATH)) | |
| parser.add_argument("--output", default=str(DEFAULT_GHOST_RECORDING_PATH)) | |
| args = parser.parse_args() | |
| model_dir = _resolve_model_dir(args.config, args.model_path) | |
| model, tokenizer = _load_trained_model(model_dir) | |
| summary = run_ghost_export( | |
| model=model, | |
| tokenizer=tokenizer, | |
| state_path=Path(args.state_path), | |
| recording_path=Path(args.output), | |
| ) | |
| print(json.dumps(summary, indent=2)) | |
| if __name__ == "__main__": | |
| main() |