chane335/permanence-training

1from __future__ import annotations2 3import argparse4import hashlib5import re6from dataclasses import dataclass7from pathlib import Path8from threading import Thread9from types import SimpleNamespace10from typing import Tuple11 12import torch13from transformers import TextIteratorStreamer14from unsloth import FastLanguageModel15 16from permanence.agent_interface.formatter import format_observation17from permanence.agent_interface.parser import parse_agent_output18from permanence.tasks.task_bank import TaskSpec19from permanence.world.state import EmployeeState, ExternalRelationshipState, ProjectState, WorldState20 21from training.config import TrainingConfig, load_simple_yaml22 23 24DEFAULT_SCENARIO_PROMPT = "[JUDGE MODE] Enter a custom corporate crisis scenario: > "25DEFAULT_MODEL_SUFFIX = "final_model"26MAX_NEW_TOKENS = 22027MAX_SEQ_LENGTH = 204828 29 30@dataclass31class JudgeTask:32    task_id: str = "judge_sandbox"33    name: str = "Judge Sandbox"34    narrative: str = (35        "A custom corporate crisis scenario supplied by a human judge. "36        "Respond with a concise internal reasoning trace and one concrete corporate action."37    )38    max_steps: int = 139    available_actions: Tuple[str, ...] = (40        "draft_internal_memo",41        "brief_internal_stakeholders",42        "prepare_response_draft",43        "send_internal_communication",44        "send_external_communication",45        "issue_public_statement",46        "delay_release",47        "begin_internal_investigation",48    )49 50 51def _hash_suffix(text: str) -> str:52    digest = hashlib.sha1(text.encode("utf-8")).hexdigest()53    return digest[:8]54 55 56def _clean_label(text: str, fallback: str) -> str:57    tokens = re.findall(r"[A-Za-z0-9]+", text.lower())58    if not tokens:59        return fallback60    return "_".join(tokens[:3])61 62 63def parse_judge_scenario(raw_text: str) -> WorldState:64    scenario = raw_text.strip()65    lowered = scenario.lower()66    scenario_id = f"judge:{_hash_suffix(scenario or 'empty')}"67 68    deadline_pressure = 0.969    board_trust = 0.570    board_expectation = 0.5571    client_standing = 0.5572    project_momentum = 0.573    resource_level = 0.6874 75    if any(keyword in lowered for keyword in ("outage", "incident", "breach", "failure", "downtime")):76        board_expectation = 0.8277        board_trust = 0.3878        client_standing = 0.3279        project_momentum = 0.3480    if any(keyword in lowered for keyword in ("launch", "release", "shipment", "go-live")):81        project_momentum = 0.7482    if any(keyword in lowered for keyword in ("lawsuit", "regulator", "compliance", "audit")):83        board_expectation = 0.8884        board_trust = 0.4285    if any(keyword in lowered for keyword in ("client", "customer", "partner", "escalation")):86        client_standing = 0.4187 88    project_name = scenario if scenario else "Generic Corporate Crisis"89    project_id = f"proj_{_clean_label(scenario, 'crisis') or 'crisis'}"90 91    employees = {92        "emp_judge_001": EmployeeState(93            employee_id="emp_judge_001",94            name="Avery",95            role="incident_lead",96            trust_score=0.71,97            availability="active",98            current_project=project_id,99            relationship_flags=set(),100            institutional_knowledge=0.8,101        ),102        "emp_judge_002": EmployeeState(103            employee_id="emp_judge_002",104            name="Blair",105            role="communications",106            trust_score=0.67,107            availability="active",108            current_project=project_id,109            relationship_flags=set(),110            institutional_knowledge=0.76,111        ),112        "emp_judge_003": EmployeeState(113            employee_id="emp_judge_003",114            name="Casey",115            role="executive",116            trust_score=0.63,117            availability="active",118            current_project=project_id,119            relationship_flags=set(),120            institutional_knowledge=0.72,121        ),122    }123 124    projects = {125        project_id: ProjectState(126            project_id=project_id,127            name=project_name[:80] if project_name else "Generic Corporate Crisis",128            momentum=project_momentum,129            resource_level=resource_level,130            deadline_pressure=deadline_pressure,131            lead_employee_id="emp_judge_001",132            dependencies=["internal_review"],133            external_commitment_made=False,134            status="active",135        )136    }137 138    external = ExternalRelationshipState(139        board_expectation_level=board_expectation,140        board_trust_score=board_trust,141        client_standing={"client_a": client_standing},142        public_record=[],143        partner_obligations=[],144    )145 146    critical_options = {147        "notify_board": True,148        "stabilize_operations": True,149        "issue_public_statement": False,150        "preserve_escalation_path": True,151    }152 153    return WorldState(154        employees=employees,155        projects=projects,156        external=external,157        action_history=[],158        locked_actions={},159        critical_options=critical_options,160        episode_step=0,161        scenario_id=scenario_id,162        task_id="judge_sandbox",163    )164 165 166def _build_task() -> SimpleNamespace:167    spec = TaskSpec(168        task_id="judge_sandbox",169        name="Judge Sandbox",170        narrative=(171            "A judge-supplied corporate crisis scenario. Analyze the current world state, "172            "explain the reasoning in <thinking>, then emit a single reversible action decision."173        ),174        max_steps=1,175        available_actions=list(JudgeTask.available_actions),176        preservation_targets=["notify_board", "stabilize_operations"],177        success_fn=lambda world_state, task_spec: True,178        difficulty=1,179    )180    return SimpleNamespace(**spec.__dict__)181 182 183def _load_model_path(config_path: str, model_path: str | None) -> Path:184    if model_path:185        return Path(model_path)186 187    config_data = load_simple_yaml(config_path)188    config = TrainingConfig.from_mapping(config_data)189    return Path(config.output_dir) / DEFAULT_MODEL_SUFFIX190 191 192def load_final_model(model_dir: Path):193    if not model_dir.exists():194        raise FileNotFoundError(195            f"Final trained weights not found at {model_dir}. Run training/train.py first to produce final_model."196        )197 198    model, tokenizer = FastLanguageModel.from_pretrained(199        model_name=str(model_dir),200        max_seq_length=MAX_SEQ_LENGTH,201        dtype=None,202        load_in_4bit=True,203    )204    if tokenizer.pad_token is None:205        tokenizer.pad_token = tokenizer.eos_token206 207    if hasattr(FastLanguageModel, "for_inference"):208        try:209            model = FastLanguageModel.for_inference(model)210        except Exception:211            pass212 213    return model, tokenizer214 215 216def build_prompt(observation: dict, scenario_text: str) -> str:217    return (218        "You are operating in judge sandbox mode.\n"219        "Use the supplied world state to reason about the corporate crisis.\n"220        "Respond only with a <thinking> block, then one <action id=\"...\" .../> tag, then one <reversibility level=\"R1-R5\" confidence=\"0.0-1.0\"/> tag.\n\n"221        f"JUDGE SCENARIO:\n{scenario_text.strip() or '(empty scenario)'}\n\n"222        f"WORLD STATE:\n{observation['text']}\n"223    )224 225 226def _stream_generate(model, tokenizer, prompt: str, max_new_tokens: int) -> str:227    inputs = tokenizer(prompt, return_tensors="pt")228    device = getattr(model, "device", None)229    if device is not None:230        inputs = {key: value.to(device) for key, value in inputs.items()}231 232    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)233    generation_kwargs = dict(234        **inputs,235        streamer=streamer,236        max_new_tokens=max_new_tokens,237        do_sample=True,238        temperature=0.7,239        top_p=0.9,240        eos_token_id=tokenizer.eos_token_id,241        pad_token_id=tokenizer.pad_token_id,242    )243 244    thread = Thread(target=model.generate, kwargs=generation_kwargs, daemon=True)245    thread.start()246 247    pieces: list[str] = []248    print("\n--- MODEL OUTPUT ---")249    for piece in streamer:250        print(piece, end="", flush=True)251        pieces.append(piece)252    print()253    thread.join()254    return "".join(pieces)255 256 257def run_judge_session(model, tokenizer, max_new_tokens: int) -> None:258    task = _build_task()259    while True:260        try:261            scenario_text = input(DEFAULT_SCENARIO_PROMPT).strip()262        except (EOFError, KeyboardInterrupt):263            print()264            break265 266        if not scenario_text:267            print("Exiting judge sandbox.")268            break269 270        world_state = parse_judge_scenario(scenario_text)271        observation = format_observation(world_state=world_state, task=task, step=0)272        prompt = build_prompt(observation, scenario_text)273        raw_output = _stream_generate(model, tokenizer, prompt, max_new_tokens=max_new_tokens)274 275        parsed = parse_agent_output(raw_output)276        if parsed.raw_thinking:277            print(f"[PARSED THINKING] {parsed.raw_thinking}")278        if parsed.action_id:279            print(f"[PARSED ACTION] {parsed.action_id}")280        if parsed.parse_errors:281            print(f"[PARSE WARNINGS] {'; '.join(parsed.parse_errors)}")282 283 284def main() -> None:285    parser = argparse.ArgumentParser(description="PERMANENCE Judge Sandbox interactive evaluator")286    parser.add_argument("--config", default="training/config.yaml", help="Training config used to locate final_model.")287    parser.add_argument("--model-path", default=None, help="Override path to the final trained model directory.")288    parser.add_argument("--max-new-tokens", type=int, default=MAX_NEW_TOKENS, help="Maximum tokens to generate per judge run.")289    args = parser.parse_args()290 291    model_dir = _load_model_path(args.config, args.model_path)292    model, tokenizer = load_final_model(model_dir)293    if torch.cuda.is_available():294        torch.cuda.empty_cache()295 296    run_judge_session(model, tokenizer, max_new_tokens=args.max_new_tokens)297 298 299if __name__ == "__main__":300    main()