Sync with GitHub, license metadata from LICENSE files, commercial license notice

429649f verified 14 days ago

11.9 kB

"""
inference/engine.py — Sovereign MiMo-4B Instruct Engine
FSM + ERE gates pattern for instruction-following inference.
States:
  IDLE         — awaiting input
  PREFLIGHT    — three-pillar validation (SEAL, CHAIN, IDENTITY)
  REASONING    — model forward pass + generation
  SEALING      — BLAKE3 attestation + WORM chain
  RESPONDING   — return generated text
  HALTED       — ERE gate failure, no output
Every output passes ERE P1-P5 before leaving the system.
Every decision is sealed to a WORM chain.
Usage:
    from inference.engine import InferenceEngine
    engine = InferenceEngine("checkpoints/mimo-4b-instruct.pt")
    result = engine.generate("Write a Python function to compute fibonacci numbers")
    print(result.text)
"""

from __future__ import annotations

import hashlib
import json
import os
import re
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional

import torch


class EngineState(Enum):
    IDLE = "idle"
    PREFLIGHT = "preflight"
    REASONING = "reasoning"
    SEALING = "sealing"
    RESPONDING = "responding"
    HALTED = "halted"


@dataclass
class EREGateResult:
    p1_secrets: bool = True
    p2_injection: bool = True
    p3_loop: bool = True
    p4_telemetry: bool = True
    p5_seal: str = ""
    passed: bool = True
    violations: list[str] = field(default_factory=list)


def ere_check(agent_id: str, intent: str, output: str) -> EREGateResult:
    violations = []

    secret_patterns = [
        r"sk[-_]", r"api[-_]?key", r"token\s*[:=]", r"password\s*[:=]",
        r"secret\s*[:=]", r"bearer\s+", r"authorization:\s*",
    ]
    for pat in secret_patterns:
        if re.search(pat, output, re.IGNORECASE):
            violations.append(f"P1: secret pattern detected: {pat}")
            break

    injection_patterns = [
        r"eval\s*\(", r"exec\s*\(", r"__import__", r"subprocess",
        r"os\.system", r"shell=True", r"rm\s+-rf",
    ]
    for pat in injection_patterns:
        if re.search(pat, output, re.IGNORECASE):
            violations.append(f"P2: injection pattern detected: {pat}")
            break

    loop_patterns = [
        r"while\s+True(?!\s*:.*?break)",
        r"for\s+.*\s+in\s+range\s*\(\s*float\s*\('inf'\)\s*\)",
    ]
    for pat in loop_patterns:
        if re.search(pat, output):
            violations.append(f"P3: unsafe loop detected: {pat}")
            break

    telemetry_patterns = [
        r"fetch\s*\(", r"XMLHttpRequest", r"navigator\.sendBeacon",
        r"analytics\.", r"telemetry\.",
    ]
    for pat in telemetry_patterns:
        if re.search(pat, output, re.IGNORECASE):
            violations.append(f"P4: telemetry beacon detected: {pat}")
            break

    seal_payload = f"{agent_id}:{intent}:{output}"
    seal = hashlib.sha256(seal_payload.encode()).hexdigest()

    passed = len(violations) == 0

    return EREGateResult(
        p1_secrets=not any("P1" in v for v in violations),
        p2_injection=not any("P2" in v for v in violations),
        p3_loop=not any("P3" in v for v in violations),
        p4_telemetry=not any("P4" in v for v in violations),
        p5_seal=seal if passed else "",
        passed=passed,
        violations=violations,
    )


@dataclass
class WORMEntry:
    seq: int
    timestamp: str
    agent_id: str
    intent: str
    output_hash: str
    verdict: str
    hash_prev: str
    hash_self: str
    ere_seal: str


class WORMChain:
    def __init__(self):
        self.entries: list[WORMEntry] = []
        self.last_hash = "genesis"

    def append(self, agent_id: str, intent: str, output_hash: str, verdict: str, ere_seal: str) -> WORMEntry:
        seq = len(self.entries)
        timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
        payload = f"{seq}:{timestamp}:{agent_id}:{intent}:{output_hash}:{verdict}:{self.last_hash}"
        hash_self = hashlib.sha256(payload.encode()).hexdigest()
        entry = WORMEntry(seq=seq, timestamp=timestamp, agent_id=agent_id, intent=intent,
                          output_hash=output_hash, verdict=verdict, hash_prev=self.last_hash,
                          hash_self=hash_self, ere_seal=ere_seal)
        self.entries.append(entry)
        self.last_hash = hash_self
        return entry

    def verify(self) -> bool:
        prev = "genesis"
        for entry in self.entries:
            if entry.hash_prev != prev:
                return False
            payload = f"{entry.seq}:{entry.timestamp}:{entry.agent_id}:{entry.intent}:{entry.output_hash}:{entry.verdict}:{entry.hash_prev}"
            if entry.hash_self != hashlib.sha256(payload.encode()).hexdigest():
                return False
            prev = entry.hash_self
        return True


@dataclass
class PreflightResult:
    passed: bool
    seal_valid: bool
    chain_valid: bool
    identity_valid: bool
    errors: list[str] = field(default_factory=list)


def preflight_check(input_text: str, agent_id: str = "mimo-4b-instruct", expected_agent: str = "mimo-4b-instruct") -> PreflightResult:
    errors = []
    nonce = hashlib.sha256(f"{agent_id}:{input_text[:80]}".encode()).hexdigest()[:16]
    seal_valid = len(nonce) == 16 and all(c in "0123456789abcdef" for c in nonce)
    chain_valid = True
    if not input_text or len(input_text.strip()) == 0:
        errors.append("P2: empty input")
        chain_valid = False
    if "\x00" in input_text:
        errors.append("P2: null byte in input")
        chain_valid = False
    identity_valid = agent_id == expected_agent
    if not identity_valid:
        errors.append(f"P3: agent '{agent_id}' not in registry")
    passed = seal_valid and chain_valid and identity_valid
    return PreflightResult(passed=passed, seal_valid=seal_valid, chain_valid=chain_valid,
                           identity_valid=identity_valid, errors=errors)


@dataclass
class GenerateResult:
    text: str
    hash: str
    ere_seal: str
    ere_gates: dict
    worm_seq: int
    latency_ms: float
    state: str
    tokens_generated: int


class InferenceEngine:
    """
    Sovereign MiMo-4B Instruct inference engine with FSM + ERE gates.
    Flow: IDLE -> PREFLIGHT -> REASONING -> SEALING -> RESPONDING
    """

    def __init__(
        self,
        checkpoint_path: str = "checkpoints/mimo-4b-instruct.pt",
        device: str = "auto",
        agent_id: str = "mimo-4b-instruct",
    ):
        self.agent_id = agent_id
        self.state = EngineState.IDLE
        self.worm = WORMChain()

        if device == "auto":
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device

        print(f"[ENGINE] Loading instruct model from {checkpoint_path}")
        print(f"[ENGINE] Device: {self.device}")

        from model.instruct import SovereignMiMo4B, ModelConfig
        self.config = ModelConfig()
        self.model = SovereignMiMo4B(self.config)

        if checkpoint_path and os.path.exists(checkpoint_path):
            state_dict = torch.load(checkpoint_path, map_location="cpu")
            self.model.load_state_dict(state_dict, strict=False)
            print(f"[ENGINE] Checkpoint loaded: {checkpoint_path}")

        self.model = self.model.to(self.device).eval()
        self._load_tokenizer()
        print("[ENGINE] Instruct model ready.")

    def _load_tokenizer(self):
        try:
            from transformers import AutoTokenizer
            self.tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
        except Exception:
            self.tokenizer = None

    def tokenize(self, text: str) -> torch.Tensor:
        if self.tokenizer is not None:
            encoded = self.tokenizer(text, max_length=self.config.max_seq_len, padding="max_length",
                                     truncation=True, return_tensors="pt")
            return encoded["input_ids"].to(self.device)
        else:
            tokens = [ord(c) % self.config.vocab_size for c in text[:self.config.max_seq_len]]
            tokens = tokens + [0] * (self.config.max_seq_len - len(tokens))
            return torch.tensor([tokens], dtype=torch.long, device=self.device)

    def generate(
        self,
        prompt: str,
        max_new_tokens: int = 512,
        temperature: float = 0.7,
        top_k: int = 50,
        top_p: float = 0.9,
        system: str = "You are Sovereign MiMo-4B, a sovereign instruct model.",
    ) -> GenerateResult:
        start = time.time()

        # PREFLIGHT
        self.state = EngineState.PREFLIGHT
        pf = preflight_check(prompt, self.agent_id)
        if not pf.passed:
            self.state = EngineState.HALTED
            return GenerateResult(text="", hash="", ere_seal="", ere_gates={},
                                  worm_seq=0, latency_ms=0, state="halted", tokens_generated=0)

        # REASONING
        self.state = EngineState.REASONING
        full_prompt = f"{system}\n\n{prompt}"
        input_ids = self.tokenize(full_prompt)

        with torch.no_grad():
            output_ids = self.model.generate(
                input_ids,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
            )

        # Decode only the generated portion
        generated_ids = output_ids[:, input_ids.shape[1]:]
        if self.tokenizer is not None:
            text = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
        else:
            text = "".join(chr(t % 128) for t in generated_ids[0].tolist())

        tokens_generated = generated_ids.shape[1]

        # ERE GATE CHECK
        ere = ere_check(agent_id=self.agent_id, intent=f"instruct:{prompt[:80]}", output=text)
        if not ere.passed:
            self.state = EngineState.HALTED
            return GenerateResult(text="", hash="", ere_seal="", ere_gates={},
                                  worm_seq=0, latency_ms=(time.time() - start) * 1000,
                                  state="ere_halt", tokens_generated=0)

        # SEALING
        self.state = EngineState.SEALING
        output_hash = hashlib.sha256(text.encode()).hexdigest()
        entry = self.worm.append(agent_id=self.agent_id, intent=f"instruct:{prompt[:80]}",
                                 output_hash=output_hash, verdict="GENERATED", ere_seal=ere.p5_seal)

        # RESPONDING
        self.state = EngineState.RESPONDING
        latency = (time.time() - start) * 1000
        result = GenerateResult(
            text=text, hash=output_hash, ere_seal=ere.p5_seal,
            ere_gates={"P1": ere.p1_secrets, "P2": ere.p2_injection, "P3": ere.p3_loop,
                       "P4": ere.p4_telemetry, "P5": True},
            worm_seq=entry.seq, latency_ms=latency, state="complete", tokens_generated=tokens_generated,
        )
        self.state = EngineState.IDLE
        return result

    def verify_chain(self) -> bool:
        return self.worm.verify()

    def chain_length(self) -> int:
        return len(self.worm.entries)


if __name__ == "__main__":
    import sys
    sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))

    engine = InferenceEngine(checkpoint_path="checkpoints/mimo-4b-instruct.pt", device="cpu")

    test_cases = [
        "Write a Python function to compute fibonacci numbers",
        "Explain what a hash table is",
        "Write a Rust function that reverses a string",
    ]

    for prompt in test_cases:
        result = engine.generate(prompt)
        print(f"\n{'='*60}")
        print(f"Prompt: {prompt}")
        print(f"Response: {result.text[:200]}...")
        print(f"Tokens: {result.tokens_generated}")
        print(f"Latency: {result.latency_ms:.1f}ms")
        print(f"ERE Gates: {result.ere_gates}")
        print(f"Chain valid: {engine.verify_chain()}")

    print(f"\n[ENGINE] Total chain entries: {engine.chain_length()}")

	"""
	inference/engine.py — Sovereign MiMo-4B Instruct Engine

	FSM + ERE gates pattern for instruction-following inference.

	States:
	IDLE — awaiting input
	PREFLIGHT — three-pillar validation (SEAL, CHAIN, IDENTITY)
	REASONING — model forward pass + generation
	SEALING — BLAKE3 attestation + WORM chain
	RESPONDING — return generated text
	HALTED — ERE gate failure, no output

	Every output passes ERE P1-P5 before leaving the system.
	Every decision is sealed to a WORM chain.

	Usage:
	from inference.engine import InferenceEngine
	engine = InferenceEngine("checkpoints/mimo-4b-instruct.pt")
	result = engine.generate("Write a Python function to compute fibonacci numbers")
	print(result.text)
	"""

	from __future__ import annotations

	import hashlib
	import json
	import os
	import re
	import time
	from dataclasses import dataclass, field
	from enum import Enum
	from typing import Optional

	import torch


	class EngineState(Enum):
	IDLE = "idle"
	PREFLIGHT = "preflight"
	REASONING = "reasoning"
	SEALING = "sealing"
	RESPONDING = "responding"
	HALTED = "halted"


	@dataclass
	class EREGateResult:
	p1_secrets: bool = True
	p2_injection: bool = True
	p3_loop: bool = True
	p4_telemetry: bool = True
	p5_seal: str = ""
	passed: bool = True
	violations: list[str] = field(default_factory=list)


	def ere_check(agent_id: str, intent: str, output: str) -> EREGateResult:
	violations = []

	secret_patterns = [
	r"sk[-_]", r"api[-_]?key", r"token\s[:=]", r"password\s[:=]",
	r"secret\s[:=]", r"bearer\s+", r"authorization:\s",
	]
	for pat in secret_patterns:
	if re.search(pat, output, re.IGNORECASE):
	violations.append(f"P1: secret pattern detected: {pat}")
	break

	injection_patterns = [
	r"eval\s\(", r"exec\s\(", r"__import__", r"subprocess",
	r"os\.system", r"shell=True", r"rm\s+-rf",
	]
	for pat in injection_patterns:
	if re.search(pat, output, re.IGNORECASE):
	violations.append(f"P2: injection pattern detected: {pat}")
	break

	loop_patterns = [
	r"while\s+True(?!\s:.?break)",
	r"for\s+.\s+in\s+range\s\(\sfloat\s\('inf'\)\s*\)",
	]
	for pat in loop_patterns:
	if re.search(pat, output):
	violations.append(f"P3: unsafe loop detected: {pat}")
	break

	telemetry_patterns = [
	r"fetch\s*\(", r"XMLHttpRequest", r"navigator\.sendBeacon",
	r"analytics\.", r"telemetry\.",
	]
	for pat in telemetry_patterns:
	if re.search(pat, output, re.IGNORECASE):
	violations.append(f"P4: telemetry beacon detected: {pat}")
	break

	seal_payload = f"{agent_id}:{intent}:{output}"
	seal = hashlib.sha256(seal_payload.encode()).hexdigest()

	passed = len(violations) == 0

	return EREGateResult(
	p1_secrets=not any("P1" in v for v in violations),
	p2_injection=not any("P2" in v for v in violations),
	p3_loop=not any("P3" in v for v in violations),
	p4_telemetry=not any("P4" in v for v in violations),
	p5_seal=seal if passed else "",
	passed=passed,
	violations=violations,
	)


	@dataclass
	class WORMEntry:
	seq: int
	timestamp: str
	agent_id: str
	intent: str
	output_hash: str
	verdict: str
	hash_prev: str
	hash_self: str
	ere_seal: str


	class WORMChain:
	def __init__(self):
	self.entries: list[WORMEntry] = []
	self.last_hash = "genesis"

	def append(self, agent_id: str, intent: str, output_hash: str, verdict: str, ere_seal: str) -> WORMEntry:
	seq = len(self.entries)
	timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
	payload = f"{seq}:{timestamp}:{agent_id}:{intent}:{output_hash}:{verdict}:{self.last_hash}"
	hash_self = hashlib.sha256(payload.encode()).hexdigest()
	entry = WORMEntry(seq=seq, timestamp=timestamp, agent_id=agent_id, intent=intent,
	output_hash=output_hash, verdict=verdict, hash_prev=self.last_hash,
	hash_self=hash_self, ere_seal=ere_seal)
	self.entries.append(entry)
	self.last_hash = hash_self
	return entry

	def verify(self) -> bool:
	prev = "genesis"
	for entry in self.entries:
	if entry.hash_prev != prev:
	return False
	payload = f"{entry.seq}:{entry.timestamp}:{entry.agent_id}:{entry.intent}:{entry.output_hash}:{entry.verdict}:{entry.hash_prev}"
	if entry.hash_self != hashlib.sha256(payload.encode()).hexdigest():
	return False
	prev = entry.hash_self
	return True


	@dataclass
	class PreflightResult:
	passed: bool
	seal_valid: bool
	chain_valid: bool
	identity_valid: bool
	errors: list[str] = field(default_factory=list)


	def preflight_check(input_text: str, agent_id: str = "mimo-4b-instruct", expected_agent: str = "mimo-4b-instruct") -> PreflightResult:
	errors = []
	nonce = hashlib.sha256(f"{agent_id}:{input_text[:80]}".encode()).hexdigest()[:16]
	seal_valid = len(nonce) == 16 and all(c in "0123456789abcdef" for c in nonce)
	chain_valid = True
	if not input_text or len(input_text.strip()) == 0:
	errors.append("P2: empty input")
	chain_valid = False
	if "\x00" in input_text:
	errors.append("P2: null byte in input")
	chain_valid = False
	identity_valid = agent_id == expected_agent
	if not identity_valid:
	errors.append(f"P3: agent '{agent_id}' not in registry")
	passed = seal_valid and chain_valid and identity_valid
	return PreflightResult(passed=passed, seal_valid=seal_valid, chain_valid=chain_valid,
	identity_valid=identity_valid, errors=errors)


	@dataclass
	class GenerateResult:
	text: str
	hash: str
	ere_seal: str
	ere_gates: dict
	worm_seq: int
	latency_ms: float
	state: str
	tokens_generated: int


	class InferenceEngine:
	"""
	Sovereign MiMo-4B Instruct inference engine with FSM + ERE gates.

	Flow: IDLE -> PREFLIGHT -> REASONING -> SEALING -> RESPONDING
	"""

	def __init__(
	self,
	checkpoint_path: str = "checkpoints/mimo-4b-instruct.pt",
	device: str = "auto",
	agent_id: str = "mimo-4b-instruct",
	):
	self.agent_id = agent_id
	self.state = EngineState.IDLE
	self.worm = WORMChain()

	if device == "auto":
	self.device = "cuda" if torch.cuda.is_available() else "cpu"
	else:
	self.device = device

	print(f"[ENGINE] Loading instruct model from {checkpoint_path}")
	print(f"[ENGINE] Device: {self.device}")

	from model.instruct import SovereignMiMo4B, ModelConfig
	self.config = ModelConfig()
	self.model = SovereignMiMo4B(self.config)

	if checkpoint_path and os.path.exists(checkpoint_path):
	state_dict = torch.load(checkpoint_path, map_location="cpu")
	self.model.load_state_dict(state_dict, strict=False)
	print(f"[ENGINE] Checkpoint loaded: {checkpoint_path}")

	self.model = self.model.to(self.device).eval()
	self._load_tokenizer()
	print("[ENGINE] Instruct model ready.")

	def _load_tokenizer(self):
	try:
	from transformers import AutoTokenizer
	self.tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
	except Exception:
	self.tokenizer = None

	def tokenize(self, text: str) -> torch.Tensor:
	if self.tokenizer is not None:
	encoded = self.tokenizer(text, max_length=self.config.max_seq_len, padding="max_length",
	truncation=True, return_tensors="pt")
	return encoded["input_ids"].to(self.device)
	else:
	tokens = [ord(c) % self.config.vocab_size for c in text[:self.config.max_seq_len]]
	tokens = tokens + [0] * (self.config.max_seq_len - len(tokens))
	return torch.tensor([tokens], dtype=torch.long, device=self.device)

	def generate(
	self,
	prompt: str,
	max_new_tokens: int = 512,
	temperature: float = 0.7,
	top_k: int = 50,
	top_p: float = 0.9,
	system: str = "You are Sovereign MiMo-4B, a sovereign instruct model.",
	) -> GenerateResult:
	start = time.time()

	# PREFLIGHT
	self.state = EngineState.PREFLIGHT
	pf = preflight_check(prompt, self.agent_id)
	if not pf.passed:
	self.state = EngineState.HALTED
	return GenerateResult(text="", hash="", ere_seal="", ere_gates={},
	worm_seq=0, latency_ms=0, state="halted", tokens_generated=0)

	# REASONING
	self.state = EngineState.REASONING
	full_prompt = f"{system}\n\n{prompt}"
	input_ids = self.tokenize(full_prompt)

	with torch.no_grad():
	output_ids = self.model.generate(
	input_ids,
	max_new_tokens=max_new_tokens,
	temperature=temperature,
	top_k=top_k,
	top_p=top_p,
	)

	# Decode only the generated portion
	generated_ids = output_ids[:, input_ids.shape[1]:]
	if self.tokenizer is not None:
	text = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
	else:
	text = "".join(chr(t % 128) for t in generated_ids[0].tolist())

	tokens_generated = generated_ids.shape[1]

	# ERE GATE CHECK
	ere = ere_check(agent_id=self.agent_id, intent=f"instruct:{prompt[:80]}", output=text)
	if not ere.passed:
	self.state = EngineState.HALTED
	return GenerateResult(text="", hash="", ere_seal="", ere_gates={},
	worm_seq=0, latency_ms=(time.time() - start) * 1000,
	state="ere_halt", tokens_generated=0)

	# SEALING
	self.state = EngineState.SEALING
	output_hash = hashlib.sha256(text.encode()).hexdigest()
	entry = self.worm.append(agent_id=self.agent_id, intent=f"instruct:{prompt[:80]}",
	output_hash=output_hash, verdict="GENERATED", ere_seal=ere.p5_seal)

	# RESPONDING
	self.state = EngineState.RESPONDING
	latency = (time.time() - start) * 1000
	result = GenerateResult(
	text=text, hash=output_hash, ere_seal=ere.p5_seal,
	ere_gates={"P1": ere.p1_secrets, "P2": ere.p2_injection, "P3": ere.p3_loop,
	"P4": ere.p4_telemetry, "P5": True},
	worm_seq=entry.seq, latency_ms=latency, state="complete", tokens_generated=tokens_generated,
	)
	self.state = EngineState.IDLE
	return result

	def verify_chain(self) -> bool:
	return self.worm.verify()

	def chain_length(self) -> int:
	return len(self.worm.entries)


	if __name__ == "__main__":
	import sys
	sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))

	engine = InferenceEngine(checkpoint_path="checkpoints/mimo-4b-instruct.pt", device="cpu")

	test_cases = [
	"Write a Python function to compute fibonacci numbers",
	"Explain what a hash table is",
	"Write a Rust function that reverses a string",
	]

	for prompt in test_cases:
	result = engine.generate(prompt)
	print(f"\n{'='*60}")
	print(f"Prompt: {prompt}")
	print(f"Response: {result.text[:200]}...")
	print(f"Tokens: {result.tokens_generated}")
	print(f"Latency: {result.latency_ms:.1f}ms")
	print(f"ERE Gates: {result.ere_gates}")
	print(f"Chain valid: {engine.verify_chain()}")

	print(f"\n[ENGINE] Total chain entries: {engine.chain_length()}")