| """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 |
|
|
| |
| |
| _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", |
| ] |
|
|
