Upload folder using huggingface_hub

__init__.py

@@ -0,0 +1,16 @@
+"""
+Hermes Edge — Package Init
+"""
+
+__version__ = "0.2.0"
+__author__ = "Barry Clerjuste"
+__email__ = "bclerjuste@gmail.com"
+
+from hermes.config import HermesConfig, get_config, PRESETS
+from hermes.chat_template import build_prompt, Message
+from hermes.litert_model import LiteRTModel
+
+try:
+    from hermes.agent import HermesAgent, AgentConfig
+except ImportError:
+    pass

agent.py

@@ -0,0 +1,196 @@
+"""
+Hermes Edge Agent — On-Device AI Agent Framework
+
+Combines DeepSeek-style reasoning + Hermes tool calling + LiteRT-LM runtime
+into a coherent agent loop for on-device inference.
+
+Usage:
+    from hermes.agent import HermesAgent
+    from hermes.tools import ToolRegistry
+    from hermes.litert_model import LiteRTModel
+
+    model = LiteRTModel("/path/to/model.litertlm")
+    agent = HermesAgent(model)
+    response = agent.run("What's the weather?")
+"""
+
+import logging
+import time
+from dataclasses import dataclass, field
+
+from hermes.chat_template import build_prompt, Message
+from scripts.deepseek_reasoning_template import ReasoningPipeline, ReasoningResult
+from scripts.hermes_tool_format import ToolRegistry, HermesToolFormatter
+from scripts.dspark_draft import DSparkDraftEngine, DSparkConfig, NGramDraftModel
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class AgentConfig:
+    max_tool_rounds: int = 5
+    max_tokens: int = 512
+    temperature: float = 0.7
+    top_k: int = 40
+    use_reasoning: bool = True
+    use_speculative_decoding: bool = True
+    draft_k: int = 4
+    system_prompt: str = ""
+
+
+DEFAULT_SYSTEM = (
+    "You are Hermes Edge, an on-device AI agent powered by Raven AI ecosystem. "
+    "You run fully offline via LiteRT-LM on iPhone 16 / Android. "
+    "You have access to tools and can reason step by step. "
+    "Always prefer local computation. Be helpful, concise, and accurate."
+)
+
+
+@dataclass
+class AgentTurn:
+    user_input: str = ""
+    assistant_response: str = ""
+    thinking: str = ""
+    tool_calls: list[dict] = field(default_factory=list)
+    tool_results: list[dict] = field(default_factory=list)
+    latency_ms: float = 0.0
+    tokens_used: int = 0
+
+
+@dataclass
+class Conversation:
+    messages: list[Message] = field(default_factory=list)
+    turns: list[AgentTurn] = field(default_factory=list)
+
+    def add_user(self, text: str) -> None:
+        self.messages.append(Message(role="user", content=text))
+
+    def add_assistant(self, text: str) -> None:
+        self.messages.append(Message(role="assistant", content=text))
+
+    def add_tool_result(self, name: str, content: str) -> None:
+        self.messages.append(Message(role="tool", content=f"<tool_response>{name}: {content}</tool_response>"))
+
+
+class HermesAgent:
+    """Full agent loop combining reasoning, tool calling, and speculative decoding."""
+
+    def __init__(
+        self,
+        model=None,
+        tool_registry: ToolRegistry | None = None,
+        config: AgentConfig | None = None,
+    ):
+        self.model = model
+        self.config = config or AgentConfig()
+        self.tools = tool_registry or ToolRegistry()
+        self.conversation = Conversation()
+        self.reasoning = ReasoningPipeline(use_reasoning=self.config.use_reasoning)
+        self.tool_formatter = HermesToolFormatter()
+        self.draft_engine: DSparkDraftEngine | None = None
+        self._init_draft_engine()
+
+    def _init_draft_engine(self) -> None:
+        if self.config.use_speculative_decoding and self.model is not None:
+            vocab_size = getattr(self.model, "vocab_size", 32000)
+            draft = NGramDraftModel(vocab_size=vocab_size, max_order=3)
+            dconfig = DSparkConfig(
+                draft_k=self.config.draft_k,
+                temperature=self.config.temperature,
+                top_k=self.config.top_k,
+            )
+            self.draft_engine = DSparkDraftEngine(self.model, draft, dconfig)
+
+    def set_model(self, model) -> None:
+        self.model = model
+        self._init_draft_engine()
+
+    def register_tool(self, name: str, description: str, func, parameters: dict | None = None) -> None:
+        self.tools.register(name, description, func, parameters)
+
+    def run(self, user_input: str, context: str | None = None) -> str:
+        """Process a user input through the full agent pipeline."""
+        if not self.model:
+            return "Error: No model loaded."
+
+        turn = AgentTurn(user_input=user_input)
+        start = time.perf_counter()
+
+        if self.config.use_reasoning:
+            prompt = self.reasoning.build_reasoning_prompt(user_input, context)
+        else:
+            tool_defs = self.tools.get_defs()
+            self.tool_formatter.set_tools(tool_defs)
+            prompt = self.tool_formatter.build_tool_prompt(user_input, context=context)
+
+        raw_output = self._generate(prompt)
+        turn.tokens_used = len(raw_output) // 4
+
+        parsed = self.reasoning.parse_response(raw_output)
+        turn.thinking = parsed.thinking
+        turn.assistant_response = parsed.answer
+        turn.tool_calls = parsed.tool_calls
+
+        tool_round = 0
+        while parsed.tool_calls and tool_round < self.config.max_tool_rounds:
+            tool_round += 1
+            for call in parsed.tool_calls:
+                name = call.get("name", "")
+                args = call.get("arguments", {})
+                result = self.tools.execute(name, args)
+                turn.tool_results.append({"name": name, "content": result.content, "success": result.success})
+                self.conversation.add_tool_result(name, result.content)
+
+            tool_prompt = self.reasoning.build_tool_result_prompt(
+                tool_name=name if parsed.tool_calls else "unknown",
+                tool_content=result.content if parsed.tool_calls else "",
+                original_prompt=prompt,
+            )
+            raw_output = self._generate(tool_prompt)
+            parsed = self.reasoning.parse_response(raw_output)
+            turn.assistant_response += "\n" + parsed.answer
+            turn.tool_calls.extend(parsed.tool_calls)
+
+        turn.latency_ms = (time.perf_counter() - start) * 1000
+        self.conversation.turns.append(turn)
+        self.conversation.add_user(user_input)
+        self.conversation.add_assistant(turn.assistant_response)
+
+        log.info(
+            "Agent turn: %d ms, %d tokens, %d tool calls, reasoning=%s",
+            turn.latency_ms,
+            turn.tokens_used,
+            len(turn.tool_calls),
+            bool(turn.thinking),
+        )
+        return turn.assistant_response
+
+    def _generate(self, prompt: str) -> str:
+        """Generate text using the model, optionally with speculative decoding."""
+        try:
+            if self.draft_engine and self.model:
+                prompt_ids = self._encode(prompt)
+                result = self.draft_engine.speculative_generate(
+                    prompt_ids=prompt_ids,
+                    max_tokens=self.config.max_tokens,
+                    tokenizer=getattr(self.model, "tokenizer", None),
+                )
+                if result.text:
+                    return result.text
+        except Exception as exc:
+            log.warning("Speculative decoding failed, falling back: %s", exc)
+
+        if hasattr(self.model, "generate"):
+            return self.model.generate(prompt, max_tokens=self.config.max_tokens)
+        return f"[Model would generate response for: {prompt[:50]}...]"
+
+    @staticmethod
+    def _encode(text: str) -> list[int]:
+        return list(text.encode("utf-8")[:256])
+
+    def get_conversation_summary(self) -> str:
+        """Get a summary of the conversation."""
+        turns = len(self.conversation.turns)
+        total_tokens = sum(t.tokens_used for t in self.conversation.turns)
+        total_latency = sum(t.latency_ms for t in self.conversation.turns)
+        return f"{turns} turns, ~{total_tokens} tokens, ~{total_latency:.0f}ms total"

chat_template.py

@@ -0,0 +1,102 @@
+"""Hermes chat + tool-calling prompt format.
+
+The format follows the ChatML-style convention used by the original Hermes
+models (`<|im_start|>role ... <|im_end|>`) and adds explicit tool-call markers
+so the on-device model can emit structured function calls that the Google AI
+Edge Gallery Agent Skills runtime can parse and dispatch.
+
+A tool call is emitted as::
+
+    <tool_call>{"name": "calculator", "arguments": {"expression": "2+2"}}</tool_call>
+
+Constrained decoding in LiteRT-LM can be anchored on the ``<tool_call>`` /
+``</tool_call>`` sentinels to guarantee well-formed JSON.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+IM_START = "<|im_start|>"
+IM_END = "<|im_end|>"
+TOOL_CALL_START = "<tool_call>"
+TOOL_CALL_END = "</tool_call>"
+TOOL_RESPONSE_START = "<tool_response>"
+TOOL_RESPONSE_END = "</tool_response>"
+
+DEFAULT_SYSTEM_PROMPT = (
+    "You are Hermes, a helpful on-device AI agent. You can call tools when "
+    "they help answer the user. To call a tool, respond ONLY with a "
+    "<tool_call> block containing JSON: "
+    '{"name": <tool_name>, "arguments": <json_args>}. '
+    "After receiving a <tool_response>, use it to answer the user."
+)
+
+
+@dataclass
+class Message:
+    role: str  # "system" | "user" | "assistant" | "tool"
+    content: str
+
+
+def render_tools(tools: List[Dict[str, Any]]) -> str:
+    """Render available tool schemas into the system context."""
+    if not tools:
+        return ""
+    lines = ["You have access to the following tools:"]
+    for tool in tools:
+        lines.append(json.dumps(tool, ensure_ascii=False))
+    return "\n".join(lines)
+
+
+def build_prompt(
+    messages: List[Message],
+    tools: Optional[List[Dict[str, Any]]] = None,
+    system_prompt: str = DEFAULT_SYSTEM_PROMPT,
+    add_generation_prompt: bool = True,
+) -> str:
+    """Render a list of messages into the Hermes ChatML training/inference string."""
+    parts: List[str] = []
+
+    system_content = system_prompt
+    tool_block = render_tools(tools or [])
+    if tool_block:
+        system_content = f"{system_prompt}\n\n{tool_block}"
+    parts.append(f"{IM_START}system\n{system_content}{IM_END}")
+
+    for msg in messages:
+        if msg.role == "tool":
+            body = f"{TOOL_RESPONSE_START}\n{msg.content}\n{TOOL_RESPONSE_END}"
+            parts.append(f"{IM_START}tool\n{body}{IM_END}")
+        else:
+            parts.append(f"{IM_START}{msg.role}\n{msg.content}{IM_END}")
+
+    if add_generation_prompt:
+        parts.append(f"{IM_START}assistant\n")
+    return "\n".join(parts)
+
+
+def format_tool_call(name: str, arguments: Dict[str, Any]) -> str:
+    """Serialize a tool call into the sentinel-wrapped JSON the model emits."""
+    payload = json.dumps({"name": name, "arguments": arguments}, ensure_ascii=False)
+    return f"{TOOL_CALL_START}{payload}{TOOL_CALL_END}"
+
+
+def parse_tool_call(text: str) -> Optional[Dict[str, Any]]:
+    """Extract a tool call from model output, or None if absent/malformed."""
+    start = text.find(TOOL_CALL_START)
+    if start == -1:
+        return None
+    start += len(TOOL_CALL_START)
+    end = text.find(TOOL_CALL_END, start)
+    snippet = text[start:end] if end != -1 else text[start:]
+    try:
+        call = json.loads(snippet.strip())
+    except json.JSONDecodeError:
+        return None
+    if not isinstance(call, dict) or "name" not in call:
+        return None
+    call.setdefault("arguments", {})
+    return call

config.py

@@ -0,0 +1,210 @@
+"""Architecture configuration for the Hermes mobile transformer.
+
+The configuration intentionally mirrors the knobs exposed by
+``ai_edge_torch.generative.layers.model_config`` so that the same numbers can
+drive both the reference PyTorch implementation (used for training) and the
+LiteRT conversion path. Keeping a single source of truth avoids the classic
+"the converted graph does not match the trained weights" failure mode.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+
+@dataclass
+class HermesConfig:
+    """Hyper-parameters for a decoder-only, grouped-query-attention model.
+
+    Attributes:
+        vocab_size: SentencePiece vocabulary size (must match the tokenizer).
+        hidden_size: Model / embedding dimension.
+        intermediate_size: Feed-forward (MLP) inner dimension.
+        num_layers: Number of transformer decoder blocks.
+        num_heads: Number of query attention heads.
+        num_kv_heads: Number of key/value heads (GQA). Must divide num_heads.
+        head_dim: Dimension per attention head.
+        max_seq_len: Maximum context window (KV-cache length) in tokens.
+        rope_theta: RoPE base frequency.
+        rms_norm_eps: Epsilon for RMSNorm numerical stability.
+        tie_embeddings: Share input embedding and output projection weights.
+        pad_token_id / bos_token_id / eos_token_id: Special token ids.
+    """
+
+    vocab_size: int = 32000
+    hidden_size: int = 2048
+    intermediate_size: int = 5632
+    num_layers: int = 22
+    num_heads: int = 32
+    num_kv_heads: int = 4
+    head_dim: int = 64
+    max_seq_len: int = 4096
+    rope_theta: float = 10000.0
+    rms_norm_eps: float = 1e-6
+    tie_embeddings: bool = True
+    pad_token_id: int = 0
+    bos_token_id: int = 1
+    eos_token_id: int = 2
+    # Tool-call sentinel tokens reserved in the tokenizer for constrained
+    # decoding of function calls (see scripts/train.py chat template).
+    tool_call_start_id: Optional[int] = 3
+    tool_call_end_id: Optional[int] = 4
+
+    def __post_init__(self) -> None:
+        if self.num_heads % self.num_kv_heads != 0:
+            raise ValueError(
+                f"num_heads ({self.num_heads}) must be divisible by "
+                f"num_kv_heads ({self.num_kv_heads}) for grouped-query attention."
+            )
+        if self.hidden_size != self.num_heads * self.head_dim:
+            raise ValueError(
+                f"hidden_size ({self.hidden_size}) must equal "
+                f"num_heads * head_dim ({self.num_heads * self.head_dim})."
+            )
+
+    @property
+    def num_query_groups(self) -> int:
+        """Heads per KV group (the GQA sharing factor)."""
+        return self.num_heads // self.num_kv_heads
+
+    def estimated_parameters(self) -> int:
+        """Rough parameter count (weights only, ignoring norms/biases)."""
+        emb = self.vocab_size * self.hidden_size
+        q = self.hidden_size * self.num_heads * self.head_dim
+        kv = 2 * self.hidden_size * self.num_kv_heads * self.head_dim
+        o = self.num_heads * self.head_dim * self.hidden_size
+        attn = q + kv + o
+        mlp = 3 * self.hidden_size * self.intermediate_size  # gate, up, down
+        per_layer = attn + mlp
+        total = emb + self.num_layers * per_layer
+        if not self.tie_embeddings:
+            total += emb  # separate lm_head
+        return total
+
+
+def hermes_1b_config() -> HermesConfig:
+    """~1B parameter variant — the default mobile target (~600 MB at INT4)."""
+    return HermesConfig(
+        vocab_size=32000,
+        hidden_size=2048,
+        intermediate_size=5632,
+        num_layers=22,
+        num_heads=32,
+        num_kv_heads=4,
+        head_dim=64,
+        max_seq_len=4096,
+    )
+
+
+def hermes_500m_config() -> HermesConfig:
+    """~500M parameter variant — quality/speed sweet spot (~280 MB at INT4)."""
+    return HermesConfig(
+        vocab_size=32000,
+        hidden_size=1536,
+        intermediate_size=4096,
+        num_layers=24,
+        num_heads=24,
+        num_kv_heads=6,
+        head_dim=64,
+        max_seq_len=4096,
+    )
+
+
+def hermes_270m_config() -> HermesConfig:
+    """~270M parameter variant — smallest, FunctionGemma-class footprint."""
+    return HermesConfig(
+        vocab_size=32000,
+        hidden_size=1024,
+        intermediate_size=2816,
+        num_layers=21,
+        num_heads=16,
+        num_kv_heads=4,
+        head_dim=64,
+        max_seq_len=4096,
+    )
+
+
+# ── Gemma-inspired presets (optimized for iPhone 16 A18 Pro / ANE) ──────────
+
+def gemma_3_1b_config() -> HermesConfig:
+    """Gemma 3 1B — Google's latest small model architecture.
+
+    Optimized for on-device inference with Apple Neural Engine:
+    - 32k vocab, 26 layers, 2048 hidden dim
+    - 16 heads, 8 KV heads (GQA ratio 2:1 — efficient for ANE)
+    - 8192 context window for longer conversations
+    - Ideal for iPhone 16 A18 Pro at INT4 (~250 MB on disk)
+    """
+    return HermesConfig(
+        vocab_size=32768,
+        hidden_size=2048,
+        intermediate_size=8192,
+        num_layers=26,
+        num_heads=16,
+        num_kv_heads=8,
+        head_dim=128,
+        max_seq_len=8192,
+        rope_theta=10000.0,
+    )
+
+
+def gemma_2_2b_config() -> HermesConfig:
+    """Gemma 2 2B — higher quality with shared KV-heads.
+
+    Uses deeper GQA (2 KV heads shared across 16 query heads) for
+    memory-efficient inference on iPhone 16 Pro / Pro Max.
+    ~1.1 GB at INT4.
+    """
+    return HermesConfig(
+        vocab_size=32768,
+        hidden_size=2560,
+        intermediate_size=9216,
+        num_layers=26,
+        num_heads=16,
+        num_kv_heads=2,
+        head_dim=160,
+        max_seq_len=8192,
+        rope_theta=10000.0,
+    )
+
+
+# ── DeepSeek-inspired distilled presets ─────────────────────────────────────
+
+def hermes_distilled_1b_config() -> HermesConfig:
+    """Distilled 1B model using DeepSeek-R1 reasoning principles.
+
+    Knowledge distilled from Gemma 3 1B teacher. Maintains the same
+    architecture as hermes-1b but with extended context and tuned
+    for step-by-step reasoning before tool calls.
+    """
+    return HermesConfig(
+        vocab_size=32000,
+        hidden_size=2048,
+        intermediate_size=5632,
+        num_layers=22,
+        num_heads=32,
+        num_kv_heads=4,
+        head_dim=64,
+        max_seq_len=8192,
+        rope_theta=10000.0,
+    )
+
+
+PRESETS = {
+    "hermes-1b": hermes_1b_config,
+    "hermes-500m": hermes_500m_config,
+    "hermes-270m": hermes_270m_config,
+    "gemma-3-1b": gemma_3_1b_config,
+    "gemma-2-2b": gemma_2_2b_config,
+    "hermes-distilled-1b": hermes_distilled_1b_config,
+}
+
+
+def get_config(name: str) -> HermesConfig:
+    """Look up a preset config by name."""
+    if name not in PRESETS:
+        raise KeyError(
+            f"Unknown preset '{name}'. Available: {sorted(PRESETS)}"
+        )
+    return PRESETS[name]()

inference.py

@@ -0,0 +1,334 @@
+"""Streaming inference engine for the Hermes mobile transformer.
+
+:class:`HermesInference` wraps a :class:`~hermes.model.HermesForCausalLM` and a
+SentencePiece tokenizer into a single object with three entry points:
+
+* :meth:`generate` — low-level text completion with nucleus (top-p), top-k, and
+  repetition-penalty sampling, optionally streaming token strings as they decode.
+* :meth:`chat` — renders a message list through the Hermes ChatML template, then
+  generates an assistant turn.
+* :meth:`tool_call_loop` — the agentic loop: generate, parse any ``<tool_call>``,
+  dispatch it to a Python callable, feed the ``<tool_response>`` back, and repeat
+  until the model produces a plain answer (or ``max_rounds`` is hit).
+
+Decoding reuses the existing :class:`~hermes.model.Attention` KV-cache: the
+prompt is run once to prime per-layer caches, then each new token is decoded with
+a single-position forward pass, so cost is linear in generated length rather than
+quadratic.
+
+The tokenizer is duck-typed: anything exposing ``encode(str) -> list[int]`` and
+``decode(list[int]) -> str`` works, which covers ``sentencepiece`` and the tiny
+byte-level stub used in tests.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+
+from hermes.chat_template import (
+    Message,
+    build_prompt,
+    parse_tool_call,
+)
+from hermes.config import HermesConfig
+from hermes.model import HermesForCausalLM, build_model
+
+KVList = List[Optional[Tuple[torch.Tensor, torch.Tensor]]]
+ToolFn = Callable[..., Any]
+
+
+class HermesInference:
+    """Load a Hermes checkpoint + tokenizer and run streaming generation."""
+
+    def __init__(
+        self,
+        model: HermesForCausalLM,
+        tokenizer: Any,
+        device: Union[str, torch.device] = "cpu",
+        preset_name: str = "custom",
+    ) -> None:
+        self.device = torch.device(device)
+        self.model = model.to(self.device).eval()
+        self.tokenizer = tokenizer
+        self.config: HermesConfig = model.config
+        self.preset_name = preset_name
+
+    # ------------------------------------------------------------------ #
+    # Construction helpers
+    # ------------------------------------------------------------------ #
+    @classmethod
+    def from_checkpoint(
+        cls,
+        config: HermesConfig,
+        checkpoint_path: Optional[str],
+        tokenizer: Any,
+        device: Union[str, torch.device] = "cpu",
+        preset_name: str = "custom",
+    ) -> "HermesInference":
+        """Build a model from ``config``, optionally load weights, and wrap it.
+
+        If ``checkpoint_path`` is None the model keeps its random init — handy for
+        CI and shape tests that don't need a trained checkpoint.
+        """
+        model = build_model(config)
+        if checkpoint_path is not None:
+            ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+            state_dict = ckpt.get("model", ckpt) if isinstance(ckpt, dict) else ckpt
+            model.load_state_dict(state_dict, strict=False)
+        return cls(model, tokenizer, device=device, preset_name=preset_name)
+
+    def __repr__(self) -> str:
+        n_params = sum(p.numel() for p in self.model.parameters())
+        return (
+            f"HermesInference(preset={self.preset_name!r}, "
+            f"params={n_params / 1e6:.1f}M, "
+            f"layers={self.config.num_layers}, ctx={self.config.max_seq_len}, "
+            f"device={self.device.type})"
+        )
+
+    # ------------------------------------------------------------------ #
+    # Sampling
+    # ------------------------------------------------------------------ #
+    @staticmethod
+    def _apply_repetition_penalty(
+        logits: torch.Tensor, generated: List[int], penalty: float
+    ) -> torch.Tensor:
+        """Divide logits of already-seen tokens by ``penalty`` (CTRL-style)."""
+        if penalty == 1.0 or not generated:
+            return logits
+        idx = torch.tensor(sorted(set(generated)), device=logits.device)
+        selected = logits.index_select(-1, idx)
+        # Positive logits are divided, negative are multiplied (push both down).
+        selected = torch.where(selected > 0, selected / penalty, selected * penalty)
+        logits = logits.index_copy(-1, idx, selected)
+        return logits
+
+    @staticmethod
+    def _sample(
+        logits: torch.Tensor,
+        temperature: float,
+        top_p: float,
+        top_k: int,
+    ) -> int:
+        """Sample a single token id from ``logits`` with top-k + nucleus filtering."""
+        if temperature <= 0.0:
+            return int(logits.argmax(dim=-1))
+
+        logits = logits / temperature
+
+        if top_k and top_k > 0:
+            k = min(top_k, logits.size(-1))
+            kth = torch.topk(logits, k).values[..., -1, None]
+            logits = torch.where(
+                logits < kth, torch.full_like(logits, float("-inf")), logits
+            )
+
+        if top_p and 0.0 < top_p < 1.0:
+            sorted_logits, sorted_idx = torch.sort(logits, descending=True)
+            probs = F.softmax(sorted_logits, dim=-1)
+            cumulative = torch.cumsum(probs, dim=-1)
+            # Keep tokens up to and including the one that crosses top_p.
+            remove = cumulative - probs > top_p
+            sorted_logits = sorted_logits.masked_fill(remove, float("-inf"))
+            logits = torch.full_like(logits, float("-inf")).scatter(
+                -1, sorted_idx, sorted_logits
+            )
+
+        probs = F.softmax(logits, dim=-1)
+        return int(torch.multinomial(probs, num_samples=1))
+
+    # ------------------------------------------------------------------ #
+    # KV-cache primed decode
+    # ------------------------------------------------------------------ #
+    def _forward_with_cache(
+        self,
+        input_ids: torch.Tensor,
+        caches: KVList,
+        start_pos: int,
+    ) -> Tuple[torch.Tensor, KVList]:
+        """Run the model for ``input_ids`` reusing/extending per-layer KV caches.
+
+        Returns the last-position logits and the updated cache list. This bypasses
+        ``HermesForCausalLM.forward`` so it can thread the per-layer cache tuples
+        through the existing :class:`Attention` ``kv_cache`` argument.
+        """
+        model = self.model
+        b, t = input_ids.shape
+        x = model.embed_tokens(input_ids)
+        cos = model.rope_cos[start_pos : start_pos + t].to(x.device)
+        sin = model.rope_sin[start_pos : start_pos + t].to(x.device)
+
+        # Causal mask over the *full* attended length (past + current).
+        total = start_pos + t
+        full_mask = torch.full((t, total), float("-inf"), device=x.device)
+        full_mask = torch.triu(full_mask, diagonal=1 + start_pos)
+
+        new_caches: KVList = [None] * len(model.layers)
+        for i, layer in enumerate(model.layers):
+            h, new_cache = layer.self_attn(
+                layer.input_layernorm(x), cos, sin, full_mask, caches[i]
+            )
+            x = x + h
+            x = x + layer.mlp(layer.post_attention_layernorm(x))
+            new_caches[i] = new_cache
+
+        x = model.norm(x)
+        logits = model.lm_head(x[:, -1, :])
+        return logits, new_caches
+
+    @torch.no_grad()
+    def _generate_ids(
+        self,
+        prompt_ids: List[int],
+        max_new_tokens: int,
+        temperature: float,
+        top_p: float,
+        top_k: int,
+        repetition_penalty: float,
+    ) -> Iterator[int]:
+        """Yield newly generated token ids one at a time, using a KV cache."""
+        self.model.eval()
+        eos = self.config.eos_token_id
+        caches: KVList = [None] * len(self.model.layers)
+
+        # Keep room for at least one generated token: truncate the prompt to the
+        # most recent (max_seq_len - 1) tokens if it would otherwise fill context.
+        max_prompt = max(1, self.config.max_seq_len - 1)
+        if len(prompt_ids) > max_prompt:
+            prompt_ids = prompt_ids[-max_prompt:]
+
+        # Prime the cache on the full prompt in one prefill pass.
+        ids = torch.tensor([prompt_ids], dtype=torch.long, device=self.device)
+        logits, caches = self._forward_with_cache(ids, caches, start_pos=0)
+        pos = len(prompt_ids)
+
+        generated: List[int] = []
+        for _ in range(max_new_tokens):
+            step_logits = self._apply_repetition_penalty(
+                logits.clone(), prompt_ids + generated, repetition_penalty
+            )
+            next_id = self._sample(step_logits.squeeze(0), temperature, top_p, top_k)
+            if next_id == eos:
+                break
+            generated.append(next_id)
+            yield next_id
+
+            if pos >= self.config.max_seq_len:
+                break
+            step = torch.tensor([[next_id]], dtype=torch.long, device=self.device)
+            logits, caches = self._forward_with_cache(step, caches, start_pos=pos)
+            pos += 1
+
+    # ------------------------------------------------------------------ #
+    # Public generation API
+    # ------------------------------------------------------------------ #
+    def generate(
+        self,
+        prompt: str,
+        max_new_tokens: int = 128,
+        temperature: float = 0.8,
+        top_p: float = 0.95,
+        top_k: int = 50,
+        repetition_penalty: float = 1.1,
+        stream: bool = False,
+    ) -> Union[str, Iterator[str]]:
+        """Generate text from ``prompt``.
+
+        Returns the full completion string, or — if ``stream=True`` — a generator
+        that yields incremental token strings as they are produced.
+        """
+        prompt_ids = self.tokenizer.encode(prompt)
+
+        def _token_strings() -> Iterator[str]:
+            prev_text = ""
+            buffer: List[int] = []
+            for tok in self._generate_ids(
+                prompt_ids, max_new_tokens, temperature, top_p, top_k, repetition_penalty
+            ):
+                buffer.append(tok)
+                # Decode incrementally so multi-token characters render correctly.
+                text = self.tokenizer.decode(buffer)
+                delta = text[len(prev_text) :]
+                if delta:
+                    prev_text = text
+                    yield delta
+
+        if stream:
+            return _token_strings()
+        return "".join(_token_strings())
+
+    def chat(
+        self,
+        messages: List[Message],
+        tools: Optional[List[Dict[str, Any]]] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Render ``messages`` (+ optional ``tools``) and generate a reply string."""
+        prompt = build_prompt(messages, tools=tools)
+        result = self.generate(prompt, stream=False, **kwargs)
+        assert isinstance(result, str)
+        return result
+
+    def tool_call_loop(
+        self,
+        messages: List[Message],
+        tools: List[Dict[str, Any]],
+        tool_functions: Dict[str, ToolFn],
+        max_rounds: int = 5,
+        **kwargs: Any,
+    ) -> List[Message]:
+        """Agentic loop: generate → parse tool call → dispatch → feed back.
+
+        Each round generates an assistant turn. If it contains a parseable
+        ``<tool_call>``, the named callable in ``tool_functions`` is invoked with
+        the parsed arguments and its result is appended as a ``tool`` message;
+        otherwise the loop ends. Returns the full conversation including all
+        assistant and tool turns appended.
+
+        Args:
+            messages: Seed conversation (mutated copy is returned).
+            tools: Tool schemas advertised to the model in the system prompt.
+            tool_functions: Maps tool ``name`` → Python callable.
+            max_rounds: Hard cap on generate/dispatch cycles.
+        """
+        convo = list(messages)
+        for _ in range(max_rounds):
+            reply = self.chat(convo, tools=tools, **kwargs)
+            convo.append(Message("assistant", reply))
+
+            call = parse_tool_call(reply)
+            if call is None:
+                break
+
+            fn = tool_functions.get(call["name"])
+            if fn is None:
+                convo.append(
+                    Message("tool", f'{{"error": "unknown tool: {call["name"]}"}}')
+                )
+                continue
+            try:
+                result = fn(**call.get("arguments", {}))
+            except Exception as exc:  # surface tool errors back to the model
+                result = {"error": str(exc)}
+            convo.append(Message("tool", str(result)))
+        return convo
+
+
+if __name__ == "__main__":  # pragma: no cover - manual smoke check
+    from hermes.config import hermes_270m_config
+
+    class _ByteTokenizer:
+        def encode(self, text: str) -> List[int]:
+            return [b % 32000 for b in text.encode("utf-8")] or [1]
+
+        def decode(self, ids: List[int]) -> str:
+            return bytes(i % 256 for i in ids).decode("utf-8", errors="replace")
+
+    engine = HermesInference.from_checkpoint(
+        hermes_270m_config(), None, _ByteTokenizer(), preset_name="hermes-270m"
+    )
+    print(engine)
+    print(engine.generate("Hello", max_new_tokens=8, temperature=0.0))

kv_cache.py

@@ -0,0 +1,311 @@
+"""Key/value cache managers for Hermes incremental decoding.
+
+Three cache strategies are provided, trading off memory, context length, and
+multi-session serving:
+
+* :class:`StaticKVCache` — a single pre-allocated, fixed-size cache. This is the
+  shape LiteRT-LM exports on device: the converted TFLite ``decode`` signature
+  writes into a buffer sized to ``max_seq_len`` and never reallocates.
+* :class:`SlidingWindowKVCache` — a :class:`StaticKVCache` subclass that keeps a
+  rolling window of the most recent ``window_size`` tokens, evicting the oldest
+  when full. Lets a 4096-ctx model hold an arbitrarily long conversation at a
+  bounded memory cost (at the price of forgetting distant context).
+* :class:`PagedKVCache` — block-level paging à la vLLM. The cache is carved into
+  fixed-size blocks that are allocated on demand and freed per sequence, which
+  makes it suitable for serving many concurrent sessions from one pool.
+
+All caches expose ``to(device)`` for device migration and
+``state_dict``/``load_state_dict`` for clean (de)serialization.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, List, Tuple
+
+import torch
+
+KVTuple = Tuple[torch.Tensor, torch.Tensor]
+
+
+class StaticKVCache:
+    """Pre-allocated fixed-size KV cache (the LiteRT-LM on-device shape).
+
+    Each layer owns a ``[1, num_kv_heads, max_seq_len, head_dim]`` key and value
+    buffer. :meth:`update` writes new entries at ``position`` and returns the
+    valid prefix; :meth:`get` returns the currently-filled slice.
+    """
+
+    def __init__(
+        self,
+        num_layers: int,
+        num_kv_heads: int,
+        max_seq_len: int,
+        head_dim: int,
+        dtype: torch.dtype = torch.float32,
+        device: torch.device | str = "cpu",
+        batch_size: int = 1,
+    ) -> None:
+        self.num_layers = num_layers
+        self.num_kv_heads = num_kv_heads
+        self.max_seq_len = max_seq_len
+        self.head_dim = head_dim
+        self.dtype = dtype
+        self.device = torch.device(device)
+        self.batch_size = batch_size
+        self._len = 0
+        self._alloc()
+
+    def _alloc(self) -> None:
+        shape = (self.batch_size, self.num_kv_heads, self.max_seq_len, self.head_dim)
+        self.keys: List[torch.Tensor] = [
+            torch.zeros(shape, dtype=self.dtype, device=self.device)
+            for _ in range(self.num_layers)
+        ]
+        self.values: List[torch.Tensor] = [
+            torch.zeros(shape, dtype=self.dtype, device=self.device)
+            for _ in range(self.num_layers)
+        ]
+
+    @property
+    def current_len(self) -> int:
+        """Number of valid (written) timesteps in the cache."""
+        return self._len
+
+    def reset(self) -> None:
+        """Zero the fill pointer (buffers are reused, not reallocated)."""
+        self._len = 0
+
+    def update(
+        self, layer_idx: int, k: torch.Tensor, v: torch.Tensor, position: int
+    ) -> KVTuple:
+        """Write ``k``/``v`` at ``position`` and return the valid prefix.
+
+        Args:
+            layer_idx: Which decoder layer this cache slot belongs to.
+            k: Keys ``[B, num_kv_heads, T_new, head_dim]``.
+            v: Values with the same shape.
+            position: Start timestep to write at (the current sequence length).
+
+        Returns:
+            ``(keys, values)`` covering timesteps ``[0, position + T_new)``.
+
+        Raises:
+            ValueError: If the write would exceed ``max_seq_len``.
+        """
+        t_new = k.shape[2]
+        end = position + t_new
+        if end > self.max_seq_len:
+            raise ValueError(
+                f"KV cache overflow: writing {t_new} tokens at position "
+                f"{position} exceeds max_seq_len={self.max_seq_len}."
+            )
+        self.keys[layer_idx][:, :, position:end, :] = k.to(self.dtype)
+        self.values[layer_idx][:, :, position:end, :] = v.to(self.dtype)
+        # The fill pointer tracks the furthest layer write of the current step.
+        self._len = max(self._len, end)
+        return self.get(layer_idx)
+
+    def get(self, layer_idx: int) -> KVTuple:
+        """Return the valid ``(keys, values)`` prefix for ``layer_idx``."""
+        return (
+            self.keys[layer_idx][:, :, : self._len, :],
+            self.values[layer_idx][:, :, : self._len, :],
+        )
+
+    def to(self, device: torch.device | str) -> "StaticKVCache":
+        """Move all buffers to ``device`` in place and return self."""
+        self.device = torch.device(device)
+        self.keys = [k.to(self.device) for k in self.keys]
+        self.values = [v.to(self.device) for v in self.values]
+        return self
+
+    def state_dict(self) -> Dict[str, object]:
+        """Serialize cache contents + metadata into a plain dict."""
+        return {
+            "class": type(self).__name__,
+            "num_layers": self.num_layers,
+            "num_kv_heads": self.num_kv_heads,
+            "max_seq_len": self.max_seq_len,
+            "head_dim": self.head_dim,
+            "batch_size": self.batch_size,
+            "len": self._len,
+            "keys": [k.cpu() for k in self.keys],
+            "values": [v.cpu() for v in self.values],
+        }
+
+    def load_state_dict(self, state: Dict[str, object]) -> None:
+        """Restore cache contents from :meth:`state_dict` output."""
+        self._len = int(state["len"])  # type: ignore[arg-type]
+        self.keys = [k.to(self.device) for k in state["keys"]]  # type: ignore[union-attr]
+        self.values = [v.to(self.device) for v in state["values"]]  # type: ignore[union-attr]
+
+
+class SlidingWindowKVCache(StaticKVCache):
+    """Rolling-window KV cache that evicts the oldest tokens when full.
+
+    Keeps at most ``window_size`` timesteps. Once full, each new write shifts the
+    buffer left (dropping the oldest entries) so memory stays bounded — handy for
+    long-running chats inside the 4096-token model.
+    """
+
+    def __init__(
+        self,
+        num_layers: int,
+        num_kv_heads: int,
+        max_seq_len: int,
+        head_dim: int,
+        dtype: torch.dtype = torch.float32,
+        device: torch.device | str = "cpu",
+        batch_size: int = 1,
+        window_size: int = 1024,
+    ) -> None:
+        self.window_size = min(window_size, max_seq_len)
+        super().__init__(
+            num_layers, num_kv_heads, max_seq_len, head_dim, dtype, device, batch_size
+        )
+
+    def update(
+        self, layer_idx: int, k: torch.Tensor, v: torch.Tensor, position: int
+    ) -> KVTuple:
+        t_new = k.shape[2]
+        if t_new > self.window_size:
+            # Only the most recent window_size tokens can be retained.
+            k = k[:, :, -self.window_size :, :]
+            v = v[:, :, -self.window_size :, :]
+            t_new = self.window_size
+
+        if self._len + t_new > self.window_size:
+            evict = self._len + t_new - self.window_size
+            kept = self._len - evict
+            if kept > 0:
+                self.keys[layer_idx][:, :, :kept, :] = self.keys[layer_idx][
+                    :, :, evict : self._len, :
+                ].clone()
+                self.values[layer_idx][:, :, :kept, :] = self.values[layer_idx][
+                    :, :, evict : self._len, :
+                ].clone()
+            write_at = kept
+        else:
+            write_at = self._len
+
+        end = write_at + t_new
+        self.keys[layer_idx][:, :, write_at:end, :] = k.to(self.dtype)
+        self.values[layer_idx][:, :, write_at:end, :] = v.to(self.dtype)
+        # current_len is shared across layers; cap it at the window.
+        self._len = min(end, self.window_size)
+        return self.get(layer_idx)
+
+
+class PagedKVCache:
+    """Block-paged KV cache for multi-session serving (vLLM-style).
+
+    The KV pool is divided into ``num_blocks`` blocks of ``block_size`` tokens.
+    Sequences request blocks via :meth:`allocate_block`; :meth:`free_sequence`
+    returns a sequence's blocks to the free list. :meth:`get_page_table` exposes
+    the per-sequence block mapping used to gather KV during attention.
+    """
+
+    def __init__(
+        self,
+        num_layers: int,
+        num_kv_heads: int,
+        head_dim: int,
+        num_blocks: int = 256,
+        block_size: int = 16,
+        dtype: torch.dtype = torch.float32,
+        device: torch.device | str = "cpu",
+    ) -> None:
+        self.num_layers = num_layers
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.num_blocks = num_blocks
+        self.block_size = block_size
+        self.dtype = dtype
+        self.device = torch.device(device)
+        self._page_table: Dict[int, List[int]] = {}
+        self._free: List[int] = list(range(num_blocks))
+        self._alloc()
+
+    def _alloc(self) -> None:
+        # Pool shape: [num_layers, num_blocks, num_kv_heads, block_size, head_dim].
+        shape = (
+            self.num_layers,
+            self.num_blocks,
+            self.num_kv_heads,
+            self.block_size,
+            self.head_dim,
+        )
+        self.key_pool = torch.zeros(shape, dtype=self.dtype, device=self.device)
+        self.value_pool = torch.zeros(shape, dtype=self.dtype, device=self.device)
+
+    @property
+    def num_free_blocks(self) -> int:
+        """Count of blocks currently available for allocation."""
+        return len(self._free)
+
+    @property
+    def num_used_blocks(self) -> int:
+        """Count of blocks currently assigned to some sequence."""
+        return self.num_blocks - len(self._free)
+
+    def allocate_block(self, seq_id: int = 0) -> int:
+        """Pop a free block, assign it to ``seq_id``, and return its index.
+
+        Raises:
+            RuntimeError: If the pool is exhausted.
+        """
+        if not self._free:
+            raise RuntimeError("PagedKVCache out of blocks; free a sequence first.")
+        block = self._free.pop(0)
+        self._page_table.setdefault(seq_id, []).append(block)
+        return block
+
+    def free_sequence(self, seq_id: int) -> List[int]:
+        """Return all of ``seq_id``'s blocks to the free list.
+
+        Returns the freed block indices (empty if the sequence was unknown).
+        """
+        blocks = self._page_table.pop(seq_id, [])
+        self._free.extend(blocks)
+        self._free.sort()
+        return blocks
+
+    def get_page_table(self) -> Dict[int, List[int]]:
+        """Return a copy of the per-sequence ``seq_id -> [block_idx]`` mapping."""
+        return {seq: list(blocks) for seq, blocks in self._page_table.items()}
+
+    def reset(self) -> None:
+        """Free every block and clear the page table."""
+        self._page_table.clear()
+        self._free = list(range(self.num_blocks))
+
+    def to(self, device: torch.device | str) -> "PagedKVCache":
+        """Move the KV pool to ``device`` in place and return self."""
+        self.device = torch.device(device)
+        self.key_pool = self.key_pool.to(self.device)
+        self.value_pool = self.value_pool.to(self.device)
+        return self
+
+    def state_dict(self) -> Dict[str, object]:
+        """Serialize pool contents + allocation state."""
+        return {
+            "class": type(self).__name__,
+            "num_layers": self.num_layers,
+            "num_kv_heads": self.num_kv_heads,
+            "head_dim": self.head_dim,
+            "num_blocks": self.num_blocks,
+            "block_size": self.block_size,
+            "page_table": self.get_page_table(),
+            "free": list(self._free),
+            "key_pool": self.key_pool.cpu(),
+            "value_pool": self.value_pool.cpu(),
+        }
+
+    def load_state_dict(self, state: Dict[str, object]) -> None:
+        """Restore pool contents + allocation state from :meth:`state_dict`."""
+        self._page_table = {
+            int(k): list(v) for k, v in state["page_table"].items()  # type: ignore[union-attr]
+        }
+        self._free = list(state["free"])  # type: ignore[arg-type]
+        self.key_pool = state["key_pool"].to(self.device)  # type: ignore[union-attr]
+        self.value_pool = state["value_pool"].to(self.device)  # type: ignore[union-attr]

litert_model.py

@@ -0,0 +1,176 @@
+"""
+LiteRT-LM Model Wrapper — Python interface for .litertlm models
+
+Wraps the LiteRT-LM C++ runtime via ctypes, providing a Pythonic
+interface for inference, tokenization, and agent integration.
+
+On actual devices, this is replaced by the Swift/Kotlin SDK.
+This Python wrapper is used for:
+  - Desktop testing and debugging
+  - HF Space demos (via Python backend)
+  - CI validation of model bundles
+
+Usage:
+    from hermes.litert_model import LiteRTModel
+
+    model = LiteRTModel("dist/hermes-mobile.litertlm")
+    model.load()
+    response = model.generate("Hello!", max_tokens=128)
+    print(response)
+"""
+
+import json
+import logging
+import os
+import subprocess
+import tempfile
+from pathlib import Path
+
+log = logging.getLogger(__name__)
+
+
+class LiteRTModel:
+    """
+    Wrapper around a .litertlm model bundle.
+
+    Uses the `litert-lm` CLI tool for inference (since the Python C++
+    binding requires libvulkan which isn't available in all environments).
+
+    On iOS/Android, the native SDK replaces this class entirely.
+    """
+
+    def __init__(self, model_path: str, cli_path: str = "litert-lm"):
+        self.model_path = Path(model_path).resolve()
+        self.cli_path = cli_path
+        self.vocab_size = 32000
+        self.tokenizer = None
+        self._loaded = False
+        self._metadata: dict = {}
+
+    def load(self) -> bool:
+        """Validate the model file and extract metadata."""
+        if not self.model_path.exists():
+            log.error("Model not found: %s", self.model_path)
+            return False
+
+        with open(self.model_path, "rb") as f:
+            header = f.read(16)
+            if header[:8] != b"LITERTLM":
+                log.error("Invalid model file (bad magic): %s", self.model_path)
+                return False
+
+        self._loaded = True
+        mb = self.model_path.stat().st_size / 1024 / 1024
+        log.info("Model loaded: %s (%.1f MB)", self.model_path.name, mb)
+        return True
+
+    def generate(
+        self,
+        prompt: str,
+        max_tokens: int = 256,
+        temperature: float = 0.7,
+        top_k: int = 40,
+    ) -> str:
+        """Generate text using the litert-lm CLI."""
+        if not self._loaded:
+            return "Error: Model not loaded."
+
+        try:
+            result = subprocess.run(
+                [
+                    self.cli_path,
+                    "run",
+                    str(self.model_path),
+                    "--prompt",
+                    prompt,
+                    "--max_tokens",
+                    str(max_tokens),
+                ],
+                capture_output=True,
+                text=True,
+                timeout=60,
+            )
+            if result.returncode == 0 and result.stdout.strip():
+                return result.stdout.strip()
+
+            if result.stderr:
+                log.warning("CLI stderr: %s", result.stderr[:200])
+
+        except FileNotFoundError:
+            log.warning("litert-lm CLI not available, using simulated response")
+        except subprocess.TimeoutExpired:
+            log.warning("Model inference timed out")
+        except Exception as exc:
+            log.warning("Model inference error: %s", exc)
+
+        return self._simulate_response(prompt)
+
+    def predict_next_token(self, context: list[int]) -> int:
+        """Predict the most likely next token (used by DSpark draft engine)."""
+        if not self._loaded:
+            return 0
+        try:
+            text = self._decode_tokens(context)
+            result = subprocess.run(
+                [
+                    self.cli_path,
+                    "run",
+                    str(self.model_path),
+                    "--prompt",
+                    text[-200:],
+                    "--max_tokens",
+                    "1",
+                    "--temperature",
+                    "0.0",
+                ],
+                capture_output=True,
+                text=True,
+                timeout=30,
+            )
+            if result.returncode == 0 and result.stdout.strip():
+                return hash(result.stdout.strip()) % self.vocab_size
+        except Exception:
+            pass
+        return context[-1] if context else 0
+
+    @staticmethod
+    def _decode_tokens(token_ids: list[int]) -> str:
+        return "".join(chr(max(32, min(126, t % 128))) for t in token_ids[-50:])
+
+    def _simulate_response(self, prompt: str) -> str:
+        """Simulated response when CLI is unavailable (for demo/dev only)."""
+        prompt_lower = prompt.lower()
+        if "hello" in prompt_lower or "hi" in prompt_lower:
+            return "Hello! I'm Hermes Edge, running on-device. How can I help?"
+        if "tool" in prompt_lower or "function" in prompt_lower:
+            return (
+                "<think>The user is asking about tool calling. "
+                "I can use calculator, web search, memory, and timer tools.</think>\n\n"
+                "I support function calling. Available tools:\n"
+                "- calculator: evaluate math expressions\n"
+                "- web_search: search the web (requires network)\n"
+                "- memory: store and recall information\n"
+                "- timer: set timers"
+            )
+        if "reason" in prompt_lower or "deep" in prompt_lower:
+            return (
+                "<think>Applying DeepSeek-style reasoning. "
+                "Breaking down the question step by step. "
+                "Verifying each step.</think>\n\n"
+                "Based on my reasoning, here's my answer."
+            )
+        return (
+            f"<think>Processing query using {self.model_path.name} "
+            f"on LiteRT-LM runtime.</think>\n\n"
+            f"I received your message. I'm running fully offline as a {self.model_path.stem} model."
+        )
+
+    def get_metadata(self) -> dict:
+        """Get model metadata."""
+        return {
+            "path": str(self.model_path),
+            "size_mb": round(self.model_path.stat().st_size / 1024 / 1024, 1),
+            "loaded": self._loaded,
+            "format": "LITERTLM",
+            "vocab_size": self.vocab_size,
+        }

model.py

@@ -0,0 +1,248 @@
+"""Reference PyTorch implementation of the Hermes mobile transformer.
+
+This is the *training* model. It is intentionally written with plain,
+conversion-friendly PyTorch ops (no custom CUDA kernels, no flash-attention
+calls) so that the same ``state_dict`` can be loaded by the LiteRT builder in
+``scripts/convert_to_litertlm.py`` and traced by ``ai_edge_torch``.
+
+Architecture: decoder-only, RMSNorm (pre-norm), rotary position embeddings,
+grouped-query attention, and a SwiGLU feed-forward block — the same family as
+Gemma / Llama, sized for on-device inference.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from hermes.config import HermesConfig
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x = x.float()
+        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return (x.to(dtype)) * self.weight
+
+
+def build_rope_cache(
+    seq_len: int, head_dim: int, theta: float, device: torch.device
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Precompute cos/sin tables for rotary position embeddings."""
+    inv_freq = 1.0 / (
+        theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)
+    )
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos(), emb.sin()
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rope(
+    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    # q, k: [B, H, T, D]; cos/sin: [T, D]
+    cos = cos.unsqueeze(0).unsqueeze(0)
+    sin = sin.unsqueeze(0).unsqueeze(0)
+    q_out = (q * cos) + (rotate_half(q) * sin)
+    k_out = (k * cos) + (rotate_half(k) * sin)
+    return q_out, k_out
+
+
+class Attention(nn.Module):
+    """Grouped-query attention with an optional incremental KV-cache."""
+
+    def __init__(self, config: HermesConfig) -> None:
+        super().__init__()
+        self.num_heads = config.num_heads
+        self.num_kv_heads = config.num_kv_heads
+        self.head_dim = config.head_dim
+        self.num_query_groups = config.num_query_groups
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, self.num_heads * self.head_dim, bias=False
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, self.num_kv_heads * self.head_dim, bias=False
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, self.num_kv_heads * self.head_dim, bias=False
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, config.hidden_size, bias=False
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        kv_cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        b, t, _ = x.shape
+
+        q = self.q_proj(x).view(b, t, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(b, t, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(b, t, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        q, k = apply_rope(q, k, cos, sin)
+
+        if kv_cache is not None:
+            past_k, past_v = kv_cache
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        new_cache = (k, v)
+
+        # Expand KV heads to match query heads (GQA).
+        if self.num_query_groups > 1:
+            k = k.repeat_interleave(self.num_query_groups, dim=1)
+            v = v.repeat_interleave(self.num_query_groups, dim=1)
+
+        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        if mask is not None:
+            attn = attn + mask
+        attn = F.softmax(attn, dim=-1, dtype=torch.float32).to(q.dtype)
+        out = torch.matmul(attn, v)
+
+        out = out.transpose(1, 2).contiguous().view(b, t, -1)
+        return self.o_proj(out), new_cache
+
+
+class FeedForward(nn.Module):
+    """SwiGLU MLP: down(silu(gate(x)) * up(x))."""
+
+    def __init__(self, config: HermesConfig) -> None:
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, config: HermesConfig) -> None:
+        super().__init__()
+        self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.self_attn = Attention(config)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.mlp = FeedForward(config)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        kv_cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        h, new_cache = self.self_attn(
+            self.input_layernorm(x), cos, sin, mask, kv_cache
+        )
+        x = x + h
+        x = x + self.mlp(self.post_attention_layernorm(x))
+        return x, new_cache
+
+
+class HermesForCausalLM(nn.Module):
+    """Full Hermes decoder-only language model with a causal LM head."""
+
+    def __init__(self, config: HermesConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            [DecoderBlock(config) for _ in range(config.num_layers)]
+        )
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        if config.tie_embeddings:
+            self.lm_head.weight = self.embed_tokens.weight
+
+        cos, sin = build_rope_cache(
+            config.max_seq_len, config.head_dim, config.rope_theta, torch.device("cpu")
+        )
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        labels: Optional[torch.Tensor] = None,
+        start_pos: int = 0,
+    ) -> dict:
+        b, t = input_ids.shape
+        x = self.embed_tokens(input_ids)
+
+        cos = self.rope_cos[start_pos : start_pos + t].to(x.device)
+        sin = self.rope_sin[start_pos : start_pos + t].to(x.device)
+
+        mask = torch.full((t, t), float("-inf"), device=x.device)
+        mask = torch.triu(mask, diagonal=1)
+
+        for layer in self.layers:
+            x, _ = layer(x, cos, sin, mask)
+
+        x = self.norm(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=self.config.pad_token_id,
+            )
+        return {"logits": logits, "loss": loss}
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        max_new_tokens: int = 64,
+        temperature: float = 0.8,
+        top_k: int = 50,
+        eos_token_id: Optional[int] = None,
+    ) -> torch.Tensor:
+        """Minimal greedy/sampling loop — sanity check for trained weights."""
+        self.eval()
+        eos_token_id = eos_token_id if eos_token_id is not None else self.config.eos_token_id
+        for _ in range(max_new_tokens):
+            ids = input_ids[:, -self.config.max_seq_len :]
+            logits = self.forward(ids)["logits"][:, -1, :]
+            if temperature > 0:
+                logits = logits / temperature
+                if top_k:
+                    v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                    logits[logits < v[:, [-1]]] = float("-inf")
+                probs = F.softmax(logits, dim=-1)
+                next_id = torch.multinomial(probs, num_samples=1)
+            else:
+                next_id = logits.argmax(dim=-1, keepdim=True)
+            input_ids = torch.cat([input_ids, next_id], dim=1)
+            if (next_id == eos_token_id).all():
+                break
+        return input_ids
+
+
+def build_model(config: HermesConfig) -> HermesForCausalLM:
+    return HermesForCausalLM(config)

quantization.py

@@ -0,0 +1,301 @@
+"""Post-training quantization (PTQ) analysis + fake-quant utilities.
+
+These helpers are deliberately **standalone** — they have no ``ai_edge_torch``
+dependency. They serve two purposes:
+
+1. **Pre-conversion analysis.** :func:`collect_calibration_stats` and
+   :func:`quantization_error_report` let you measure activation ranges and the
+   weight/perplexity error a given bit-width would introduce, *before* you spend
+   minutes lowering the model through the LiteRT stack. Use them to sanity-check
+   that INT4 is viable for a checkpoint, or to pick which layers are sensitive.
+
+2. **Training-time fake quantization.** :func:`apply_weight_only_int4` and
+   :func:`apply_weight_only_int8` replace each ``nn.Linear`` weight with its
+   quantized-then-dequantized value using a straight-through estimator (STE) so
+   gradients still flow. This is the quantization-aware-training (QAT) path: fine
+   tune with fake-quant on to recover accuracy the real INT4 graph would lose.
+
+Relationship to ``scripts/convert_to_litertlm.py``
+--------------------------------------------------
+The *real* mobile INT4 graph is produced by ``convert_to_litertlm.py`` via
+``ai_edge_torch``'s ``full_int4_dynamic_recipe`` — that is what actually ships in
+the ``.litertlm`` bundle. The functions here do **not** replace that conversion:
+they approximate the same symmetric per-group INT4 scheme in pure PyTorch so you
+can (a) estimate the error offline and (b) QAT-finetune to minimize it. Numbers
+from here are guidance; the converter's output is ground truth.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Dict, Iterable, Optional
+
+import torch
+import torch.nn as nn
+
+
+# --------------------------------------------------------------------------- #
+# Symmetric per-group quantization core
+# --------------------------------------------------------------------------- #
+def _quant_levels(bits: int) -> tuple[int, int]:
+    """Return ``(qmin, qmax)`` for a signed ``bits``-bit integer."""
+    qmax = 2 ** (bits - 1) - 1
+    qmin = -(2 ** (bits - 1))
+    return qmin, qmax
+
+
+def fake_quantize_per_group(
+    weight: torch.Tensor, bits: int, group_size: int
+) -> torch.Tensor:
+    """Symmetric per-group fake quantization of a 2-D weight matrix.
+
+    The weight ``[out_features, in_features]`` is split along ``in_features`` into
+    groups of ``group_size``; each group gets its own scale ``max(|w|) / qmax``.
+    The result is quantized to the integer grid and dequantized back to float, so
+    the returned tensor has the same dtype/shape but only takes representable
+    values. Used by both the analysis and STE paths.
+    """
+    qmin, qmax = _quant_levels(bits)
+    out_features, in_features = weight.shape
+    gs = group_size if group_size > 0 else in_features
+    pad = (gs - in_features % gs) % gs
+    w = weight
+    if pad:
+        w = torch.nn.functional.pad(w, (0, pad))
+    w = w.reshape(out_features, -1, gs)
+
+    max_abs = w.abs().amax(dim=-1, keepdim=True)
+    scale = (max_abs / qmax).clamp(min=1e-8)
+    q = torch.clamp(torch.round(w / scale), qmin, qmax)
+    deq = (q * scale).reshape(out_features, -1)
+    if pad:
+        deq = deq[:, :in_features]
+    return deq.to(weight.dtype)
+
+
+class _STEFakeQuant(torch.autograd.Function):
+    """Straight-through estimator: quantize on forward, identity on backward."""
+
+    @staticmethod
+    def forward(ctx, weight: torch.Tensor, bits: int, group_size: int) -> torch.Tensor:  # type: ignore[override]
+        return fake_quantize_per_group(weight, bits, group_size)
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):  # type: ignore[override]
+        # Identity gradient w.r.t. the weight; None for the int hyper-params.
+        return grad_output, None, None
+
+
+def _apply_weight_only(model: nn.Module, bits: int, group_size: int) -> nn.Module:
+    """In-place STE fake-quant of every ``nn.Linear`` weight in ``model``."""
+    for module in model.modules():
+        if isinstance(module, nn.Linear):
+            with torch.no_grad():
+                quantized = _STEFakeQuant.apply(module.weight, bits, group_size)
+                module.weight.copy_(quantized)
+    return model
+
+
+def apply_weight_only_int4(model: nn.Module, group_size: int = 128) -> nn.Module:
+    """Fake-quantize all ``nn.Linear`` weights to symmetric per-group INT4.
+
+    Each weight is mapped onto the signed 4-bit grid ``[-8, 7]`` (per group of
+    ``group_size`` input channels) and dequantized in place. Uses a
+    straight-through estimator so the operation is differentiable for QAT.
+
+    This mirrors the per-group INT4 scheme that
+    ``ai_edge_torch``'s ``full_int4_dynamic_recipe`` applies during the real
+    conversion in ``scripts/convert_to_litertlm.py`` — call this to QAT-finetune
+    or to estimate INT4 error offline; the converter produces the shipped graph.
+
+    Returns the same model (mutated in place).
+    """
+    return _apply_weight_only(model, bits=4, group_size=group_size)
+
+
+def apply_weight_only_int8(model: nn.Module, group_size: int = 0) -> nn.Module:
+    """Fake-quantize all ``nn.Linear`` weights to symmetric INT8 (``[-128, 127]``).
+
+    Per-channel by default (``group_size=0`` → one scale per output row). Same STE
+    semantics as :func:`apply_weight_only_int4`; useful as the higher-quality
+    fallback recipe when INT4 degrades a sensitive checkpoint too much.
+
+    Returns the same model (mutated in place).
+    """
+    return _apply_weight_only(model, bits=8, group_size=group_size)
+
+
+# --------------------------------------------------------------------------- #
+# Calibration + error analysis
+# --------------------------------------------------------------------------- #
+@torch.no_grad()
+def collect_calibration_stats(
+    model: nn.Module,
+    dataloader: Iterable,
+    num_batches: int = 64,
+) -> Dict[str, Dict[str, float]]:
+    """Run forward passes and collect per-layer activation statistics.
+
+    Forward hooks on every ``nn.Linear`` record the running min/max and a coarse
+    99th-percentile estimate of the *output* activations across up to
+    ``num_batches`` batches. These ranges are what an activation-quantization
+    scheme (or a converter calibration pass) would use to pick scales.
+
+    Args:
+        model: The model to profile (set to eval).
+        dataloader: Yields either tensors of ``input_ids`` or ``(inputs, _)``
+            tuples / dicts with an ``input_ids`` key.
+        num_batches: Max number of batches to run.
+
+    Returns:
+        ``{layer_name: {"min", "max", "abs_max", "p99", "mean", "num_samples"}}``.
+    """
+    model.eval()
+    stats: Dict[str, Dict[str, float]] = {}
+    handles = []
+
+    def make_hook(name: str):
+        def hook(_module, _inp, out):
+            t = out.detach()
+            if not torch.is_floating_point(t):
+                return
+            flat = t.float().reshape(-1)
+            entry = stats.setdefault(
+                name,
+                {
+                    "min": math.inf,
+                    "max": -math.inf,
+                    "abs_max": 0.0,
+                    "p99": 0.0,
+                    "mean": 0.0,
+                    "num_samples": 0.0,
+                },
+            )
+            entry["min"] = min(entry["min"], float(flat.min()))
+            entry["max"] = max(entry["max"], float(flat.max()))
+            entry["abs_max"] = max(entry["abs_max"], float(flat.abs().max()))
+            # Running mean + percentile (cheap quantile on a subsample).
+            n_prev = entry["num_samples"]
+            n_new = flat.numel()
+            entry["mean"] = (
+                entry["mean"] * n_prev + float(flat.sum())
+            ) / max(n_prev + n_new, 1)
+            sample = flat if flat.numel() <= 16384 else flat[torch.randint(
+                0, flat.numel(), (16384,), device=flat.device)]
+            entry["p99"] = max(entry["p99"], float(torch.quantile(sample.abs(), 0.99)))
+            entry["num_samples"] = n_prev + n_new
+
+        return hook
+
+    for name, module in model.named_modules():
+        if isinstance(module, nn.Linear):
+            handles.append(module.register_forward_hook(make_hook(name)))
+
+    try:
+        for i, batch in enumerate(dataloader):
+            if i >= num_batches:
+                break
+            input_ids = _extract_input_ids(batch)
+            model(input_ids)
+    finally:
+        for h in handles:
+            h.remove()
+
+    return stats
+
+
+def _extract_input_ids(batch) -> torch.Tensor:
+    """Pull an ``input_ids`` tensor out of common dataloader batch shapes."""
+    if isinstance(batch, torch.Tensor):
+        return batch
+    if isinstance(batch, dict):
+        return batch["input_ids"]
+    if isinstance(batch, (tuple, list)):
+        return batch[0]
+    raise TypeError(f"Cannot extract input_ids from batch of type {type(batch)}.")
+
+
+@torch.no_grad()
+def _perplexity(model: nn.Module, dataloader: Iterable, num_batches: int) -> float:
+    """Mean token-level perplexity over ``num_batches`` (labels == inputs)."""
+    model.eval()
+    total_loss = 0.0
+    count = 0
+    for i, batch in enumerate(dataloader):
+        if i >= num_batches:
+            break
+        input_ids = _extract_input_ids(batch)
+        out = model(input_ids, labels=input_ids)
+        loss = out["loss"] if isinstance(out, dict) else out
+        if loss is None:
+            continue
+        total_loss += float(loss)
+        count += 1
+    if count == 0:
+        return float("nan")
+    return math.exp(total_loss / count)
+
+
+@torch.no_grad()
+def quantization_error_report(
+    original_model: nn.Module,
+    quantized_model: nn.Module,
+    dataloader: Iterable,
+    num_batches: int = 8,
+) -> Dict[str, object]:
+    """Compare a model against its quantized copy.
+
+    Computes, per ``nn.Linear`` layer, the relative L2 error between the original
+    and quantized weights, and the model-level perplexity delta on ``dataloader``.
+
+    Returns:
+        ``{"per_layer_l2": {name: rel_l2}, "max_layer_l2": float,
+           "perplexity_original": float, "perplexity_quantized": float,
+           "perplexity_delta": float}``.
+    """
+    orig_linears = dict(_named_linears(original_model))
+    quant_linears = dict(_named_linears(quantized_model))
+
+    per_layer: Dict[str, float] = {}
+    for name, orig in orig_linears.items():
+        if name not in quant_linears:
+            continue
+        diff = (orig.weight - quant_linears[name].weight).float()
+        denom = orig.weight.float().norm().clamp(min=1e-8)
+        per_layer[name] = float(diff.norm() / denom)
+
+    ppl_orig = _perplexity(original_model, dataloader, num_batches)
+    ppl_quant = _perplexity(quantized_model, dataloader, num_batches)
+
+    return {
+        "per_layer_l2": per_layer,
+        "max_layer_l2": max(per_layer.values()) if per_layer else 0.0,
+        "perplexity_original": ppl_orig,
+        "perplexity_quantized": ppl_quant,
+        "perplexity_delta": ppl_quant - ppl_orig,
+    }
+
+
+def _named_linears(model: nn.Module):
+    """Yield ``(name, module)`` for every ``nn.Linear`` in ``model``."""
+    for name, module in model.named_modules():
+        if isinstance(module, nn.Linear):
+            yield name, module
+
+
+if __name__ == "__main__":  # pragma: no cover - manual smoke check
+    import copy
+
+    from hermes.config import HermesConfig
+    from hermes.model import build_model
+
+    cfg = HermesConfig(
+        vocab_size=128, hidden_size=64, intermediate_size=128, num_layers=2,
+        num_heads=4, num_kv_heads=2, head_dim=16, max_seq_len=32,
+    )
+    fp_model = build_model(cfg)
+    q_model = apply_weight_only_int4(copy.deepcopy(fp_model))
+    data = [torch.randint(0, cfg.vocab_size, (1, 8)) for _ in range(4)]
+    report = quantization_error_report(fp_model, q_model, data, num_batches=4)
+    print("max layer L2 error:", round(report["max_layer_l2"], 4))
+    print("perplexity delta:", round(report["perplexity_delta"], 4))