proprioceptive-lens 0.1.0__tar.gz

proprioceptive_lens-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Proprioceptive AI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

proprioceptive_lens-0.1.0/PKG-INFO

@@ -0,0 +1,85 @@
+Metadata-Version: 2.4
+Name: proprioceptive-lens
+Version: 0.1.0
+Summary: See inside your local model: live confidence, hesitations, long-term memory and internal-state behaviour for local LLMs.
+License-Expression: MIT
+Keywords: llm,ollama,interpretability,local-ai,memory,hidden-states
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Provides-Extra: deep
+Requires-Dist: torch>=2.1; extra == "deep"
+Requires-Dist: transformers>=4.43; extra == "deep"
+Dynamic: license-file
+
+# Proprioceptive Lens
+
+**See inside your local model.** A free app for anyone running local AI with [Ollama](https://ollama.com) or Hugging Face models.
+
+![Lens showing a confidently wrong answer: every word at near-100% certainty](docs/demo-confidence.gif)
+
+*qwen3:4b says the 17th US president was Garfield, with 100% certainty. (It was Andrew Johnson.) Token confidence shows the surface; Lens also lets you look underneath.*
+
+- **Confidence you can see.** Every word is tinted by how sure the model was. Hover any word to see what it almost said.
+- **Where it hesitated.** A live certainty trace, plus the exact words where the model weighed alternatives.
+- **It remembers you.** Private memory of your conversations, on your computer, for any model.
+- **Inside the model** (deep mode). A live layer-by-layer activity map of the network while it writes.
+- **Behaviour** (optional). Internal-state probes score what the model is doing: writing code or prose, formal or casual
+  register, following your instructions, creative or literal, assertive or hedged, repeating itself, asking or stating.
+  These probes were validated on the original model, then transferred to each supported model. There they were
+  re-tested against length and random-direction controls, and with a live test on the model's own answers.
+- **Terminal mode and an API** for wiring Lens into your own stack.
+- **Private.** Everything runs locally. Behaviour is opt-in and sends only 16 scrambled numbers per word, never your text.
+
+![Deep mode: layer activity map and live Behaviour scores](docs/demo-deep-behaviour.gif)
+
+## Install and run
+
+```bash
+pip install proprioceptive-lens      # no dependencies; Python 3.9+
+ollama pull qwen3:4b                 # any chat model works
+lens                                 # opens the app in your browser
+```
+
+Options: `lens --port 7777 --ollama http://127.0.0.1:11434 --hindsight http://localhost:8888`.
+
+## Deep mode and Behaviour
+
+```bash
+pip install "proprioceptive-lens[deep]"
+lens --deep Qwen/Qwen2.5-3B-Instruct
+```
+
+Deep mode runs a Hugging Face model locally and reads its hidden states. It works with Llama, Qwen, Mistral, Gemma,
+Phi, Falcon, Granite, OLMo, SmolLM, DeepSeek and other families. Lens tells Ollama users which command matches the
+model they already run. Switch on **Behaviour** in the app to see internal-state probe scores. The app lists which
+models have a Behaviour pack; so does [docs/MODELS.md](docs/MODELS.md).
+
+## Proprioceptive Terminal Mode
+
+```bash
+lens term                                   # interactive: confidence-coloured output, /why, /memory, /behaviour
+lens term -m qwen3:4b "your question"       # one-shot
+echo "your question" | lens term --json     # JSON-lines events for scripts and pipelines
+```
+
+Events: `memory`, `token` (`text`, `p`, `alternatives`, plus `layers` in deep mode), `behaviour` (score per probe),
+`done`. Any OpenAI-compatible client pointed at `http://127.0.0.1:7777/v1` also gets Lens memory.
+
+## Share a prompt
+
+`http://127.0.0.1:7777/?ask=Why%20is%20the%20sky%20blue` opens Lens with the prompt ready to send.
+
+## Privacy
+
+Lens listens only on `127.0.0.1` and rejects requests from other websites. Memories live in
+`~/.proprioceptive-lens/memory.sqlite3`. You can view or forget them in the app, or delete the file. Behaviour is off
+until you turn it on.
+
+## License
+
+MIT for this client. The Behaviour probes and the scoring service are proprietary to Proprioceptive AI, are not part
+of this package, and no licence to them (or to any Proprioceptive AI patent) is granted by this licence.

proprioceptive_lens-0.1.0/README.md

@@ -0,0 +1,68 @@
+# Proprioceptive Lens
+
+**See inside your local model.** A free app for anyone running local AI with [Ollama](https://ollama.com) or Hugging Face models.
+
+![Lens showing a confidently wrong answer: every word at near-100% certainty](docs/demo-confidence.gif)
+
+*qwen3:4b says the 17th US president was Garfield, with 100% certainty. (It was Andrew Johnson.) Token confidence shows the surface; Lens also lets you look underneath.*
+
+- **Confidence you can see.** Every word is tinted by how sure the model was. Hover any word to see what it almost said.
+- **Where it hesitated.** A live certainty trace, plus the exact words where the model weighed alternatives.
+- **It remembers you.** Private memory of your conversations, on your computer, for any model.
+- **Inside the model** (deep mode). A live layer-by-layer activity map of the network while it writes.
+- **Behaviour** (optional). Internal-state probes score what the model is doing: writing code or prose, formal or casual
+  register, following your instructions, creative or literal, assertive or hedged, repeating itself, asking or stating.
+  These probes were validated on the original model, then transferred to each supported model. There they were
+  re-tested against length and random-direction controls, and with a live test on the model's own answers.
+- **Terminal mode and an API** for wiring Lens into your own stack.
+- **Private.** Everything runs locally. Behaviour is opt-in and sends only 16 scrambled numbers per word, never your text.
+
+![Deep mode: layer activity map and live Behaviour scores](docs/demo-deep-behaviour.gif)
+
+## Install and run
+
+```bash
+pip install proprioceptive-lens      # no dependencies; Python 3.9+
+ollama pull qwen3:4b                 # any chat model works
+lens                                 # opens the app in your browser
+```
+
+Options: `lens --port 7777 --ollama http://127.0.0.1:11434 --hindsight http://localhost:8888`.
+
+## Deep mode and Behaviour
+
+```bash
+pip install "proprioceptive-lens[deep]"
+lens --deep Qwen/Qwen2.5-3B-Instruct
+```
+
+Deep mode runs a Hugging Face model locally and reads its hidden states. It works with Llama, Qwen, Mistral, Gemma,
+Phi, Falcon, Granite, OLMo, SmolLM, DeepSeek and other families. Lens tells Ollama users which command matches the
+model they already run. Switch on **Behaviour** in the app to see internal-state probe scores. The app lists which
+models have a Behaviour pack; so does [docs/MODELS.md](docs/MODELS.md).
+
+## Proprioceptive Terminal Mode
+
+```bash
+lens term                                   # interactive: confidence-coloured output, /why, /memory, /behaviour
+lens term -m qwen3:4b "your question"       # one-shot
+echo "your question" | lens term --json     # JSON-lines events for scripts and pipelines
+```
+
+Events: `memory`, `token` (`text`, `p`, `alternatives`, plus `layers` in deep mode), `behaviour` (score per probe),
+`done`. Any OpenAI-compatible client pointed at `http://127.0.0.1:7777/v1` also gets Lens memory.
+
+## Share a prompt
+
+`http://127.0.0.1:7777/?ask=Why%20is%20the%20sky%20blue` opens Lens with the prompt ready to send.
+
+## Privacy
+
+Lens listens only on `127.0.0.1` and rejects requests from other websites. Memories live in
+`~/.proprioceptive-lens/memory.sqlite3`. You can view or forget them in the app, or delete the file. Behaviour is off
+until you turn it on.
+
+## License
+
+MIT for this client. The Behaviour probes and the scoring service are proprietary to Proprioceptive AI, are not part
+of this package, and no licence to them (or to any Proprioceptive AI patent) is granted by this licence.

proprioceptive_lens-0.1.0/lens/__init__.py

@@ -0,0 +1,2 @@
+"""Proprioceptive Lens — see inside your local model."""
+__version__ = "0.1.0"

proprioceptive_lens-0.1.0/lens/__main__.py

@@ -0,0 +1,63 @@
+"""`lens` command: start the local app and open it in the browser."""
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import time
+import webbrowser
+from pathlib import Path
+
+from . import __version__
+from .server import LensApp, serve
+
+
+def main(argv=None) -> int:
+    argv = sys.argv[1:] if argv is None else argv
+    if argv and argv[0] in ("term", "terminal"):
+        from .terminal import main as terminal_main
+        return terminal_main(argv[1:])
+    parser = argparse.ArgumentParser(prog="lens", description="Proprioceptive Lens — see inside your local model. "
+                                     "Run `lens term` for Proprioceptive Terminal Mode.")
+    parser.add_argument("--port", type=int, default=int(os.environ.get("LENS_PORT", 7777)))
+    parser.add_argument("--ollama", default=os.environ.get("OLLAMA_HOST", "http://127.0.0.1:11434"),
+                        help="Ollama URL (default http://127.0.0.1:11434)")
+    parser.add_argument("--hindsight", default=os.environ.get("LENS_HINDSIGHT_URL"),
+                        help="optional Hindsight server URL to mirror memories into")
+    parser.add_argument("--data", default=os.environ.get("LENS_DATA", str(Path.home() / ".proprioceptive-lens")))
+    parser.add_argument("--deep", default=os.environ.get("LENS_DEEP_MODEL"),
+                        help="Hugging Face model to run in deep mode, e.g. Qwen/Qwen2.5-3B-Instruct "
+                             "(needs: pip install proprioceptive-lens[deep])")
+    parser.add_argument("--cloud", default=os.environ.get("LENS_CLOUD_URL"), help=argparse.SUPPRESS)
+    parser.add_argument("--no-browser", action="store_true")
+    parser.add_argument("--version", action="version", version=f"lens {__version__}")
+    args = parser.parse_args(argv)
+    ollama = args.ollama if args.ollama.startswith("http") else "http://" + args.ollama
+    if args.deep:
+        print(f"Loading deep model {args.deep} (first run downloads it)…")
+    app = LensApp(Path(args.data), ollama, args.hindsight, deep_model=args.deep, cloud_url=args.cloud)
+    if app.deep_error:
+        print(f"  Deep mode unavailable: {app.deep_error}")
+    try:
+        serve(app, args.port)
+    except OSError as exc:
+        print(f"Could not start on port {args.port}: {exc}", file=sys.stderr)
+        return 1
+    url = f"http://127.0.0.1:{args.port}/"
+    status = app.status()
+    print(f"Proprioceptive Lens {__version__} running at {url}")
+    print(f"  Ollama: {'connected' if status['backend_ok'] else 'NOT FOUND at ' + ollama} · "
+          f"{len(status['models'])} models · memory: {status['memory_count']} saved")
+    print(f"  Use it from other apps: OpenAI-compatible endpoint {url}v1 (adds memory to any chat app)")
+    if not args.no_browser:
+        webbrowser.open(url)
+    try:
+        while True:
+            time.sleep(3600)
+    except KeyboardInterrupt:
+        print("\nStopped.")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

proprioceptive_lens-0.1.0/lens/catalog.py

@@ -0,0 +1,69 @@
+"""Map the Ollama model names people use to their Hugging Face equivalents (for deep mode).
+
+Ollama does not expose hidden states, so deep view and Behaviour run the same model's Hugging Face weights.
+This table lets Lens tell an Ollama user the exact command for the model they already use.
+"""
+from __future__ import annotations
+
+OLLAMA_TO_HF = {
+    # Qwen
+    "qwen2.5:0.5b": "Qwen/Qwen2.5-0.5B-Instruct", "qwen2.5:1.5b": "Qwen/Qwen2.5-1.5B-Instruct",
+    "qwen2.5:3b": "Qwen/Qwen2.5-3B-Instruct", "qwen2.5:7b": "Qwen/Qwen2.5-7B-Instruct",
+    "qwen2.5-coder:0.5b": "Qwen/Qwen2.5-Coder-0.5B-Instruct", "qwen2.5-coder:1.5b": "Qwen/Qwen2.5-Coder-1.5B-Instruct",
+    "qwen2.5-coder:3b": "Qwen/Qwen2.5-Coder-3B-Instruct", "qwen2.5-coder:7b": "Qwen/Qwen2.5-Coder-7B-Instruct",
+    "qwen2:0.5b": "Qwen/Qwen2-0.5B-Instruct", "qwen2:7b": "Qwen/Qwen2-7B-Instruct",
+    "qwen3:0.6b": "Qwen/Qwen3-0.6B", "qwen3:1.7b": "Qwen/Qwen3-1.7B", "qwen3:4b": "Qwen/Qwen3-4B", "qwen3:8b": "Qwen/Qwen3-8B",
+    # Llama (ungated copies of Meta's weights, so no Hugging Face account is needed)
+    "llama3.2:1b": "unsloth/Llama-3.2-1B-Instruct", "llama3.2:3b": "unsloth/Llama-3.2-3B-Instruct",
+    "llama3.1:8b": "unsloth/Llama-3.1-8B-Instruct", "llama3:8b": "NousResearch/Meta-Llama-3-8B-Instruct",
+    "hermes3:8b": "NousResearch/Hermes-3-Llama-3.1-8B", "tinyllama:1.1b": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+    # Mistral, Gemma
+    "mistral:7b": "mistralai/Mistral-7B-Instruct-v0.3", "mistral-nemo:12b": "m8than/Mistral-Nemo-Instruct-2407-lenient-chatfix",
+    "gemma2:2b": "Efficient-Large-Model/gemma-2-2b-it", "gemma3:1b": "unsloth/gemma-3-1b-it",
+    # Microsoft, DeepSeek
+    "phi3:3.8b": "microsoft/Phi-3-mini-4k-instruct", "phi3.5:3.8b": "microsoft/Phi-3.5-mini-instruct",
+    "phi4-mini:3.8b": "microsoft/Phi-4-mini-instruct", "phi4-mini-reasoning:3.8b": "microsoft/Phi-4-mini-reasoning",
+    "deepseek-r1:1.5b": "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B", "deepseek-r1:7b": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
+    "deepseek-r1:8b": "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B", "deepseek-coder:1.3b": "deepseek-ai/deepseek-coder-1.3b-instruct",
+    "deepseek-coder:6.7b": "deepseek-ai/deepseek-coder-6.7b-instruct",
+    # Others
+    "smollm2:1.7b": "HuggingFaceTB/SmolLM2-1.7B-Instruct", "smollm2:360m": "HuggingFaceTB/SmolLM2-360M-Instruct",
+    "smollm3:3b": "HuggingFaceTB/SmolLM3-3B", "granite3.3:2b": "ibm-granite/granite-3.3-2b-instruct",
+    "granite3.3:8b": "ibm-granite/granite-3.3-8b-instruct", "granite3.2:8b": "ibm-granite/granite-3.2-8b-instruct",
+    "granite3.1-dense:2b": "ibm-granite/granite-3.1-2b-instruct", "granite3.1-dense:8b": "ibm-granite/granite-3.1-8b-instruct",
+    "falcon3:3b": "tiiuae/Falcon3-3B-Instruct", "falcon3:7b": "tiiuae/Falcon3-7B-Instruct",
+    "olmo2:1b": "allenai/OLMo-2-0425-1B-Instruct", "olmo2:7b": "allenai/OLMo-2-1124-7B-Instruct",
+    "exaone3.5:2.4b": "LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct", "exaone3.5:7.8b": "LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct",
+    "internlm2:7b": "internlm/internlm2-chat-7b", "stablelm-zephyr:3b": "stabilityai/stablelm-zephyr-3b",
+    "qwen2:1.5b": "Qwen/Qwen2-1.5B-Instruct", "qwen:0.5b": "Qwen/Qwen1.5-0.5B-Chat", "qwen:1.8b": "Qwen/Qwen1.5-1.8B-Chat",
+    "zephyr:7b": "HuggingFaceH4/zephyr-7b-beta", "openchat:7b": "openchat/openchat-3.5-1210",
+    "solar:10.7b": "upstage/SOLAR-10.7B-Instruct-v1.0", "llama2:7b": "unsloth/llama-2-7b-chat",
+    "llama2:13b": "Featherless-Chat-Models/llama2-13b-chat-hf", "falcon:7b": "tiiuae/falcon-7b-instruct",
+    "foundation-sec:8b": "fdtn-ai/Foundation-Sec-8B-Instruct", "danube3:500m": "h2oai/h2o-danube3-500m-chat",
+}
+# More specific Ollama tags that point at a different checkpoint than the base name.
+TAGGED = {
+    "qwen3:4b-instruct-2507": "Qwen/Qwen3-4B-Instruct-2507", "qwen3:4b-thinking-2507": "Qwen/Qwen3-4B-Thinking-2507",
+    "deepseek-r1:8b-llama-distill": "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
+    "phi3:3.8b-mini-128k-instruct": "microsoft/Phi-3-mini-128k-instruct",
+    "mistral:7b-instruct-v0.3": "mistralai/Mistral-7B-Instruct-v0.3", "mistral:7b-instruct-v0.1": "mistralai/Mistral-7B-Instruct-v0.1",
+    "deepseek-coder:6.7b-instruct": "deepseek-ai/deepseek-coder-6.7b-instruct",
+}
+# What Ollama pulls when no tag (or ":latest") is given.
+DEFAULT_SIZE = {
+    "qwen3": "8b", "qwen2.5": "7b", "qwen2.5-coder": "7b", "qwen2": "7b", "llama3.2": "3b", "llama3.1": "8b", "llama3": "8b",
+    "llama2": "7b", "mistral": "7b", "phi3": "3.8b", "phi3.5": "3.8b", "phi4-mini": "3.8b", "deepseek-r1": "8b",
+    "smollm2": "1.7b", "granite3.3": "8b", "falcon3": "7b", "olmo2": "7b", "tinyllama": "1.1b", "zephyr": "7b",
+    "openchat": "7b", "solar": "10.7b", "hermes3": "8b", "exaone3.5": "7.8b",
+}
+
+
+def hf_for(ollama_name: str) -> str | None:
+    """Best Hugging Face match for an Ollama model name ('qwen2.5:3b-instruct-q4_K_M' -> Qwen/Qwen2.5-3B-Instruct)."""
+    name = ollama_name.lower().split("/")[-1]
+    base, _, tag = name.partition(":")
+    for full, hf in TAGGED.items():
+        if name == full or name.startswith(full + "-"):
+            return hf
+    size = tag.split("-")[0] if tag and tag != "latest" else DEFAULT_SIZE.get(base, "")
+    return OLLAMA_TO_HF.get(f"{base}:{size}")

proprioceptive_lens-0.1.0/lens/cloud.py

@@ -0,0 +1,98 @@
+"""Proprioceptive Cloud client (opt-in): behaviour scores from internal state.
+
+What is sent: for each generated token, 16 numbers per probe layer (the hidden
+state multiplied by this device's scrambled matrix) and the hidden state's
+length. No text, no prompts and no raw hidden states leave the machine.
+What comes back: behaviour scores. The probes themselves stay in the cloud.
+"""
+from __future__ import annotations
+
+import base64
+import json
+import struct
+import urllib.error
+import urllib.parse
+import urllib.request
+from pathlib import Path
+
+from . import __version__
+
+DEFAULT_URL = "https://api.proprioceptiveai.com"
+
+
+class CloudError(RuntimeError):
+    pass
+
+
+class Cloud:
+    def __init__(self, data_dir: Path, url: str | None = None):
+        self.url = (url or DEFAULT_URL).rstrip("/")
+        self.creds_path = Path(data_dir) / "cloud.json"
+        self.creds = json.loads(self.creds_path.read_text()) if self.creds_path.exists() else {}
+        self.config: dict | None = None
+
+    @property
+    def enabled(self) -> bool:
+        return bool(self.creds.get("token")) and self.creds.get("url") == self.url
+
+    def _call(self, method: str, path: str, body: dict | None = None, timeout: float = 15.0) -> dict:
+        headers = {"Content-Type": "application/json", "User-Agent": f"proprioceptive-lens/{__version__}"}
+        if self.creds.get("token"):
+            headers["Authorization"] = "Bearer " + self.creds["token"]
+        req = urllib.request.Request(self.url + path, method=method, headers=headers,
+                                     data=json.dumps(body).encode() if body is not None else None)
+        try:
+            with urllib.request.urlopen(req, timeout=timeout) as response:
+                return json.load(response)
+        except urllib.error.HTTPError as exc:
+            try:
+                message = json.load(exc).get("error", str(exc))
+            except (ValueError, AttributeError):
+                message = str(exc)
+            raise CloudError(message) from None
+        except (urllib.error.URLError, OSError) as exc:
+            raise CloudError(f"Proprioceptive Cloud unreachable ({exc})") from None
+
+    def enable(self) -> dict:
+        data = self._call("POST", "/v1/lens/devices", {"client": f"lens/{__version__}"})
+        self.creds = {"url": self.url, "device_id": data["device_id"], "token": data["token"]}
+        self.creds_path.parent.mkdir(parents=True, exist_ok=True)
+        self.creds_path.write_text(json.dumps(self.creds))
+        self.creds_path.chmod(0o600)
+        return data
+
+    def disable(self) -> None:
+        self.creds = {}
+        self.config = None
+        self.creds_path.unlink(missing_ok=True)
+
+    def load_config(self, model_id: str) -> dict:
+        """Fetch this device's scrambled matrices for a model. Returns {layer: rows (list of lists)}."""
+        data = self._call("GET", "/v1/lens/packs/" + urllib.parse.quote(model_id, safe="") + "/config")
+        matrices = {}
+        for layer in data["layers"]:
+            raw = base64.b64decode(layer["matrix_b64"])
+            values = struct.unpack(f"<{layer['rows'] * layer['cols']}f", raw)
+            cols = layer["cols"]
+            matrices[int(layer["layer"])] = [list(values[r * cols:(r + 1) * cols]) for r in range(layer["rows"])]
+        self.config = data
+        return matrices
+
+    def supported_models(self) -> list[str]:
+        """Models with a Behaviour pack (public list, cached for 10 minutes)."""
+        import time
+        cached = getattr(self, "_supported", None)
+        if cached and time.time() - cached[0] < 600:
+            return cached[1]
+        try:
+            models = self._call("GET", "/v1/lens/health", timeout=5).get("packs", [])
+        except CloudError:
+            models = cached[1] if cached else []
+        self._supported = (time.time(), models)
+        return models
+
+    def score(self, steps: list[dict]) -> dict:
+        if not self.config:
+            raise CloudError("no behaviour pack loaded")
+        return self._call("POST", "/v1/lens/score", {"pack_id": self.config["pack_id"],
+                                                    "key_version": self.config["key_version"], "steps": steps})

proprioceptive_lens-0.1.0/lens/hf.py

@@ -0,0 +1,141 @@
+"""Deep mode: run a Hugging Face model locally and read its internal state.
+
+Needs `pip install proprioceptive-lens[deep]` (torch + transformers).
+Per generated token this yields the same events as the Ollama backend plus
+  layers: activity (RMS of the hidden state) for every layer, for the deep view
+  sigma/norm: for opted-in cloud scoring, M_L @ h and |h| at the layers the
+              cloud requested, where M_L is the device's scrambled matrix.
+This module contains no proprietary components: it multiplies hidden states by
+whatever matrix the cloud sent and nothing else.
+"""
+from __future__ import annotations
+
+import threading
+from typing import Iterator
+
+
+# Where decoder blocks live, by architecture family. Covers Llama, Qwen 1-3, Mistral, Mixtral, Gemma 1-3, Phi 1-4,
+# Granite, OLMo, DeepSeek, InternLM, StableLM, Falcon, GPT-2/J/NeoX, BLOOM, MPT, OPT, Mamba/Falcon-Mamba, RWKV and more.
+LAYER_PATHS = ("model.layers", "model.language_model.layers", "language_model.model.layers", "transformer.h",
+               "gpt_neox.layers", "model.decoder.layers", "transformer.blocks", "backbone.layers", "rwkv.blocks",
+               "transformer.layers", "model.transformer.h")
+
+
+def find_layers(model):
+    """Return the model's list of decoder blocks, or raise if the architecture is unknown."""
+    for path in LAYER_PATHS:
+        obj = model
+        for part in path.split("."):
+            obj = getattr(obj, part, None)
+            if obj is None:
+                break
+        if obj is not None and hasattr(obj, "__len__") and len(obj) > 0:
+            return obj
+    raise ValueError(f"deep mode does not recognise the layer layout of {type(model).__name__}")
+
+
+class DeepModel:
+    def __init__(self, model_id: str, device: str | None = None):
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        self.torch = torch
+        # "Org/Name@/local/path" loads local weights while reporting the public model ID.
+        model_id, _, local_path = model_id.partition("@")
+        source = local_path or model_id
+        self.model_id = model_id
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        dtype = torch.bfloat16 if self.device == "cuda" else torch.float32
+        self.tokenizer = AutoTokenizer.from_pretrained(source)
+        try:
+            model = AutoModelForCausalLM.from_pretrained(source, dtype=dtype)
+        except TypeError:   # transformers < 4.56 uses torch_dtype
+            model = AutoModelForCausalLM.from_pretrained(source, torch_dtype=dtype)
+        self.model = model.to(self.device).eval()
+        self.layers = find_layers(self.model)
+        self.n_layers = len(self.layers)
+        self.lock = threading.Lock()
+        self.projections: dict[int, "torch.Tensor"] = {}   # layer -> (16, d) scrambled matrix from the cloud
+
+    def set_projections(self, matrices: dict[int, list]) -> None:
+        torch = self.torch
+        self.projections = {int(layer): torch.tensor(m, dtype=torch.float32, device=self.device)
+                            for layer, m in matrices.items()}
+
+    def chat_stream(self, messages: list[dict], max_new_tokens: int = 512, temperature: float = 0.6,
+                    top_alternatives: int = 3) -> Iterator[dict]:
+        torch = self.torch
+        captured: dict[int, "torch.Tensor"] = {}
+
+        def make_hook(index):
+            def hook(_module, _inputs, output):
+                hidden = output[0] if isinstance(output, tuple) else output
+                captured[index] = hidden[0, -1, :].detach().float()   # last token, like the validated board
+            return hook
+
+        # Running mean of each projected layer's hidden state over the answer so far. The probes were
+        # validated on mean-pooled states, so Behaviour scores the mean (generic arithmetic, no IP).
+        self._sums = {}
+        self._count = 0
+        with self.lock:
+            handles = [layer.register_forward_hook(make_hook(i)) for i, layer in enumerate(self.layers)]
+            try:
+                prompt = self.tokenizer.apply_chat_template(messages, add_generation_prompt=True,
+                                                            return_tensors="pt", return_dict=True)
+                input_ids = prompt["input_ids"].to(self.device)   # dict form works on transformers 4.4x and 5.x
+                past = None
+                generated = 0
+                with torch.no_grad():
+                    while generated < max_new_tokens:
+                        out = self.model(input_ids=input_ids, past_key_values=past, use_cache=True)
+                        past = out.past_key_values
+                        logits = out.logits[0, -1].float()
+                        probs = torch.softmax(logits, dim=-1)
+                        if temperature > 0:
+                            choice = int(torch.multinomial(torch.softmax(logits / temperature, -1), 1))
+                        else:
+                            choice = int(torch.argmax(probs))
+                        if generated == 0:
+                            captured.clear()      # the prompt pass is not a generated token
+                        else:
+                            yield self._step_event(previous_token, previous_probs, captured, top_alternatives)
+                        if choice in (self.tokenizer.eos_token_id, *self._stop_ids()):
+                            break
+                        previous_token, previous_probs = choice, probs
+                        input_ids = torch.tensor([[choice]], device=self.device)
+                        generated += 1
+                    else:
+                        out = self.model(input_ids=input_ids, past_key_values=past, use_cache=True)
+                        yield self._step_event(previous_token, previous_probs, captured, top_alternatives)
+            finally:
+                for handle in handles:
+                    handle.remove()
+        yield {"type": "done", "tokens": generated, "seconds": None}
+
+    def _stop_ids(self) -> tuple:
+        ids = []
+        for token in ("<|im_end|>", "<|endoftext|>", "<|eot_id|>"):
+            tid = self.tokenizer.convert_tokens_to_ids(token)
+            if isinstance(tid, int) and tid >= 0 and tid != self.tokenizer.unk_token_id:
+                ids.append(tid)
+        return tuple(ids)
+
+    def _step_event(self, token_id: int, probs, captured: dict, top_k: int) -> dict:
+        """Event for one generated token; hidden states are those produced while reading that token."""
+        torch = self.torch
+        top = torch.topk(probs, top_k + 1)
+        alternatives = [{"text": self.tokenizer.decode([int(i)]), "p": float(p)}
+                        for p, i in zip(top.values, top.indices) if int(i) != token_id][:top_k]
+        alternatives = [a for a in alternatives if a["text"].strip()]
+        event = {"type": "token", "text": self.tokenizer.decode([token_id]), "p": float(probs[token_id]),
+                 "alternatives": alternatives,
+                 "layers": [round(float(captured[i].pow(2).mean().sqrt()), 4) for i in range(self.n_layers)
+                            if i in captured]}
+        if self.projections:
+            self._count += 1
+            for L in self.projections:
+                if L in captured:
+                    self._sums[L] = self._sums.get(L, 0) + captured[L]
+            means = {L: self._sums[L] / self._count for L in self._sums}
+            event["sigma"] = {str(L): (M @ means[L]).tolist() for L, M in self.projections.items() if L in means}
+            event["norm"] = {str(L): float(means[L].norm()) for L in self.projections if L in means}
+        return event