add: A CUDA path for the demo Space, behind the existing engine contract

MLX runs on Apple Silicon and nowhere else, so the public Space -- Linux/NVIDIA
on ZeroGPU -- has had no way to run the current code since the backends were
collapsed. This adds one alternative Engine implementation behind the same
contract, not an orchestrator: CONTROLAI_BACKEND=torch makes app.py build
TorchEngine instead of LocalEngine, and Embedder switches on the same variable.
Two branches. The agent loop, the registry and every tool are untouched.

engine_torch.py mirrors engine.py rather than calling model.generate, because
generate cannot express the two things that matter here: DynamicCache.crop() for
prefix reuse across tool steps, and injecting </think> to close an overrunning
reasoning block. Verified against Qwen2.5-0.5B on CPU: prefix reuse (39 of 56
prompt tokens served from cache on a follow-up turn), stop sequences, greedy
determinism, top-k/top-p/presence-penalty sampling, and the think-budget
injection all behave as the MLX engine does.

Loading is 4-bit NF4 -- Qwen3-14B is ~28GB in bf16, ~9GB quantised -- with
device_map={"": 0}. Not "auto": it inspects free VRAM at load time, which on
ZeroGPU happens before real hardware is attached to the process, and silently
offloads layers to CPU. The device is logged at startup so that is checkable.

app_space.py carries the ZeroGPU platform workarounds, each of which was
previously learned by having the Space fail to start: the @spaces.GPU function
must be module-level and wired to a real Gradio event handler or startup
validation does not see it; the probe demo must launch on its own port rather
than being mounted into the FastAPI app; and the agent must be reached through a
module global, never passed as an argument, or ZeroGPU tries to share its CUDA
tensors across a process boundary and hangs with no output.

Embedder gains the same torch path, since the Space also has to embed queries.
The index was built with the MLX 4-bit checkpoint and the Space queries with
bf16 Qwen/Qwen3-Embedding-0.6B; measured, the two agree at cosine 0.96, and
against the real 80,370-chunk index the cosine gate behaves the same -- in-domain
worst best-match 0.712 vs 0.708, off-domain best 0.528 vs 0.541. MIN_COSINE=0.62
holds unchanged, and is now overridable via CONTROLAI_MIN_COSINE if it drifts.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

CLAUDE.md

@@ -11,9 +11,13 @@ never from the model's own arithmetic. A BM25 + dense hybrid retriever over a lo
 corpus grounds conceptual answers. A FastAPI + vanilla-JS console and a terminal CLI are the two
 front ends. Nothing leaves the machine.
 
-Apple Silicon only. There is no CUDA, GGUF, Ollama, or Hugging Face Spaces path — an earlier
-version carried all four, and maintaining four backends for one machine was most of the complexity
-in the codebase.
+Apple Silicon only for local use. There is no GGUF or Ollama path — an earlier version carried
+both plus CUDA and a Spaces deployment, and maintaining four backends for one machine was most of
+the complexity in the codebase.
+
+One CUDA path came back, narrowly, for the public demo Space: see **Deployment** below. It is a
+second `Engine` implementation behind the same contract, selected by one env var. It does not touch
+the agent loop, the registry, or any tool.
 
 ## Commands
 
@@ -175,6 +179,40 @@ formats correctly. `CONTROLAI_ADAPTER=<path>` loads one anyway for A/B work.
 discovery → extraction → dataset generation → training → evaluation). That pipeline is independent
 of the serving path and uses `requirements-training.txt`/`requirements-corpus.txt`.
 
+### Deployment (`app_space.py`, `engine_torch.py`, `requirements-space.txt`)
+The demo Space (huggingface.co/spaces/atakankahya/ControlAI-Agent) runs Linux/NVIDIA on ZeroGPU,
+where MLX does not exist. `CONTROLAI_BACKEND=torch` makes `app.py::_make_engine` build
+`TorchEngine` instead of `LocalEngine`; `Embedder` switches on the same variable. That is the whole
+switch — two branches, no orchestrator.
+
+`engine_torch.py` mirrors `engine.py` rather than calling `model.generate`, because `generate`
+cannot express either of the two things that matter: `DynamicCache.crop()` for prefix reuse across
+tool steps, and injecting `</think>` to close an overrunning reasoning block. Loading is 4-bit NF4
+(Qwen3-14B is ~28GB bf16, ~9GB quantised) with `device_map={"": 0}` — **never `"auto"`**, which
+inspects free VRAM at load time, before ZeroGPU has attached hardware, and silently offloads to CPU.
+
+**ZeroGPU platform gotchas, each learned by having the Space fail:**
+- A `@spaces.GPU` function is only detected if wired to a real Gradio event handler. One called
+  solely from a FastAPI route fails startup with "No @spaces.GPU function detected". Hence the
+  hidden probe button.
+- That function must be a module-level `def`; nested inside `with gr.Blocks():` the detection fails.
+- Don't `mount_gradio_app()` the probe and then run uvicorn on the same port — Gradio's own server
+  setup collides ("address already in use"). The probe launches on `port + 1`, non-blocking.
+- **Never pass the model-holding object as an argument** to a `@spaces.GPU` function. ZeroGPU
+  marshals arguments across a process boundary and tries to share the model's CUDA tensors, failing
+  with `_share_cuda_: only available on CUDA` after emitting nothing — which reads exactly like
+  "just slow". Reach the agent through the module global.
+
+The Space needs `HF_TOKEN` as a secret: the retrieval index is in a private dataset repo and
+`app_space.py::_fetch_index` pulls it at startup. Without it the Space still boots, and answers
+from model knowledge alone.
+
+**The Space's queries are embedded with the bf16 `Qwen/Qwen3-Embedding-0.6B` while the index was
+built with the MLX 4-bit checkpoint.** Measured, the two produce vectors agreeing at cosine 0.96,
+and against the real 80,370-chunk index the gate behaves the same: in-domain worst best-match 0.712
+(MLX 0.708), off-domain best 0.528 (MLX 0.541). `MIN_COSINE = 0.62` therefore holds unchanged on
+both. `CONTROLAI_MIN_COSINE` overrides it if that ever drifts.
+
 ### Benchmark (`benchmarks/`)
 `controlbench_v1.jsonl` is the eval set; `SCOPE.md` defines the taxonomy and `README.md` the
 data-hygiene rules (never let benchmark prompts leak into training data; split by `family`, not by

app.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 
 import asyncio
 import json
+import os
 import queue
 import shutil
 import sys
@@ -56,10 +57,26 @@ inference_lock = threading.Lock()
 _agent: ControlAgent | None = None
 
 
+def _make_engine():
+    """The engine for this process. MLX unless the Space asked for CUDA.
+
+    This is the *only* backend branch left in the serving path, and it exists
+    for one reason: the public demo Space runs on Linux/NVIDIA, where MLX does
+    not exist. `TorchEngine` is imported lazily so a normal Apple Silicon run
+    never needs torch installed. Everything downstream -- the agent loop, the
+    registry, every tool -- is identical either way.
+    """
+    if os.environ.get("CONTROLAI_BACKEND", "mlx").lower() == "torch":
+        from controlai_agent.engine_torch import TorchEngine
+
+        return TorchEngine()
+    return None  # ControlAgent's own default, LocalEngine
+
+
 def get_agent() -> ControlAgent:
     global _agent
     if _agent is None:
-        _agent = ControlAgent()
+        _agent = ControlAgent(engine=_make_engine())
     return _agent
 
 

app_space.py

@@ -0,0 +1,89 @@
+"""Hugging Face Space entry point: the same FastAPI app, on ZeroGPU.
+
+Not used for local runs -- `./run.sh` runs `app.py` directly. This module exists
+because the Space's platform imposes requirements that have nothing to do with
+ControlAI, and each of the four below was learned by having the Space fail:
+
+1.  A `@spaces.GPU`-decorated function is only recognised by the platform's
+    startup validation if it is wired to an actual Gradio event handler. A
+    function called only from a FastAPI route fails startup with "No @spaces.GPU
+    function detected". Hence the hidden probe button.
+2.  That function must be a module-level `def`. Nested inside a `with
+    gr.Blocks():` block, the same detection fails.
+3.  Do not `gr.mount_gradio_app()` the probe into the FastAPI app and then run
+    uvicorn on the same port -- Gradio's own server setup collides with it
+    ("address already in use"). The probe launches on its own port; FastAPI is
+    the only thing bound to the public one.
+4.  Never pass the model-holding object as an *argument* to a `@spaces.GPU`
+    function. ZeroGPU marshals arguments across a process boundary and will try
+    to share the model's CUDA tensors, failing with `_share_cuda_: only
+    available on CUDA` after producing no output at all -- which reads exactly
+    like "just slow". Reach the agent through the module-level global instead.
+
+The retrieval index is not in the repo (see CLAUDE.md), so it is pulled from the
+private Hub dataset repo at startup using the Space's HF_TOKEN secret. Without
+it the retriever has no corpus and every answer falls back to model knowledge.
+"""
+
+from __future__ import annotations
+
+import os
+
+# Must be set before app.py is imported: it decides which engine gets built.
+os.environ.setdefault("CONTROLAI_BACKEND", "torch")
+# Qwen3-14B in 4-bit NF4. Override with CONTROLAI_MODEL_TORCH in Space settings.
+
+import gradio as gr
+import spaces
+import uvicorn
+
+from app import app, get_agent
+
+
+def _fetch_index() -> None:
+    """Pull the retrieval index into data/rag_index/ before the agent loads."""
+    from controlai_rag.fetch_index import fetch
+
+    if not os.environ.get("HF_TOKEN"):
+        print("[space] HF_TOKEN not set -- skipping index fetch, retrieval disabled")
+        return
+    try:
+        fetch()
+        print("[space] retrieval index ready")
+    except Exception as exc:  # noqa: BLE001 - a missing index must not stop boot
+        print(f"[space] could not fetch the index ({exc}); retrieval disabled")
+
+
+@spaces.GPU(duration=120)
+def _gpu_probe(text: str) -> str:
+    """Satisfies ZeroGPU's startup validation, and warms the model on first use.
+
+    Takes only a plain string. The agent is reached through the module global --
+    see note 4 above; passing it in is what hangs the whole Space.
+    """
+    agent = get_agent()
+    return agent.run(text).answer if text else "ready"
+
+
+with gr.Blocks() as _gpu_demo:
+    _probe_in = gr.Textbox(visible=False)
+    _probe_out = gr.Textbox(visible=False)
+    _probe_btn = gr.Button("probe", visible=False)
+    _probe_btn.click(fn=_gpu_probe, inputs=_probe_in, outputs=_probe_out)
+
+
+def main() -> None:
+    _fetch_index()
+    port = int(os.environ.get("PORT", 7860))
+    # Separate port, non-blocking: see note 3.
+    _gpu_demo.launch(
+        server_name="0.0.0.0",
+        server_port=port + 1,
+        prevent_thread_lock=True,
+        share=False,
+    )
+    uvicorn.run(app, host="0.0.0.0", port=port, log_level="info")
+
+
+if __name__ == "__main__":
+    main()

controlai_agent/engine_torch.py

@@ -0,0 +1,348 @@
+"""CUDA inference engine: the same contract as `LocalEngine`, on transformers.
+
+`LocalEngine` is MLX, so it runs on Apple Silicon and nowhere else. The public
+demo Space runs on Linux/NVIDIA (ZeroGPU), which needs a second path. This is
+that path and nothing more -- it is not a supported way to run ControlAI locally,
+and `agent.py`, `registry.py` and everything under `tools/` are untouched by it.
+The seam is `ControlAgent(engine=...)`.
+
+It mirrors `LocalEngine`'s design decisions rather than reaching for
+`model.generate`:
+
+*   **Prefix-reusing KV cache.** `DynamicCache.crop()` trims the cache to the
+    longest prefix the incoming prompt shares with it, so the ~8k-token
+    system-prompt-plus-tool-schema prefix is prefilled once per process, not
+    once per tool step. `model.generate` cannot express that.
+*   **A hand-written decode loop.** `think_budget` closes an overrunning
+    `<think>` block by *injecting* the closing token, which no `generate`
+    callback can do, and `presence_penalty` (not `repetition_penalty`) is
+    applied for the reason spelled out in `engine.py`: a flat repetition penalty
+    punishes the `[`, `0`, `,` that matrices and JSON are made of.
+
+Loading is 4-bit NF4 via bitsandbytes: Qwen3-14B is ~28GB in bf16 and ~9GB
+quantised. `device_map={"": 0}` pins every layer to GPU 0 explicitly. Do not use
+`device_map="auto"` here -- it inspects free VRAM at load time, which on ZeroGPU
+happens *before* real hardware is attached to the process, and silently offloads
+layers to CPU.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from typing import Any, Generator, Iterable, Sequence
+
+from controlai_agent.engine import Chunk, SamplingConfig, Stats
+
+# The MLX default is a 4-bit MLX conversion, which transformers cannot read.
+# This is the same weights in a format it can.
+DEFAULT_TORCH_MODEL = os.environ.get("CONTROLAI_MODEL_TORCH", "Qwen/Qwen3-14B")
+
+
+class TorchEngine:
+    """Streaming generation against a CUDA-resident transformers model."""
+
+    def __init__(
+        self,
+        model_id: str = DEFAULT_TORCH_MODEL,
+        adapter_path: str | None = None,
+        sampling: SamplingConfig | None = None,
+        max_cache_tokens: int = 32768,
+        load_in_4bit: bool = True,
+    ) -> None:
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        self.model_id = model_id
+        self.adapter_path = adapter_path
+        self.sampling = sampling or SamplingConfig()
+        self.max_cache_tokens = max_cache_tokens
+
+        t0 = time.time()
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+        # `{"": 0}` pins every layer to GPU 0 explicitly -- never "auto", see the
+        # module docstring. The CPU branch exists only so the decode loop can be
+        # exercised on a small model off a GPU box; it is far too slow to serve.
+        cuda = torch.cuda.is_available()
+        # `torch_dtype`, not `dtype`: the newer spelling only exists from
+        # transformers 4.56, and requirements-space.txt pins older than that.
+        kwargs: dict[str, Any] = {
+            "torch_dtype": torch.bfloat16 if cuda else torch.float32,
+            "device_map": {"": 0} if cuda else "cpu",
+        }
+        if load_in_4bit and not cuda:
+            load_in_4bit = False  # bitsandbytes is CUDA-only
+        if load_in_4bit:
+            from transformers import BitsAndBytesConfig
+
+            kwargs["quantization_config"] = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_compute_dtype=torch.bfloat16,
+                # Quantising the quantisation constants too; ~0.4GB saved on a
+                # 14B model for no measurable quality cost.
+                bnb_4bit_use_double_quant=True,
+            )
+        self.model = AutoModelForCausalLM.from_pretrained(model_id, **kwargs)
+        if adapter_path:
+            from peft import PeftModel
+
+            self.model = PeftModel.from_pretrained(self.model, adapter_path)
+        self.model.eval()
+        self.load_seconds = time.time() - t0
+        # Verifiable in the Space logs: if this says cpu, device_map silently
+        # offloaded and every request will be minutes rather than seconds.
+        print(f"[torch] {model_id} on {next(self.model.parameters()).device} "
+              f"in {self.load_seconds:.1f}s")
+
+        self._cache: Any | None = None
+        self._cache_tokens: list[int] = []
+        self.last_stats = Stats()
+
+        self.supports_thinking = self._probe_thinking_support()
+        self._think_open = self._single_token("<think>")
+        self._think_close = self._single_token("</think>")
+        self._tool_open = self._single_token("<tool_call>")
+        self._tool_close = self._single_token("</tool_call>")
+
+    # ------------------------------------------------------------------ setup
+
+    def _token_ids(self, text: str) -> list[int]:
+        try:
+            return self.tokenizer.encode(text, add_special_tokens=False)
+        except TypeError:
+            return self.tokenizer.encode(text)
+
+    def _single_token(self, text: str) -> int | None:
+        ids = self._token_ids(text)
+        return ids[0] if len(ids) == 1 else None
+
+    def _probe_thinking_support(self) -> bool:
+        try:
+            self.tokenizer.apply_chat_template(
+                [{"role": "user", "content": "x"}],
+                tokenize=False,
+                add_generation_prompt=True,
+                enable_thinking=False,
+            )
+            return True
+        except (TypeError, ValueError):
+            return False
+
+    def render(
+        self,
+        messages: Sequence[dict[str, Any]],
+        tools: Sequence[dict[str, Any]] | None = None,
+        enable_thinking: bool = False,
+    ) -> str:
+        kwargs: dict[str, Any] = {"tokenize": False, "add_generation_prompt": True}
+        if tools:
+            kwargs["tools"] = list(tools)
+        if self.supports_thinking:
+            kwargs["enable_thinking"] = enable_thinking
+        return self.tokenizer.apply_chat_template(list(messages), **kwargs)
+
+    def encode(self, text: str) -> list[int]:
+        return self.tokenizer.encode(text)
+
+    def count_tokens(self, text: str) -> int:
+        return len(self.tokenizer.encode(text))
+
+    # ------------------------------------------------------------------ cache
+
+    def reset_cache(self) -> None:
+        self._cache = None
+        self._cache_tokens = []
+
+    def _align_cache(self, tokens: list[int]) -> list[int]:
+        """Trim the cache to the longest prefix it shares with `tokens`.
+
+        Returns the suffix that still has to be fed to the model. Mirrors
+        `LocalEngine._align_cache`; see that docstring for why this exists.
+        """
+        from transformers import DynamicCache
+
+        if self._cache is None or not self._cache_tokens:
+            self._cache = DynamicCache()
+            self._cache_tokens = []
+            return list(tokens)
+
+        shared = 0
+        for a, b in zip(self._cache_tokens, tokens):
+            if a != b:
+                break
+            shared += 1
+
+        # Never keep the whole prompt: the model needs at least one token to
+        # run forward on, or there are no logits to sample from.
+        if shared >= len(tokens):
+            shared = len(tokens) - 1
+        if shared > self.max_cache_tokens:
+            shared = 0
+
+        if shared == 0:
+            self._cache = DynamicCache()
+            self._cache_tokens = []
+            return list(tokens)
+
+        if shared < len(self._cache_tokens):
+            self._cache.crop(shared)
+        self._cache_tokens = list(tokens[:shared])
+        return list(tokens[shared:])
+
+    def prewarm(self, text: str) -> int:
+        """Prefill a prompt prefix so the first real question doesn't pay for it."""
+        import torch
+
+        tokens = self.encode(text)
+        to_feed = self._align_cache(tokens)
+        if to_feed:
+            with torch.inference_mode():
+                self.model(
+                    input_ids=torch.tensor([to_feed], device=self.model.device),
+                    past_key_values=self._cache,
+                    use_cache=True,
+                )
+        self._cache_tokens = list(tokens)
+        return len(tokens)
+
+    # ------------------------------------------------------------- generation
+
+    def _sample(self, logits: Any, cfg: SamplingConfig, seen: set[int]) -> int:
+        import torch
+
+        logits = logits.float()
+        if cfg.presence_penalty and seen:
+            idx = torch.tensor(sorted(seen), device=logits.device)
+            logits[idx] -= cfg.presence_penalty
+        if cfg.temperature <= 0:
+            return int(torch.argmax(logits).item())
+        logits = logits / cfg.temperature
+
+        if cfg.top_k and cfg.top_k > 0:
+            kth = torch.topk(logits, min(cfg.top_k, logits.numel())).values[-1]
+            logits = logits.masked_fill(logits < kth, float("-inf"))
+
+        probs = torch.softmax(logits, dim=-1)
+        if cfg.top_p and 0 < cfg.top_p < 1:
+            ordered, order = torch.sort(probs, descending=True)
+            cumulative = torch.cumsum(ordered, dim=-1)
+            # Keep the first token that crosses top_p, so the mask is never
+            # empty even when one token already carries more than top_p mass.
+            drop = cumulative - ordered > cfg.top_p
+            ordered[drop] = 0.0
+            probs = torch.zeros_like(probs).scatter_(0, order, ordered)
+            probs = probs / probs.sum()
+
+        return int(torch.multinomial(probs, 1).item())
+
+    def stream(
+        self,
+        prompt: str | list[int],
+        sampling: SamplingConfig | None = None,
+        stop: Iterable[str] = (),
+        think_budget: int | None = None,
+    ) -> Generator[Chunk, None, None]:
+        """Yield output chunks as they are generated. See `LocalEngine.stream`."""
+        import torch
+
+        cfg = sampling or self.sampling
+        tokens = self.encode(prompt) if isinstance(prompt, str) else list(prompt)
+
+        t0 = time.time()
+        to_feed = self._align_cache(tokens)
+        self.last_stats = Stats(
+            prompt_tokens=len(tokens),
+            cached_tokens=len(tokens) - len(to_feed),
+        )
+
+        stop = tuple(s for s in stop if s)
+        stop_ids = {self._tool_close} if "</tool_call>" in stop and self._tool_close else set()
+        text_stops = tuple(s for s in stop if not (s == "</tool_call>" and self._tool_close))
+        window = max((len(s) for s in text_stops), default=0) + 8
+
+        eos_ids = {self.tokenizer.eos_token_id}
+        for extra in ("<|im_end|>", "<|endoftext|>"):
+            tid = self._single_token(extra)
+            if tid is not None:
+                eos_ids.add(tid)
+        eos_ids.discard(None)
+
+        emitted: list[int] = []
+        seen: set[int] = set()
+        tail = ""
+        thinking = False
+        think_tokens = 0
+        prefill_done = False
+
+        with torch.inference_mode():
+            step_input = to_feed
+            while len(emitted) < cfg.max_tokens:
+                out = self.model(
+                    input_ids=torch.tensor([step_input], device=self.model.device),
+                    past_key_values=self._cache,
+                    use_cache=True,
+                )
+                self._cache = out.past_key_values
+                self._cache_tokens.extend(step_input)
+                if not prefill_done:
+                    self.last_stats.prefill_seconds = time.time() - t0
+                    t1 = time.time()
+                    prefill_done = True
+
+                token = self._sample(out.logits[0, -1, :], cfg, seen)
+                if token in eos_ids:
+                    break
+
+                seen.add(token)
+                emitted.append(token)
+
+                if token == self._think_open:
+                    thinking, think_tokens = True, 0
+                elif token == self._think_close:
+                    thinking = False
+                elif thinking:
+                    think_tokens += 1
+
+                text = self.tokenizer.decode([token], skip_special_tokens=False)
+                yield Chunk(
+                    text=text,
+                    token=token,
+                    thinking=thinking,
+                    tool_call=token == self._tool_open,
+                )
+
+                step_input = [token]
+
+                # Overran the reasoning budget: close the block by hand and make
+                # the model answer. Bounds worst-case latency on a reasoning model.
+                if (
+                    thinking
+                    and think_budget
+                    and think_tokens >= think_budget
+                    and self._think_close is not None
+                ):
+                    thinking = False
+                    emitted.append(self._think_close)
+                    yield Chunk(text="</think>", token=self._think_close, thinking=False)
+                    step_input = [token, self._think_close]
+
+                if token in stop_ids:
+                    break
+                if text_stops:
+                    tail = (tail + text)[-window:]
+                    if any(s in tail for s in text_stops):
+                        break
+
+        self.last_stats.generated_tokens = len(emitted)
+        self.last_stats.decode_seconds = time.time() - (t1 if prefill_done else t0)
+
+    def generate(
+        self,
+        prompt: str | list[int],
+        sampling: SamplingConfig | None = None,
+        stop: Iterable[str] = (),
+        think_budget: int | None = None,
+    ) -> str:
+        return "".join(c.text for c in self.stream(prompt, sampling, stop, think_budget))

controlai_rag/embeddings.py

@@ -15,6 +15,9 @@ from pathlib import Path
 import numpy as np
 
 MODEL_ID = os.environ.get("CONTROLAI_EMBED_MODEL", "mlx-community/Qwen3-Embedding-0.6B-4bit-DWQ")
+# The Space runs on Linux, where MLX does not exist, so the same embedder has a
+# transformers path. It is the same weights unquantised; see `Embedder._backend`.
+TORCH_MODEL_ID = os.environ.get("CONTROLAI_EMBED_MODEL_TORCH", "Qwen/Qwen3-Embedding-0.6B")
 MAX_TOKENS = 512
 # Qwen3-Embedding pools the hidden state at the final position, and it was
 # trained with an explicit end-of-text token in that position. Omitting it is
@@ -34,31 +37,58 @@ QUERY_INSTRUCTION = (
 class Embedder:
     """Lazily-loaded sentence embedder producing L2-normalised float32 vectors."""
 
-    def __init__(self, model_id: str = MODEL_ID) -> None:
-        self.model_id = model_id
+    def __init__(self, model_id: str | None = None, backend: str | None = None) -> None:
+        # "mlx" locally, "torch" on the Space. The vectors in embeddings.npz were
+        # produced by the MLX 4-bit checkpoint; the bf16 transformers weights are
+        # the same model, so the two agree closely but not bit-exactly. If
+        # retrieval on the Space looks over- or under-eager, MIN_COSINE is the
+        # knob (CONTROLAI_MIN_COSINE), not this.
+        self._backend = (backend or os.environ.get("CONTROLAI_BACKEND", "mlx")).lower()
+        if self._backend != "torch":
+            self._backend = "mlx"
+        self.model_id = model_id or (TORCH_MODEL_ID if self._backend == "torch" else MODEL_ID)
         self._model = None
         self._tokenizer = None
         self._eos_id: int | None = None
         self._pad_id: int | None = None
 
     def _ensure_loaded(self) -> None:
-        if self._model is None:
+        if self._model is not None:
+            return
+        if self._backend == "torch":
+            import torch
+            from transformers import AutoModel, AutoTokenizer
+
+            self._tokenizer = AutoTokenizer.from_pretrained(self.model_id)
+            self._model = AutoModel.from_pretrained(
+                self.model_id,
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            )
+            self._model = self._model.to("cuda" if torch.cuda.is_available() else "cpu")
+            self._model.eval()
+        else:
             from mlx_lm import load
 
             self._model, self._tokenizer = load(self.model_id)
-            ids = self._tokenizer.encode(EOS_TOKEN)
-            self._eos_id = ids[-1] if ids else self._tokenizer.eos_token_id
-            self._pad_id = self._tokenizer.pad_token_id or self._eos_id
+        ids = self._tokenizer.encode(EOS_TOKEN)
+        self._eos_id = ids[-1] if ids else self._tokenizer.eos_token_id
+        self._pad_id = self._tokenizer.pad_token_id or self._eos_id
 
     @property
     def dim(self) -> int:
         self._ensure_loaded()
+        if self._backend == "torch":
+            return int(self._model.config.hidden_size)
         return int(self._model.args.hidden_size)
 
     def _tokens_for(self, text: str) -> list[int]:
         return self._tokenizer.encode(text)[: MAX_TOKENS - 1] + [self._eos_id]
 
     def _encode_one(self, text: str) -> np.ndarray:
+        self._ensure_loaded()
+        if self._backend == "torch":
+            return self._encode_batch([self._tokens_for(text)])[0]
+
         import mlx.core as mx
 
         ids = self._tokens_for(text)
@@ -78,6 +108,10 @@ class Embedder:
         positions <= i, so tokens appended after the real end cannot influence
         the hidden state being pooled.
         """
+        self._ensure_loaded()
+        if self._backend == "torch":
+            return self._encode_batch_torch(batch)
+
         import mlx.core as mx
 
         lengths = [len(ids) for ids in batch]
@@ -89,6 +123,34 @@ class Embedder:
         picked = picked / (mx.linalg.norm(picked, axis=-1, keepdims=True) + 1e-9)
         return np.array(picked, copy=True)
 
+    def _encode_batch_torch(self, batch: list[list[int]]) -> np.ndarray:
+        """`_encode_batch` on transformers. Same right-padding and same pooling.
+
+        An explicit attention mask is passed even though right-padding a causal
+        backbone is already safe, because transformers otherwise warns on every
+        call and the mask costs nothing.
+        """
+        import torch
+
+        lengths = [len(ids) for ids in batch]
+        width = max(lengths)
+        pad = self._pad_id
+        device = self._model.device
+        ids = torch.tensor(
+            [row + [pad] * (width - len(row)) for row in batch], device=device
+        )
+        mask = torch.zeros_like(ids)
+        for i, n in enumerate(lengths):
+            mask[i, :n] = 1
+
+        with torch.inference_mode():
+            hidden = self._model(input_ids=ids, attention_mask=mask).last_hidden_state
+        picked = torch.stack(
+            [hidden[i, n - 1] for i, n in enumerate(lengths)]
+        ).float()
+        picked = picked / (picked.norm(dim=-1, keepdim=True) + 1e-9)
+        return picked.cpu().numpy().astype(np.float32)
+
     def encode_documents(
         self,
         texts: list[str],

controlai_rag/retriever.py

@@ -18,6 +18,7 @@ Without it, retrieval degrades to lexical-only rather than failing.
 from __future__ import annotations
 
 import argparse
+import os
 import re
 from pathlib import Path
 from typing import Any
@@ -39,7 +40,10 @@ EMBEDDINGS_PATH = INDEX_DIR / "embeddings.npz"
 # 0.62 sits in that 0.108-wide gap. The off-domain outlier is a pollen-allergy
 # query at 0.571, pulled up by biomedical material in the open_books tier --
 # everything genuinely unrelated lands near 0.42.
-MIN_COSINE = 0.62
+# Overridable because the Space embeds queries with the bf16 transformers
+# checkpoint rather than the MLX 4-bit one the index was built with; the two
+# agree closely, but this is the knob if the gate turns out mis-set there.
+MIN_COSINE = float(os.environ.get("CONTROLAI_MIN_COSINE", "0.62"))
 # Standard RRF constant; damps the influence of any single ranker's tail.
 RRF_K = 60
 

requirements-space.txt

@@ -0,0 +1,31 @@
+# Dependencies for the Hugging Face Space only -- see app_space.py.
+#
+# The Space runs on Linux/NVIDIA, where MLX does not exist, so inference goes
+# through controlai_agent/engine_torch.py instead. Do NOT install this file
+# locally: requirements.txt is the Apple Silicon runtime and is MLX-only.
+
+torch>=2.4.0
+# <4.56: engine_torch passes `torch_dtype=`, which the newer releases renamed.
+transformers>=4.51.0,<4.56
+accelerate>=0.30.0
+bitsandbytes>=0.43.0          # 4-bit NF4; CUDA-only
+sentencepiece>=0.2.0
+
+gradio>=4.44.0                # only for the ZeroGPU probe
+spaces>=0.30.0
+huggingface-hub>=0.23.0
+
+# Everything below is shared with requirements.txt (minus mlx-lm).
+numpy>=1.24.0
+scipy>=1.11.0
+control>=0.9.4
+cvxpy>=1.4.0
+matplotlib>=3.7.0
+jsonschema>=4.20.0
+rank-bm25>=0.2.2
+pypdf>=3.17.0
+pymupdf>=1.23.0
+fastapi>=0.110.0
+uvicorn>=0.28.0
+python-multipart>=0.0.9
+pydantic>=2.0.0