rev-decision 0.2.0__py3-none-any.whl

rev/__init__.py

@@ -0,0 +1,141 @@
+"""
+rev: Unified System 1 Decision Engine with TypeSafe API and Sub-Microsecond Multi-Backbone Routing.
+
+Unifies Pretrained Encoders:
+  - ModernBERT-large (421M): English backbone with long-context, deep reasoning.
+  - mmBERT-base (322M): Multilingual backbone natively supporting 100+ languages.
+  - Typed Decisions (421M): Fine-tuned on specialized enterprise decision schemas.
+  - Causal LM Backbones: Document Prefix KV-Caching (<5ms repeated queries).
+
+Features:
+  - Sub-microsecond Unicode script routing (rev.lang).
+  - RLCD strictly proper scoring rules (log score + spherical + RPS) for calibrated probabilities.
+  - High-cardinality candidate shortlisting (rev.shortlist).
+  - Production presets for triage, email, guardrails, moderation, invoice, security, traces.
+"""
+
+__version__ = "0.2.0"
+
+# Language and Script Detection
+from . import lang
+
+# Common Utilities & Scoring Rules
+from .common import (
+    QTYPES,
+    QTYPE_NAMES,
+    serialize_state,
+    render_criterion,
+    render_options,
+    build_sequence,
+    collate_items,
+    proper_reward,
+    confidence_from_probs,
+    clamp_temperature,
+    ece_score,
+    amp_dtype,
+)
+
+# Encoder Model Architecture
+from .encoder_model import (
+    EncoderDecisionModel,
+    build_encoder_model,
+)
+
+# Inference Runtime
+from .agent import Agent
+
+# Unified 'One Model' Router
+from .router import (
+    UnifiedModel,
+    Router,
+    RouteDecision,
+    predict,
+    get_default_model,
+)
+
+# Primary 'Model' alias pointing to the unified engine
+Model = UnifiedModel
+
+# Candidate Shortlisting
+from .shortlist import (
+    shortlist_choice,
+    predict_shortlist,
+    embed_fn_from_agent,
+)
+
+# Workflow Presets
+from . import presets
+
+# Email Utilities
+from . import email
+from .email import email_state, clean_email_body
+
+# TypeSafe API Models (for backward compatibility and server endpoints)
+from .api import (
+    Noul,
+    Choice,
+    Score,
+    SystemOneRequest,
+    to_record,
+    to_answers,
+    render,
+    choice_confidence,
+    score_confidence,
+)
+
+# Optional Causal LM Backbone and Pointer Head
+try:
+    from .model import (
+        DecisionModel as CausalDecisionModel,
+        PointerHead,
+        encode,
+        user_tokens,
+        branch_mask_batch,
+        encode_state,
+        encode_question_branches,
+    )
+except ImportError:
+    CausalDecisionModel = None
+    PointerHead = None
+
+__all__ = [
+    # Unified Model
+    "Model",
+    "UnifiedModel",
+    "Router",
+    "Agent",
+    "EncoderDecisionModel",
+    "predict",
+    "get_default_model",
+    "RouteDecision",
+    # Presets & Modules
+    "presets",
+    "lang",
+    "email",
+    "email_state",
+    "clean_email_body",
+    # Shortlisting
+    "shortlist_choice",
+    "predict_shortlist",
+    "embed_fn_from_agent",
+    # Scoring & Calibration
+    "proper_reward",
+    "confidence_from_probs",
+    "clamp_temperature",
+    "ece_score",
+    "render_options",
+    "build_sequence",
+    # TypeSafe API
+    "Noul",
+    "Choice",
+    "Score",
+    "SystemOneRequest",
+    "to_record",
+    "to_answers",
+    "render",
+    "choice_confidence",
+    "score_confidence",
+    # Causal
+    "CausalDecisionModel",
+    "PointerHead",
+]

rev/agent.py

@@ -0,0 +1,256 @@
+"""
+High-level inference runtime for rev encoder decision models.
+Provides calibrated, non-autoregressive, sub-10ms decision evaluations.
+"""
+
+import json
+import os
+import warnings
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+
+from .common import (
+    QTYPES,
+    TEMP_MAX,
+    TEMP_MIN,
+    amp_dtype,
+    build_sequence,
+    clamp_temperature,
+    collate_items,
+    confidence_from_probs,
+    render_options,
+    temp_bucket,
+)
+from .encoder_model import EncoderDecisionModel, build_encoder_model
+
+
+def _fix_tokenizer_config(path: str):
+    """Ensure tokenizer_config.json loads cleanly across all transformers versions."""
+    cfg_file = os.path.join(path, "tokenizer", "tokenizer_config.json")
+    if not os.path.exists(cfg_file):
+        cfg_file = os.path.join(path, "tokenizer_config.json")
+        if not os.path.exists(cfg_file):
+            return
+    try:
+        with open(cfg_file) as f:
+            tcfg = json.load(f)
+        changed = False
+        if tcfg.get("tokenizer_class") in (None, "TokenizersBackend"):
+            tcfg["tokenizer_class"] = "PreTrainedTokenizerFast"
+            tcfg.pop("backend", None)
+            tcfg.pop("is_local", None)
+            changed = True
+        extra = tcfg.get("extra_special_tokens")
+        if isinstance(extra, list):
+            tcfg["extra_special_tokens"] = {"extra_%d" % i: t for i, t in enumerate(extra)}
+            changed = True
+        if changed:
+            with open(cfg_file, "w") as f:
+                json.dump(tcfg, f, indent=2)
+    except Exception:
+        pass
+
+
+class Agent:
+    """System 1 encoder decision runtime: fast, non-autoregressive, calibrated decisions."""
+
+    def __init__(
+        self,
+        model_id_or_path: str = "convaiinnovations/laya",
+        device: Optional[str] = None,
+        token: Optional[str] = None,
+        subfolder: Optional[str] = None,
+    ):
+        """
+        Load an encoder checkpoint.
+        `subfolder` can select specific checkpoints (e.g. 'multilingual' or 'typed-decisions').
+        """
+        from safetensors.torch import load_file
+        from transformers import AutoTokenizer
+
+        model_dir = model_id_or_path
+        if not os.path.exists(model_dir):
+            if model_id_or_path.startswith(("/", "./", "../")) or os.path.isabs(model_id_or_path):
+                raise FileNotFoundError(f"Local model path not found: {model_id_or_path!r}")
+            from huggingface_hub import snapshot_download
+
+            prefix = f"{subfolder}/" if subfolder else ""
+            kw = {
+                "token": token or os.environ.get("HF_TOKEN"),
+                "allow_patterns": [
+                    prefix + name
+                    for name in (
+                        "rl_agent_config.json",
+                        "model.safetensors",
+                        "tokenizer/*",
+                        "tokenizer_config.json",
+                        "vocab.json",
+                        "merges.txt",
+                        "encoder/*",
+                    )
+                ],
+            }
+            model_dir = snapshot_download(model_id_or_path, **kw)
+
+        if subfolder:
+            model_dir = os.path.join(model_dir, subfolder)
+            if not os.path.isdir(model_dir):
+                raise FileNotFoundError(f"Subfolder {subfolder!r} not found in {model_id_or_path!r}.")
+
+        _fix_tokenizer_config(model_dir)
+
+        cfg_path = os.path.join(model_dir, "rl_agent_config.json")
+        if os.path.exists(cfg_path):
+            with open(cfg_path) as f:
+                self.cfg = json.load(f)
+        else:
+            self.cfg = {
+                "encoder": "answerdotai/ModernBERT-large",
+                "head_layers": 2,
+                "max_len": 512,
+                "head_max_len": 192,
+            }
+
+        weights_path = os.path.join(model_dir, "model.safetensors")
+        weights = load_file(weights_path) if os.path.exists(weights_path) else None
+
+        tok_dir = os.path.join(model_dir, "tokenizer")
+        if not os.path.exists(tok_dir):
+            tok_dir = model_dir
+        self.tok = AutoTokenizer.from_pretrained(tok_dir, fix_markdown=False)
+
+        enc_dir = os.path.join(model_dir, "encoder")
+        self.model = build_encoder_model(self.cfg, enc_dir if os.path.exists(enc_dir) else None)
+
+        if weights is not None:
+            self.model.load_state_dict(weights, strict=False)
+
+        self.model.eval()
+
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = torch.device(device)
+        self.model.to(self.device)
+
+        self.max_len = int(self.cfg.get("max_len", 512))
+        self.head_max_len = int(self.cfg.get("head_max_len", 192))
+        self.temperatures = self.cfg.get("temperatures", {})
+
+    @torch.no_grad()
+    def system_one(
+        self,
+        state: Union[str, dict, list],
+        questions: Dict[str, Any],
+        temperatures: Optional[Dict[str, float]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Evaluate all questions for state in a single non-autoregressive forward pass.
+        Returns dict containing answers, probabilities, and confidence scores.
+        """
+        if not questions:
+            return {"answers": {}, "actions": {}}
+
+        items = []
+        q_order = []
+        for qid, q in questions.items():
+            q_norm = {
+                "t": q.get("type") or q.get("t", "choice"),
+                "ins": q.get("instructions") or q.get("ins", ""),
+                "crit": q.get("criteria") or q.get("crit", {}),
+            }
+            qtype_id = QTYPES[q_norm["t"]]
+            ids, markers = build_sequence(
+                self.tok,
+                state,
+                q_norm,
+                max_len=self.max_len,
+                head_max_len=self.head_max_len,
+            )
+            items.append({
+                "ids": ids,
+                "markers": markers,
+                "qtype": qtype_id,
+                "qid": qid,
+                "q": q_norm,
+            })
+            q_order.append((qid, q_norm))
+
+        pad_id = getattr(self.tok, "pad_token_id", None) or 0
+        batch = collate_items(items, pad_id=pad_id)
+
+        input_ids = batch["input_ids"].to(self.device)
+        attention_mask = batch["attention_mask"].to(self.device)
+        marker_pos = batch["marker_pos"].to(self.device)
+        marker_mask = batch["marker_mask"].to(self.device)
+        qtype = batch["qtype"].to(self.device)
+
+        logits, act_logits = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            marker_pos=marker_pos,
+            marker_mask=marker_mask,
+            qtype=qtype,
+        )
+
+        answers = {}
+        for i, (qid, q_norm) in enumerate(q_order):
+            k = int(marker_mask[i].sum().item())
+            row_logits = logits[i, :k].cpu().numpy()
+
+            # Temperature calibration
+            bucket = temp_bucket(QTYPES[q_norm["t"]], k)
+            t_val = 1.0
+            if temperatures and bucket in temperatures:
+                t_val = clamp_temperature(temperatures[bucket])
+            elif bucket in self.temperatures:
+                t_val = clamp_temperature(self.temperatures[bucket])
+
+            calibrated_logits = row_logits / t_val
+            exp_l = np.exp(calibrated_logits - np.max(calibrated_logits))
+            probs = exp_l / np.sum(exp_l)
+
+            conf = confidence_from_probs(probs, k)
+            t = q_norm["t"]
+
+            if t == "choice":
+                opts = render_options(q_norm)
+                crit = q_norm["crit"]
+                keys = list(crit.keys()) if isinstance(crit, dict) else [str(j) for j in range(len(opts))]
+                best_idx = int(np.argmax(probs))
+                answers[qid] = {
+                    "type": "choice",
+                    "choice": keys[best_idx] if best_idx < len(keys) else str(best_idx),
+                    "confidence": round(conf, 4),
+                    "probabilities": {keys[j]: round(float(probs[j]), 4) for j in range(min(len(keys), k))},
+                }
+            elif t == "score":
+                expected_score = float(np.sum(np.arange(k) * probs))
+                answers[qid] = {
+                    "type": "score",
+                    "score": round(expected_score, 3),
+                    "confidence": round(conf, 4),
+                    "probabilities": {str(j): round(float(probs[j]), 4) for j in range(k)},
+                }
+            else:  # noul
+                # probs[0] is false, probs[1] is true
+                p_true = float(probs[1]) if k >= 2 else 0.0
+                answers[qid] = {
+                    "type": "noul",
+                    "noul": round(p_true, 4),
+                    "confidence": round(conf, 4),
+                }
+
+        # Action head / deferral
+        act_probs = torch.softmax(act_logits, dim=-1).cpu().numpy()
+        actions = {}
+        for i, (qid, _) in enumerate(q_order):
+            actions[qid] = {
+                "action": int(np.argmax(act_probs[i])),
+                "probabilities": [round(float(p), 4) for p in act_probs[i]],
+            }
+
+        return {"answers": answers, "actions": actions}
+
+    predict = system_one

rev/api.py

@@ -0,0 +1,120 @@
+"""
+TypeSafe-compatible request/response shapes (POST /v1/systemone) mapped onto the single pointer primitive.
+
+Noul   -> 2 options [false, true];              answer = p(true)
+Choice -> options 'name' or 'name: desc';       answer = argmax, probabilities by name, confidence
+Score  -> options = ordered level descriptions; answer = expected level, legend, probabilities by index
+"""
+
+from typing import Any, Literal, Union, Dict, List
+from pydantic import BaseModel, Field, model_validator
+
+JSONContent = Union[str, dict, list, int, float, bool, None]
+MAX_OPTIONS = 255
+
+
+class Noul(BaseModel):
+    type: Literal["noul"] = "noul"
+    instructions: JSONContent
+    criteria: dict[str, JSONContent] | None = None
+
+
+class Choice(BaseModel):
+    type: Literal["choice"] = "choice"
+    instructions: JSONContent
+    criteria: dict[str, JSONContent]
+
+    @model_validator(mode="after")
+    def _check(self):
+        if not 1 <= len(self.criteria) <= MAX_OPTIONS:
+            raise ValueError(f"criteria must have 1..{MAX_OPTIONS} options")
+        return self
+
+
+class Score(BaseModel):
+    type: Literal["score"] = "score"
+    instructions: JSONContent
+    criteria: list[JSONContent] = Field(min_length=2, max_length=MAX_OPTIONS)
+
+
+Question = Union[Noul, Choice, Score]
+
+
+class SystemOneRequest(BaseModel):
+    state: JSONContent
+    model: str = "rev-latest"
+    questions: dict[str, Question] = Field(min_length=1)
+
+
+def render(v: JSONContent, indent: int = 0) -> str:
+    """Flatten str | object | array into text the model sees."""
+    pad = "  " * indent
+    if v is None:
+        return ""
+    if isinstance(v, (str, int, float, bool)):
+        return str(v)
+    if isinstance(v, list):
+        return "\n".join(f"{pad}- {render(x, indent + 1).lstrip()}" for x in v)
+    return "\n".join(
+        f"{pad}{k}:\n{render(x, indent + 1)}" if isinstance(x, (dict, list)) else f"{pad}{k}: {render(x)}"
+        for k, x in v.items()
+    )
+
+
+def option_text(name: str, desc: JSONContent) -> str:
+    return name if desc is None or desc == "" else f"{name}: {render(desc)}"
+
+
+def to_record(req: SystemOneRequest):
+    """-> internal record for encode(), plus per-question metadata to map probabilities back."""
+    qs, meta = [], []
+    for qid, q in req.questions.items():
+        instr = render(q.instructions)
+        if q.type == "noul":
+            c = q.criteria or {}
+            opts = [option_text("no", c.get("false")), option_text("yes", c.get("true"))]
+            meta.append({"id": qid, "type": "noul"})
+        elif q.type == "choice":
+            opts = [option_text(k, v) for k, v in q.criteria.items()]
+            meta.append({"id": qid, "type": "choice", "keys": list(q.criteria.keys())})
+        else:
+            opts = [render(x) for x in q.criteria]
+            meta.append({"id": qid, "type": "score", "legend": {str(i): render(x) for i, x in enumerate(q.criteria)}})
+        qs.append({"instr": instr, "options": opts, "label": 0})
+    return {"state": render(req.state), "questions": qs}, meta
+
+
+def choice_confidence(p: list[float]) -> float:
+    K = len(p)
+    return 1.0 if K == 1 else (max(p) - 1.0 / K) / (1.0 - 1.0 / K)
+
+
+def score_confidence(p: list[float]) -> float:
+    L = len(p)
+    mode = max(range(L), key=lambda i: p[i])
+    return 1.0 - sum(pi * abs(i - mode) for i, pi in enumerate(p)) / (L - 1)
+
+
+def to_answers(probs: list[list[float]], meta: list[dict]) -> dict[str, Any]:
+    out = {}
+    for p, m in zip(probs, meta):
+        if m["type"] == "noul":
+            out[m["id"]] = {"type": "noul", "noul": round(p[1], 3)}
+        elif m["type"] == "choice":
+            best_idx = max(range(len(p)), key=lambda i: p[i])
+            out[m["id"]] = {
+                "type": "choice",
+                "choice": m["keys"][best_idx],
+                "confidence": round(choice_confidence(p), 3),
+                "probabilities": {k: round(v, 4) for k, v in zip(m["keys"], p)},
+            }
+        else:
+            score = sum(i * pi for i, pi in enumerate(p))
+            out[m["id"]] = {
+                "type": "score",
+                "score": round(score, 2),
+                "confidence": round(score_confidence(p), 3),
+                "legend": m["legend"],
+                "probabilities": {str(i): round(v, 4) for i, v in enumerate(p)},
+            }
+    return out

rev/cache.py

@@ -0,0 +1,110 @@
+"""
+Thread-safe LRU state prefix KV-cache manager for rev.
+Caches precomputed key/value tensors of document prefixes to enable sub-5ms decision queries.
+"""
+
+import collections
+import hashlib
+import threading
+import time
+from typing import Any, Optional
+
+
+class CachedStateEntry:
+    def __init__(self, key: str, past_key_values: Any, state_len: int, token_ids: list[int] = None):
+        self.key = key
+        self.past_key_values = past_key_values
+        self.state_len = state_len
+        self.token_ids = token_ids or []
+        self.created_at = time.time()
+        self.last_accessed = time.time()
+        self.hits = 0
+
+    def touch(self):
+        self.last_accessed = time.time()
+        self.hits += 1
+
+
+class StateKVCacheManager:
+    """
+    LRU cache for prefix key-value states.
+    Thread-safe and eviction-ready.
+    """
+    def __init__(self, max_entries: int = 64):
+        self.max_entries = max_entries
+        self._cache: collections.OrderedDict[str, CachedStateEntry] = collections.OrderedDict()
+        self._lock = threading.Lock()
+        self._total_hits = 0
+        self._total_misses = 0
+
+    @staticmethod
+    def hash_state(state_text: str) -> str:
+        return hashlib.sha256(state_text.strip().encode("utf-8")).hexdigest()
+
+    def get(self, state_text: str) -> Optional[CachedStateEntry]:
+        key = self.hash_state(state_text)
+        return self.get_by_hash(key)
+
+    def get_by_hash(self, key: str) -> Optional[CachedStateEntry]:
+        with self._lock:
+            if key in self._cache:
+                entry = self._cache[key]
+                entry.touch()
+                self._cache.move_to_end(key)
+                self._total_hits += 1
+                return entry
+            else:
+                self._total_misses += 1
+                return None
+
+    def put(self, state_text: str, past_key_values: Any, state_len: int, token_ids: list[int] = None) -> str:
+        key = self.hash_state(state_text)
+        with self._lock:
+            if key in self._cache:
+                entry = self._cache[key]
+                entry.past_key_values = past_key_values
+                entry.state_len = state_len
+                entry.token_ids = token_ids or []
+                entry.touch()
+                self._cache.move_to_end(key)
+            else:
+                if len(self._cache) >= self.max_entries:
+                    self._cache.popitem(last=False)
+                entry = CachedStateEntry(key, past_key_values, state_len, token_ids or [])
+                self._cache[key] = entry
+        return key
+
+    def evict(self, key: str) -> bool:
+        with self._lock:
+            if key in self._cache:
+                del self._cache[key]
+                return True
+            return False
+
+    def clear(self):
+        with self._lock:
+            self._cache.clear()
+            self._total_hits = 0
+            self._total_misses = 0
+
+    def stats(self) -> dict:
+        with self._lock:
+            total = self._total_hits + self._total_misses
+            hit_rate = round(self._total_hits / total, 4) if total > 0 else 0.0
+            entries_info = [
+                {
+                    "hash": k,
+                    "state_len": entry.state_len,
+                    "hits": entry.hits,
+                    "age_seconds": round(time.time() - entry.created_at, 1),
+                }
+                for k, entry in self._cache.items()
+            ]
+            return {
+                "total_entries": len(self._cache),
+                "max_entries": self.max_entries,
+                "hits": self._total_hits,
+                "misses": self._total_misses,
+                "hit_rate": hit_rate,
+                "entries": entries_info,
+            }