tinyjev 0.1.0__py3-none-any.whl

tinyjev/__init__.py

@@ -0,0 +1,13 @@
+"""tinyjev: run tiny Jev-style decision models locally, on any machine.
+
+    import tinyjev
+    agent = tinyjev.load("AnkitAI/tinyjev-kev-0.6b")          # MLX on Apple Silicon, torch elsewhere
+    agent.predict({"state": "...", "questions": {...}})       # System One request shape
+    agent.predict({"states": [...]})                          # NanoJev's native shape
+"""
+from .agent import Agent, load, normalize_request
+from .convert import convert
+from .registry import MODELS
+
+__all__ = ["Agent", "load", "convert", "normalize_request", "MODELS", "__version__"]
+__version__ = "0.1.0"

tinyjev/agent.py

@@ -0,0 +1,172 @@
+"""One loaded model: a family (prompt + head) on a backend (MLX or torch)."""
+from __future__ import annotations
+
+import json
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from . import backends, families
+from .families import softmax
+
+SCHEMA_VERSION = "tinyjev-v1"
+TYPE_ALIASES = {"noul": "boolean", "boolean": "boolean", "choice": "choice", "score": "score"}
+
+
+def _resolve(model_path, subfolder: Optional[str] = None) -> Path:
+    from .registry import resolve
+    repo, sub = resolve(str(model_path))
+    sub = subfolder or sub
+    p = Path(repo).expanduser()
+    if p.exists():
+        root = p.resolve()
+    else:
+        from huggingface_hub import snapshot_download
+        patterns = [f"{sub}/*"] if sub else ["*"]
+        root = Path(snapshot_download(repo, allow_patterns=patterns))
+    return root / sub if sub else root
+
+
+def normalize_request(payload: Dict[str, Any]) -> List[dict]:
+    """Accept NanoJev's {"states": [...]} or a System One {"state", "questions"} body.
+
+    Returns records: {"id", "state", "questions": [{"id","type","instructions","criteria"?}]}
+    with `noul` folded into `boolean`."""
+    if not isinstance(payload, dict):
+        raise ValueError("request must be a JSON object")
+    if "states" in payload:
+        if set(payload) != {"states"} or not isinstance(payload["states"], list) or not payload["states"]:
+            raise ValueError('native request must be exactly {"states": [...]} with at least one state')
+        raw = payload["states"]
+    elif "state" in payload and "questions" in payload:
+        raw = [{"id": payload.get("id") or "request", "state": payload["state"],
+                "questions": payload["questions"]}]
+    else:
+        raise ValueError('request needs either {"states": [...]} or {"state", "questions"}')
+
+    records, seen = [], set()
+    for s in raw:
+        if not isinstance(s, dict) or not {"id", "state", "questions"} <= set(s):
+            raise ValueError("each state needs id, state and questions")
+        sid = s["id"]
+        if not isinstance(sid, str) or not sid.strip() or sid in seen:
+            raise ValueError("state id must be a unique non-empty string")
+        seen.add(sid)
+        qs = s["questions"]
+        if not isinstance(qs, dict) or not qs:
+            raise ValueError(f"{sid}: questions must be a non-empty object")
+        out = []
+        for qid, q in qs.items():
+            if not isinstance(qid, str) or not qid.strip() or not isinstance(q, dict):
+                raise ValueError(f"{sid}: question ids must be non-empty strings mapping to objects")
+            typ = TYPE_ALIASES.get(q.get("type"))
+            if typ is None:
+                raise ValueError(f"{sid}:{qid}: unsupported question type {q.get('type')!r}")
+            item = {"id": qid, "type": typ, "instructions": q.get("instructions")}
+            if "criteria" in q and q["criteria"] is not None:
+                item["criteria"] = q["criteria"]
+            out.append(item)
+        records.append({"id": sid, "state": s["state"], "questions": out})
+    return records
+
+
+class Agent:
+    def __init__(self, model_path, backend: Optional[str] = None, device: Optional[str] = None,
+                 subfolder: Optional[str] = None, quantize: int = 0):
+        root = _resolve(model_path, subfolder)
+        manifest_path = root / "tinyjev.json"
+        if not manifest_path.exists():
+            raise FileNotFoundError(f"{root} is not a tinyjev checkpoint (no tinyjev.json); "
+                                    f"run `tinyjev convert` first")
+        self.manifest = json.loads(manifest_path.read_text())
+        if self.manifest.get("format") != "tinyjev-v2":
+            raise ValueError(f"{root}: layout {self.manifest.get('format')!r} is not tinyjev-v2; re-run `tinyjev convert`")
+        self.manifest["backbone_config"] = json.loads((root / "config.json").read_text())
+        self.root = root
+        self.family = families.make(self.manifest["family"], root, self.manifest)
+        name = backend or backends.default_backend()
+        kw = {"device": device} if (device and name == "torch") else {}
+        if name == "torch":
+            if quantize:
+                raise ValueError("quantize is only implemented on the mlx backend for now")
+            from .backends.torch_backend import Qwen3Backbone
+            self.backbone = Qwen3Backbone(self.manifest["backbone_config"], str(root / "model.safetensors"), **kw)
+        else:
+            self.backbone = backends.make(name, self.manifest["backbone_config"], str(root / "model.safetensors"),
+                                          quantize=quantize)
+        self.backend = self.backbone.name
+        self.quantize = int(quantize)
+
+    @property
+    def name(self) -> str:
+        return self.manifest.get("name", self.manifest["family"])
+
+    def _run(self, records: List[dict], temperature: float):
+        results, paths, elapsed = {}, 0, 0.0
+        for rec in records:
+            enc = self.family.encode(rec)
+            t = time.perf_counter()
+            hs = self.backbone.hidden_rows(enc.prefix, enc.rows, self.family.pad_token_id)
+            logits = self.family.logits(hs, enc)
+            elapsed += time.perf_counter() - t
+            paths += len(enc.rows)
+            for q in enc.questions:
+                z = logits[q["id"]]
+                probs = softmax(z / temperature).tolist()
+                results[f"{rec['id']}:{q['id']}"] = {
+                    "state_id": rec["id"], "qid": q["id"], "type": q["type"], "keys": q["keys"],
+                    "logits": [float(v) for v in z], "probabilities": probs,
+                    "answer": self.family.answer(q, probs),
+                    "path_token_counts": [len(enc.prefix) + len(enc.rows[i]) for i in
+                                          (q["rows"] if "rows" in q else [q["row"]])]}
+        return results, paths, elapsed * 1000.0
+
+    def logits(self, payload: Dict[str, Any]) -> Dict[str, dict]:
+        """Raw per-candidate logits/probabilities keyed by `<state id>:<question id>`."""
+        results, _, _ = self._run(normalize_request(payload), 1.0)
+        return {k: {kk: v[kk] for kk in ("type", "logits", "probabilities", "answer", "path_token_counts")}
+                | {"candidate_ids": v["keys"]} for k, v in results.items()}
+
+    def predict(self, payload: Dict[str, Any], temperature: float = 1.0) -> Dict[str, Any]:
+        if not isinstance(temperature, (int, float)) or isinstance(temperature, bool) \
+                or temperature <= 0 or temperature != temperature:
+            raise ValueError("temperature must be a finite positive number")
+        records = normalize_request(payload)
+        results, paths, ms = self._run(records, float(temperature))
+        states = {r["id"]: {"id": r["id"], "answers": {}} for r in records}
+        for v in results.values():
+            states[v["state_id"]]["answers"][v["qid"]] = v["answer"]
+        return {
+            "schema_version": SCHEMA_VERSION,
+            "model": {"name": self.name, "family": self.manifest["family"], "backend": self.backend,
+                      "quantize_bits": self.quantize or None,
+                      "directory": str(self.root), "upstream": self.manifest.get("upstream", {})},
+            "temperature": {"value": float(temperature)},
+            "execution": {"states": len(records), "questions": len(results), "candidate_paths": paths,
+                          "autoregressive_decode_steps": 0, "model_ms": round(ms, 2)},
+            "states": list(states.values()),
+        }
+
+    def systemone(self, body: Dict[str, Any]) -> Dict[str, Any]:
+        """TypeSafe System One response for a System One request."""
+        res = self.predict(body)
+        answers = {}
+        for qid, a in res["states"][0]["answers"].items():
+            if a["type"] in ("boolean", "noul"):
+                answers[qid] = {"type": "noul", "noul": round(float(a["p_true"]), 4)}
+            elif a["type"] == "choice":
+                answers[qid] = {"type": "choice", "choice": a["choice"],
+                                "confidence": a.get("confidence"), "probabilities": a["probabilities"]}
+            else:
+                answers[qid] = {"type": "score", "score": round(float(a["score"]), 4),
+                                "legend": a.get("legend"), "probabilities": a["probabilities"],
+                                "confidence": a.get("confidence")}
+        return {"model": body.get("model") or self.name, "answers": answers,
+                "latency_ms": res["execution"]["model_ms"]}
+
+
+def load(model_path, backend: Optional[str] = None, device: Optional[str] = None,
+         subfolder: Optional[str] = None, quantize: int = 0) -> Agent:
+    """`load("kev-0.6b")` (alias), `load("/path/to/dir")`, or `load("org/repo", subfolder="name")`.
+    quantize=8 or 4 quantizes the backbone's Linear layers at load time (mlx backend)."""
+    return Agent(model_path, backend=backend, device=device, subfolder=subfolder, quantize=quantize)

tinyjev/backends/__init__.py

@@ -0,0 +1,46 @@
+"""Backbone backends. Each exposes `Qwen3Backbone(config, weights_path)` with
+
+    hidden_rows(prefix: list[int], suffixes: list[list[int]], pad_token: int) -> list[np.ndarray]
+
+returning the fp32 hidden states [len(prefix)+len(suffix_i), d] of every row prefix+suffix_i.
+Attention is causal, so a shared prefix can be run once and its KV broadcast to every suffix.
+"""
+from __future__ import annotations
+
+import platform
+
+
+def available() -> list[str]:
+    out = []
+    try:
+        import mlx.core  # noqa: F401
+        out.append("mlx")
+    except Exception:
+        pass
+    try:
+        import torch  # noqa: F401
+        out.append("torch")
+    except Exception:
+        pass
+    return out
+
+
+def default_backend() -> str:
+    have = available()
+    if "mlx" in have and platform.machine() == "arm64" and platform.system() == "Darwin":
+        return "mlx"
+    if have:
+        return have[0]
+    raise RuntimeError("install either mlx (Apple Silicon) or torch")
+
+
+def make(name: str, config: dict, weights_path: str, prefix_min_tokens: int = 96, quantize: int = 0):
+    if name == "mlx":
+        from .mlx_backend import Qwen3Backbone
+        return Qwen3Backbone(config, weights_path, prefix_min_tokens=prefix_min_tokens, quantize=quantize)
+    if name == "torch":
+        if quantize:
+            raise ValueError("quantize is only implemented on the mlx backend for now")
+        from .torch_backend import Qwen3Backbone
+        return Qwen3Backbone(config, weights_path, prefix_min_tokens=prefix_min_tokens)
+    raise ValueError(f"unknown backend {name!r}; choose mlx or torch")

tinyjev/backends/mlx_backend.py

@@ -0,0 +1,74 @@
+"""Qwen3 backbone on MLX with shared-prefix KV reuse."""
+from __future__ import annotations
+
+from typing import List, Sequence
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+from mlx_lm.models.qwen3 import ModelArgs, Qwen3Model
+
+
+def qwen3_args(cfg: dict) -> ModelArgs:
+    rope = cfg.get("rope_parameters") or {}
+    return ModelArgs(
+        model_type="qwen3", hidden_size=cfg["hidden_size"],
+        num_hidden_layers=cfg["num_hidden_layers"], intermediate_size=cfg["intermediate_size"],
+        num_attention_heads=cfg["num_attention_heads"], rms_norm_eps=cfg["rms_norm_eps"],
+        vocab_size=cfg["vocab_size"], num_key_value_heads=cfg["num_key_value_heads"],
+        head_dim=cfg["head_dim"], max_position_embeddings=cfg.get("max_position_embeddings", 40960),
+        rope_theta=rope.get("rope_theta", cfg.get("rope_theta", 1000000)),
+        tie_word_embeddings=cfg.get("tie_word_embeddings", True))
+
+
+class Qwen3Backbone:
+    name = "mlx"
+
+    def __init__(self, config: dict, weights_path: str, prefix_min_tokens: int = 96,
+                 quantize: int = 0, group_size: int = 64):
+        self.model = Qwen3Model(qwen3_args(config))
+        weights = mx.load(weights_path)          # model.safetensors: standard Qwen3Model keys
+        self.model.load_weights(list(weights.items()))
+        if quantize:
+            # Backbone Linear layers only (the decision head runs fp32 in numpy). Embeddings stay
+            # unquantized: they are gathered, not multiplied, and small models lose most at 4-bit.
+            nn.quantize(self.model, group_size=group_size, bits=int(quantize),
+                        class_predicate=lambda _p, m: isinstance(m, nn.Linear))
+        self.quantize = int(quantize)
+        self.model.eval()
+        mx.eval(self.model.parameters())
+        self.prefix_min_tokens = prefix_min_tokens
+
+    def _rows_plain(self, rows: Sequence[Sequence[int]], pad: int) -> List[np.ndarray]:
+        lengths = [len(r) for r in rows]
+        width = max(lengths)
+        tokens = mx.array([list(r) + [pad] * (width - len(r)) for r in rows])
+        h = self.model(tokens).astype(mx.float32)
+        mx.eval(h)
+        arr = np.array(h)
+        return [arr[i, :n] for i, n in enumerate(lengths)]
+
+    def _rows_shared(self, prefix: Sequence[int], suffixes: Sequence[Sequence[int]], pad: int):
+        from mlx_lm.models.cache import KVCache
+        cache = [KVCache() for _ in self.model.layers]
+        h_prefix = self.model(mx.array([list(prefix)]), cache=cache).astype(mx.float32)
+        k = len(suffixes)
+        shared = []
+        for c in cache:
+            keys, values = c.state
+            bc = KVCache()
+            bc.keys, bc.values = mx.repeat(keys, k, axis=0), mx.repeat(values, k, axis=0)
+            bc.offset = keys.shape[2]
+            shared.append(bc)
+        lengths = [len(s) for s in suffixes]
+        width = max(lengths)
+        tokens = mx.array([list(s) + [pad] * (width - len(s)) for s in suffixes])
+        h = self.model(tokens, cache=shared).astype(mx.float32)
+        mx.eval(h_prefix, h)
+        hp, hs = np.array(h_prefix)[0], np.array(h)
+        return [np.concatenate([hp, hs[i, :n]], axis=0) for i, n in enumerate(lengths)]
+
+    def hidden_rows(self, prefix: Sequence[int], suffixes: Sequence[Sequence[int]], pad: int):
+        if len(suffixes) >= 2 and len(prefix) >= self.prefix_min_tokens:
+            return self._rows_shared(prefix, suffixes, pad)
+        return self._rows_plain([list(prefix) + list(s) for s in suffixes], pad)

tinyjev/backends/torch_backend.py

@@ -0,0 +1,71 @@
+"""Qwen3 backbone on PyTorch (CPU, CUDA or MPS) with shared-prefix KV reuse."""
+from __future__ import annotations
+
+from typing import List, Sequence
+
+import numpy as np
+
+
+class Qwen3Backbone:
+    name = "torch"
+
+    def __init__(self, config: dict, weights_path: str, prefix_min_tokens: int = 96, device=None):
+        """`weights_path` is <root>/model.safetensors; the root is a standard transformers model dir."""
+        import torch
+        from pathlib import Path
+        from transformers import AutoConfig, AutoModel
+
+        self.torch = torch
+        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else
+                                              "mps" if torch.backends.mps.is_available() else "cpu"))
+        root = str(Path(weights_path).parent)
+        cfg = AutoConfig.from_pretrained(root)
+        # transformers 4.x ignores `rope_parameters`; the converter also writes rope_theta at the top
+        # level, but force it here too in case the config came from elsewhere
+        rope = config.get("rope_parameters") or {}
+        cfg.rope_theta = rope.get("rope_theta", config.get("rope_theta", getattr(cfg, "rope_theta", None)))
+        cfg.use_cache = True
+        dtype = torch.float16 if self.device.type != "cpu" else torch.float32
+        self.model = AutoModel.from_pretrained(root, config=cfg, dtype=dtype, attn_implementation="sdpa")
+        self.model.to(self.device).eval()
+        self.prefix_min_tokens = prefix_min_tokens
+
+    def _pad(self, rows, pad):
+        torch = self.torch
+        lengths = [len(r) for r in rows]
+        width = max(lengths)
+        ids = torch.full((len(rows), width), pad, dtype=torch.long)
+        att = torch.zeros((len(rows), width), dtype=torch.long)
+        for i, r in enumerate(rows):
+            ids[i, :len(r)] = torch.tensor(r)
+            att[i, :len(r)] = 1
+        return ids.to(self.device), att.to(self.device), lengths
+
+    def _rows_plain(self, rows, pad) -> List[np.ndarray]:
+        torch = self.torch
+        ids, att, lengths = self._pad(rows, pad)
+        with torch.inference_mode():
+            h = self.model(input_ids=ids, attention_mask=att, use_cache=False).last_hidden_state.float().cpu().numpy()
+        return [h[i, :n] for i, n in enumerate(lengths)]
+
+    def _rows_shared(self, prefix, suffixes, pad):
+        torch = self.torch
+        from transformers import DynamicCache
+        k = len(suffixes)
+        with torch.inference_mode():
+            p = torch.tensor([list(prefix)], device=self.device)
+            out = self.model(input_ids=p, past_key_values=DynamicCache(), use_cache=True)
+            hp = out.last_hidden_state[0].float().cpu().numpy()
+            cache = out.past_key_values
+            cache.reorder_cache(torch.zeros(k, dtype=torch.long, device=self.device))
+            ids, att, lengths = self._pad(suffixes, pad)
+            full_att = torch.cat([torch.ones((k, len(prefix)), dtype=torch.long, device=self.device), att], 1)
+            pos = torch.arange(len(prefix), len(prefix) + ids.shape[1], device=self.device)[None].expand(k, -1)
+            h = self.model(input_ids=ids, attention_mask=full_att, position_ids=pos,
+                           past_key_values=cache, use_cache=True).last_hidden_state.float().cpu().numpy()
+        return [np.concatenate([hp, h[i, :n]], axis=0) for i, n in enumerate(lengths)]
+
+    def hidden_rows(self, prefix, suffixes, pad):
+        if len(suffixes) >= 2 and len(prefix) >= self.prefix_min_tokens:
+            return self._rows_shared(prefix, suffixes, pad)
+        return self._rows_plain([list(prefix) + list(s) for s in suffixes], pad)

tinyjev/cli.py

@@ -0,0 +1,64 @@
+"""tinyjev command line."""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+
+
+def main(argv=None) -> int:
+    parser = argparse.ArgumentParser(prog="tinyjev", description="tiny Jev-style decision models, locally")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    c = sub.add_parser("convert", help="convert an upstream checkpoint into the tinyjev layout")
+    c.add_argument("family", choices=["nanojev", "kev"])
+    c.add_argument("dest")
+    c.add_argument("--source", help="nanojev: upstream checkpoint dir")
+    c.add_argument("--adapter", help="kev: adapter dir (adapter_model.safetensors + head.pt)")
+    c.add_argument("--base", help="kev: Qwen3 base model dir")
+    c.add_argument("--name", default=None)
+    c.add_argument("--dtype", default="float16", choices=["float16", "float32"])
+
+    def model_args(p):
+        p.add_argument("--backend", choices=["mlx", "torch"], default=None)
+        p.add_argument("--device", default=None, help="torch backend: cpu, mps or cuda")
+        p.add_argument("--quantize", type=int, default=0, choices=[0, 4, 8], help="mlx: quantize backbone Linear layers to 4 or 8 bits at load")
+
+    s = sub.add_parser("serve", help="serve a checkpoint over HTTP (/predict, /v1/systemone)")
+    s.add_argument("model"); s.add_argument("--host", default="127.0.0.1")
+    s.add_argument("--port", type=int, default=8077); model_args(s)
+
+    a = sub.add_parser("ask", help="answer one request (System One or native shape) from a file or stdin")
+    a.add_argument("model"); a.add_argument("request", nargs="?", default="-"); model_args(a)
+
+    ls = sub.add_parser("models", help="list known checkpoints")
+
+    args = parser.parse_args(argv)
+
+    if args.command == "models":
+        from .registry import MODELS
+        for k, v in MODELS.items():
+            print(f"{k:<13} {v['params']:<5} {v['family']:<8} {v['repo']:<22} {v['what']}")
+        print("\nOther projects' models: convert them locally, e.g.\n"
+              "  tinyjev convert nanojev <dest> --source <C-Tianyu/NanoJev checkout>\n"
+              "  tinyjev convert kev <dest> --adapter <kev run or hub dir> --base <Qwen3 base dir>")
+        return 0
+    if args.command == "convert":
+        from .convert import convert
+        kw = {"source": args.source, "adapter": args.adapter, "base": args.base, "name": args.name or args.family}
+        convert(args.family, args.dest, dtype=args.dtype, **{k: v for k, v in kw.items() if v})
+        return 0
+
+    from . import load
+    if args.command == "serve":
+        from .serve import serve
+        serve(load(args.model, backend=args.backend, device=args.device, quantize=args.quantize), host=args.host, port=args.port)
+        return 0
+    payload = json.load(sys.stdin) if args.request == "-" else json.load(open(args.request))
+    agent = load(args.model, backend=args.backend, device=args.device, quantize=args.quantize)
+    print(json.dumps(agent.predict(payload), indent=2, ensure_ascii=False))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())