mindxtrain/data/scripts.py

"""Author a training *script* for an *actor*.

The mindXtrain mental model: a **model is an actor**; an actor has a **persona**
(identity / voice) and a **script** (the training examples — the "impression" left
on the actor when it trains). This module is the clean-room primitive for building
a script from a persona + a handful of exchanges, written as the OpenAI-chat JSONL
that `data.source: local` ingests (`{"messages": [{role, content}, ...]}`).

Clean-room: the Codephreak persona is *loaded* at runtime from
`MINDXTRAIN_PERSONA_PATH` (or a caller-supplied path); we never copy mindX bytes —
we read recognised fields and ignore the rest.

Pure stdlib + pydantic; importable on a base install (no `--extra ml`).
"""

from __future__ import annotations

import json
import os
from pathlib import Path

from pydantic import BaseModel, ConfigDict, Field

# Recognised keys for a persona's identity/voice in an mindX-style persona JSON.
# We map these defensively — any other keys are ignored (clean-room read).
_NAME_KEYS = ("name", "persona", "id", "title")
_SYSTEM_KEYS = ("system_prompt", "system", "description", "bio", "summary", "prompt")
_VOICE_KEYS = ("voice_examples", "examples", "utterances", "samples", "voice")


class Persona(BaseModel):
    """The identity to imprint onto an actor."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    name: str = "actor"
    system_prompt: str = ""
    voice_examples: list[str] = Field(
        default_factory=list,
        description="Example in-voice utterances; seed rows + the imprint baseline.",
    )


class Exchange(BaseModel):
    """One user→assistant turn in a script."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    user: str
    assistant: str


def load_persona(path: str | Path | None = None) -> Persona:
    """Load a persona, clean-room, from JSON.

    Resolution: explicit `path` → `MINDXTRAIN_PERSONA_PATH` → a built-in minimal
    default. Reads only recognised fields; unknown keys are ignored so an
    arbitrary mindX persona file maps cleanly without copying its schema.
    """
    resolved = path or os.environ.get("MINDXTRAIN_PERSONA_PATH")
    if not resolved:
        return _default_persona()
    p = Path(resolved).expanduser()
    if not p.is_file():
        return _default_persona()
    try:
        raw = json.loads(p.read_text())
    except (json.JSONDecodeError, OSError):
        return _default_persona()
    if not isinstance(raw, dict):
        return _default_persona()
    return persona_from_dict(raw)


def persona_from_dict(raw: dict) -> Persona:
    """Build a Persona from a loosely-shaped dict (recognised keys only)."""
    name = next((str(raw[k]) for k in _NAME_KEYS if raw.get(k)), "actor")
    system = next((str(raw[k]) for k in _SYSTEM_KEYS if raw.get(k)), "")
    voice: list[str] = []
    for k in _VOICE_KEYS:
        v = raw.get(k)
        if isinstance(v, list):
            voice.extend(str(x) for x in v if isinstance(x, (str, int, float)))
        elif isinstance(v, str):
            voice.append(v)
    return Persona(name=name, system_prompt=system, voice_examples=voice)


def _default_persona() -> Persona:
    return Persona(
        name="actor",
        system_prompt="You are a helpful, concise assistant.",
        voice_examples=[],
    )


def persona_system_prompt(persona: Persona) -> str:
    """The system message that fronts every row of the script.

    Uses the persona's own system prompt when present, otherwise synthesises a
    minimal one from the name so the actor still has an identity to imprint.
    """
    if persona.system_prompt.strip():
        return persona.system_prompt.strip()
    return f"You are {persona.name}. Stay in character and answer in your own voice."


def build_script_rows(
    persona: Persona,
    exchanges: list[Exchange],
    *,
    seed_voice: bool = True,
) -> list[dict]:
    """Turn a persona + exchanges into OpenAI-chat rows for `source: local`.

    Each row carries the persona system prompt + one user→assistant turn. When
    `seed_voice` is set, the persona's voice examples are added as extra
    assistant-only demonstrations so a tiny model has voice to imprint even from
    very few exchanges.
    """
    system = persona_system_prompt(persona)
    rows: list[dict] = []
    for ex in exchanges:
        rows.append(
            {
                "messages": [
                    {"role": "system", "content": system},
                    {"role": "user", "content": ex.user},
                    {"role": "assistant", "content": ex.assistant},
                ],
            },
        )
    if seed_voice:
        for sample in persona.voice_examples:
            rows.append(
                {
                    "messages": [
                        {"role": "system", "content": system},
                        {"role": "user", "content": f"Say something as {persona.name}."},
                        {"role": "assistant", "content": sample},
                    ],
                },
            )
    return rows


def write_script_jsonl(rows: list[dict], out_path: str | Path) -> Path:
    """Write script rows as JSONL; returns the path. Parent dirs are created."""
    out = Path(out_path).expanduser()
    out.parent.mkdir(parents=True, exist_ok=True)
    with out.open("w", encoding="utf-8") as fh:
        for row in rows:
            fh.write(json.dumps(row, ensure_ascii=False) + "\n")
    return out


def derive_training_params(num_rows: int) -> dict[str, int]:
    """Derive CPU-imprint training params from the dataset size.

    A small persona/skill script must *overfit* to imprint (many epochs, grad_accum 1
    so a few-row script still does many optimizer steps); larger datasets taper toward
    ordinary SFT. Returns `{epochs, grad_accum, per_device}` the imprint lane can apply.
    """
    n = max(1, num_rows)
    if n <= 8:
        epochs, grad_accum = 24, 1
    elif n <= 32:
        epochs, grad_accum = 16, 1
    elif n <= 128:
        epochs, grad_accum = 8, 1
    elif n <= 512:
        epochs, grad_accum = 4, 2
    else:
        epochs, grad_accum = 2, 4
    return {"epochs": epochs, "grad_accum": grad_accum, "per_device": 1}


def author_script(
    *,
    out_path: str | Path,
    exchanges: list[Exchange],
    persona: Persona | None = None,
    persona_path: str | Path | None = None,
    seed_voice: bool = True,
) -> tuple[Path, int]:
    """One-call script authoring: persona + exchanges → JSONL on disk.

    Returns (path, row_count). The persona is taken as-given, else loaded
    clean-room from `persona_path` / `MINDXTRAIN_PERSONA_PATH` / the default.
    """
    actor_persona = persona or load_persona(persona_path)
    rows = build_script_rows(actor_persona, exchanges, seed_voice=seed_voice)
    path = write_script_jsonl(rows, out_path)
    return path, len(rows)


__all__ = [
    "Exchange",
    "Persona",
    "author_script",
    "build_script_rows",
    "derive_training_params",
    "load_persona",
    "persona_from_dict",
    "persona_system_prompt",
    "write_script_jsonl",
]