weightslab 1.2.4__tar.gz → 1.2.6__tar.gz

{weightslab-1.2.4 → weightslab-1.2.6}/.github/workflows/ci.yml

@@ -2,19 +2,52 @@ name: Code CI
 
 on:
   push:
-    branches: [ main, dev ]
+    # Fire on every branch so the `gate` job below can decide whether to run.
+    # GitHub cannot filter on commit message at the event level, so we trigger
+    # broadly and gate execution: main/dev always run, other branches run only
+    # when the head commit message contains "[force ci]".
+    branches: [ '**' ]
     tags: [ 'v*' ]
   pull_request:
     types: [ opened, reopened, ready_for_review ]
   workflow_dispatch:
 
 jobs:
-  # Required gate: a lint violation here FAILS the CI run. This job has no
-  # `needs` and nothing `needs` it, so it runs in parallel with install/test/
-  # build — its failure turns the run red but does NOT stop the other jobs from
-  # running (GitHub Actions does not cancel sibling jobs on failure).
+  # Single source of truth for "should the core CI jobs run?". Reproduces the
+  # previous trigger rules (PR / manual dispatch / main / dev run; tags skip)
+  # and adds a manual escape hatch for custom branches via "[force ci]".
+  gate:
+    runs-on: ubuntu-latest
+    outputs:
+      run_ci: ${{ steps.decide.outputs.run_ci }}
+    steps:
+      - name: Decide whether to run CI
+        id: decide
+        env:
+          # Evaluated by GitHub (not the shell) so the raw commit message is
+          # never interpolated into bash — avoids quoting/injection issues.
+          FORCE_CI: ${{ contains(github.event.head_commit.message, '[force ci]') }}
+        run: |
+          if [ "${{ github.event_name }}" != "push" ]; then
+            run_ci=true            # PRs and manual dispatch always run
+          elif [[ "${{ github.ref }}" == refs/tags/* ]]; then
+            run_ci=false           # tag pushes are handled by the release workflow
+          elif [ "${{ github.ref }}" = "refs/heads/main" ] || [ "${{ github.ref }}" = "refs/heads/dev" ]; then
+            run_ci=true            # default branches always run
+          elif [ "${FORCE_CI}" = "true" ]; then
+            run_ci=true            # custom branch opted in via [force ci]
+          else
+            run_ci=false           # custom branch without [force ci] → skip
+          fi
+          echo "run_ci=${run_ci}" >> "$GITHUB_OUTPUT"
+          echo "Decision: run_ci=${run_ci} (event=${{ github.event_name }}, ref=${{ github.ref }}, force_ci=${FORCE_CI})"
+  # Required gate: a lint violation here FAILS the CI run. Nothing `needs` this
+  # job, so it runs in parallel with install/test/build — its failure turns the
+  # run red but does NOT stop the other jobs from running (GitHub Actions does
+  # not cancel sibling jobs on failure).
   code-quality:
-    if: ${{ github.event_name != 'push' || !startsWith(github.ref, 'refs/tags/') }}
+    needs: gate
+    if: ${{ needs.gate.outputs.run_ci == 'true' }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
@@ -99,7 +132,8 @@ jobs:
   # Single-version install on 3.11 — always runs so `test` has a gate to
   # depend on regardless of branch.
   install:
-    if: ${{ github.event_name != 'push' || !startsWith(github.ref, 'refs/tags/') }}
+    needs: gate
+    if: ${{ needs.gate.outputs.run_ci == 'true' }}
     runs-on: ubuntu-latest
     name: install (Python 3.11)
     steps:
@@ -160,10 +194,11 @@ jobs:
           python -c "import weightslab; print(f'weightslab imported successfully on Python ${{ matrix.python-version }}')"
 
   test:
-    if: ${{ github.event_name != 'push' || !startsWith(github.ref, 'refs/tags/') }}
+    if: ${{ needs.gate.outputs.run_ci == 'true' }}
     runs-on: ubuntu-latest
-    # Depends on the fast 3.11 gate only — the matrix runs independently in parallel.
-    needs: install
+    # Depends on the fast 3.11 install gate; `gate` is also a direct need so this
+    # job can read run_ci (the matrix and main-only jobs run independently).
+    needs: [ gate, install ]
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4

weightslab-1.2.6/AGENTS.md

@@ -0,0 +1,268 @@
+# WeightsLab — agent context for users & debugging
+
+This file is a **portable context for AI coding agents** (Claude Code, etc.) and
+the humans driving them. Its job is to let you — or an agent helping you —
+**install, configure, run, and debug WeightsLab and Weights Studio** without
+having to reverse-engineer the system first.
+
+It deliberately covers only the two shipped repositories:
+
+- **weightslab** — the Python backend / core (training instrumentation, data
+  ledger, gRPC service, the shared proto).
+- **weights_studio** — the browser frontend (the studio UI that inspects and
+  edits a *running* experiment).
+
+> File/line references drift as the code evolves — treat them as starting points
+> and verify against the current source before relying on them. Environment
+> variable names and defaults are the most stable thing here; when in doubt the
+> authoritative reference is `weightslab/docs/configuration.rst`.
+
+---
+
+## 0. How to load this guide into Claude Code
+
+So an agent actually *has* this context when you ask it for help:
+
+- **Working inside a checkout of the repo** (`git clone`): this guide is
+  committed as `AGENTS.md`; the repo keeps a gitignored `CLAUDE.md` copy of it at
+  the root so Claude Code auto-loads it every session. Nothing to do. (Claude
+  Code also loads `~/.claude/CLAUDE.md` global memory and any parent-dir
+  `CLAUDE.md`.)
+- **You only ran `pip install weightslab`** (no checkout — the package lives in
+  `site-packages`): absolute `@import` paths are fragile because the path
+  changes per venv/OS. The robust pattern is a small **skill** that locates the
+  installed file at runtime. Create `~/.claude/skills/weightslab/SKILL.md`:
+
+  ```yaml
+  ---
+  name: weightslab
+  description: Load the WeightsLab debugging & configuration guide when helping with weightslab or weights_studio problems (connection, TLS, env vars, training hangs, rendering).
+  ---
+  !`python -c "import weightslab, os; print(open(os.path.join(os.path.dirname(weightslab.__file__), 'AGENTS.md')).read())"`
+
+  Use the guide above to diagnose the user's weightslab / weights_studio issue.
+  ```
+
+  Then run `/weightslab` (or let Claude auto-invoke it). This requires the guide
+  to be **shipped as package data** inside the installed package (see §7); the
+  copy at the repo root is for contributors working in a checkout.
+- **Quick-and-dirty:** copy this file to `~/.claude/WEIGHTSLAB.md` and add
+  `@~/.claude/WEIGHTSLAB.md` to your `~/.claude/CLAUDE.md`.
+
+---
+
+## 1. What it is and how the pieces connect
+
+A user wraps their own PyTorch training script with WeightsLab so a running
+experiment becomes inspectable/editable; Weights Studio is the UI for that.
+
+**Wire path (the thing that breaks most often):**
+
+```
+Browser (Vite app)  →  Envoy :8080 (grpc-web ↔ grpc)  →  Python gRPC servicer  →  training loop
+```
+
+- The browser cannot speak raw gRPC, so **Envoy** transcodes grpc-web ↔ gRPC.
+  If Envoy is down or misconfigured, the UI loads but no data appears.
+- The gRPC servicer and the training loop run in the **same process, different
+  threads**, coordinated by locks in
+  `weightslab/weightslab/components/global_monitoring.py`.
+- One proto is the single source of truth:
+  `weightslab/weightslab/proto/experiment_service.proto`.
+
+---
+
+## 2. Install & run (the happy path)
+
+```bash
+pip install weightslab
+```
+
+In your training script:
+
+```python
+import weightslab as wl
+# wrap your objects so the studio can see/edit them (see §3), then:
+wl.serve(serving_grpc=True, serving_cli=True)   # background threads, same process
+# ... your training loop ...
+wl.keep_serving()                                # keep the process alive for the UI
+```
+
+Then start the studio stack (Envoy + frontend) and open it in a browser.
+Working starting points live in
+`weightslab/weightslab/examples/{PyTorch,Lightning,Usecases}/<usecase>/`
+(each is a `main.py` + `config.yaml`) — find the closest example and mirror it.
+
+Studio deployment details (Docker compose, Envoy, ports, certs) are in
+`weights_studio/docker/` and documented in `weightslab/docs/weights_studio.rst`.
+
+---
+
+## 3. The integration API (`import weightslab as wl`)
+
+How a user's script plugs in. Wrap each training object with
+`wl.watch_or_edit(obj, flag=...)`; the returned tracked proxy is registered in
+the global ledger (`weightslab/weightslab/backend/ledgers.py`,
+`GLOBAL_LEDGER` — the hub everything reads/mutates through).
+
+- `flag="hyperparameters"` (dict), `flag="model"` (nn.Module, `device=…`),
+  `flag="optimizer"`, `flag="data"` (Dataset → tracked DataLoader: `loader_name`,
+  `batch_size`, `is_training`, `collate_fn`, …), `flag="loss"` (a
+  `reduction="none"` criterion, called with `(preds_raw, targets, batch_ids=ids,
+  preds=preds)`), `flag="metric"`.
+
+Conventions that matter for correctness:
+
+- Wrap the train step in `with guard_training_context:` and eval in
+  `with guard_testing_context:` (from
+  `weightslab.components.global_monitoring`). This is how pause/resume and
+  train/test separation work — **skip it and pause/resume or stats will misbehave.**
+- Use `model.get_age()` (steps actually trained; survives checkpoint reloads),
+  not the raw loop counter.
+- `task_type` on the dataset/model selects rendering: `classification`,
+  `segmentation`, `detection`, `detection_pointcloud`.
+- **Hyperparameter handle access:** the registered hyperparameters proxy
+  supports both `hp.get("lr")` and `hp["lr"]` (subscript == `.get`), and stays
+  live — reads reflect in-place updates and re-registration.
+
+---
+
+## 4. Configuration (environment variables)
+
+WeightsLab and Weights Studio are configured almost entirely through env vars.
+**Authoritative reference: `weightslab/docs/configuration.rst`.** The high-signal
+ones when debugging:
+
+**Backend (Python):**
+
+| Variable | Default | Why you touch it |
+|---|---|---|
+| `WEIGHTSLAB_LOG_LEVEL` | `INFO` | Set `DEBUG` to see what's happening. (`WATCHDOG` level sits between WARNING/ERROR.) |
+| `GRPC_BACKEND_HOST` / `GRPC_BACKEND_PORT` | `0.0.0.0` / `50051` | Backend must listen where Envoy expects it. |
+| `GRPC_TLS_ENABLED` | `1` | TLS on the gRPC socket. Set `0` **only** for isolated local debugging. |
+| `GRPC_TLS_REQUIRE_CLIENT_AUTH` | `1` | mTLS. Must match what Envoy presents. |
+| `GRPC_TLS_CERT_DIR` | `~/certs` | Where default cert files are looked up. |
+| `GRPC_AUTH_TOKEN` | *(unset)* | Optional metadata-token auth on top of mTLS. |
+| `GRPC_MAX_MESSAGE_BYTES` | `268435456` (256 MB) | Raise it if large tensors/image batches fail. |
+| `WEIGHTSLAB_DISABLE_WATCHDOGS` | `0` | Set `1` when debugging with breakpoints (see §5). |
+| `GRPC_WATCHDOG_STUCK_SECONDS` | `60` | Lock/RPC stuck threshold + lock-acquire timeout. |
+
+**Frontend (Weights Studio) — runtime-injected `window.*` globals:**
+
+| Variable | Default | Why you touch it |
+|---|---|---|
+| `WS_SERVER_HOST` / `WS_SERVER_PORT` / `WS_SERVER_PROTOCOL` | `localhost` / `8080` / `https` | How the browser reaches the backend (via Envoy). The #1 connection-issue knob. |
+| `WS_HISTOGRAM_MAX_BINS` | `512` | Cap on metadata histogram bars. |
+| `BB_THUMB_RENDER` | `10` | Max bounding boxes drawn per **thumbnail**, per overlay (GT and PRED capped independently). |
+| `BB_MODAL_RENDER` | `100` | Max bounding boxes drawn per **modal** image, per overlay. A `?` button in the modal shows the active limit. |
+
+> **VITE_ vs WS_/BB_:** `VITE_*` variables are baked at **build time** (changing
+> them needs a rebuild). `WS_*` / `BB_*` are injected at **container start** into
+> `config.js` and read as `window.*` globals — changing them needs only a
+> container restart + browser reload (see the caching note in §5).
+
+---
+
+## 5. Troubleshooting — symptom → cause → fix
+
+This is the core of the guide. Each entry is a real failure mode (several are
+distilled from issues hit in development).
+
+**UI loads but the sample grid is empty / "failed to fetch" / gRPC errors.**
+The wire path (§1) is broken somewhere. Check in order: (1) backend actually
+serving on `0.0.0.0:50051`; (2) Envoy running and reachable on `:8080`;
+(3) frontend `WS_SERVER_HOST/PORT/PROTOCOL` point at Envoy, not the raw backend;
+(4) **TLS mismatch** — `WS_SERVER_PROTOCOL=https` vs `http`, Envoy server certs,
+and Envoy→backend mTLS certs all consistent. For local debugging you can drop
+TLS end-to-end (`GRPC_TLS_ENABLED=0` + `VITE_SERVER_PROTOCOL=http`).
+
+**Changed an env var, restarted, but the UI still uses the old value.**
+- `VITE_*` is build-time → you must **rebuild** the frontend, not just restart.
+- `WS_*` / `BB_*` are read once per page load → you must **reload the tab**.
+- Historically `config.js` was served `Cache-Control: immutable` so even a
+  restart needed a **hard refresh**; current builds serve `/config.js` with
+  `no-store`, so a container restart + normal reload is enough. On an older
+  deployment, hard-refresh (Ctrl+Shift+R) or clear cache.
+
+**Sample grid flashes empty cells when auto-refresh fires.**
+An auto-refresh (timer or manual) that lands while a `GetDataSamples` grid fetch
+is still in flight used to clear the cache mid-render. The fix in
+`weights_studio/src/grid_data/gridDataManager.ts` is `isFetchInProgress()`:
+refreshes are skipped while a grid fetch is ongoing. If you see this, confirm
+you're on a build that has that guard.
+
+**Detection overlays are slow or unreadably cluttered.**
+Dense detection samples can carry hundreds of boxes. Cap rendering with
+`BB_THUMB_RENDER` (thumbnails) and `BB_MODAL_RENDER` (modal); each is applied
+separately to GT and to predictions. Render-only — no sample data is dropped.
+
+**Training appears hung; RPCs return `RESOURCE_EXHAUSTED`; server "restarts".**
+A watchdog monitors the global rlock and in-flight RPCs. If a lock/RPC is held
+longer than `GRPC_WATCHDOG_STUCK_SECONDS` (60s) it's flagged; locks get
+interrupted, and after `GRPC_WATCHDOG_RESTART_THRESHOLD` unhealthy polls the
+gRPC server restarts. When **debugging with breakpoints** that intentionally
+pause longer than that, set `WEIGHTSLAB_DISABLE_WATCHDOGS=1`. If RPCs fail with
+`RESOURCE_EXHAUSTED`, a handler couldn't acquire the lock within the window —
+something else is holding it; check for a long/blocking train or eval step.
+
+**Pause/resume doesn't work, or train vs test stats are mixed up.**
+The train step isn't wrapped in `guard_training_context` (or eval in
+`guard_testing_context`). See §3 — these context managers are how the system
+gates and separates phases.
+
+**Large weights/images fail to transfer.** Raise `GRPC_MAX_MESSAGE_BYTES`.
+
+**The agent bar says it's unconfigured.** The LLM agent needs a provider: a
+local **Ollama** server (`provider: ollama`, available immediately) or **cloud
+OpenRouter** initialized from the UI via `/init` (then `/model` to switch,
+`/reset` to clear). See `weightslab/docs/weights_studio.rst`.
+
+---
+
+## 6. Where things live (for deeper digging)
+
+**Backend (`weightslab/weightslab/`):**
+- `src.py` — the public verbs (`watch_or_edit`, `serve`, `keep_serving`,
+  `tag_samples`, `query_*`, decorators) re-exported from `__init__.py`.
+- `trainer/services/` — `experiment_service.py` (gRPC servicer) delegating to
+  `{model,data,agent}_service.py`; `data_image_utils.py` (preview/mask encoding).
+- `components/` — `global_monitoring.py` (locks, `guard_*` contexts, pause),
+  `checkpoint_manager.py`, `evaluation_controller.py`.
+- `data/` — `dataframe_manager.py`, `data_samples_with_ops.py`, `sample_stats.py`,
+  H5 storage (`h5_dataframe_store.py`, `h5_array_store.py`, `array_proxy.py`).
+- `backend/` — `ledgers.py` (`GLOBAL_LEDGER`), `logger.py`, `audit_logger.py`, `cli.py`.
+- `security/` (`CertAuthManager`), `proto/`, `examples/`, `docs/`.
+
+**Frontend (`weights_studio/src/`):**
+- `main.ts` — bootstrap; builds the grpc-web transport from `WS_SERVER_*`.
+- `experiment_service.client.ts` / `experiment_service.ts` — generated client
+  (regenerate with `npm run generate-proto:data`; do not hand-edit).
+- `grid_data/` — grid + modal rendering (`GridCell.ts`, `DataImageService.ts`,
+  `gridDataManager.ts`, `BboxRenderer.ts`, `SegmentationRenderer.ts`,
+  `PointCloudViewer.ts`).
+- `docker/` — compose, `nginx-entrypoint.sh` (injects `config.js`), Envoy assets.
+
+**Docs:** `weightslab/docs/` (Sphinx) — `configuration.rst` (all env vars),
+`weights_studio.rst` (studio deploy + agent), `quickstart.rst`, `grpc/`.
+
+---
+
+## 7. For contributors (working in a checkout)
+
+- **The two repos must sit side by side** (`…/weightslab`, `…/weights_studio`);
+  codegen and Envoy configs reach across by relative path.
+- **Editing the proto is cross-repo** — do all of: edit
+  `experiment_service.proto`; regenerate Python stubs from the repo root; run
+  `npm run generate-proto:data` in weights_studio. Editing one side only leaves
+  the build broken.
+- **Tests:** backend `python -m pytest weightslab/tests/...`; frontend unit
+  `npm run test` (vitest); E2E/user-simulation Playwright lives in
+  **weights_studio** (`test:realtime:*`, `test:e2e:*`), not here.
+- **CI on a custom branch:** pushes to non-`main`/`dev` branches only run CI when
+  the commit message contains `[force ci]` (both repos).
+- **TLS/auth in the bundled UI** is decided by cert presence under
+  `WEIGHTSLAB_CERTS_DIR` (single source of truth) — don't hardcode secure/insecure.
+- **To make this guide available to pip users**, ship it as package data inside
+  the installed package (e.g. as `weightslab/weightslab/AGENTS.md`) so the §0
+  skill can locate it; keep the root `AGENTS.md` (mirrored as the gitignored
+  `CLAUDE.md`) as the contributor-facing source.

weightslab-1.2.6/CHANGELOG.md

@@ -0,0 +1 @@
+# Changelog - 2026-06-18 v1.2.6 (1)

{weightslab-1.2.4/weightslab.egg-info → weightslab-1.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weightslab
-Version: 1.2.4
+Version: 1.2.6
 Summary: Paving the way between black-box and white-box modeling.
 Author-email: Alexandru-Andrei Rotaru <alexandru@graybx.com>, Guillaue Pelluet <guillaue@graybx.com>
 License: BSD 2-clause
@@ -120,23 +120,16 @@ The granular statistics and interactive paradigm enable powerful workflows:
 - **Docker Desktop v4.77 or newer** — required to deploy the Weights Studio UI (`weightslab ui launch`).
 - **Docker Compose v2** (the `docker compose` CLI plugin, bundled with Docker Desktop) — *recommended*. The legacy v1 standalone binary (`docker-compose`, **≥ 1.27**) also works: `weightslab ui launch` auto-detects whichever is installed and uses it. Compose v1 below 1.27 is **not** supported.
 - **Python >=3.10, <3.15** — to install and run the `weightslab` framework.
-> [!OPTIONAL]
-> We are fully compatible with Conda.
-
-> [!TIP]
-> Check your Compose version with `docker compose version` (v2) or `docker-compose --version` (v1). If `docker compose version` prints a version string, you have v2.
 
 
 ### Installation
-
-> [!OPTIONAL]
-> Seting a clean environment:
-> Define a Python environment:
-```bash
-python -m venv weightslab_venv
-./weightslab_venv/Scripts/activate
-```
-> Or install directly on your machine.
+Install directly on your machine.
+> [!TIP]
+> Setting a clean Python environment:
+> ```bash
+> python -m venv weightslab_venv
+> ./weightslab_venv/Scripts/activate
+> ```
 
 Install our framework:
 ```bash
@@ -148,17 +141,16 @@ Deploy our interface:
 weightslab ui launch
 ```
 
-`weightslab ui launch` removes any stale weightslab/weights_studio Docker resources that could break the launch, then starts the UI stack. **By default it runs unsecured (HTTP, no gRPC auth) — no certificates are generated.**
-
-To run secured (HTTPS + gRPC auth), pass `--certs`:
-```bash
-weightslab ui launch --certs   # generates TLS certs + a gRPC auth token if missing, then launches secured
-```
-
-> [!IMPORTANT]
+The command `weightslab ui launch` removes any stale weightslab/weights_studio Docker resources that could break the launch, then starts the UI stack. **By default, it runs unsecured (HTTP, no gRPC auth) — no certificates are generated. However, communication are not safe.**
+> [!TIP]
+> To run secured communication, pass the arguments `--certs`:
+> ```bash
+> weightslab ui launch --certs   # generates TLS certs + a gRPC auth token if missing, then launches secured
+> ```
 > When using certs, set `WEIGHTSLAB_CERTS_DIR` so the training backend and any new terminal use the **same** certificates (it is the single source of truth). `weightslab se` and `weightslab ui launch --certs` print the exact export/`setx` command for your shell. You can also generate certs up front with `weightslab se`.
+
 <!--
-Want to see it working end to end without writing any code? Start a bundled demo (it installs the example's requirements first — no prompts — then trains and serves until you stop it with `Ctrl+C`). It defaults to classification; pick another with a flag:
+Want to see it working end-to-end without writing any code? Start a bundled demo (it installs the example's requirements first — no prompts — then trains and serves until you stop it with `Ctrl+C`). It defaults to classification; pick another with a flag:
 ```bash
 weightslab start example          # classification (default)
 weightslab start example --cls    # classification
@@ -186,27 +178,29 @@ weightslab start example --gen    # generation
 
 2. **Wrap your parameters** with WeightsLab tracking:
    ```python
-   model = wl.watch_or_edit(parameters, ...)  # ← Now WeightsLab monitors your parameters and allow you to update them from your UI
+   model = wl.watch_or_edit(parameters, flag='hp', ...)  # ← Now WeightsLab monitors your parameters and allow you to update them from your UI
    ```
 
 3. **Wrap your model** with WeightsLab tracking:
    ```python
-   model = wl.watch_or_edit(SimpleModel(...), ...)  # ← Now WeightsLab monitors your model state
+   model = wl.watch_or_edit(SimpleModel(...), flag='model', ...)  # ← Now WeightsLab monitors your model state
    ```
 
 4. **Wrap your optimizer** with WeightsLab tracking:
    ```python
-   optimizer = wl.watch_or_edit(optim.Adam(...), ...)  # ← Tracks optimizer state and update optimizer learning rate from your UI
+   optimizer = wl.watch_or_edit(optim.Adam(...), flag='opt', ...)  # ← Tracks optimizer state and update optimizer learning rate from your UI
    ```
 
 5. **Wrap your signal** with WeightsLab tracking:
    ```python
-    criterion = wl.watch_or_edit(nn.CrossEntropyLoss(reduction="none"), ...)  # ← Tracks this signal and others (metrics, ..etc) from your UI
+    train_criterion = wl.watch_or_edit(nn.CrossEntropyLoss(reduction="none"), flag='signal', name="train_loss/sample", per_sample=True, log=True)  # ← Tracks this signal and others (metrics, ..etc) from your UI
+    test_criterion = wl.watch_or_edit(nn.CrossEntropyLoss(reduction="none"), flag='signal', name="test_loss/sample", per_sample=True, log=False)  # ← Tracks this signal and others (metrics, ..etc) from your UI - Plot is disabled, only per sample signal
    ```
 
 6. **Wrap your dataset** with WeightsLab tracking:
    ```python
-    train_loader = wl.watch_or_edit(dataset, ...)  # ← Tracks this dataset and others (validation, test) from your UI
+    train_loader = wl.watch_or_edit(train_dataset, flag='data', loader_name="train_loader", ...)  # ← Tracks this dataset and others (validation, test) from your UI
+    val_loader = wl.watch_or_edit(val_dataset, flag='data', loader_name="val_loader", ...)  # ← Tracks this dataset and others (validation, test) from your UI
    ```
 
 7. **Run your training script** as usual:
@@ -451,9 +445,10 @@ if __name__ == "__main__":
 ## Examples
 
 ### Local examples
+
 After starting the UI, launch a local experiment with the command:
 ```bash
-weightslab start example          # classification (default)
+weightslab start example            # classification (default)
 # weightslab start example --cls    # classification
 # weightslab start example --seg    # segmentation
 # weightslab start example --det    # detection
@@ -461,66 +456,9 @@ weightslab start example          # classification (default)
 # weightslab start example --gen    # generation
 ```
 
-###  Cloud examples
-Find our demos online. The password is **graybx**.
-<a href="https://preview.graybx.com/">
-  <p style="text-indent:20px;">DEMOS</p>
-</a>
-
-
-## AI Agent
-
-WeightsLab can run its data agent in two modes:
-
-- Local provider with Ollama
-- Cloud provider with OpenRouter
-
-Use local Ollama when you want a fully local setup and do not need cloud-hosted models (see more details installation from the documentation).
-Use OpenRouter when you want larger hosted models and model selection directly from Weightslab UI.
-
-
-### Cloud OpenRouter
-You can either preconfigure OpenRouter in `agent_config.yaml` / `.env`, or initialize it interactively from Weightslab UI.
-
-#### SDK Configuration
-
-Example static configuration:
-
-```yaml
-agent:
-  provider: openrouter
-  openrouter_model: meta-llama/llama-3.3-70b-instruct
-  fallback_to_local: false
-  # openrouter_api_key: ${OPENROUTER_API_KEY}
-```
-
-Environment variable:
-
-```bash
-export OPENROUTER_API_KEY=your_key_here
-```
-
-#### Interactive setup from Weightslab UI
 
-OpenRouter models can be initialized and set directly from the UI:
-1. Click in the agent bar or double-click to expand the agent window.
-2. Type `/init`.
-3. Choose either:
-   - `A` Enter your OpenRouter API key manually
-   - `B` Use the OpenRouter OAuth flow
-4. Select a model from the fetched list, then confirm.
-
-The default OpenRouter model, as recommended by Graybx, is `meta-llama/llama-3.3-70b-instruct`.
-
-### Typical Usage Flow
-
-1. Start your WeightsLab backend (e.g., "main.py").
-2. Start Weightslab UI.
-3. If you use Ollama, query the agent directly.
-4. If you use OpenRouter and the agent is not configured yet, type `/init`.
-5. Ask natural-language data operations such as sorting, filtering, slicing, and inspection requests. You can also ask questions about the data.
-6. Use `/model` to try another cloud model without re-entering the key.
-7. Use `/reset` if you want to clear the current connection and start over.
+###  Cloud examples
+Find our [sandbox](https://preview.graybx.com) online. The password is **graybx**.
 
 
 ## Documentation (API + SDK)