tinymind-native-8b-remote-handoff/bundle/model/axiom_lang.py

"""AxiomLang: a tiny AI-native language for building growing models.

The language is intentionally small. It describes a model seed, memory policy,
tool/retrieval gates, and self-growth rules, then compiles into TinyMind config
objects and an auditable growth plan.
"""

from __future__ import annotations

from dataclasses import asdict
from datetime import datetime, timezone
import json
from pathlib import Path
import re

from .config import OmegaConfig, axiomweave_config


TOKEN_RE = re.compile(r'"[^"]*"|[{}]|[A-Za-z_][\w-]*|[0-9]+(?:\.[0-9]+)?|[><=]+|x')


class AxiomLangError(ValueError):
    pass


class AxiomLangParser:
    def __init__(self, source: str):
        self.tokens = [tok for tok in TOKEN_RE.findall(source) if tok.strip()]
        self.i = 0

    def parse(self) -> dict:
        self._expect("ai")
        name = self._identifier()
        self._expect("{")
        spec = {
            "name": name,
            "size": "tiny",
            "dim": None,
            "layers": None,
            "heads": None,
            "head_dim": None,
            "memory": {},
            "tools": {},
            "gates": [],
            "growth": [],
            "precision": "bf16_quality",
        }
        while not self._peek_is("}"):
            key = self._identifier()
            if key == "size":
                spec["size"] = self._identifier()
            elif key in {"dim", "layers", "heads", "head_dim"}:
                spec[key] = int(self._number())
            elif key == "memory":
                spec["memory"].update(self._memory_clause())
            elif key == "tool":
                tool_name = self._identifier()
                state = self._identifier()
                payload = {"enabled": state == "on"}
                while not self._at_clause_boundary():
                    opt = self._identifier()
                    value = self._value()
                    payload[opt] = value
                spec["tools"][tool_name] = payload
            elif key == "gate":
                name = self._identifier()
                policy = self._identifier()
                spec["gates"].append({"name": name, "policy": policy})
            elif key == "grow":
                spec["growth"].append(self._growth_clause())
            elif key == "precision":
                spec["precision"] = self._identifier()
            else:
                raise AxiomLangError(f"unknown clause '{key}'")
        self._expect("}")
        if self.i != len(self.tokens):
            raise AxiomLangError(f"unexpected trailing token '{self.tokens[self.i]}'")
        return spec

    def _memory_clause(self) -> dict:
        data = {}
        while not self._at_clause_boundary():
            key = self._identifier()
            if key in {"slots", "ranks", "window", "timescales"}:
                data[key] = int(self._number())
            elif key == "persistent":
                data[key] = int(self._number())
            else:
                raise AxiomLangError(f"unknown memory option '{key}'")
        return data

    def _growth_clause(self) -> dict:
        self._expect("when")
        metric = self._identifier()
        op = self._operator()
        threshold = float(self._number())
        action = self._identifier()
        amount = int(self._number())
        limit_name = self._identifier()
        limit = int(self._number())
        if limit_name not in {"max", "max_rank", "max_dim", "max_layers"}:
            raise AxiomLangError("growth rule must end with max/max_rank/max_dim/max_layers")
        return {
            "metric": metric,
            "operator": op,
            "threshold": threshold,
            "action": action,
            "amount": amount,
            "limit_name": limit_name,
            "limit": limit,
        }

    def _value(self) -> int | float | str | bool:
        tok = self._take()
        if tok in {"on", "true"}:
            return True
        if tok in {"off", "false"}:
            return False
        if tok.startswith('"'):
            return tok[1:-1]
        if re.fullmatch(r"[0-9]+", tok):
            return int(tok)
        if re.fullmatch(r"[0-9]+\.[0-9]+", tok):
            return float(tok)
        return tok

    def _at_clause_boundary(self) -> bool:
        return self._peek_is("}") or self._peek() in {"size", "dim", "layers", "heads", "head_dim", "memory", "tool", "gate", "grow", "precision"}

    def _identifier(self) -> str:
        tok = self._take()
        if not re.fullmatch(r"[A-Za-z_][\w-]*", tok):
            raise AxiomLangError(f"expected identifier, got '{tok}'")
        return tok

    def _number(self) -> str:
        tok = self._take()
        if not re.fullmatch(r"[0-9]+(?:\.[0-9]+)?", tok):
            raise AxiomLangError(f"expected number, got '{tok}'")
        return tok

    def _operator(self) -> str:
        tok = self._take()
        if tok not in {">", "<", ">=", "<=", "=", "=="}:
            raise AxiomLangError(f"expected comparison operator, got '{tok}'")
        return "==" if tok == "=" else tok

    def _expect(self, expected: str) -> None:
        tok = self._take()
        if tok != expected:
            raise AxiomLangError(f"expected '{expected}', got '{tok}'")

    def _peek(self) -> str | None:
        return self.tokens[self.i] if self.i < len(self.tokens) else None

    def _peek_is(self, value: str) -> bool:
        return self._peek() == value

    def _take(self) -> str:
        if self.i >= len(self.tokens):
            raise AxiomLangError("unexpected end of source")
        tok = self.tokens[self.i]
        self.i += 1
        return tok


def parse_axiomlang(source: str) -> dict:
    return AxiomLangParser(source).parse()


def compile_axiomlang(source: str) -> dict:
    spec = parse_axiomlang(source)
    cfg = axiomweave_config(str(spec["size"]))
    overrides = {
        "dim": "dim",
        "layers": "n_layers",
        "heads": "n_heads",
        "head_dim": "head_dim",
    }
    for source_key, cfg_key in overrides.items():
        if spec[source_key] is not None:
            setattr(cfg, cfg_key, int(spec[source_key]))
    memory = spec["memory"]
    if "slots" in memory:
        cfg.memory_slots = int(memory["slots"])
    if "ranks" in memory:
        cfg.memory_ranks = int(memory["ranks"])
    if "window" in memory:
        cfg.local_window = int(memory["window"])
    if "timescales" in memory:
        cfg.timescale_count = int(memory["timescales"])
    if "persistent" in memory:
        cfg.max_persistent_tokens = int(memory["persistent"])
    cfg.precision_mode = "auto" if spec["precision"] == "auto" else str(spec["precision"])
    retrieval = spec["tools"].get("retrieval")
    if retrieval:
        cfg.retrieval_top_k = int(retrieval.get("top_k", cfg.retrieval_top_k))
    compiled = {
        "schema_version": "tinymind-axiomlang-compile-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "language": "AxiomLang",
        "spec": spec,
        "config": asdict(cfg),
        "growth_plan": _growth_plan(spec, cfg),
        "claim_gate": {
            "world_first_language_claim_allowed": False,
            "world_best_model_claim_allowed": False,
            "reason": "AxiomLang compiles locally; superiority requires public external benchmarks and ablations.",
        },
    }
    return compiled


def _growth_plan(spec: dict, cfg: OmegaConfig) -> list[dict]:
    plan = []
    for rule in spec["growth"]:
        action = rule["action"]
        target = {
            "add_rank": "memory_ranks",
            "add_dim": "dim",
            "add_layer": "n_layers",
            "add_layers": "n_layers",
            "add_window": "local_window",
        }.get(action)
        if target is None:
            raise AxiomLangError(f"unknown growth action '{action}'")
        current = int(getattr(cfg, target))
        proposed = min(current + int(rule["amount"]), int(rule["limit"]))
        plan.append({**rule, "target": target, "current": current, "proposed": proposed, "resource_guard": "apply only if eval improves per added parameter"})
    return plan


def write_axiomlang_compile(source_path: str | Path, out_path: str | Path) -> dict:
    source = Path(source_path).read_text(encoding="utf-8")
    compiled = compile_axiomlang(source)
    out = Path(out_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(compiled, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    compiled["out_path"] = str(out)
    return compiled