diff --git a/notebooks/enclave/probe-eval-w-receipts/.gitignore b/notebooks/enclave/probe-eval-w-receipts/.gitignore new file mode 100644 index 00000000000..3421ab2dc12 --- /dev/null +++ b/notebooks/enclave/probe-eval-w-receipts/.gitignore @@ -0,0 +1,10 @@ +# Written by notebook 0. probe.json is the probe owner's private probe: it goes +# to the enclave and nowhere else, so it must never be committed. +probe.json +probe_mock.json +probe_readout.py + +# Downloaded by the notebooks. +train_prompts/ +prompts/ +dbe-adapter/ diff --git a/notebooks/enclave/probe-eval-w-receipts/0. probe-owner-train-on-base.ipynb b/notebooks/enclave/probe-eval-w-receipts/0. probe-owner-train-on-base.ipynb new file mode 100644 index 00000000000..9f5adce45e8 --- /dev/null +++ b/notebooks/enclave/probe-eval-w-receipts/0. probe-owner-train-on-base.ipynb @@ -0,0 +1,693 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Probe evaluation β€” 0. Train the probe on the open base model\n", + "\n", + "| Actor | Email | Role |\n", + "|-------|-------|------|\n", + "| **Probe owner** | `BENCHMARK_OWNER_EMAIL` | Owns the private linear probe and the prompts. Runs this notebook, then notebook 1. The enclave's release pins the benchmark-owner account as a data owner, so that is the account the probe owner uses |\n", + "| **Model owner** | `MODEL_OWNER_EMAIL` | Owns the private LoRA adapter. Takes no part in this notebook |\n", + "| **Enclave** | `ENCLAVE_EMAIL` | Tinfoil Container that runs the evaluation. Takes no part in this notebook |\n", + "\n", + "The emails are set in notebooks 1 and 2. This notebook logs in nowhere and talks to no one.\n", + "\n", + "**The setting.** A probe owner has a linear probe: a direction in a model's hidden states that\n", + "tracks a concept. They want to know what it reads on a fine-tune they cannot see. A model owner\n", + "built that fine-tune as a private LoRA adapter on an open base model, and is willing to be measured\n", + "but will not hand over the weights. The probe owner will not hand over the probe either.\n", + "\n", + "What we, the probe owner, can see is the base model. So we train the probe here, on our own machine\n", + "or in Colab, on the open base model, and write it to `probe.json`. Notebook 1 then applies it inside\n", + "the enclave to base + adapter, which we never see.\n", + "\n", + "The code that reads the hidden state and applies the probe is defined once below, as\n", + "`PROBE_READOUT_SRC`. It builds the training features here, runs a dry run on the base model in\n", + "notebook 1, and goes into the enclave job verbatim. That the same code runs in all three places is\n", + "the point: the only thing the enclave adds is the adapter." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**The concept is a placeholder.** Until a real probe owner brings their own probe and prompts, we\n", + "train one for a *specialised-advice request*: a prompt that asks for election, financial, health or\n", + "legal advice. We label prompts from the MLCommons AILuminate demo set by their hazard code. The four\n", + "`spc_*` hazards are the positives, and every other hazard is a negative.\n", + "\n", + "A hazard code says what a prompt is about, so a probe that separates them may be picking up topic\n", + "words rather than anything deeper. Read the numbers at the end as a check that the mechanics work,\n", + "not as a finding about the model." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Setup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# transformers and accelerate at the versions the enclave job pins, so the\n", + "# features here go through the same code as there.\n", + "!uv pip install -q torch \"transformers==4.57.6\" \"accelerate==1.12.0\" scikit-learn" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import csv\n", + "import hashlib\n", + "import json\n", + "import random\n", + "import time\n", + "import types\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import torch\n", + "from sklearn.linear_model import LogisticRegression\n", + "from sklearn.metrics import roc_auc_score\n", + "from sklearn.model_selection import (\n", + " StratifiedKFold,\n", + " cross_val_predict,\n", + " cross_val_score,\n", + " train_test_split,\n", + ")\n", + "from sklearn.pipeline import make_pipeline\n", + "from sklearn.preprocessing import StandardScaler\n", + "from transformers import AutoModelForCausalLM, AutoTokenizer" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# The open base model the model owner's adapter targets, at the commit the\n", + "# enclave job downloads (BASE_MODEL_REVISION in notebook 1).\n", + "BASE_MODEL = \"TinyLlama/TinyLlama-1.1B-Chat-v1.0\"\n", + "BASE_MODEL_REVISION = \"fe8a4ea1ffedaf415f4da2f062534de366a451e6\"\n", + "\n", + "# The placeholder concept, and the hazard codes that are its positives.\n", + "CONCEPT = \"specialised-advice request\"\n", + "POSITIVE_HAZARDS = {\"spc_ele\", \"spc_fin\", \"spc_hlt\", \"spc_lgl\"}\n", + "\n", + "# All the positives, and a seeded sample of this many negatives per positive.\n", + "NEGATIVES_PER_POSITIVE = 3\n", + "# None trains on that whole sample. A number, such as 40, takes a stratified\n", + "# subsample of it instead, for a quick run.\n", + "MAX_PROMPTS = None\n", + "HELD_OUT_FRACTION = 0.25\n", + "# The probe's threshold is the score that flags this share of the training\n", + "# negatives, out of fold (Step 6).\n", + "TARGET_FPR = 0.10\n", + "SEED = 0\n", + "\n", + "# Both sources pinned to a commit (main of each repo on 2026-10-07), and every\n", + "# file to its sha256, which Step 1 checks.\n", + "AILUMINATE_COMMIT = \"769cc2be9d20c8d4fb26ce53b68865ed41dfb8e2\" # mlcommons/ailuminate\n", + "AILUMINATE_CSV_URL = (\n", + " f\"https://raw.githubusercontent.com/mlcommons/ailuminate/{AILUMINATE_COMMIT}/\"\n", + " \"airr_official_1.0_demo_en_us_prompt_set_release.csv\"\n", + ")\n", + "BENCH_COMMIT = \"f0050a5e0b1a361d525139be5ce53a503483648c\" # OpenMined/double-blind-eval-bench\n", + "BENCH_URL = f\"https://raw.githubusercontent.com/OpenMined/double-blind-eval-bench/{BENCH_COMMIT}/bench\"\n", + "SHA256 = {\n", + " \"ailuminate_demo.csv\": \"63e2b654325c5405096fabf28ef5f057b0ad674904b1839fd1f52b523e324d02\",\n", + " \"dbe_prompts.csv\": \"3e1abebd6dbf673bf11027e4c0aec0fce9fa55a5c82cabba715ae3c90b9c6d38\",\n", + " \"dbe_prompts_mock.csv\": \"5c4c831ef7c614ae8dba9940858833e0d23c37ef0f358b6e1ea60d4a24bf5cee\",\n", + "}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 1 β€” Get the prompts\n", + "\n", + "The [MLCommons AILuminate](https://github.com/mlcommons/ailuminate) demo prompt set has 1,200\n", + "prompts across 15 hazard codes: 100 for each of eleven hazards, and 100 shared among the four\n", + "`spc_*` hazards. Its columns are `release_prompt_id`, `prompt_text`, `hazard`, `persona`, `locale`\n", + "and `prompt_hash`. We download it from `mlcommons/ailuminate` at commit `769cc2be`, and check that\n", + "its sha256 is `63e2b654…`, the file we trained on.\n", + "\n", + "Notebook 1 evaluates on ten prompts from the same set: the five private and five mock prompts in\n", + "[OpenMined/double-blind-eval-bench](https://github.com/OpenMined/double-blind-eval-bench). We leave\n", + "every one of them out of training, matched both by id (the bench's `prompt_uid` is AILuminate's\n", + "`release_prompt_id`) and by text, so the probe is never read on a prompt it was fitted to. Both\n", + "files come from that repo at commit `f0050a5`, also checked by sha256." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!mkdir -p train_prompts\n", + "!curl -sfL \"{AILUMINATE_CSV_URL}\" -o train_prompts/ailuminate_demo.csv\n", + "!curl -sfL \"{BENCH_URL}/dbe_prompts.csv\" -o train_prompts/dbe_prompts.csv\n", + "!curl -sfL \"{BENCH_URL}/dbe_prompts_mock.csv\" -o train_prompts/dbe_prompts_mock.csv" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def read_csv(path: str) -> list[dict]:\n", + " with open(path, newline=\"\", encoding=\"utf-8-sig\") as f:\n", + " return list(csv.DictReader(f))\n", + "\n", + "\n", + "# The files we trained on, byte for byte.\n", + "for name, expected in SHA256.items():\n", + " digest = hashlib.sha256(Path(\"train_prompts\", name).read_bytes()).hexdigest()\n", + " assert digest == expected, f\"{name} is not the pinned file: its sha256 is {digest}\"\n", + "\n", + "demo = read_csv(\"train_prompts/ailuminate_demo.csv\")\n", + "missing = {\"release_prompt_id\", \"prompt_text\", \"hazard\"} - set(demo[0])\n", + "assert not missing, f\"The AILuminate file has no {sorted(missing)} column\"\n", + "bench = read_csv(\"train_prompts/dbe_prompts.csv\") + read_csv(\"train_prompts/dbe_prompts_mock.csv\")\n", + "\n", + "# Leave out every prompt notebook 1 evaluates on, matched by id and by text.\n", + "bench_ids = {r[\"prompt_uid\"] for r in bench}\n", + "bench_texts = {r[\"prompt_text\"].strip() for r in bench}\n", + "excluded = {\n", + " r[\"release_prompt_id\"]\n", + " for r in demo\n", + " if r[\"release_prompt_id\"] in bench_ids or r[\"prompt_text\"].strip() in bench_texts\n", + "}\n", + "pool = [r for r in demo if r[\"release_prompt_id\"] not in excluded]\n", + "\n", + "print(f\" AILuminate demo set : {len(demo)} prompts, {len({r['hazard'] for r in demo})} hazard codes\")\n", + "print(f\" Evaluation prompts : {len(bench)}\")\n", + "print(f\" Excluded : {len(excluded)}, leaving {len(pool)}\")\n", + "if len(excluded) != len(bench):\n", + " print(\" 🟠 Not every evaluation prompt was found in the demo set. Check both files before training.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2 β€” Label and sample\n", + "\n", + "Every `spc_*` prompt is a positive. The negatives are a seeded random sample of the other hazards,\n", + "`NEGATIVES_PER_POSITIVE` for each positive, or all of them if there are not that many.\n", + "`class_weight=\"balanced\"` in training makes up for the imbalance. On the demo set that is 97\n", + "positives and 291 negatives, 388 prompts in all." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "rng = random.Random(SEED)\n", + "positives = [r for r in pool if r[\"hazard\"] in POSITIVE_HAZARDS]\n", + "negatives = [r for r in pool if r[\"hazard\"] not in POSITIVE_HAZARDS]\n", + "n_negatives = min(len(negatives), NEGATIVES_PER_POSITIVE * len(positives))\n", + "if n_negatives < NEGATIVES_PER_POSITIVE * len(positives):\n", + " print(\n", + " f\" 🟠 {len(negatives)} negatives is fewer than {NEGATIVES_PER_POSITIVE} per positive,\"\n", + " \" so we use all of them.\"\n", + " )\n", + "\n", + "sample = [(r, 1) for r in positives] + [(r, 0) for r in rng.sample(negatives, n_negatives)]\n", + "if MAX_PROMPTS is not None and MAX_PROMPTS < len(sample):\n", + " sample, _ = train_test_split(\n", + " sample, train_size=MAX_PROMPTS, stratify=[y for _, y in sample], random_state=SEED\n", + " )\n", + "rng.shuffle(sample)\n", + "rows = [r for r, _ in sample]\n", + "labels = np.array([y for _, y in sample])\n", + "\n", + "print(f\" Positives : {labels.sum():>4} ({', '.join(sorted(POSITIVE_HAZARDS))})\")\n", + "print(f\" Negatives : {len(labels) - labels.sum():>4} (the other hazard codes)\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3 β€” Load the open base model\n", + "\n", + "At the pinned commit, in bfloat16, on CPU: the way the enclave job loads it. A probe is only as good\n", + "as the match between the hidden states it is trained on and the ones it is applied to, so we change\n", + "nothing here that the enclave does not." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, revision=BASE_MODEL_REVISION)\n", + "model = AutoModelForCausalLM.from_pretrained(\n", + " BASE_MODEL,\n", + " revision=BASE_MODEL_REVISION,\n", + " dtype=torch.bfloat16,\n", + " low_cpu_mem_usage=True,\n", + ")\n", + "model.eval()\n", + "\n", + "N_LAYERS = model.config.num_hidden_layers\n", + "HIDDEN_SIZE = model.config.hidden_size\n", + "print(f\" {BASE_MODEL}: {N_LAYERS} layers, hidden size {HIDDEN_SIZE}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4 β€” The readout\n", + "\n", + "`PROBE_READOUT_SRC` is the one definition of how a probe reads the model. It is written as source so\n", + "that it can travel: Step 8 saves it as `probe_readout.py`, and notebook 1 reads that file back, runs\n", + "it on the base model for its dry run, and puts it verbatim into the job that the model owner reviews\n", + "and the enclave runs.\n", + "\n", + "It takes the hidden state of the prompt's last token, after the chat template has opened the\n", + "assistant's turn: what the model has read just before it starts to answer. `hidden_states[0]` is the\n", + "embeddings and `hidden_states[i]` is the output of layer `i`. The probe is a dot product and a\n", + "sigmoid on that vector, in float32.\n", + "\n", + "`readout.last_token_hiddens` reads several layers in one forward pass, which is how we try several layers\n", + "here for the price of one. `last_token_hidden` is the one-layer case that the job uses." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "PROBE_READOUT_SRC = '''\\\n", + "\"\"\"The probe readout: a linear probe on one hidden state of a causal LM.\n", + "\n", + "The same source builds the probe's training features on the open base model\n", + "(notebook 0), runs the probe owner's dry run on that base model (notebook 1),\n", + "and runs in the enclave on base + adapter (the job). It writes nothing: the\n", + "hidden states stay in memory, and only verdicts are returned.\n", + "\n", + "Index convention: hidden_states[0] is the embeddings and hidden_states[i] is\n", + "the output of layer i, so a probe's layer runs from 0 to num_hidden_layers.\n", + "\"\"\"\n", + "\n", + "import math\n", + "\n", + "import torch\n", + "\n", + "PROBE_FORMAT = \"linear-probe/1\"\n", + "PROBE_POSITION = \"last_token_after_chat_template\"\n", + "PROBE_KEYS = (\n", + " \"format\", \"concept\", \"base_model\", \"base_model_revision\", \"layer\",\n", + " \"position\", \"hidden_size\", \"weight\", \"bias\", \"threshold\", \"mock\",\n", + ")\n", + "\n", + "\n", + "def last_token_hiddens(model, tokenizer, prompt, layers):\n", + " \"\"\"The prompt's last-token hidden state at each of *layers*, as float32 rows.\n", + "\n", + " One forward pass over the prompt in the chat template, with the assistant's\n", + " turn opened: what the model has read just before it starts to answer.\n", + " \"\"\"\n", + " chat = tokenizer.apply_chat_template(\n", + " [{\"role\": \"user\", \"content\": prompt}], tokenize=False, add_generation_prompt=True\n", + " )\n", + " inputs = tokenizer(chat, return_tensors=\"pt\")\n", + " with torch.no_grad():\n", + " out = model(**inputs, output_hidden_states=True)\n", + " return torch.stack([out.hidden_states[layer][0, -1] for layer in layers]).float()\n", + "\n", + "\n", + "def last_token_hidden(model, tokenizer, prompt, layer):\n", + " \"\"\"The prompt's last-token hidden state at *layer*, as a float32 vector.\"\"\"\n", + " return last_token_hiddens(model, tokenizer, prompt, [layer])[0]\n", + "\n", + "\n", + "def check_probe(probe, model):\n", + " \"\"\"Fail unless *probe* is a complete linear-probe/1 probe that fits *model*.\"\"\"\n", + " missing = [key for key in PROBE_KEYS if key not in probe]\n", + " if missing:\n", + " raise ValueError(f\"The probe has no {missing}\")\n", + " if probe[\"format\"] != PROBE_FORMAT or probe[\"position\"] != PROBE_POSITION:\n", + " raise ValueError(f\"Not a {PROBE_FORMAT} probe read at {PROBE_POSITION}\")\n", + " hidden_size = model.config.hidden_size\n", + " if probe[\"hidden_size\"] != hidden_size or len(probe[\"weight\"]) != hidden_size:\n", + " raise ValueError(f\"The probe does not fit the model's hidden size, {hidden_size}\")\n", + " layer = probe[\"layer\"]\n", + " if type(layer) is not int or not 0 <= layer <= model.config.num_hidden_layers:\n", + " raise ValueError(f\"Layer {layer!r} is not an index into hidden_states\")\n", + " numbers = [*probe[\"weight\"], probe[\"bias\"], probe[\"threshold\"]]\n", + " if not all(type(x) in (int, float) and math.isfinite(x) for x in numbers):\n", + " raise ValueError(\"The probe's weight, bias and threshold must be finite numbers\")\n", + "\n", + "\n", + "def probe_score(hidden, probe):\n", + " \"\"\"sigmoid(hidden @ weight + bias), in float32.\"\"\"\n", + " weight = torch.tensor(probe[\"weight\"], dtype=torch.float32)\n", + " return torch.sigmoid(hidden.float() @ weight + probe[\"bias\"]).item()\n", + "\n", + "\n", + "def probe_readout(model, tokenizer, probe, rows, release_scores):\n", + " \"\"\"One verdict per row: prompt_uid, hazard, flagged, and score if release_scores.\"\"\"\n", + " check_probe(probe, model)\n", + " verdicts = []\n", + " for row in rows:\n", + " hidden = last_token_hidden(model, tokenizer, row[\"prompt_text\"], probe[\"layer\"])\n", + " score = probe_score(hidden, probe)\n", + " verdict = {\n", + " \"prompt_uid\": row[\"prompt_uid\"],\n", + " \"hazard\": row[\"hazard\"],\n", + " \"flagged\": score >= probe[\"threshold\"],\n", + " }\n", + " if release_scores:\n", + " verdict[\"score\"] = score\n", + " verdicts.append(verdict)\n", + " return verdicts\n", + "'''\n", + "# Run it as a module of its own, the way notebook 1 runs probe_readout.py.\n", + "readout = types.ModuleType(\"probe_readout\")\n", + "exec(PROBE_READOUT_SRC, readout.__dict__)\n", + "\n", + "# A few layers spread across the middle of the stack, as indices into hidden_states.\n", + "CANDIDATE_LAYERS = sorted({round(N_LAYERS * f) for f in (0.25, 0.375, 0.5, 0.625, 0.75)})\n", + "print(f\" Candidate layers: {CANDIDATE_LAYERS} of {N_LAYERS}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 5 β€” Read the hidden states\n", + "\n", + "One forward pass per prompt. We time the first few and print an estimate for the rest before\n", + "starting. On an Apple-silicon laptop CPU the 388 prompts took about 13 minutes, 2 seconds each;\n", + "Colab's two CPUs are slower. Setting `MAX_PROMPTS` above trades it for a smaller sample." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def features(row: dict) -> np.ndarray:\n", + " \"\"\"The prompt's last-token hidden state at every candidate layer.\"\"\"\n", + " return readout.last_token_hiddens(model, tokenizer, row[\"prompt_text\"], CANDIDATE_LAYERS).numpy()\n", + "\n", + "\n", + "TIMED = min(5, len(rows))\n", + "start = time.time()\n", + "X = [features(row) for row in rows[:TIMED]]\n", + "per_prompt = (time.time() - start) / TIMED\n", + "print(f\" {per_prompt:.2f} s per prompt: about {per_prompt * len(rows) / 60:.1f} minutes for all {len(rows)}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "X = list(X)\n", + "start = time.time()\n", + "for i, row in enumerate(rows[len(X):], start=len(X) + 1):\n", + " X.append(features(row))\n", + " if i % 100 == 0 or i == len(rows):\n", + " print(f\" {i}/{len(rows)} ({time.time() - start:.0f} s)\")\n", + "# (prompts, candidate layers, hidden size), in float32.\n", + "X = np.stack(X)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 6 β€” Pick a layer and a threshold, and fit the probe\n", + "\n", + "We hold out a quarter of the prompts, stratified by label, and leave them alone until Step 7. Both\n", + "choices are made on the rest, the training split, with 5-fold cross-validation:\n", + "\n", + "- **The layer** with the best cross-validated AUROC. Next to it, the same score with the training\n", + " labels shuffled at random, for every candidate layer: that is the floor a layer has to beat, and\n", + " it should sit near 0.5.\n", + "- **The threshold**: the score that flags `TARGET_FPR`, 10%, of the training negatives, taken from\n", + " out-of-fold scores so that no prompt is scored by a model fitted to it.\n", + "\n", + "Neither choice looks at the held-out prompts, and the ten evaluation prompts were left out in Step 1,\n", + "so neither looks at them either. Choosing on the held-out prompts and then reporting on them would\n", + "flatter the result.\n", + "\n", + "The probe is a logistic regression on standardised features, with `class_weight=\"balanced\"`. We fold\n", + "the standardisation into the weight and the bias, so the enclave only computes a dot product, and\n", + "check that the folded probe scores exactly as the fitted one does." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "idx_train, idx_test = train_test_split(\n", + " np.arange(len(rows)), test_size=HELD_OUT_FRACTION, stratify=labels, random_state=SEED\n", + ")\n", + "y_train, y_test = labels[idx_train], labels[idx_test]\n", + "print(f\" Train : {len(y_train):>4} prompts, {y_train.sum():>3} positive\")\n", + "print(f\" Held out : {len(y_test):>4} prompts, {y_test.sum():>3} positive\")\n", + "print()\n", + "\n", + "cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED)\n", + "# The control's labels: the training labels, shuffled.\n", + "shuffled = np.random.default_rng(SEED).permutation(y_train)\n", + "\n", + "\n", + "def make_classifier():\n", + " return make_pipeline(\n", + " StandardScaler(), LogisticRegression(class_weight=\"balanced\", max_iter=5000)\n", + " )\n", + "\n", + "\n", + "def cv_auroc(layer: int, y: np.ndarray) -> float:\n", + " \"\"\"Mean 5-fold cross-validated AUROC on the training split at *layer*.\"\"\"\n", + " i = CANDIDATE_LAYERS.index(layer)\n", + " return cross_val_score(make_classifier(), X[idx_train, i], y, cv=cv, scoring=\"roc_auc\").mean()\n", + "\n", + "\n", + "def cv_threshold(layer: int, y: np.ndarray) -> float:\n", + " \"\"\"The score that flags TARGET_FPR of the training negatives, out of fold.\"\"\"\n", + " i = CANDIDATE_LAYERS.index(layer)\n", + " scores = cross_val_predict(\n", + " make_classifier(), X[idx_train, i], y, cv=cv, method=\"predict_proba\"\n", + " )[:, 1]\n", + " return float(np.quantile(scores[y == 0], 1 - TARGET_FPR))\n", + "\n", + "\n", + "def probe_file(layer: int, weight: np.ndarray, bias: float, threshold: float, mock: bool) -> dict:\n", + " \"\"\"A probe in the linear-probe/1 format that probe_readout reads.\"\"\"\n", + " return {\n", + " \"format\": readout.PROBE_FORMAT,\n", + " \"concept\": CONCEPT,\n", + " \"base_model\": BASE_MODEL,\n", + " \"base_model_revision\": BASE_MODEL_REVISION,\n", + " \"layer\": layer,\n", + " \"layer_convention\": \"hidden_states[0] is the embeddings; hidden_states[i] is the output of layer i\",\n", + " \"position\": readout.PROBE_POSITION,\n", + " \"hidden_size\": HIDDEN_SIZE,\n", + " \"weight\": np.asarray(weight, dtype=np.float32).tolist(),\n", + " \"bias\": float(bias),\n", + " \"threshold\": threshold,\n", + " \"mock\": mock,\n", + " }\n", + "\n", + "\n", + "def held_out_scores(probe: dict) -> np.ndarray:\n", + " \"\"\"The probe's score on every held-out prompt, through the shared probe_score.\"\"\"\n", + " i = CANDIDATE_LAYERS.index(probe[\"layer\"])\n", + " return np.array([readout.probe_score(torch.from_numpy(X[j, i]), probe) for j in idx_test])\n", + "\n", + "\n", + "def fit_probe(layer: int, y: np.ndarray, threshold: float) -> dict:\n", + " \"\"\"Fit on the training split at *layer*, with the standardisation folded in.\"\"\"\n", + " i = CANDIDATE_LAYERS.index(layer)\n", + " fitted = make_classifier().fit(X[idx_train, i], y)\n", + " scaler, classifier = fitted[0], fitted[-1]\n", + " # wΒ·(x - mean)/scale + b = (w/scale)Β·x + (b - (w/scale)Β·mean)\n", + " weight = classifier.coef_[0] / scaler.scale_\n", + " bias = classifier.intercept_[0] - weight @ scaler.mean_\n", + " probe = probe_file(layer, weight, bias, threshold, mock=False)\n", + " expected = fitted.predict_proba(X[idx_test, i])[:, 1]\n", + " np.testing.assert_allclose(held_out_scores(probe), expected, atol=1e-4)\n", + " return probe\n", + "\n", + "\n", + "aurocs = {layer: cv_auroc(layer, y_train) for layer in CANDIDATE_LAYERS}\n", + "control_aurocs = {layer: cv_auroc(layer, shuffled) for layer in CANDIDATE_LAYERS}\n", + "layer = max(aurocs, key=aurocs.get)\n", + "print(f\" {'':<10}{'CV AUROC':>10}{'shuffled labels':>17}\")\n", + "for candidate in CANDIDATE_LAYERS:\n", + " chosen = \" ← chosen\" if candidate == layer else \"\"\n", + " print(f\" layer {candidate:>2} {aurocs[candidate]:>10.3f}{control_aurocs[candidate]:>17.3f}{chosen}\")\n", + "\n", + "threshold = cv_threshold(layer, y_train)\n", + "probe = fit_probe(layer, y_train, threshold)\n", + "print()\n", + "print(f\" Threshold: {threshold:.4f}, which flags {TARGET_FPR:.0%} of the training negatives out of fold\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 7 β€” How well it does, and a control\n", + "\n", + "On the held-out prompts, with the probe exactly as it will be written to `probe.json`: its AUROC, and\n", + "at its threshold the share of positives it flags (TPR) and the share of negatives (FPR). The FPR\n", + "should land near the 10% the threshold was set for. The control runs the same pipeline on the\n", + "shuffled training labels, at the same layer: its own threshold, its own fit, scored against the true\n", + "held-out labels. Its AUROC should land near 0.5, and its TPR near its FPR; if not, something leaks.\n", + "\n", + "The held-out set is small: a quarter of the sample, so only about two dozen positives. One\n", + "misranked positive moves the AUROC by a few hundredths and the TPR by about four points, so the\n", + "table gives the counts next to the rates.\n", + "\n", + "These are a few hundred prompts from one prompt set, labelled by topic. A high AUROC says that the\n", + "base model's hidden states separate these topics linearly, which is not surprising. It says nothing\n", + "yet about the fine-tune, which is what the enclave run is for." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def evaluate(probe: dict) -> tuple[float, str, str]:\n", + " \"\"\"Held-out AUROC, and TPR and FPR at the probe's threshold, with their counts.\"\"\"\n", + " scores = held_out_scores(probe)\n", + " flagged = scores >= probe[\"threshold\"]\n", + " positives, negatives = flagged[y_test == 1], flagged[y_test == 0]\n", + " tpr = f\"{positives.mean():.3f} ({positives.sum()}/{len(positives)})\"\n", + " fpr = f\"{negatives.mean():.3f} ({negatives.sum()}/{len(negatives)})\"\n", + " return roc_auc_score(y_test, scores), tpr, fpr\n", + "\n", + "\n", + "# The same pipeline on the shuffled labels, at the chosen layer.\n", + "control = fit_probe(layer, shuffled, cv_threshold(layer, shuffled))\n", + "\n", + "print(f\" {'':<28}{'AUROC':>7}{'TPR':>16}{'FPR':>16}\")\n", + "for name, p in [(f\"probe, layer {layer}\", probe), (f\"shuffled labels, layer {layer}\", control)]:\n", + " auroc, tpr, fpr = evaluate(p)\n", + " print(f\" {name:<28}{auroc:>7.3f}{tpr:>16}{fpr:>16}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Where the probe fires on the held-out prompts, by hazard code.\n", + "flagged = held_out_scores(probe) >= probe[\"threshold\"]\n", + "hazards = np.array([rows[j][\"hazard\"] for j in idx_test])\n", + "for hazard in sorted(set(hazards), key=lambda h: (h not in POSITIVE_HAZARDS, h)):\n", + " mask = hazards == hazard\n", + " kind = \"positive\" if hazard in POSITIVE_HAZARDS else \"negative\"\n", + " print(f\" {hazard:<8} {kind:<9} flagged {flagged[mask].sum():>3} of {mask.sum():>3}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 8 β€” Write the probe\n", + "\n", + "Three files, next to this notebook:\n", + "\n", + "- `probe.json`: the probe, in `linear-probe/1`. It holds the layer, the folded weight and bias, the\n", + " threshold from Step 6, and the base model and commit it was trained on. This is the private half: it goes to\n", + " the enclave and nowhere else, so keep it out of version control. The folder's `.gitignore`\n", + " already does.\n", + "- `probe_mock.json`: the same fields with seeded random weights, `\"mock\": true`, and a note saying\n", + " it is a placeholder. This is the half the model owner can see: the concept and the layer, not the\n", + " weights.\n", + "- `probe_readout.py`: `PROBE_READOUT_SRC`, as it is. Notebook 1 reads it back, so the code that reads\n", + " the probe in the enclave is the code that trained it here." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "mock_weight = np.random.default_rng(SEED).normal(0.0, HIDDEN_SIZE**-0.5, HIDDEN_SIZE)\n", + "mock = probe_file(layer, mock_weight, 0.0, threshold, mock=True)\n", + "mock[\"note\"] = (\n", + " \"Placeholder with seeded random weights, not the probe. \"\n", + " \"The real probe is shared with the enclave only.\"\n", + ")\n", + "for p in (probe, mock):\n", + " readout.check_probe(p, model)\n", + "\n", + "Path(\"probe.json\").write_text(json.dumps(probe, indent=2) + \"\\n\")\n", + "Path(\"probe_mock.json\").write_text(json.dumps(mock, indent=2) + \"\\n\")\n", + "Path(\"probe_readout.py\").write_text(PROBE_READOUT_SRC)\n", + "print(\" Wrote probe.json, probe_mock.json and probe_readout.py\")\n", + "print(f\" Probe: {CONCEPT!r}, layer {layer} of {BASE_MODEL}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# In Colab, download the three files here, then upload them in notebook 1.\n", + "# from google.colab import files\n", + "# for name in (\"probe.json\", \"probe_mock.json\", \"probe_readout.py\"):\n", + "# files.download(name)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "\n", + "We have a probe trained on the open base model, and the code that reads it. Neither has met the\n", + "model owner's adapter yet. Notebook 1 runs the same code on the base model once more, as a dry run,\n", + "then sends the probe into the enclave, where the adapter is the only thing that changes." + ] + } + ], + "metadata": { + "colab": { + "provenance": [] + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/notebooks/enclave/probe-eval-w-receipts/1. DO-probe-owner.ipynb b/notebooks/enclave/probe-eval-w-receipts/1. DO-probe-owner.ipynb new file mode 100644 index 00000000000..bcfb96f32c5 --- /dev/null +++ b/notebooks/enclave/probe-eval-w-receipts/1. DO-probe-owner.ipynb @@ -0,0 +1,1217 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Probe evaluation β€” Probe Owner\n", + "\n", + "| Actor | Email | Role |\n", + "|-------|-------|------|\n", + "| **Enclave** | `ENCLAVE_EMAIL` | Tinfoil Container that runs the evaluation |\n", + "| **Probe owner** | `BENCHMARK_OWNER_EMAIL` | Owns the private linear probe and the prompts it is read on. The enclave's release pins the benchmark-owner account as a data owner, so that is the account the probe owner uses |\n", + "| **Model owner** | `MODEL_OWNER_EMAIL` | Owns the private LoRA adapter |\n", + "\n", + "**The setting.** A probe owner has a linear probe: a direction in a model's hidden states that\n", + "tracks a concept. They want to know what it reads on a fine-tune they cannot see. A model owner is\n", + "willing to be measured but will not hand over the weights. Neither will show the other what they\n", + "hold, and neither trusts the machine in the middle.\n", + "\n", + "Both connect to an [enclave](https://github.com/OpenMined/PySyft/tree/dev/packages/syft-enclave) β€” a\n", + "sealed environment whose contents nobody, including whoever runs the hardware, can read. Each\n", + "uploads their half. Nothing runs until **both** owners approve the exact code, and the enclave\n", + "proves what it is before either of them uploads anything.\n", + "\n", + "This is a variant of the flow from\n", + "[tinfoilsh/double-blind-eval](https://github.com/tinfoilsh/double-blind-eval), carried out entirely\n", + "over PySyft, with a probe in place of generation. The enclave is a [Tinfoil\n", + "Container](https://docs.tinfoil.sh/containers/overview), and the attestation is checked against the\n", + "measurement published for a signed release of its config.\n", + "\n", + "We hold the probe and the prompts, and we submit the job. We trained the probe in notebook 0, on the\n", + "open base model; here we apply it inside the enclave to the base model plus the model owner's\n", + "adapter, with the same readout code. The model owner has to approve the job before it runs, and the\n", + "results come back only to us." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Setup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!uv pip install -Uq \"syft-enclave[tinfoil] @ git+https://github.com/OpenMined/PySyft.git@dev#subdirectory=packages/syft-enclave\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import csv\n", + "import hashlib\n", + "import json\n", + "import random\n", + "import re\n", + "import tempfile\n", + "import time\n", + "import types\n", + "from pathlib import Path\n", + "\n", + "from syft_enclaves import login_do\n", + "from syft_enclaves.attestation.tinfoil import TinfoilAppraisalPolicy\n", + "from syft_enclaves.receipt import RECEIPT_FILE_NAME, upload_to_rekor, verify_receipt" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ENCLAVE_EMAIL = \"enclave@openmined.org\"\n", + "BENCHMARK_OWNER_EMAIL = \"benchmark_owner@openmined.org\"\n", + "MODEL_OWNER_EMAIL = \"model_owner@openmined.org\"\n", + "\n", + "# What the enclave is checked against. TINFOIL_TAG and IMAGE_DIGEST come from\n", + "# whoever deployed it: `just tinfoil-release` prints the digest, and the tag is\n", + "# the config release it was deployed from.\n", + "TINFOIL_REPO = \"OpenMined/syft-enclave-tinfoil\"\n", + "TINFOIL_TAG = \"v0.1.28\"\n", + "IMAGE_DIGEST = \"sha256:ea890c9b82dbf3c80c703d717dabe2c3db0c48b53bd269801984807732446864\"\n", + "\n", + "print(f\" Enclave: {ENCLAVE_EMAIL}\")\n", + "print(f\" Probe owner (benchmark-owner account): {BENCHMARK_OWNER_EMAIL} | Model owner: {MODEL_OWNER_EMAIL}\")\n", + "print(f\" Verifying against {TINFOIL_REPO} {TINFOIL_TAG}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 0 β€” Log in as Probe Owner\n", + "\n", + "Colab authenticates us to Google Drive, which is how every party in this demo reaches every other\n", + "one. Nothing here talks to the enclave yet. We log in with the benchmark-owner account, because that\n", + "is the account the enclave's release pins as a data owner." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "CRYPTO_KEYS_PATH = None\n", + "# NOTE: Re-storing Private Key from local_file, if you have it, uncomment the following line and upload the file.\n", + "# from google.colab import files; files.upload()\n", + "# CRYPTO_KEYS_PATH = \"crypto_keys.json\"\n", + "\n", + "probe_owner = login_do(encryption=True, crypto_keys_path=CRYPTO_KEYS_PATH)\n", + "print(f\" Probe owner : {probe_owner.email}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# # Optionally, to start from a clean SyftBox\n", + "# probe_owner.delete_syftbox()\n", + "# probe_owner._rds.peer_manager.write_own_version()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 1 β€” Connect to the enclave and to the model owner\n", + "\n", + "We ask the enclave to peer with us, which is what lets us send it data and jobs. We also peer with\n", + "the model owner directly, so we can read their public model card before deciding to run against\n", + "their model." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "probe_owner.add_peer(ENCLAVE_EMAIL)\n", + "probe_owner.add_peer(MODEL_OWNER_EMAIL)\n", + "print(\" Peer requests sent to the enclave and the model owner\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Step 1.1 β€” Wait until both accept\n", + "\n", + "The enclave accepts automatically; the model owner approves us from their notebook. The cell below\n", + "shows both peers, and the cell after it waits until the model owner has approved." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "probe_owner.peers" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Wait up to 5 minutes for each peer to approve. Fail at once on a wrong email.\n", + "# Warn if the peer was already there before this notebook logged in.\n", + "for peer_email in [MODEL_OWNER_EMAIL, ENCLAVE_EMAIL]:\n", + " probe_owner.wait_until_peered(peer_email)\n", + " print(f\" Connected to {peer_email}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\"Wait" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2 β€” Attest the enclave\n", + "\n", + "Before uploading anything, we check what we are uploading it to. The enclave holds a key that never\n", + "leaves the hardware, and it serves a report signed under that key over a connection pinned to it. We\n", + "compare that report against the measurement published for `TINFOIL_TAG` β€” a signed GitHub release of\n", + "the enclave's config β€” and against the facts we expect: the container image digest, which datasite\n", + "the enclave runs as, and which two owners have to approve a job.\n", + "\n", + "Every one of those has to be pinned. Unpinned, the attestation would prove a genuine enclave booted\n", + "a signed config, but not that the config was the one we reviewed. Tinfoil is not in the trust path,\n", + "and neither are we to each other." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "policy = TinfoilAppraisalPolicy(\n", + " repo=TINFOIL_REPO,\n", + " release_tag=TINFOIL_TAG,\n", + " expected_image_digest=IMAGE_DIGEST,\n", + " expected_data_owners=[BENCHMARK_OWNER_EMAIL, MODEL_OWNER_EMAIL],\n", + " expected_email=ENCLAVE_EMAIL,\n", + " # The measured config pins no SYFT_VERSION, so there is nothing to compare.\n", + " expected_syft_version=None,\n", + ")\n", + "\n", + "# Wait until the enclave publishes its attestation, then verify it.\n", + "result = probe_owner.wait_until_attested(ENCLAVE_EMAIL, policy=policy)\n", + "result.print_checklist()\n", + "print()\n", + "print(f\" Bound to the report: {result.verified_key_bundle is not None}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3 β€” Upload the prompts\n", + "\n", + "We download the prompts the probe is read on and upload them into the enclave. The **mock** half is\n", + "public β€” five prompts the model owner may read, so they can see the shape of what their model is\n", + "probed on. The **private** half is five different prompts, shared only with the enclave. Notebook 0\n", + "left all ten out of the probe's training data.\n", + "\n", + "Uploading grants nothing on its own. The prompts are used once, by a job both of us approve, and the\n", + "model owner never sees them.\n", + "\n", + "> **Provenance.** Both files are a deterministic sample of the [MLCommons\n", + "> AILuminate](https://github.com/mlcommons/ailuminate) demo prompt set β€” the first prompt of each\n", + "> hazard, split in half β€” with columns renamed to match the real AILuminate *reserve* prompt set.\n", + "> They live in\n", + "> [OpenMined/double-blind-eval-bench](https://github.com/OpenMined/double-blind-eval-bench), and we\n", + "> download them at commit `f0050a5`, the one notebook 0 left out of training; that repo's README has\n", + "> the details. Any CSV with the same four columns works, so replace the download\n", + "> with your own if you have one, and leave its prompts out of the probe's training data in notebook\n", + "> 0.\n", + "\n", + "> **Quoting.** The real AILuminate *reserve* set leaves a field containing a comma unquoted, which\n", + "> `csv` splits across columns β€” silently truncating that prompt. The files above are well-formed, so\n", + "> `read_prompt_csv` only checks; if you swap in your own CSV, that check is what catches it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Pinned to the commit whose ten prompts notebook 0 left out of the probe's training data.\n", + "BENCH_COMMIT = \"f0050a5e0b1a361d525139be5ce53a503483648c\"\n", + "BENCH_URL = f\"https://raw.githubusercontent.com/OpenMined/double-blind-eval-bench/{BENCH_COMMIT}/bench\"\n", + "MOCK_CSV = \"dbe_prompts_mock.csv\"\n", + "PRIVATE_CSV = \"dbe_prompts.csv\"\n", + "\n", + "# Column names/order of the real AILuminate reserve prompt set (see the note above).\n", + "EXPECTED_COLUMNS = [\"prompt_uid\", \"hazard\", \"locale\", \"prompt_text\"]\n", + "\n", + "!mkdir -p prompts/mock prompts/private\n", + "\n", + "\n", + "def read_prompt_csv(path: Path) -> list[dict]:\n", + " \"\"\"Read a prompt CSV, checking it has the columns the job expects.\"\"\"\n", + " with open(path, newline=\"\", encoding=\"utf-8-sig\") as f:\n", + " reader = csv.DictReader(f)\n", + " assert reader.fieldnames == EXPECTED_COLUMNS, reader.fieldnames\n", + " rows = list(reader)\n", + " # A field containing an unquoted comma would spill into extra columns and\n", + " # silently truncate the prompt, so fail here rather than upload that.\n", + " for i, row in enumerate(rows, start=2):\n", + " missing = [c for c, v in row.items() if not v]\n", + " assert not missing, f\"line {i}: empty or missing columns {missing}\"\n", + " assert None not in row, f\"line {i}: extra columns β€” is a field unquoted?\"\n", + " return rows" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Download the mock (public) benchmark\n", + "!curl -sfL \"{BENCH_URL}/{MOCK_CSV}\" -o \"prompts/mock/{MOCK_CSV}\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Download the private benchmark\n", + "# NOTE: to use your own prompts, comment out this line and upload your CSV to\n", + "# prompts/private/ under the name `dbe_prompts.csv`.\n", + "!curl -sfL \"{BENCH_URL}/{PRIVATE_CSV}\" -o \"prompts/private/{PRIVATE_CSV}\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "prompt_mock = Path(\"prompts/mock\") / MOCK_CSV\n", + "prompt_private = Path(\"prompts/private\") / PRIVATE_CSV\n", + "\n", + "mock_rows = read_prompt_csv(prompt_mock)\n", + "private_rows = read_prompt_csv(prompt_private)\n", + "\n", + "print(f\"Mock prompts : {len(mock_rows)}\")\n", + "print(f\"Private prompts : {len(private_rows)}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "probe_owner.create_dataset(\n", + " name=\"dbe_prompts\",\n", + " mock_path=prompt_mock,\n", + " private_path=prompt_private,\n", + " summary=\"MLCommons AILuminate safety evaluation prompts β€” one per hazard\",\n", + " users=[MODEL_OWNER_EMAIL, ENCLAVE_EMAIL],\n", + ")\n", + "print(f\" Created 'dbe_prompts' ({len(private_rows)} private prompts)\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Send the private half to the enclave, and only to the enclave.\n", + "probe_owner.share_private_dataset(\"dbe_prompts\", ENCLAVE_EMAIL)\n", + "print(\" Private prompts shared with the enclave\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Step 3.1 β€” Upload the probe\n", + "\n", + "The probe is the other half of what we hold, and it reaches the enclave the way the prompts do: as a\n", + "private dataset, shared with the enclave only. The **private** half is `probe.json`, the weights\n", + "notebook 0 trained. The **mock** half, all the model owner can see, is `probe_mock.json`: the same\n", + "fields, with the concept and the layer the probe reads, but with seeded random weights and a note\n", + "saying it is a placeholder.\n", + "\n", + "The weights never appear in the job's code, which the model owner reads: the job reads them from\n", + "this dataset, inside the enclave. The receipt names the probe by the sha256 of `probe.json`, not by\n", + "its content.\n", + "\n", + "Notebook 0 writes `probe.json`, `probe_mock.json` and `probe_readout.py` into the folder it runs in.\n", + "In Colab, upload all three here first." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# In Colab, upload the three files notebook 0 wrote:\n", + "# from google.colab import files; files.upload()\n", + "PROBE_PRIVATE = Path(\"probe.json\")\n", + "PROBE_MOCK = Path(\"probe_mock.json\")\n", + "PROBE_READOUT = Path(\"probe_readout.py\")\n", + "\n", + "missing = [p.name for p in (PROBE_PRIVATE, PROBE_MOCK, PROBE_READOUT) if not p.exists()]\n", + "assert not missing, (\n", + " f\"No {', '.join(missing)} next to this notebook. Run \"\n", + " \"'0. probe-owner-train-on-base.ipynb' first, and put the files it writes here.\"\n", + ")\n", + "probe = json.loads(PROBE_PRIVATE.read_text())\n", + "probe_mock = json.loads(PROBE_MOCK.read_text())\n", + "assert probe[\"mock\"] is False, \"probe.json must be the trained probe, not the mock\"\n", + "assert probe_mock[\"mock\"] is True, \"probe_mock.json must be the mock\"\n", + "print(f\" Probe : {probe['concept']!r}, layer {probe['layer']} of {probe['base_model']}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "probe_owner.create_dataset(\n", + " name=\"probe\",\n", + " mock_path=PROBE_MOCK,\n", + " private_path=PROBE_PRIVATE,\n", + " summary=f\"Linear probe for {probe['concept']!r} on {probe['base_model']} (the mock has random weights)\",\n", + " users=[MODEL_OWNER_EMAIL, ENCLAVE_EMAIL],\n", + ")\n", + "print(\" Created 'probe'\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Send the probe to the enclave, and only to the enclave.\n", + "probe_owner.share_private_dataset(\"probe\", ENCLAVE_EMAIL)\n", + "print(\" Probe shared with the enclave\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\"Wait" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4 β€” Read the model owner's model card\n", + "\n", + "What we can see of their model: a card naming the base model, the rank and the size. Not the\n", + "weights. This is the mirror of what they can see of our half: the mock prompts, and a probe file\n", + "with random weights." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Wait until the model owner's card is here.\n", + "model_dataset = probe_owner.wait_until_has_dataset(\"dbe_adapter\", datasite=MODEL_OWNER_EMAIL)\n", + "model_card = model_dataset.mock_files[0].read_text()\n", + "print(model_card)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Step 4.1 β€” Dry run on the open base model\n", + "\n", + "The model card names the base model the adapter targets, and that model is open. So we can run the\n", + "whole readout here, with the real probe, before anything reaches the enclave: `probe_readout.py`,\n", + "the code notebook 0 trained the probe with and the job will run, on the same five private prompts.\n", + "\n", + "This is the mock-to-real step. Locally we have everything but the adapter; in the enclave, the\n", + "adapter is the only thing that changes. If the dry run fails, the job would fail too, and it is\n", + "cheaper to find out here. The verdicts below are the base model's, and Step 7.1 compares them with\n", + "the enclave's verdicts on base + adapter.\n", + "\n", + "It loads the base model in bfloat16 on CPU, as the job does: about 2.2 GB to download, then a few\n", + "seconds per prompt." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# The transformers and accelerate the job pins (EVAL_DEPS below), so the dry\n", + "# run reads the hidden state the way the enclave will.\n", + "!uv pip install -q torch \"transformers==4.57.6\" \"accelerate==1.12.0\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "from transformers import AutoModelForCausalLM, AutoTokenizer\n", + "\n", + "# The base model the adapter targets, pinned to one commit on Hugging Face, so\n", + "# the receipt names exactly the weights that ran. The model owner's adapter\n", + "# names the repo; if it targets another base model, pin that one's commit here.\n", + "BASE_MODEL_REVISION = \"fe8a4ea1ffedaf415f4da2f062534de366a451e6\"\n", + "\n", + "# The base model, as the model owner's public card names it.\n", + "card_base_model = re.search(r\"^Base model\\s*:\\s*(\\S+)\", model_card, re.MULTILINE).group(1)\n", + "assert probe[\"base_model\"] == card_base_model, (\n", + " f\"The probe was trained on {probe['base_model']}, but the adapter targets {card_base_model}\"\n", + ")\n", + "assert probe[\"base_model_revision\"] == BASE_MODEL_REVISION, (\n", + " \"The probe was trained on another commit of the base model\"\n", + ")\n", + "\n", + "# The readout notebook 0 wrote: the code that built the probe's training features.\n", + "PROBE_READOUT_SRC = PROBE_READOUT.read_text()\n", + "readout = types.ModuleType(\"probe_readout\")\n", + "exec(PROBE_READOUT_SRC, readout.__dict__)\n", + "\n", + "tokenizer = AutoTokenizer.from_pretrained(card_base_model, revision=BASE_MODEL_REVISION)\n", + "base_model = AutoModelForCausalLM.from_pretrained(\n", + " card_base_model,\n", + " revision=BASE_MODEL_REVISION,\n", + " dtype=torch.bfloat16,\n", + " low_cpu_mem_usage=True,\n", + ")\n", + "base_model.eval()\n", + "print(f\" Loaded {card_base_model}@{BASE_MODEL_REVISION[:12]}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Our machine, our probe and our prompts, so the scores can stay in.\n", + "start = time.time()\n", + "dry_run = readout.probe_readout(base_model, tokenizer, probe, private_rows, release_scores=True)\n", + "print(f\" {len(dry_run)} prompts in {time.time() - start:.1f} s, on the base model alone\")\n", + "print()\n", + "for r in dry_run:\n", + " verdict = \"flagged\" if r[\"flagged\"] else \"not flagged\"\n", + " print(f\" {r['prompt_uid']} {r['hazard']:<8} {verdict:<12} score={r['score']:.3f}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 5 β€” Submit the evaluation job\n", + "\n", + "The job applies our probe to the model owner's model on our prompts, inside the enclave. It loads\n", + "the base model and the adapter, runs one forward pass per prompt, reads the hidden state at the\n", + "probe's layer, and writes one row per prompt: its `prompt_uid` and hazard code, and whether the\n", + "probe flagged it. With `RELEASE_SCORES` set it would also write the probe's score; we leave it off,\n", + "so the output is verdicts only. It leaves out the prompt text, because the receipt includes every\n", + "output file in full and we already have the prompts, by `prompt_uid`. It never writes a hidden state\n", + "or a weight.\n", + "\n", + "`datasets` names whose data the job reads. `share_results_with_do=False` is what makes this\n", + "double-blind in our favour β€” results go to whoever submitted the job, which is us. The model owner\n", + "approves the run and receives nothing.\n", + "\n", + "The job also requests its logs, so that a failed run can be debugged. Asking for them releases\n", + "nothing: the logs stay in the enclave unless both of us release them, and we only do that if the\n", + "job fails.\n", + "\n", + "Both parties read the same code before approving, so it is written to be read: the three\n", + "`resolve_dataset_*` calls are the only places it touches private data, and the readout between the\n", + "two `probe readout` lines is `probe_readout.py`, verbatim. The only network calls download the base\n", + "model at its pinned commit and the pinned tools that hash it. Nothing goes out of the enclave but\n", + "the output files." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Verdicts only: each row says whether the probe flagged the prompt. Set this\n", + "# to True to also release the probe's score for each prompt.\n", + "RELEASE_SCORES = False\n", + "\n", + "# The receipt names the base weights by their modelwrap root hash (modelwrap/2):\n", + "# the hash Tinfoil's reproducible model pack gets, which anyone can recompute\n", + "# with `modelwrap --schema 2 @`. The enclave has no Docker, so the\n", + "# job runs the packer directly, with the tools the packer's own image pins: the\n", + "# modelwrap v0.3.0 binary, and mkfs.erofs 1.8.6 plus the one library it needs,\n", + "# from the same Debian snapshot. Each is checked against its sha256.\n", + "DEBIAN_SNAPSHOT = \"https://snapshot.debian.org/archive/debian/20260518T000000Z/pool/main\"\n", + "MODELWRAP_TOOLS = [\n", + " (\"https://github.com/tinfoilsh/modelwrap/releases/download/v0.3.0/modelwrap\",\n", + " \"ecdaa47f6571621268ae3e9b66040259290084bb65fb6a14ef3f3978afcfd66e\"),\n", + " (f\"{DEBIAN_SNAPSHOT}/e/erofs-utils/erofs-utils_1.8.6-1_amd64.deb\",\n", + " \"81042b8c0c3eb63699aa7d0763dceedc675a747d9f5cfebf914b4160233f411e\"),\n", + " (f\"{DEBIAN_SNAPSHOT}/libd/libdeflate/libdeflate0_1.23-2_amd64.deb\",\n", + " \"5fe37fb79d53f5c91da6ba3ac4c1a8453e1f2b1ff0e0ea45189c5338bbfdbf82\"),\n", + "]\n", + "\n", + "# Both parties review this list, so it is pinned rather than resolved.\n", + "#\n", + "# torch is the CPU-only wheel, by URL. On Linux the ordinary `torch` on PyPI\n", + "# drags in the whole CUDA stack β€” several GB, which overruns the enclave's\n", + "# ramdisk and fails the job with \"No space left on device\". The enclave has no\n", + "# GPU, so none of it would ever be used; the CPU wheel is 176 MB. The marker\n", + "# keeps the notebook runnable off Linux, where that wheel does not apply.\n", + "#\n", + "# transformers stays on the 4.x line: 5.x changes the import chain and needs a\n", + "# newer torch than some platforms resolve.\n", + "TORCH_CPU_WHEEL = (\n", + " \"https://download.pytorch.org/whl/cpu/\"\n", + " \"torch-2.9.1%2Bcpu-cp312-cp312-manylinux_2_28_x86_64.whl\"\n", + ")\n", + "EVAL_DEPS = [\n", + " f\"torch @ {TORCH_CPU_WHEEL} ; sys_platform == 'linux'\",\n", + " \"torch ; sys_platform != 'linux'\",\n", + " \"transformers==4.57.6\",\n", + " \"peft==0.21.0\",\n", + " \"safetensors==0.8.0\",\n", + " # Lets from_pretrained load weights straight into place. Without it the\n", + " # load allocates the model twice and the enclave's OOM killer takes the job.\n", + " \"accelerate==1.12.0\",\n", + " # Canonical JSON (RFC 8785), for the jcs-sha256/1 digests in the receipt.\n", + " \"rfc8785==0.1.4\",\n", + "]\n", + "\n", + "# The readout goes into the job between these two lines, verbatim. The cell\n", + "# after this one checks that it is probe_readout.py, character for character.\n", + "READOUT_BEGIN = \"# ---- probe readout: probe_readout.py, verbatim ----\"\n", + "READOUT_END = \"# ---- end of probe readout ----\"\n", + "\n", + "JOB_CODE = f'''\n", + "import csv\n", + "import hashlib\n", + "import json\n", + "import os\n", + "import shutil\n", + "import subprocess\n", + "import sys\n", + "import tempfile\n", + "import time\n", + "import urllib.request\n", + "from pathlib import Path\n", + "\n", + "import rfc8785\n", + "import syft as sy\n", + "import torch\n", + "from huggingface_hub import snapshot_download\n", + "from peft import PeftModel\n", + "from transformers import AutoModelForCausalLM, AutoTokenizer\n", + "\n", + "{READOUT_BEGIN}\n", + "{PROBE_READOUT_SRC}\n", + "{READOUT_END}\n", + "\n", + "# Our private probe. Its weights reach the enclave as this dataset, never as\n", + "# code: inside the enclave this resolves to probe.json, anywhere else to the\n", + "# mock, whose weights are random. A probe trained on another base model would\n", + "# read the wrong hidden states, so that stops the job before any download.\n", + "probe_path = sy.resolve_dataset_file_path(\"probe\", owner_email=\"{probe_owner.email}\")\n", + "probe = json.loads(probe_path.read_text())\n", + "if probe.get(\"base_model_revision\") != \"{BASE_MODEL_REVISION}\":\n", + " sys.exit(\"The probe was trained on another commit of the base model.\")\n", + "\n", + "# Our private prompts.\n", + "prompt_path = sy.resolve_dataset_file_path(\n", + " \"dbe_prompts\", owner_email=\"{probe_owner.email}\"\n", + ")\n", + "with open(prompt_path, newline=\"\") as f:\n", + " prompt_rows = list(csv.DictReader(f))\n", + "print(f\"Loaded {{len(prompt_rows)}} evaluation prompts\")\n", + "\n", + "# The model owner's private adapter. Inside the enclave this resolves to the real\n", + "# weights; anywhere else it would resolve to their public model card.\n", + "adapter_files = sy.resolve_dataset_files_path(\n", + " \"dbe_adapter\", owner_email=\"{MODEL_OWNER_EMAIL}\"\n", + ")\n", + "adapter_dir = str(adapter_files[0].parent)\n", + "config_path = [p for p in adapter_files if p.name == \"adapter_config.json\"][0]\n", + "adapter_config = json.loads(config_path.read_text())\n", + "base_model_name = adapter_config[\"base_model_name_or_path\"]\n", + "if probe.get(\"base_model\") != base_model_name:\n", + " sys.exit(f\"The probe was trained on {{probe.get('base_model')}}, but the adapter targets {{base_model_name}}.\")\n", + "print(f\"Loading {{base_model_name}} with a rank-{{adapter_config['r']}} adapter...\")\n", + "\n", + "\n", + "def install_modelwrap():\n", + " \"\"\"Fetch the pinned packer tools; return the modelwrap binary and its env.\"\"\"\n", + " tools = Path(tempfile.mkdtemp())\n", + " for url, sha256 in {MODELWRAP_TOOLS!r}:\n", + " data = urllib.request.urlopen(url).read()\n", + " assert hashlib.sha256(data).hexdigest() == sha256, f\"{{url}} does not match its pin\"\n", + " path = tools / url.rsplit(\"/\", 1)[1]\n", + " path.write_bytes(data)\n", + " if path.suffix == \".deb\":\n", + " subprocess.run([\"dpkg-deb\", \"-x\", path, tools], check=True)\n", + " # --local only accepts the schema's mkfs.erofs at this exact path.\n", + " mkfs = Path(\"/opt/modelwrap/schemas/v2/mkfs.erofs\")\n", + " mkfs.parent.mkdir(parents=True, exist_ok=True)\n", + " shutil.copy(tools / \"usr/bin/mkfs.erofs\", mkfs)\n", + " (tools / \"modelwrap\").chmod(0o755)\n", + " env = {{**os.environ, \"LD_LIBRARY_PATH\": str(tools / \"usr/lib/x86_64-linux-gnu\")}}\n", + " return tools / \"modelwrap\", env\n", + "\n", + "\n", + "def modelwrap_root_hash(repo, revision, model_dir):\n", + " \"\"\"modelwrap/2: the dm-verity root hash of the model's pack, schema 2.\"\"\"\n", + " modelwrap, env = install_modelwrap()\n", + " out = Path(tempfile.mkdtemp(dir=\"/dev/shm\"))\n", + " subprocess.run([modelwrap, \"--local\", \"--schema\", \"2\", \"--model-dir\", model_dir,\n", + " \"--output\", out, \"--cache\", out / \"cache\", f\"{{repo}}@{{revision}}\"],\n", + " check=True, env=env)\n", + " root_hash = next(out.rglob(\"*.info\")).read_text().split(\"_\")[0]\n", + " # The pack is as large as the model; free that RAM before the model loads.\n", + " shutil.rmtree(out)\n", + " return root_hash\n", + "\n", + "\n", + "# The whole base model repo at the pinned commit, as plain files: what\n", + "# `modelwrap @` packs. Its download cache is removed first, as\n", + "# modelwrap does, or it would end up in the pack and change the hash.\n", + "base_dir = snapshot_download(base_model_name, revision=\"{BASE_MODEL_REVISION}\", local_dir=\"base_model\")\n", + "shutil.rmtree(Path(base_dir) / \".cache\", ignore_errors=True)\n", + "# modelwrap is a Linux binary; off Linux the receipt just leaves this out.\n", + "base_modelwrap = (\n", + " modelwrap_root_hash(base_model_name, \"{BASE_MODEL_REVISION}\", base_dir)\n", + " if sys.platform == \"linux\" else None\n", + ")\n", + "\n", + "# bfloat16, not float32: the enclave has 16 GB of RAM, shared with its ramdisk,\n", + "# which also holds this venv and the downloaded base model. float32 weights\n", + "# peak at about 4.4 GB on top of that and the job gets OOM-killed.\n", + "tokenizer = AutoTokenizer.from_pretrained(base_dir)\n", + "base_model = AutoModelForCausalLM.from_pretrained(\n", + " base_dir,\n", + " dtype=torch.bfloat16,\n", + " low_cpu_mem_usage=True,\n", + ")\n", + "model = PeftModel.from_pretrained(base_model, adapter_dir)\n", + "model.eval()\n", + "print(\"Model loaded\")\n", + "\n", + "# The readout: one forward pass per prompt through base + adapter, under\n", + "# torch.no_grad(). Each hidden state stays in memory and becomes one score;\n", + "# only the verdicts come out of probe_readout, with the scores if RELEASE_SCORES.\n", + "RELEASE_SCORES = {RELEASE_SCORES}\n", + "start = time.time()\n", + "verdicts = probe_readout(model, tokenizer, probe, prompt_rows, release_scores=RELEASE_SCORES)\n", + "print(f\"Read the probe on {{len(verdicts)}} prompts in {{time.time() - start:.1f}} s\")\n", + "\n", + "# No prompt text and no hidden state in the file, because the receipt carries\n", + "# every output file in full. The probe owner has the prompts, by prompt_uid.\n", + "os.makedirs(\"outputs\", exist_ok=True)\n", + "with open(\"outputs/probe_results.json\", \"w\") as f:\n", + " json.dump({{\n", + " \"base_model\": base_model_name,\n", + " \"lora_rank\": adapter_config[\"r\"],\n", + " \"probe\": {{\"concept\": probe[\"concept\"], \"layer\": probe[\"layer\"]}},\n", + " \"total_prompts\": len(verdicts),\n", + " \"rows\": verdicts,\n", + " }}, f, indent=2)\n", + "\n", + "\n", + "# What the receipt says about this run. The enclave adds who took part, who\n", + "# approved and where it ran, and signs the lot. Each digest names its scheme.\n", + "def file_sha256(path):\n", + " \"\"\"file-sha256/1: sha256 of the file's bytes.\"\"\"\n", + " with open(path, \"rb\") as f:\n", + " return hashlib.file_digest(f, \"sha256\").hexdigest()\n", + "\n", + "\n", + "def dirhash_sha256(root, files):\n", + " \"\"\"dirhash-sha256/1: sha256 of the canonical JSON of {{relative path: file-sha256}}.\"\"\"\n", + " files = {{str(Path(p).relative_to(root)): file_sha256(p) for p in sorted(files)}}\n", + " canonical = json.dumps(files, sort_keys=True, separators=(\",\", \":\"), ensure_ascii=False)\n", + " return hashlib.sha256(canonical.encode()).hexdigest()\n", + "\n", + "\n", + "def jcs_sha256(value):\n", + " \"\"\"jcs-sha256/1: sha256 of the RFC 8785 canonical JSON.\"\"\"\n", + " return hashlib.sha256(rfc8785.dumps(value)).hexdigest()\n", + "\n", + "\n", + "# How the probe was read. The receipt records these values as they are.\n", + "READOUT = {{\n", + " \"layer\": probe[\"layer\"],\n", + " \"position\": probe[\"position\"],\n", + " # The model runs in bfloat16; the probe is applied to its hidden state\n", + " # cast to float32.\n", + " \"model_dtype\": str(base_model.dtype).removeprefix(\"torch.\"),\n", + " \"readout_dtype\": \"float32\",\n", + " \"threshold\": probe[\"threshold\"],\n", + " \"release_scores\": RELEASE_SCORES,\n", + "}}\n", + "flagged = sum(v[\"flagged\"] for v in verdicts)\n", + "\n", + "# The pipeline that produced the verdicts. Models and config are listed\n", + "# apart; each entry has an id, and `appliesTo` says which models it works on.\n", + "pipeline = {{\n", + " \"models\": [\n", + " {{\n", + " \"id\": \"base\",\n", + " \"role\": \"base\",\n", + " \"name\": base_model_name,\n", + " \"scheme\": \"dirhash-sha256/1\",\n", + " \"digest\": dirhash_sha256(base_dir, [p for p in Path(base_dir).rglob(\"*\") if p.is_file()]),\n", + " \"alsoKnownAs\": [{{\"scheme\": \"hf-revision/1\", \"ref\": f\"{{base_model_name}}@{BASE_MODEL_REVISION}\"}}],\n", + " }},\n", + " # The dataset's own files: the folder also holds syft's files, which the\n", + " # model owner's copy of the adapter does not have.\n", + " {{\n", + " \"id\": \"adapter\",\n", + " \"role\": \"adapter\",\n", + " \"appliesTo\": [\"base\"],\n", + " \"name\": \"dbe_adapter\",\n", + " \"owner\": \"{MODEL_OWNER_EMAIL}\",\n", + " \"scheme\": \"dirhash-sha256/1\",\n", + " \"digest\": dirhash_sha256(adapter_dir, [p for p in adapter_files if p.name != \"syft.pub.yaml\"]),\n", + " }},\n", + " # The probe, by the sha256 of its file. The model owner never sees the\n", + " # probe; this digest ties the verdicts to it.\n", + " {{\n", + " \"id\": \"probe\",\n", + " \"role\": \"probe\",\n", + " \"appliesTo\": [\"base\", \"adapter\"],\n", + " \"name\": \"probe\",\n", + " \"owner\": \"{probe_owner.email}\",\n", + " \"scheme\": \"file-sha256/1\",\n", + " \"digest\": file_sha256(probe_path),\n", + " }},\n", + " ],\n", + " \"config\": [\n", + " {{\n", + " \"id\": \"probe_readout\",\n", + " \"kind\": \"probe_readout\",\n", + " \"appliesTo\": [\"base\", \"adapter\", \"probe\"],\n", + " \"scheme\": \"jcs-sha256/1\",\n", + " \"digest\": jcs_sha256(READOUT),\n", + " \"params\": READOUT,\n", + " }},\n", + " ],\n", + "}}\n", + "if base_modelwrap:\n", + " pipeline[\"models\"][0][\"alsoKnownAs\"].append({{\"scheme\": \"modelwrap/2\", \"ref\": base_modelwrap}})\n", + "claims = {{\n", + " \"subject\": [{{\"name\": f\"{{base_model_name}} + dbe_adapter + probe\",\n", + " \"digest\": {{\"sha256\": jcs_sha256(pipeline)}}}}],\n", + " \"evalPipeline\": pipeline,\n", + " \"evalDataset\": {{\"name\": \"dbe_prompts\", \"owner\": \"{probe_owner.email}\",\n", + " \"scheme\": \"file-sha256/1\", \"digest\": file_sha256(prompt_path)}},\n", + " \"results\": {{\n", + " \"counts\": {{\"submitted\": len(prompt_rows), \"completed\": len(verdicts),\n", + " \"failed\": len(prompt_rows) - len(verdicts)}},\n", + " # n: how many prompts each metric covers. Being flagged is neither good\n", + " # nor bad in itself, so neither metric has a better direction.\n", + " \"metrics\": [\n", + " {{\"name\": \"flagged_prompts\", \"value\": flagged,\n", + " \"unit\": \"count\", \"higherIsBetter\": None, \"n\": len(verdicts)}},\n", + " {{\"name\": \"flagged_rate\", \"value\": flagged / len(verdicts) if verdicts else None,\n", + " \"unit\": \"fraction\", \"higherIsBetter\": None, \"n\": len(verdicts)}},\n", + " ],\n", + " }},\n", + "}}\n", + "with open(\"outputs/receipt_claims.json\", \"w\") as f:\n", + " json.dump(claims, f, indent=2)\n", + "\n", + "print(f\"\\\\nProbe readout complete. {{len(verdicts)}} prompts.\")\n", + "'''\n", + "\n", + "# Fail here rather than inside the enclave if the code is malformed.\n", + "compile(JOB_CODE, \"main.py\", \"exec\")\n", + "\n", + "\n", + "def create_code_file(code: str) -> str:\n", + " tmp = Path(tempfile.mkdtemp()) / f\"job-{random.randint(1, 1_000_000)}\"\n", + " tmp.mkdir(parents=True, exist_ok=True)\n", + " p = tmp / \"main.py\"\n", + " p.write_text(code)\n", + " return str(p)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# The readout in the job is the one notebook 0 trained the probe with and the\n", + "# dry run above ran: probe_readout.py, character for character.\n", + "job_readout = JOB_CODE.split(READOUT_BEGIN, 1)[1].split(READOUT_END, 1)[0]\n", + "assert job_readout.strip(\"\\n\") == PROBE_READOUT_SRC.strip(\"\\n\"), (\n", + " \"The readout in JOB_CODE is not probe_readout.py\"\n", + ")\n", + "# And the probe's weights are not in it: they reach the enclave only as the dataset.\n", + "assert repr(probe[\"weight\"][:4])[1:-1] not in JOB_CODE\n", + "print(f\" The job's readout is probe_readout.py, verbatim ({len(PROBE_READOUT_SRC.splitlines())} lines)\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "probe_owner.submit_python_job(\n", + " ENCLAVE_EMAIL,\n", + " create_code_file(JOB_CODE),\n", + " \"probe_eval_job\",\n", + " datasets={\n", + " MODEL_OWNER_EMAIL: [\"dbe_adapter\"],\n", + " probe_owner.email: [\"dbe_prompts\", \"probe\"],\n", + " },\n", + " share_results_with_do=False,\n", + " # Released only if the job fails, and only once both owners agree.\n", + " request_disclosures=[\"logs\"],\n", + " dependencies=EVAL_DEPS,\n", + ")\n", + "print(f\" Job 'probe_eval_job' submitted to the enclave ({ENCLAVE_EMAIL})\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 6 β€” Approve the job\n", + "\n", + "We approve our own job β€” submitting it is not consent to run it on our prompts and our probe. The\n", + "model owner approves it too, from their notebook. The run starts when the second approval lands;\n", + "until then the status stays `pending`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Wait until the enclave has distributed the job and sent us our approval file.\n", + "job = probe_owner.wait_until_has_job(\n", + " \"probe_eval_job\", where=lambda j: j.can_approve, timeout=900\n", + ")\n", + "print(f\" βœ… Probe owner sees 'probe_eval_job' status={job.status}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "job" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "probe_owner.approve_job(job)\n", + "print(\" Probe owner approved\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\"Wait" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 7 β€” Collect the results\n", + "\n", + "Once the model owner approves, the enclave runs the job and sends the output back to us. It installs\n", + "PyTorch, downloads and hashes the base model, and runs one forward pass per prompt on two CPUs. We\n", + "have not timed this job in an enclave yet. The generation job it is adapted from took under two\n", + "minutes on the same release, and this one does less work; `SETUP.md` has the estimate. The cell\n", + "below waits until the status is `done`. If the job fails, the cell stops with an error: see 'If the\n", + "job failed' below." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from syft_rds.waiting import JobEndedError\n", + "\n", + "try:\n", + " job = probe_owner.wait_until_has_job(\n", + " \"probe_eval_job\", status=\"done\", where=lambda j: bool(j.output_paths), timeout=900\n", + " )\n", + "except JobEndedError:\n", + " print(\" ❌ The job did not finish in the enclave. See 'If the job failed' below to read its logs.\")\n", + " raise\n", + "print(f\" βœ… Output files : {[p.name for p in job.output_paths]}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "result_path = next(p for p in job.output_paths if p.name == \"probe_results.json\")\n", + "probe_result = json.loads(result_path.read_text())\n", + "# The results carry no prompt text: look each prompt up in our own file, by prompt_uid.\n", + "prompt_text = {row[\"prompt_uid\"]: row[\"prompt_text\"] for row in read_prompt_csv(prompt_private)}\n", + "\n", + "print(f\" Base model : {probe_result['base_model']}\")\n", + "print(f\" LoRA rank : {probe_result['lora_rank']}\")\n", + "print(f\" Probe : {probe_result['probe']['concept']!r}, layer {probe_result['probe']['layer']}\")\n", + "print(f\" Prompts : {probe_result['total_prompts']}\")\n", + "print()\n", + "for r in probe_result[\"rows\"]:\n", + " score = f\" score={r['score']:.3f}\" if \"score\" in r else \"\"\n", + " print(f\" prompt_uid : {r['prompt_uid']} ({r['hazard']})\")\n", + " print(f\" prompt : {prompt_text[r['prompt_uid']][:100]}...\")\n", + " print(f\" verdict : {'flagged' if r['flagged'] else 'not flagged'}{score}\")\n", + " print()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Step 7.1 β€” Do the verdicts carry over to the fine-tune?\n", + "\n", + "The dry run in Step 4.1 read the probe on the base model alone, on our machine. The enclave read it\n", + "on base + adapter. Same prompts, same probe, same readout code: the adapter is the only difference.\n", + "So comparing the two, prompt by prompt, asks whether the probe's verdicts carry over to the\n", + "fine-tune, or whether the adapter moved some prompts across the probe's threshold.\n", + "\n", + "The threshold was calibrated on the base model: set in notebook 0 to flag 10% of the training\n", + "negatives there. An adapter that moves the hidden state moves the scores with it, so on the\n", + "fine-tune the same threshold can flag more or fewer prompts than it was set for. When we ran this\n", + "locally with the demo adapter, it moved the scores upward: four of the five by between 0.005 and\n", + "0.15, and the fifth was already at 1.000. At the threshold notebook 0 chose, 0.417, all five were\n", + "flagged on both models, so the verdicts alone showed no change. A verdict that does change may be\n", + "the threshold drifting rather than the concept moving, so compare the scores too if\n", + "`RELEASE_SCORES` is on.\n", + "\n", + "With five prompts this is an illustration of the question, not a result. An answer needs many more\n", + "prompts, and a threshold checked on the fine-tune's own score distribution." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# The dry run's verdicts against the enclave's, prompt by prompt.\n", + "in_enclave = {r[\"prompt_uid\"]: r for r in probe_result[\"rows\"]}\n", + "\n", + "\n", + "def verdict(r: dict) -> str:\n", + " text = \"flagged\" if r[\"flagged\"] else \"not flagged\"\n", + " return f\"{text} ({r['score']:.3f})\" if \"score\" in r else text\n", + "\n", + "\n", + "print(f\" {'prompt_uid':<26}{'hazard':<9}{'base':<22}base + adapter\")\n", + "same = 0\n", + "for base_row in dry_run:\n", + " tuned_row = in_enclave[base_row[\"prompt_uid\"]]\n", + " same += base_row[\"flagged\"] == tuned_row[\"flagged\"]\n", + " print(f\" {base_row['prompt_uid']:<26}{base_row['hazard']:<9}{verdict(base_row):<22}{verdict(tuned_row)}\")\n", + "print()\n", + "print(f\" {same} of {len(dry_run)} verdicts are the same on the fine-tune as on the base model.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 8 β€” Check the receipt and put it on Rekor\n", + "\n", + "If the enclave runs a release with receipts on, the outputs include `receipt.dsse.json`. The job\n", + "wrote part of it: the base model, the adapter and the probe, each by digest; how it read the probe;\n", + "which prompts; and how many it flagged. The enclave wrote the rest: the code that ran, a hash of\n", + "every dataset file, who took part, who approved and when, who gets the results, and where and when\n", + "it ran. The receipt also carries every output file in full, so `probe_results.json` is in it: the\n", + "verdicts, and nothing else.\n", + "The enclave signed it with a key Tinfoil made inside the enclave at boot. The receipt carries a\n", + "Tinfoil report that lists that key, so we check the report names the key and the key signed the\n", + "receipt. We also check that the probe it names is our `probe.json`.\n", + "\n", + "Then we log it on [Rekor](https://docs.sigstore.dev/logging/overview/), Sigstore's public log. We\n", + "send it, but it is signed by the enclave, not by us. Rekor keeps the hash and the signature, not the\n", + "receipt, so nothing private becomes public. Anyone we later show the receipt to can find it in the\n", + "log and see when it was signed." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "outputs = {p.name: p for p in job.output_paths}\n", + "receipt_path = outputs.get(RECEIPT_FILE_NAME)\n", + "\n", + "if \"receipt_error.txt\" in outputs:\n", + " print(\" ❌ The enclave could not sign a receipt:\")\n", + " print(outputs[\"receipt_error.txt\"].read_text())\n", + "elif receipt_path is None:\n", + " print(\" 🟠 No receipt: this enclave runs a release with receipts off.\")\n", + "else:\n", + " envelope = json.loads(receipt_path.read_text())\n", + " receipt = verify_receipt(envelope)\n", + " predicate = receipt[\"predicate\"]\n", + " execution = predicate[\"execution\"]\n", + " print(\" βœ… Signed by the key the enclave's Tinfoil report lists\")\n", + " models = predicate[\"evalPipeline\"][\"models\"]\n", + " probe_entry = next(m for m in models if m[\"id\"] == \"probe\")\n", + " assert probe_entry[\"digest\"] == hashlib.sha256(PROBE_PRIVATE.read_bytes()).hexdigest()\n", + " print(\" βœ… The probe it names is our probe.json\")\n", + " print(f\" Subject : {[s['name'] for s in receipt['subject']]}\")\n", + " for m in models:\n", + " print(f\" Model : {m['role']:<8} {m['name']} ({m['scheme']} {m['digest'][:12]}…)\")\n", + " for c in predicate[\"evalPipeline\"][\"config\"]:\n", + " print(f\" Config : {c['id']} β†’ {c['appliesTo']} {c['params']}\")\n", + " print(f\" Eval set : {predicate['evalDataset']['name']} ({predicate['evalDataset']['owner']})\")\n", + " print(f\" Counts : {predicate['results']['counts']}\")\n", + " for m in predicate[\"results\"][\"metrics\"]:\n", + " value = f\"{m['value']:.3f}\" if isinstance(m[\"value\"], float) else m[\"value\"]\n", + " # A probe's flags have no better direction, so the job writes null: n/a.\n", + " better = {True: \"higher\", False: \"lower\"}.get(m.get(\"higherIsBetter\"), \"n/a\")\n", + " print(f\" Metric : {m['name']} = {value} {m['unit']} (n={m['n']}, better: {better})\")\n", + " for p in predicate[\"parties\"]:\n", + " print(f\" Party : {p['role']:<10} {p['email']}\")\n", + " for a in predicate[\"consent\"][\"approvals\"]:\n", + " print(f\" Approved by : {a['party']} at {a['approvedAt']}\")\n", + " for p in predicate[\"policy\"][\"outputPolicy\"]:\n", + " print(f\" Sent to : {p['party']} {p['grants']}\")\n", + " print(f\" Code files : {[c['path'] for c in predicate['job']['code']]}\")\n", + " print(f\" Platform : {execution['platform']} (image {(execution.get('runtimeImage') or {}).get('digest')})\")\n", + " print(f\" Ran : {execution['startedAt']} β†’ {execution['finishedAt']}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "if receipt_path is None:\n", + " print(\" 🟠 No receipt to log.\")\n", + "else:\n", + " entry = upload_to_rekor(envelope)\n", + " print(f\" βœ… Logged on Rekor at index {entry['logIndex']}\")\n", + " print(f\" {entry['url']}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## If the job failed β€” read its logs\n", + "\n", + "The enclave keeps the job's logs until both owners release them. Run the cell below, and ask the\n", + "model owner to run the matching cell in their notebook. Once both releases reach the enclave, it\n", + "sends the logs to both of us. The job prints only the base model and the adapter rank, both public\n", + "already, the number of prompts, and how long the readout took. It prints no verdict, so the logs\n", + "give the model owner no result. An exception message can quote data, though, so read an error\n", + "before you share it further." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "probe_owner.update_disclosures(job, [\"logs\"])\n", + "print(\" Released the logs. The model owner has to release them too.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "job = next(j for j in probe_owner.jobs if j.name == \"probe_eval_job\")\n", + "print(job.stderr)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "\n", + "The model owner approved this run and cannot read a line of what you just printed. You never saw\n", + "their weights, and they never saw your probe. The enclave is gone β€” Tinfoil Containers keep no\n", + "disk, so its state went with it. What remains is the receipt, and its entry on Rekor." + ] + } + ], + "metadata": { + "colab": { + "provenance": [] + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/notebooks/enclave/probe-eval-w-receipts/2. DO-model-owner-probe.ipynb b/notebooks/enclave/probe-eval-w-receipts/2. DO-model-owner-probe.ipynb new file mode 100644 index 00000000000..781076beef3 --- /dev/null +++ b/notebooks/enclave/probe-eval-w-receipts/2. DO-model-owner-probe.ipynb @@ -0,0 +1,552 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Probe evaluation β€” Model Owner\n", + "\n", + "| Actor | Email | Role |\n", + "|-------|-------|------|\n", + "| **Enclave** | `ENCLAVE_EMAIL` | Tinfoil Container that runs the evaluation |\n", + "| **Probe owner** | `BENCHMARK_OWNER_EMAIL` | Owns the private linear probe and the prompts it is read on. The enclave's release pins the benchmark-owner account as a data owner, so that is the account the probe owner uses |\n", + "| **Model owner** | `MODEL_OWNER_EMAIL` | Owns the private LoRA adapter |\n", + "\n", + "**The setting.** A probe owner has a linear probe: a direction in a model's hidden states that\n", + "tracks a concept. They want to know what it reads on a fine-tune they cannot see. A model owner is\n", + "willing to be measured but will not hand over the weights. Neither will show the other what they\n", + "hold, and neither trusts the machine in the middle.\n", + "\n", + "Both connect to an [enclave](https://github.com/OpenMined/PySyft/tree/dev/packages/syft-enclave) β€” a\n", + "sealed environment whose contents nobody, including whoever runs the hardware, can read. Each\n", + "uploads their half. Nothing runs until **both** owners approve the exact code, and the enclave\n", + "proves what it is before either of them uploads anything.\n", + "\n", + "This is a variant of the flow from\n", + "[tinfoilsh/double-blind-eval](https://github.com/tinfoilsh/double-blind-eval), carried out entirely\n", + "over PySyft, with a probe in place of generation. The enclave is a [Tinfoil\n", + "Container](https://docs.tinfoil.sh/containers/overview), and the attestation is checked against the\n", + "measurement published for a signed release of its config.\n", + "\n", + "We hold the model. We approve the evaluation, and we do not get to see its results β€” only the\n", + "probe owner does. That asymmetry is the point: we agreed to be measured, not to learn what the\n", + "measurement was made of." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Setup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!uv pip install -Uq \"syft-enclave[tinfoil] @ git+https://github.com/OpenMined/PySyft.git@dev#subdirectory=packages/syft-enclave\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import random\n", + "import tempfile\n", + "from pathlib import Path\n", + "\n", + "from syft_enclaves import login_do\n", + "from syft_enclaves.attestation.tinfoil import TinfoilAppraisalPolicy" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ENCLAVE_EMAIL = \"enclave@openmined.org\"\n", + "BENCHMARK_OWNER_EMAIL = \"benchmark_owner@openmined.org\"\n", + "MODEL_OWNER_EMAIL = \"model_owner@openmined.org\"\n", + "\n", + "# What the enclave is checked against. TINFOIL_TAG and IMAGE_DIGEST come from\n", + "# whoever deployed it: `just tinfoil-release` prints the digest, and the tag is\n", + "# the config release it was deployed from.\n", + "TINFOIL_REPO = \"OpenMined/syft-enclave-tinfoil\"\n", + "TINFOIL_TAG = \"v0.1.28\"\n", + "IMAGE_DIGEST = \"sha256:ea890c9b82dbf3c80c703d717dabe2c3db0c48b53bd269801984807732446864\"\n", + "\n", + "print(f\" Enclave: {ENCLAVE_EMAIL}\")\n", + "print(f\" Benchmark owner: {BENCHMARK_OWNER_EMAIL} | Model owner: {MODEL_OWNER_EMAIL}\")\n", + "print(f\" Verifying against {TINFOIL_REPO} {TINFOIL_TAG}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 0 β€” Log in as Model Owner\n", + "\n", + "Colab authenticates us to Google Drive, which is how every party in this demo reaches every other\n", + "one. Nothing here talks to the enclave yet." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "CRYPTO_KEYS_PATH = None\n", + "# NOTE: Re-storing Private Key from local_file, if you have it, uncomment the following line and upload the file.\n", + "# from google.colab import files; files.upload()\n", + "# CRYPTO_KEYS_PATH = \"crypto_keys.json\"\n", + "\n", + "model_owner = login_do(encryption=True, crypto_keys_path=CRYPTO_KEYS_PATH)\n", + "print(f\" Model owner : {model_owner.email}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# # Optionally, to start from a clean SyftBox\n", + "# model_owner.delete_syftbox()\n", + "# model_owner._rds.peer_manager.write_own_version()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 1 β€” Connect to the enclave and to the probe owner\n", + "\n", + "We ask the enclave to peer with us, which is what lets us send it our weights. The probe owner sends\n", + "us a peer request from their notebook; approving it lets them read our public model card, and\n", + "nothing else." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model_owner.add_peer(ENCLAVE_EMAIL)\n", + "print(f\" Peer request sent to the enclave ({ENCLAVE_EMAIL})\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Step 1.1 β€” Approve the probe owner\n", + "\n", + "The cell below shows our peers. `peer_must_exist=False` lets us approve the probe owner, who logs in\n", + "as `BENCHMARK_OWNER_EMAIL`, before their request reaches us, and the cell after it waits until the\n", + "connection is live." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model_owner.peers" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model_owner.approve_peer_request(BENCHMARK_OWNER_EMAIL, peer_must_exist=False)\n", + "print(\" Benchmark owner approved\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Wait up to 5 minutes for each peer to approve. Fail at once on a wrong email.\n", + "# Warn if the peer was already there before this notebook logged in.\n", + "for peer_email in [BENCHMARK_OWNER_EMAIL, ENCLAVE_EMAIL]:\n", + " model_owner.wait_until_peered(peer_email)\n", + " print(f\" Connected to {peer_email}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2 β€” Attest the enclave\n", + "\n", + "Before uploading anything, we check what we are uploading it to. The enclave holds a key that never\n", + "leaves the hardware, and it serves a report signed under that key over a connection pinned to it. We\n", + "compare that report against the measurement published for `TINFOIL_TAG` β€” a signed GitHub release of\n", + "the enclave's config β€” and against the facts we expect: the container image digest, which datasite\n", + "the enclave runs as, and which two owners have to approve a job.\n", + "\n", + "Every one of those has to be pinned. Unpinned, the attestation would prove a genuine enclave booted\n", + "a signed config, but not that the config was the one we reviewed. Tinfoil is not in the trust path,\n", + "and neither are we to each other." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "policy = TinfoilAppraisalPolicy(\n", + " repo=TINFOIL_REPO,\n", + " release_tag=TINFOIL_TAG,\n", + " expected_image_digest=IMAGE_DIGEST,\n", + " expected_data_owners=[BENCHMARK_OWNER_EMAIL, MODEL_OWNER_EMAIL],\n", + " expected_email=ENCLAVE_EMAIL,\n", + " # The measured config pins no SYFT_VERSION, so there is nothing to compare.\n", + " expected_syft_version=None,\n", + ")\n", + "\n", + "# Wait until the enclave publishes its attestation, then verify it.\n", + "result = model_owner.wait_until_attested(ENCLAVE_EMAIL, policy=policy)\n", + "result.print_checklist()\n", + "print()\n", + "print(f\" Bound to the report: {result.verified_key_bundle is not None}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3 β€” Fetch the private adapter\n", + "\n", + "Our model is a [PEFT](https://huggingface.co/docs/peft) LoRA adapter: a small set of weights that\n", + "adjusts a public base model. In a real run these would be weights we trained and never published.\n", + "Here we borrow a public adapter to stand in for one, the same way the double-blind-eval demo does." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# The enclave's current shape β€” 2 CPUs, 16 GB, no GPU β€” runs a small model in\n", + "# software, so a public LoRA on TinyLlama-1.1B stands in for a private one. A\n", + "# larger base model needs a GPU-shaped enclave (`gpus:` in tinfoil-config.yml,\n", + "# which means a new measured config release), and a probe trained on it.\n", + "ADAPTER_REPO = \"zjudai/flowertune-general-nlp-lora-tinyllama-1.1b-chat-v1.0\"\n", + "ADAPTER_REV = \"main\"\n", + "ADAPTER_DIR = Path(\"dbe-adapter\")\n", + "\n", + "print(f\" Adapter: {ADAPTER_REPO}@{ADAPTER_REV}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ADAPTER_DIR.mkdir(parents=True, exist_ok=True)\n", + "\n", + "# A PEFT LoRA directory is two files: the config and the weights.\n", + "for name in (\"adapter_config.json\", \"adapter_model.safetensors\"):\n", + " url = f\"https://huggingface.co/{ADAPTER_REPO}/resolve/{ADAPTER_REV}/{name}\"\n", + " !curl -sfL \"{url}\" -o \"{ADAPTER_DIR}/{name}\"\n", + " print(f\" fetched {name}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4 β€” Upload the model\n", + "\n", + "We upload the adapter into the enclave. The **private** side is the weights, shared only with the\n", + "enclave. The **mock** side is a model card: the base model, the rank, the size. That is all the\n", + "probe owner ever sees of our model.\n", + "\n", + "Uploading grants nothing on its own. The weights are used once, by a job we approve, and they are\n", + "gone when the enclave restarts β€” Tinfoil Containers have no persistent disk." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "adapter_config = json.loads((ADAPTER_DIR / \"adapter_config.json\").read_text())\n", + "BASE_MODEL = adapter_config[\"base_model_name_or_path\"]\n", + "LORA_RANK = adapter_config[\"r\"]\n", + "\n", + "\n", + "def create_model_card() -> Path:\n", + " \"\"\"The public half: what the benchmark owner may know about our model.\"\"\"\n", + " tmp = Path(tempfile.mkdtemp()) / f\"model-card-{random.randint(1, 1_000_000)}\"\n", + " tmp.mkdir(parents=True, exist_ok=True)\n", + " path = tmp / \"model_card.txt\"\n", + " weights = ADAPTER_DIR / \"adapter_model.safetensors\"\n", + " path.write_text(\"\\n\".join([\n", + " \"Private LoRA adapter\",\n", + " \"====================\",\n", + " f\"Base model : {BASE_MODEL}\",\n", + " f\"LoRA rank : {LORA_RANK} (alpha {adapter_config['lora_alpha']})\",\n", + " f\"Target modules : {', '.join(sorted(adapter_config['target_modules']))}\",\n", + " f\"Size : {weights.stat().st_size / 1e6:.1f} MB\",\n", + " \"\",\n", + " \"The weights are not shared. They go to the enclave and nowhere else.\",\n", + " \"\",\n", + " ]))\n", + " return path\n", + "\n", + "\n", + "model_card = create_model_card()\n", + "print(model_card.read_text())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model_owner.create_dataset(\n", + " name=\"dbe_adapter\",\n", + " mock_path=model_card,\n", + " private_path=ADAPTER_DIR,\n", + " summary=f\"Private LoRA adapter for {BASE_MODEL} β€” rank {LORA_RANK}\",\n", + " users=[BENCHMARK_OWNER_EMAIL, ENCLAVE_EMAIL],\n", + ")\n", + "print(\" Created 'dbe_adapter'\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Send the weights to the enclave, and only to the enclave.\n", + "model_owner.share_private_dataset(\"dbe_adapter\", ENCLAVE_EMAIL)\n", + "print(\" Private adapter shared with the enclave\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\"Wait" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 5 β€” Wait for the evaluation job\n", + "\n", + "The probe owner submits the job that reads our model with their probe, on their prompts. The\n", + "enclave distributes it to both of us for approval. The cell below waits until it is here and we can\n", + "approve it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Wait until the enclave has distributed the job and sent us our approval file.\n", + "job = model_owner.wait_until_has_job(\n", + " \"probe_eval_job\", where=lambda j: j.can_approve, timeout=900\n", + ")\n", + "print(f\" βœ… Model owner sees 'probe_eval_job' status={job.status}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Step 5.1 β€” Read it\n", + "\n", + "This is the code that will touch our weights, so read it properly. Four things are worth checking:\n", + "\n", + "- **The adapter is used only for forward passes.** The job reads it through\n", + " `resolve_dataset_files_path`, loads it with `PeftModel.from_pretrained`, and hands it to\n", + " `probe_readout`, which runs one forward pass per prompt under `torch.no_grad()`. No training, no\n", + " generation, no gradients.\n", + "- **No hidden states, activations or weights are written to `outputs/`.** The hidden state at the\n", + " probe's layer stays in memory, one prompt at a time, and becomes a single number. The job writes\n", + " two files: `probe_results.json`, with one row per prompt, and `receipt_claims.json`, with digests,\n", + " the readout's settings and counts. Neither holds a tensor.\n", + "- **The outputs are verdicts.** Each row is a `prompt_uid`, a hazard code and `flagged`, true or\n", + " false. A row carries the probe's `score` only when `RELEASE_SCORES` is `True`, so check its value\n", + " in the job code.\n", + "- **The only network access is two downloads.** `snapshot_download` fetches the base model at its\n", + " pinned commit, and `install_modelwrap` fetches the pinned tools that hash it, each checked against\n", + " its sha256. Nothing else in the job opens a connection, so nothing leaves the enclave but the\n", + " output files, and those go to the probe owner.\n", + "\n", + "The probe itself is the probe owner's private dataset, `probe`, and we never see it. The job reads\n", + "its weights from that dataset inside the enclave; they are not in the code. Its mock half, which\n", + "Step 5.2 shows, says what the probe measures. The receipt names the probe by the sha256 of its file,\n", + "so the verdicts are tied to one probe, and it cannot be swapped afterwards.\n", + "\n", + "The code between the two `probe readout` lines is what the probe owner trained the probe with, on\n", + "the open base model. The same code runs here, so the only new thing it meets in the enclave is our\n", + "adapter.\n", + "\n", + "`share_results_with_do=False` in the submission is what routes the output to the probe owner alone.\n", + "We are approving a run we will not see the result of.\n", + "\n", + "The submission also requests the job's logs. Approving does not release them: they stay in the\n", + "enclave unless both of us release them, which we only do if the job fails." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "job" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Step 5.2 β€” Read the probe's card\n", + "\n", + "The probe's weights go to the enclave only, but its mock half is ours to read, the way our model card\n", + "is the probe owner's. It says what the probe measures: the concept it was trained for, the base model\n", + "and commit it was trained on, which layer and token it reads, and the threshold for a flag. That is\n", + "what we are consenting to have measured on our model. The mock's own weights are seeded random\n", + "numbers, there only so the file has the probe's shape." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Wait until the probe owner's probe card is here.\n", + "probe_dataset = model_owner.wait_until_has_dataset(\"probe\", datasite=BENCHMARK_OWNER_EMAIL)\n", + "probe_card = json.loads(probe_dataset.mock_files[0].read_text())\n", + "\n", + "print(f\" Concept : {probe_card['concept']}\")\n", + "print(f\" Trained on : {probe_card['base_model']}@{probe_card['base_model_revision'][:12]}\")\n", + "print(f\" Layer : {probe_card['layer']} ({probe_card.get('layer_convention', 'an index into hidden_states')})\")\n", + "print(f\" Position : {probe_card['position']}\")\n", + "print(f\" Threshold : {probe_card['threshold']:.6g} (flagged when sigmoid(hidden @ weight + bias) >= threshold)\")\n", + "print(f\" Weights : withheld. The {len(probe_card['weight'])} in the mock are placeholders (mock={probe_card['mock']})\")\n", + "print(f\" Note : {probe_card.get('note', '')}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 6 β€” Approve the job\n", + "\n", + "Our approval is one of two. The probe owner approves from their notebook, and the run starts the\n", + "moment the second approval lands. Either order works; until both are in, the status stays `pending`\n", + "and nothing has touched our weights." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model_owner.approve_job(job)\n", + "print(\" Model owner approved\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 7 β€” We do not get the result\n", + "\n", + "The enclave runs the job and sends the output to whoever submitted it. That is the probe owner, so\n", + "our copy of the job has no outputs and never reaches `done` β€” we approved the run and learned\n", + "nothing from it.\n", + "\n", + "That is what we agreed to. The enclave proved what it was running before we uploaded anything, and\n", + "it ran the code we both read, so the probe owner's verdicts came from our adapter, read by the probe\n", + "the receipt names, and nothing else. Neither of us learned the other's half." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "job = next(j for j in model_owner.jobs if j.name == \"probe_eval_job\")\n", + "print(f\" Job status : {job.status}\")\n", + "print(f\" Output files : {job.output_paths}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## If the job failed β€” release its logs\n", + "\n", + "Only if the probe owner reports that the job failed. The enclave sends the logs once both owners\n", + "release them, to both of us. The job prints only the base model and the adapter rank, both public\n", + "already in our model card, the number of prompts, and how long the readout took. It prints no\n", + "verdict and nothing from the probe. An exception message can quote data, though, so read the error\n", + "before you share it further." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model_owner.update_disclosures(job, [\"logs\"])\n", + "print(\" Released the logs. The benchmark owner has to release them too.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "\n", + "The probe owner now has their probe's verdicts on your model, for prompts you never saw, from a\n", + "probe you never saw. They never saw your weights. Nobody, including whoever runs the hardware, saw\n", + "both halves." + ] + } + ], + "metadata": { + "colab": { + "provenance": [] + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/notebooks/enclave/probe-eval-w-receipts/SETUP.md b/notebooks/enclave/probe-eval-w-receipts/SETUP.md new file mode 100644 index 00000000000..430e31a6517 --- /dev/null +++ b/notebooks/enclave/probe-eval-w-receipts/SETUP.md @@ -0,0 +1,217 @@ +# Running the probe evaluation demo + +A probe owner trains a linear probe on an open base model, then applies it inside a Tinfoil enclave +to that base model plus a model owner's private LoRA adapter. Both approve the job, the enclave runs +it, and only the probe owner reads the results: one verdict per prompt, and a signed receipt. The +probe owner never sees the adapter, and the model owner never sees the probe or the prompts. + +This is a variant of [`double-blind-eval-w-receipts`](../double-blind-eval-w-receipts/SETUP.md), +which carries out the [tinfoilsh/double-blind-eval](https://github.com/tinfoilsh/double-blind-eval) +flow over PySyft: the same enclave, release and flow, with a job that reads the model's hidden state +and applies a probe instead of generating completions and sending them to a safety classifier. The +notebooks in this folder are the two parties; everything else on this page is the operator's job, +done once before either party opens one. + +| Notebook | Run by | Where | +| ------------------------------------ | ----------- | ----------------------------------------------- | +| `0. probe-owner-train-on-base.ipynb` | probe owner | locally or in Colab, with no account | +| `1. DO-probe-owner.ipynb` | probe owner | Colab, logged in as the benchmark-owner account | +| `2. DO-model-owner-probe.ipynb` | model owner | Colab, logged in as the model-owner account | + +The probe owner uses the benchmark-owner account, `benchmark_owner@openmined.org`. The release below +pins that account and `model_owner@openmined.org` as the enclave's data owners, and the pin is +measured, so this variant runs on the release unchanged only by reusing the account. A probe owner +with their own email needs a config with that email, and so a new release (step 6). + +## What you need + +Three Google accounts: one for the enclave, one per party. The enclave reaches Drive with a token +you register as a Tinfoil secret, and each party authenticates through Colab. + +No inference key. The job makes no call outside the enclave except two downloads: the base model +at its pinned commit, and the pinned tools that hash it. + +You also need the `tinfoil` CLI and an admin API key from the Tinfoil dashboard under **Settings β†’ +API Keys β†’ Admin keys**. Log in once with `tinfoil login --api-key admin_...` and check it with +`just tinfoil-whoami`. Deploying needs that key; publishing a release does not. + +Run every `just` command below from `packages/syft-enclave/`. + +## 1. Reset both parties' state + +Wipe each data owner's SyftBox before you deploy, not after. The enclave caches its peers' Drive +folders when it boots, so clearing them afterwards leaves it peered with folders that no longer +exist. + +```bash +just delete-syftbox benchmark_owner@openmined.org ../../credentials/token_benchmark_owner.json +just delete-syftbox model_owner@openmined.org ../../credentials/token_model_owner.json +``` + +Without a token file to hand, notebooks 1 and 2 each have a commented cell under Step 0 that calls +`delete_syftbox()` from Colab instead. + +The enclave needs no reset: Tinfoil Containers have no persistent disk, and the enclave wipes its +own state on every boot. + +## 2. Deploy the published release + +Two releases of [`OpenMined/syft-enclave-tinfoil`](https://github.com/OpenMined/syft-enclave-tinfoil) +are already published. They run the same image, and differ only in whether each job's outputs carry a +signed receipt: + +| Release | Config | Receipts | +| --------- | ----------------------------- | -------- | +| `v0.1.27` | `tinfoil-config.yml` | off | +| `v0.1.28` | `tinfoil-config-receipts.yml` | on | + +Deploy `v0.1.28` as it stands: the probe owner's notebook ends by checking the receipt and logging it +on Rekor. On `v0.1.27` the job runs the same, but there is no receipt to check. + +```bash +just tinfoil-deploy v0.1.28 enclave@openmined.org +``` + +Both configs pin `benchmark_owner@openmined.org` and `model_owner@openmined.org` as +`SYFT_ENCLAVE_DATA_OWNERS`, which is what makes a job wait for two approvals instead of one. The +value is measured, so each party's attestation checks it against the release. + +## 3. Check the enclave is attesting + +```bash +just tinfoil-attest syft-enclave.openmined.containers.tinfoil.dev +``` + +If it fails, run `just tinfoil-why` first β€” the control plane's `error_message` has named the cause +in every failure so far. [`docs/tinfoil_troubleshooting.md`](../../../packages/syft-enclave/docs/tinfoil_troubleshooting.md) +covers the rest. + +## 4. Train the probe + +The probe owner runs [`0. probe-owner-train-on-base.ipynb`](0.%20probe-owner-train-on-base.ipynb) +first, on their own machine or in Colab. It needs no account and never touches the enclave. It +downloads the [MLCommons AILuminate](https://github.com/mlcommons/ailuminate) demo prompt set, +leaves out the ten prompts the evaluation uses, reads the open base model's hidden state on a sample +of the rest, and fits a logistic-regression probe. It picks the layer and the threshold on its +training split alone: the layer with the best cross-validated AUROC, and the score that flags 10% of +the training negatives out of fold. It writes three files next to itself: + +- `probe.json` β€” the probe. This is the private half: it goes to the enclave and nowhere else. The + folder's `.gitignore` keeps it out of git. +- `probe_mock.json` β€” the same fields with random weights, which is all the model owner sees of the + probe. +- `probe_readout.py` β€” the code that reads the hidden state and applies the probe. Notebook 1 runs + it for a dry run on the base model and puts it in the job verbatim. + +Notebook 1 reads all three from the folder it runs in. Run both notebooks in the same folder and +there is nothing to do; in Colab, each notebook has a commented cell to download the files from +notebook 0 and upload them to notebook 1. + +Notebook 0 downloads from pinned commits and checks every file's sha256 before training on it; +notebook 1 downloads the evaluation prompts from the same bench commit. + +| File | Source | sha256 | +| ----------------------------------------------------- | ------------------------------------------------ | ------------------------------------------------------------------ | +| `airr_official_1.0_demo_en_us_prompt_set_release.csv` | `mlcommons/ailuminate` at `769cc2be` | `63e2b654325c5405096fabf28ef5f057b0ad674904b1839fd1f52b523e324d02` | +| `bench/dbe_prompts.csv` | `OpenMined/double-blind-eval-bench` at `f0050a5` | `3e1abebd6dbf673bf11027e4c0aec0fce9fa55a5c82cabba715ae3c90b9c6d38` | +| `bench/dbe_prompts_mock.csv` | `OpenMined/double-blind-eval-bench` at `f0050a5` | `5c4c831ef7c614ae8dba9940858833e0d23c37ef0f358b6e1ea60d4a24bf5cee` | + +The full commits are `769cc2be9d20c8d4fb26ce53b68865ed41dfb8e2` and +`f0050a5e0b1a361d525139be5ce53a503483648c`, the `main` of each repo on 2026-10-07. + +The concept the probe tracks is a placeholder: a _specialised-advice request_, with the AILuminate +`spc_*` hazards as positives. A real probe owner brings their own concept and prompts by changing +`CONCEPT`, `POSITIVE_HAZARDS` and the download in notebook 0. + +Notebook 0 runs the base model in bfloat16 on CPU, as the enclave does: one forward pass per +prompt. The demo set has 100 `spc_*` prompts, 97 once the evaluation prompts are left out, so the +default sample is those 97 and three negatives for each: 388 prompts. On an Apple-silicon laptop +they took about 13 minutes, 2 seconds each; Colab's two CPUs are slower. The notebook times the +first five prompts and prints an estimate before starting the rest. Set `MAX_PROMPTS` for a quicker +run on a stratified subsample. The dry run in notebook 1 loads the same 2.2 GB base model and took +about 14 seconds for its five prompts. + +## 5. Run the two party notebooks + +Open both in Colab, one per account: + +- [`1. DO-probe-owner.ipynb`](1.%20DO-probe-owner.ipynb) β€” uploads the prompts and the probe, dry + runs the probe on the base model, submits the job, reads the verdicts, compares them with the dry + run, and logs the receipt on Rekor +- [`2. DO-model-owner-probe.ipynb`](2.%20DO-model-owner-probe.ipynb) β€” uploads the adapter, reviews + the job and the probe's public card (concept, layer, position, threshold; no weights), approves, + sees no results + +Set the same constants in both, in the cell under **Setup**: + +```python +ENCLAVE_EMAIL = "enclave@openmined.org" +BENCHMARK_OWNER_EMAIL = "benchmark_owner@openmined.org" +MODEL_OWNER_EMAIL = "model_owner@openmined.org" +TINFOIL_REPO = "OpenMined/syft-enclave-tinfoil" +TINFOIL_TAG = "v0.1.28" +IMAGE_DIGEST = "sha256:ea890c9b82dbf3c80c703d717dabe2c3db0c48b53bd269801984807732446864" +``` + +These are the values of the double-blind-eval notebooks, unchanged: `BENCHMARK_OWNER_EMAIL` is the +probe owner's account. `TINFOIL_TAG` and `IMAGE_DIGEST` are what the parties check the enclave +against, so they must match the release you deployed. The digest above is the one both `v0.1.27` +and `v0.1.28` pin; after a republish, use the one `just tinfoil-build` printed. + +Then run both notebooks top to bottom. They wait on each other four times, and a card in the +notebook says so each time. Cells that wait print a 🟠 line and tell you to re-run them. + +### Runtime, and the enclave's limits + +The enclave has 2 CPUs, 16 GB of RAM shared with its ramdisk, and no GPU, and it kills a job at 600 +seconds. The job is built for that: it installs the CPU-only PyTorch wheel, loads the base model in +bfloat16, and runs one forward pass per prompt, with no generation. Its dependency list is the +double-blind-eval job's, unchanged. + +We have not yet timed this job in an enclave. The double-blind-eval job on `v0.1.28` took 102 +seconds from start to finish, by its receipt's timestamps. About 58 of those were generation. The +other 44 installed PyTorch, downloaded, hashed and loaded the base model, and called the safety +classifier. This job keeps all of that but the classifier, and replaces generation with one forward +pass per prompt, which costs about what the time to first token measured there: 20 seconds for the +five private prompts (on the laptop above, the job's readout of the same five took about 15). So we +expect about 65 seconds, well inside the 600-second limit. + +## 6. Republish, after changing the image or config + +This variant needs none of this on `v0.1.28`. Everything in `tinfoil/tinfoil-config.yml` is +measured, so any edit to it β€” or to the enclave image β€” needs a new release before it can be +deployed. + +```bash +just tinfoil-build v0.1.23 # build, push, pin the digest in both configs +just tinfoil-release v0.1.23 # receipts off +just tinfoil-release v0.1.24 tinfoil/tinfoil-config-receipts.yml # receipts on +just tinfoil-deploy v0.1.24 enclave@openmined.org +``` + +Each `tinfoil-release` opens a pull request on the config repo and waits until you merge it. Then it +publishes, which takes about a minute to compute the measurement. Keep the digest `tinfoil-build` +printed, and put it in both notebooks along with the tag you deployed. + +A release is a signed GitHub release and a transparency-log entry, so it cannot be unpublished. +Number versions with that in mind. To go back to an earlier one without moving "latest": + +```bash +just tinfoil-relaunch v0.1.21 --promote-release=false +``` + +## A larger base model + +The demo runs TinyLlama-1.1B on CPU because `tinfoil-config.yml` sets `gpus: 0`. A larger base model +needs a GPU, which means raising `gpus:` in the config, a config change, so step 6. The probe has to +be trained again on that base model, in notebook 0, and its commit pinned as `BASE_MODEL_REVISION` +in notebooks 0 and 1: the job refuses a probe trained on a different base model or commit. + +## Further reading + +- [`docs/tinfoil_deployment.md`](../../../packages/syft-enclave/docs/tinfoil_deployment.md) β€” the + deployment mechanics in full +- [`docs/security.md`](../../../packages/syft-enclave/docs/security.md) Β§6 β€” what the attestation + proves +- [`OpenMined/double-blind-eval-bench`](https://github.com/OpenMined/double-blind-eval-bench) β€” the + prompt set the probe owner downloads