gemmarouter 0.1.0__py3-none-any.whl

gemmarouter/cli.py

@@ -0,0 +1,103 @@
+"""Command-line interface for gemmarouter (`gemmarouter route` and `gemmarouter playground`)."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from typing import Sequence
+
+from .router import DEFAULT_REPO_ID, load
+
+
+def _parse_option_flags(raw_options: list[str]) -> dict[str, str]:
+    """Parse `-o 'key:description'` CLI flags into an options dictionary."""
+    parsed: dict[str, str] = {}
+    for item in raw_options:
+        if ":" in item:
+            k, v = item.split(":", 1)
+            parsed[k.strip()] = v.strip()
+        else:
+            parsed[item.strip()] = item.strip()
+    return parsed
+
+
+def _run_route_cmd(args: argparse.Namespace) -> int:
+    """Execute the `gemmarouter route` CLI subcommand."""
+    options = _parse_option_flags(args.option)
+    router = load(model_id=args.model, device=args.device, threshold=args.threshold)
+    res = router.route(args.query, options=options, threshold=args.threshold)
+    if args.json:
+        print(json.dumps(res.to_dict(), indent=2))
+        return 0
+
+    status = "ABSTAIN (None of the options fit)" if res.abstain else f"ROUTED -> {res.decision}"
+    print("=" * 68)
+    print(f"Query          : {args.query}")
+    print(f"Decision       : {status}")
+    print(f"Top Candidate  : {res.top_option} ({res.confidence * 100:.2f}%)")
+    print(f"P(null)        : {res.null_probability:.4f} (threshold={res.threshold:.2f})")
+    print(f"Latency        : {res.latency_ms:.2f} ms (Head: {res.head_latency_ms:.3f} ms)")
+    print("=" * 68)
+    return 0
+
+
+def _run_playground_cmd(args: argparse.Namespace) -> int:
+    """Execute the `gemmarouter playground` CLI subcommand."""
+    from .playground import launch_playground
+
+    launch_playground(
+        host=args.host,
+        port=args.port,
+        our_repo=args.model,
+        open_browser=not args.no_browser,
+        preload=not args.no_preload,
+    )
+    return 0
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Build the top-level `gemmarouter` CLI argument parser."""
+    parser = argparse.ArgumentParser(
+        prog="gemmarouter",
+        description="Set-Equivariant Decision & Tool Router on EmbeddingGemma 2",
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    p_route = sub.add_parser("route", help="Route a query across candidate options from the terminal")
+    p_route.add_argument("query", help="Incoming user query or agent prompt")
+    p_route.add_argument(
+        "-o",
+        "--option",
+        action="append",
+        required=True,
+        help="Candidate option in 'key:description' format (pass at least two -o flags)",
+    )
+    p_route.add_argument("--threshold", type=float, default=0.58, help="Null Gate abstention threshold (default: 0.58)")
+    p_route.add_argument("--model", default=DEFAULT_REPO_ID, help="Hugging Face model ID or local path")
+    p_route.add_argument("--device", default=None, help="PyTorch device ('cpu' or 'cuda')")
+    p_route.add_argument("--json", action="store_true", help="Output raw JSON")
+
+    p_play = sub.add_parser("playground", help="Launch the interactive 3-column Live Decision Playground web UI")
+    p_play.add_argument("--host", default="127.0.0.1", help="Bind host (default: 127.0.0.1)")
+    p_play.add_argument("--port", type=int, default=7860, help="Bind port (default: 7860)")
+    p_play.add_argument("--model", default=DEFAULT_REPO_ID, help="Hugging Face model ID or local path")
+    p_play.add_argument("--no-browser", action="store_true", help="Do not auto-open browser tab")
+    p_play.add_argument("--no-preload", action="store_true", help="Lazy-load models on first request")
+    return parser
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    """Entrypoint for the `gemmarouter` console script."""
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    if args.command == "route":
+        return _run_route_cmd(args)
+    if args.command == "playground":
+        return _run_playground_cmd(args)
+    parser.print_help()
+    return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

gemmarouter/configuration.py

@@ -0,0 +1,50 @@
+"""Configuration for GemmaDecisionRouterModel: Set-Equivariant Decision Head + Null Gate on EmbeddingGemma 2."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from transformers import AutoConfig, PretrainedConfig
+from transformers.models.auto.configuration_auto import CONFIG_MAPPING
+
+QUESTION_TYPES = ("choice", "score", "noul")
+
+
+class GemmaDecisionRouterConfig(PretrainedConfig):
+    """Configuration class for GemmaDecisionRouterModel."""
+
+    model_type = "gemma_decision_router"
+    sub_configs = {"encoder_config": AutoConfig}
+
+    def __init__(
+        self,
+        encoder_config: dict[str, Any] | PretrainedConfig | None = None,
+        embed_dim: int = 768,
+        n_heads: int = 8,
+        n_layers: int = 2,
+        decision_temperature: float = 0.05,
+        null_threshold: float = 0.58,
+        max_seq_length: int = 2048,
+        prompt: str = "task: classification | query: ",
+        option_prompt: str = "title: none | text: ",
+        question_types: list[str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        if isinstance(encoder_config, dict):
+            model_type = encoder_config.get("model_type")
+            if model_type is not None and model_type in CONFIG_MAPPING:
+                encoder_config = CONFIG_MAPPING[model_type].from_dict(encoder_config)
+        self.encoder_config = encoder_config
+        self.embed_dim = embed_dim
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.decision_temperature = decision_temperature
+        self.null_threshold = null_threshold
+        self.max_seq_length = max_seq_length
+        self.prompt = prompt
+        self.option_prompt = option_prompt
+        self.question_types = list(question_types or QUESTION_TYPES)
+        super().__init__(**kwargs)
+
+
+__all__ = ["GemmaDecisionRouterConfig", "QUESTION_TYPES"]

gemmarouter/head.py

@@ -0,0 +1,90 @@
+"""Set-equivariant decision head with cross-option competition and decoupled coverage gate."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+class SetEquivariantDecisionHead(nn.Module):
+    """Permutation-equivariant 2-layer Transformer decision head with zero option positional encodings."""
+
+    def __init__(self, d_model: int = 768, n_heads: int = 8, n_layers: int = 2) -> None:
+        super().__init__()
+        self.query_role = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
+        self.option_role = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
+        self.q_proj = nn.Linear(d_model, d_model)
+        self.o_proj = nn.Linear(d_model, d_model)
+        self.pair_proj = nn.Linear(d_model + 1, d_model)
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=n_heads,
+            dim_feedforward=d_model * 2,
+            dropout=0.1,
+            activation="gelu",
+            batch_first=True,
+            norm_first=True,
+        )
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
+        self.option_scorer = nn.Sequential(
+            nn.Linear(d_model * 2, 256),
+            nn.GELU(),
+            nn.Dropout(0.1),
+            nn.Linear(256, 1),
+        )
+        nn.init.normal_(self.option_scorer[-1].weight, std=0.01)
+        nn.init.zeros_(self.option_scorer[-1].bias)
+
+        self.log_tau = nn.Parameter(torch.tensor(-3.5))
+        self.a_null = nn.Parameter(torch.tensor(-100.0))
+        self.b_null = nn.Parameter(torch.tensor(74.0))
+        self.null_mlp = nn.Sequential(nn.Linear(6, 64), nn.GELU(), nn.Linear(64, 1))
+        nn.init.normal_(self.null_mlp[-1].weight, std=0.02)
+        nn.init.zeros_(self.null_mlp[-1].bias)
+
+    def _compute_null_logits(
+        self, raw_sims: Tensor, adapt_sims: Tensor, choice_logits: Tensor
+    ) -> Tensor:
+        """Compute 6-feature permutation-invariant set statistics and return null logits [B]."""
+        k_top = min(2, raw_sims.size(1))
+        top2_raw, _ = torch.topk(raw_sims, k=k_top, dim=-1)
+        top2_adp, _ = torch.topk(adapt_sims, k=k_top, dim=-1)
+        max_sim = top2_raw[:, 0:1]
+        raw_margin = top2_raw[:, 0:1] - (top2_raw[:, 1:2] if k_top > 1 else top2_raw[:, 0:1])
+        adp_max = top2_adp[:, 0:1]
+        adp_margin = top2_adp[:, 0:1] - (top2_adp[:, 1:2] if k_top > 1 else top2_adp[:, 0:1])
+        mean_sim = raw_sims.mean(dim=-1, keepdim=True)
+
+        probs = F.softmax(choice_logits, dim=-1)
+        entropy = -(probs * torch.log(probs.clamp(min=1e-9))).sum(dim=-1, keepdim=True)
+        stats = torch.cat(
+            [max_sim * 10.0, raw_margin * 10.0, adp_max * 10.0, adp_margin * 10.0, mean_sim * 10.0, entropy],
+            dim=-1,
+        )
+        base_null = self.a_null * max_sim.squeeze(-1) + self.b_null
+        return base_null + self.null_mlp(stats).squeeze(-1)
+
+    def forward(self, query_emb: Tensor, option_embs: Tensor) -> tuple[Tensor, Tensor]:
+        """Compute per-option choice logits [B, K] and set-level null logits [B]."""
+        raw_sims = torch.bmm(option_embs, query_emb.unsqueeze(-1)).squeeze(-1)
+        q_adapt = F.normalize(query_emb + 0.2 * self.q_proj(query_emb), p=2, dim=-1)
+        o_adapt = F.normalize(option_embs + 0.2 * self.o_proj(option_embs), p=2, dim=-1)
+        adapt_sims = torch.bmm(o_adapt, q_adapt.unsqueeze(-1)).squeeze(-1)
+
+        q_exp = q_adapt.unsqueeze(1).expand_as(o_adapt)
+        pair_feat = torch.cat([q_exp * o_adapt, adapt_sims.unsqueeze(-1)], dim=-1)
+        q_tok = q_adapt.unsqueeze(1) + self.query_role
+        o_toks = o_adapt + self.option_role + 0.2 * self.pair_proj(pair_feat)
+
+        hidden = self.transformer(torch.cat([q_tok, o_toks], dim=1))
+        q_context = hidden[:, 0:1, :].expand_as(option_embs)
+        o_context = hidden[:, 1:, :]
+
+        tau = torch.exp(self.log_tau).clamp(min=1e-3, max=1.0)
+        delta_scores = self.option_scorer(torch.cat([o_context, q_context * o_context], dim=-1)).squeeze(-1)
+        choice_logits = (0.5 * raw_sims + 0.5 * adapt_sims) / tau + delta_scores
+        null_logits = self._compute_null_logits(raw_sims, adapt_sims, choice_logits)
+        return choice_logits, null_logits

gemmarouter/modeling.py

@@ -0,0 +1,335 @@
+"""GemmaDecisionRouter: Set-Equivariant Typed Decision & Router Model on top of EmbeddingGemma 2."""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+import json
+from dataclasses import asdict, dataclass
+from typing import Any, Literal
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from transformers import AutoModel, AutoTokenizer, PreTrainedModel
+from transformers.utils import ModelOutput
+
+from .configuration import QUESTION_TYPES, GemmaDecisionRouterConfig
+from .head import SetEquivariantDecisionHead
+
+Kind = Literal["choice", "score", "noul"]
+
+
+@dataclass(frozen=True)
+class QuestionSpec:
+    """Normalized specification for a single routing or decision question."""
+
+    qid: str
+    kind: Kind
+    instructions: str
+    options: tuple[str, ...]
+    descriptions: tuple[str, ...]
+
+    @property
+    def kind_id(self) -> int:
+        return QUESTION_TYPES.index(self.kind)
+
+
+@dataclass(frozen=True)
+class TypedAnswer:
+    """Structured answer returned by GemmaDecisionRouterModel.decide()."""
+
+    kind: Kind
+    label: str
+    value: str | int | bool
+    confidence: float
+    probabilities: dict[str, float]
+    expected_level: float | None = None
+    null_probability: float | None = None
+    abstain: bool = False
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass
+class GemmaDecisionOutput(ModelOutput):
+    """Output container for GemmaDecisionRouterModel.forward()."""
+
+    loss: Tensor | None = None
+    logits: Tensor | None = None
+    null_logits: Tensor | None = None
+    probabilities: Tensor | None = None
+    text_embeds: Tensor | None = None
+
+
+def _json(x: Any) -> Any:
+    return json.loads(x) if isinstance(x, str) else x
+
+
+def _is_int(x: Any) -> bool:
+    return str(x).lstrip("-").isdigit()
+
+
+def _parse_score_criteria(crit: Any, probs: dict[str, Any]) -> tuple[tuple[str, ...], tuple[str, ...]]:
+    """Extract option keys and descriptions for a `score` question."""
+    if isinstance(crit, list):
+        return tuple(str(i) for i in range(len(crit))), tuple(str(c) for c in crit)
+    if isinstance(crit, dict) and all(_is_int(k) for k in crit):
+        keys = sorted(crit, key=int)
+        return tuple(str(k) for k in keys), tuple(str(crit[k]) for k in keys)
+    if isinstance(crit, dict):
+        return tuple(str(i) for i in range(len(crit))), tuple(f"{k}: {v}" for k, v in crit.items())
+    opts = tuple(str(k) for k in sorted(probs, key=int))
+    return opts, opts
+
+
+def parse_question(qid: str, question: dict[str, Any], gold: dict[str, Any] | None = None) -> QuestionSpec:
+    """Parse a raw question dictionary into a strongly typed QuestionSpec."""
+    raw_kind = question.get("type") or (gold or {}).get("type") or "choice"
+    kind: Kind = raw_kind if raw_kind in QUESTION_TYPES else "choice"
+    crit = question.get("criteria")
+    probs = (gold or {}).get("probabilities") or {}
+
+    if kind == "noul":
+        if isinstance(crit, dict) and "false" in crit and "true" in crit:
+            opts, desc = ("false", "true"), (str(crit["false"]), str(crit["true"]))
+        else:
+            opts, desc = ("false", "true"), ("False / No: does not hold", "True / Yes: holds")
+    elif kind == "score":
+        opts, desc = _parse_score_criteria(crit, probs)
+    else:
+        if isinstance(crit, dict):
+            opts, desc = tuple(map(str, crit)), tuple(str(v) for v in crit.values())
+        elif isinstance(crit, list):
+            opts = desc = tuple(map(str, crit))
+        else:
+            opts = desc = tuple(map(str, probs))
+
+    if not opts:
+        raise ValueError(f"question {qid!r} has no options")
+    return QuestionSpec(qid, kind, str(question.get("instructions") or qid), opts, desc)
+
+
+def render_state(workflow: str, state: Any) -> str:
+    """Flatten a string or nested dictionary state into a canonical text string."""
+    state = _json(state) if isinstance(state, str) and state.strip()[:1] in "{[" else state
+    if not isinstance(state, dict):
+        if state == "":
+            return f"workflow: {workflow}" if workflow else ""
+        return f"workflow: {workflow} | state: {state}" if workflow else str(state)
+    lines: list[str] = []
+
+    def walk(prefix: str, value: Any) -> None:
+        if isinstance(value, dict):
+            for k, v in value.items():
+                walk(f"{prefix} · {k}" if prefix else k, v)
+        else:
+            lines.append(f"{prefix.replace('_', ' ')}: {value}")
+
+    walk("", state)
+    return (f"workflow: {workflow}\n" if workflow else "") + "\n".join(lines)
+
+
+def _format_option_texts(sp: QuestionSpec) -> list[str]:
+    """Format candidate option strings according to the question kind."""
+    opt_texts: list[str] = []
+    for key, desc in zip(sp.options, sp.descriptions):
+        if sp.kind == "noul":
+            opt_texts.append(f"Question: {sp.instructions} | Answer: {key} ({desc})")
+        elif sp.kind == "score":
+            opt_texts.append(f"Question: {sp.instructions} | Score level {key} of {len(sp.options)-1}: {desc}")
+        else:
+            opt_texts.append(f"{key}: {desc}" if desc != key else str(key))
+    return opt_texts
+
+
+class GemmaDecisionRouterPreTrainedModel(PreTrainedModel):
+    """Base PreTrainedModel wrapper for GemmaDecisionRouter."""
+
+    config_class = GemmaDecisionRouterConfig
+    base_model_prefix = "gemma_decision_router"
+    _supports_sdpa = True
+
+    def _init_weights(self, module: nn.Module) -> None:
+        pass
+
+
+class GemmaDecisionRouterModel(GemmaDecisionRouterPreTrainedModel):
+    """EmbeddingGemma 2 text encoder + Set-Equivariant Decision Head + Calibrated Null Gate."""
+
+    def __init__(self, config: GemmaDecisionRouterConfig) -> None:
+        super().__init__(config)
+        self.encoder = AutoModel.from_config(config.encoder_config)
+        self.head = SetEquivariantDecisionHead(
+            d_model=config.embed_dim,
+            n_heads=config.n_heads,
+            n_layers=config.n_layers,
+        )
+        self._option_emb_cache: OrderedDict[tuple[str, ...], Tensor] = OrderedDict()
+        self._max_option_cache_size: int = 128
+        self.post_init()
+
+    def clear_option_cache(self) -> None:
+        """Clear the in-memory option embedding cache."""
+        self._option_emb_cache.clear()
+
+    def get_tokenizer(self, tokenizer: Any = None) -> Any:
+        """Return or lazy-load the tokenizer associated with this model."""
+        if tokenizer is not None:
+            return tokenizer
+        if getattr(self, "_tokenizer", None) is None:
+            source = (
+                getattr(self.config, "_name_or_path", None)
+                or getattr(self, "name_or_path", None)
+                or "Ajeet-001/embeddinggemma2-decision-router"
+            )
+            self._tokenizer = AutoTokenizer.from_pretrained(source, trust_remote_code=False)
+        return self._tokenizer
+
+    def embed(self, input_ids: Tensor, attention_mask: Tensor) -> Tensor:
+        """Mean-pool and L2-normalize token representations from the encoder."""
+        out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
+        tokens = out.last_hidden_state
+        mask = attention_mask.unsqueeze(-1).to(tokens.dtype)
+        pooled = (tokens * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
+        return F.normalize(pooled[:, : self.config.embed_dim].float(), p=2, dim=-1)
+
+    @torch.no_grad()
+    def encode(
+        self, texts: list[str], prefix: str | None = None, batch_size: int = 64, tokenizer: Any = None
+    ) -> Tensor:
+        """Encode a list of strings with the given prompt prefix into L2-normalized vectors."""
+        tok = self.get_tokenizer(tokenizer)
+        pref = self.config.prompt if prefix is None else prefix
+        prompted = [f"{pref}{t}" for t in texts]
+        all_vecs = []
+        for i in range(0, len(prompted), batch_size):
+            batch = prompted[i : i + batch_size]
+            enc = tok(
+                batch,
+                padding=True,
+                truncation=True,
+                max_length=self.config.max_seq_length,
+                return_tensors="pt",
+            ).to(self.device)
+            all_vecs.append(self.embed(enc["input_ids"], enc["attention_mask"]))
+        return torch.cat(all_vecs, dim=0)
+
+    @torch.no_grad()
+    def encode_options_cached(
+        self, opt_texts: list[str], tokenizer: Any = None
+    ) -> tuple[Tensor, bool]:
+        """Encode option strings once and reuse bounded LRU-cached L2-normalized vectors."""
+        cache_key = (str(self.config.option_prompt), *opt_texts)
+        cached = self._option_emb_cache.get(cache_key)
+        if cached is not None and cached.device == self.device:
+            self._option_emb_cache.move_to_end(cache_key)
+            return cached, True
+        encoded = self.encode(opt_texts, prefix=self.config.option_prompt, tokenizer=tokenizer)
+        if len(self._option_emb_cache) >= self._max_option_cache_size:
+            self._option_emb_cache.popitem(last=False)
+        self._option_emb_cache[cache_key] = encoded
+        return encoded, False
+
+    def forward(
+        self,
+        query_emb: Tensor,
+        option_embs: Tensor,
+        return_dict: bool = True,
+    ) -> GemmaDecisionOutput | tuple[Tensor, Tensor, Tensor]:
+        """Run the SetEquivariantDecisionHead on pre-encoded query and option embeddings."""
+        head_dtype = next(self.head.parameters()).dtype
+        q_in = query_emb.to(head_dtype)
+        o_in = option_embs.to(head_dtype)
+        sims = torch.bmm(o_in.float(), q_in.float().unsqueeze(-1)).squeeze(-1)
+        head_logits, null_logits = self.head(q_in, o_in)
+        logits = (sims / self.config.decision_temperature) + 0.15 * head_logits.float()
+        probs = F.softmax(logits, dim=-1)
+        if not return_dict:
+            return (logits, null_logits, probs)
+        return GemmaDecisionOutput(logits=logits, null_logits=null_logits.float(), probabilities=probs)
+
+    def _score_single_question(
+        self,
+        q_vec: Tensor,
+        sp: QuestionSpec,
+        effective_thresh: float,
+        tokenizer: Any = None,
+    ) -> TypedAnswer:
+        """Score a single QuestionSpec using cached option embeddings and the decision head."""
+        opt_texts = _format_option_texts(sp)
+        o_vecs, _ = self.encode_options_cached(opt_texts, tokenizer=tokenizer)
+        out = self.forward(q_vec, o_vecs.unsqueeze(0))
+        if sp.kind == "choice":
+            head_dtype = next(self.head.parameters()).dtype
+            head_logits, _ = self.head(q_vec.to(head_dtype), o_vecs.unsqueeze(0).to(head_dtype))
+            p = F.softmax(head_logits.float(), dim=-1)[0].cpu()
+        else:
+            p = out.probabilities[0].cpu()
+        j = int(p.argmax())
+        label = sp.options[j]
+        value: str | int | bool = (
+            (label == "true") if sp.kind == "noul" else int(label) if sp.kind == "score" else label
+        )
+        expected = (
+            float(sum(int(o) * float(pi) for o, pi in zip(sp.options, p))) if sp.kind == "score" else None
+        )
+        null_prob = float(torch.sigmoid(out.null_logits[0]).cpu()) if sp.kind == "choice" else None
+        abstain = (null_prob >= effective_thresh) if null_prob is not None else False
+        return TypedAnswer(
+            kind=sp.kind,
+            label=label,
+            value=value,
+            confidence=float(p[j]),
+            probabilities={k: round(float(pi), 4) for k, pi in zip(sp.options, p.tolist())},
+            expected_level=round(expected, 4) if expected is not None else None,
+            null_probability=round(null_prob, 4) if null_prob is not None else None,
+            abstain=abstain,
+        )
+
+    @torch.no_grad()
+    def decide(
+        self,
+        state: Any,
+        questions: dict[str, dict[str, Any]],
+        workflow: str = "",
+        tokenizer: Any = None,
+        abstain_threshold: float | None = None,
+        include_question: bool | None = None,
+    ) -> dict[str, TypedAnswer]:
+        """Answer every question in a schema (`choice`, `score`, `noul`) with automatic option caching."""
+        was_training = self.training
+        self.eval()
+        state_str = render_state(workflow, state)
+        q_dict = _json(questions)
+        specs = [parse_question(qid, q) for qid, q in q_dict.items()]
+        plain_text = isinstance(state, str) and state.strip() != "" and state.strip()[0] not in "{["
+
+        q_texts: list[str] = []
+        for sp in specs:
+            drop_q = include_question is False or (include_question is None and plain_text and sp.kind == "choice")
+            if drop_q and state_str:
+                q_texts.append(f"State: {state.strip() if plain_text else state_str}")
+            else:
+                q_texts.append(f"State: {state_str}\nQuestion: {sp.instructions}" if state_str else f"Question: {sp.instructions}")
+
+        q_vecs = self.encode(q_texts, prefix=self.config.prompt, tokenizer=tokenizer)
+        effective_thresh = float(
+            getattr(self.config, "null_threshold", 0.58) if abstain_threshold is None else abstain_threshold
+        )
+        answers = {
+            sp.qid: self._score_single_question(q_vecs[i : i + 1], sp, effective_thresh, tokenizer=tokenizer)
+            for i, sp in enumerate(specs)
+        }
+        self.train(was_training)
+        return answers
+
+
+__all__ = [
+    "GemmaDecisionOutput",
+    "GemmaDecisionRouterModel",
+    "GemmaDecisionRouterPreTrainedModel",
+    "QuestionSpec",
+    "TypedAnswer",
+]

gemmarouter/playground.py

@@ -0,0 +1,32 @@
+"""Self-contained launcher module for the Live Head-to-Head Decision Playground."""
+
+from __future__ import annotations
+
+import webbrowser
+
+from . import _playground_server
+
+
+def launch_playground(
+    host: str = "127.0.0.1",
+    port: int = 7860,
+    our_repo: str = "Ajeet-001/embeddinggemma2-decision-router",
+    open_browser: bool = True,
+    preload: bool = True,
+) -> None:
+    """Start the interactive 3-column Live Head-to-Head Decision Playground server."""
+    server = _playground_server.create_server(host=host, port=port, our_repo=our_repo, preload=preload)
+    actual_port = server.server_address[1]
+    url = f"http://{host}:{actual_port}"
+    print("=" * 78)
+    print("GEMMAROUTER LIVE HEAD-TO-HEAD DECISION PLAYGROUND READY")
+    print(f"URL: {url}")
+    print("=" * 78)
+    if open_browser:
+        webbrowser.open(url)
+    try:
+        server.serve_forever()
+    except KeyboardInterrupt:
+        pass
+    finally:
+        server.server_close()