Update app.py

app.py

@@ -8,28 +8,30 @@ from huggingface_hub import hf_hub_download
 
 app = Flask(__name__)
 
-# Cache per EE model name so repeated requests don't re-download
 _cache = {}
 
 
-def get_sigma(hidden_size: int, seed: int):
-    """Derive the hidden-dimension permutation from the secret seed."""
+def get_sigma(hidden_size: int, seed: int) -> np.ndarray:
+    """
+    Derive the encryption permutation from the secret seed.
+    This is the CLIENT'S secret key — it never leaves this Space.
+    The server only ever sees embeddings already scrambled by sigma.
+    """
     rng = np.random.default_rng(seed)
     return rng.permutation(hidden_size)
 
 
 def load_client_components(ee_model_name: str):
     """
-    Load and cache everything the client needs:
+    Load and cache:
       - ee_config  → hidden_size + original model name
-      - tokenizer  → from the EE model
-      - embed_layer → from the ORIGINAL (unmodified) model
-
-    Why we need the original embed layer:
-      The EE model's weights were permuted with sigma, but its embedding table was
-      NOT permuted (it maps token IDs → plain vectors). The client must embed with
-      the original model and then apply sigma to produce the scrambled vectors the
-      EE model expects.
+      - tokenizer  → from EE model
+      - embed_layer → from the ORIGINAL (untransformed) model
+
+    The original embed_layer is used to produce plain vectors from token IDs.
+    The client then applies sigma to those plain vectors before sending.
+    The server's EE model has weights permuted with sigma_inv, so:
+        EE_model(sigma(plain_embed(tokens))) == original_model(plain_embed(tokens))
     """
     if ee_model_name in _cache:
         return _cache[ee_model_name]
@@ -43,16 +45,16 @@ def load_client_components(ee_model_name: str):
 
     tokenizer = AutoTokenizer.from_pretrained(ee_model_name, trust_remote_code=True)
 
-    # We only need embed_tokens — load the full model then discard everything else
+    # Load ORIGINAL model just for its embed layer — discard everything else
     original_model = AutoModelForCausalLM.from_pretrained(
         original_model_name,
-        torch_dtype=torch.float32,  # float32 for precision on CPU
+        torch_dtype=torch.float32,
         device_map="cpu",
         trust_remote_code=True,
     )
     embed_layer = original_model.model.embed_tokens
     embed_layer.eval()
-    del original_model  # free RAM — we only keep the embed layer
+    del original_model
 
     _cache[ee_model_name] = (tokenizer, embed_layer, hidden_size)
     return tokenizer, embed_layer, hidden_size
@@ -68,57 +70,48 @@ def index():
         form_data = request.form.to_dict()
         server_url    = request.form["server_url"].rstrip("/")
         ee_model_name = request.form["ee_model_name"].strip()
-        ee_seed       = int(request.form["ee_seed"])
+        ee_seed       = int(request.form["ee_seed"])   # SECRET — client only
         prompt        = request.form["prompt"].strip()
         max_tokens    = int(request.form.get("max_tokens", 256))
 
         try:
             tokenizer, embed_layer, hidden_size = load_client_components(ee_model_name)
 
-            # --- Step 1: tokenize ---
+            # --- CLIENT-SIDE ENCRYPTION ---
+
+            # Step 1: tokenize
             inputs = tokenizer(prompt, return_tensors="pt")
-            input_ids = inputs.input_ids  # (1, seq_len)
 
-            # --- Step 2: embed with ORIGINAL model's embed layer ---
+            # Step 2: embed with ORIGINAL model embed layer → plain vectors
             with torch.no_grad():
-                plain_embeds = embed_layer(input_ids)  # (1, seq_len, hidden)
+                plain_embeds = embed_layer(inputs.input_ids)  # (1, seq_len, hidden)
 
-            # --- Step 3: ENCRYPT — permute hidden dim with secret sigma ---
-            # The EE model's weight matrices were pre-permuted with sigma,
-            # so feeding sigma-permuted embeddings is equivalent to feeding
-            # plain embeddings to the original model.
+            # Step 3: apply sigma permutation — this is the encryption
+            # The server NEVER sees plain_embeds, only the scrambled version.
+            # Without knowing the seed, the server cannot recover the original.
             sigma = get_sigma(hidden_size, ee_seed)
-            encrypted_embeds = plain_embeds[..., sigma]  # (1, seq_len, hidden)
-
-            # Match server model dtype (float16)
+            encrypted_embeds = plain_embeds[..., sigma]        # (1, seq_len, hidden)
             encrypted_embeds = encrypted_embeds.to(torch.float16)
 
-            # --- Step 4: send to server ---
+            # --- SEND TO SERVER ---
             payload = {
                 "encrypted_embeds": encrypted_embeds.tolist(),
-                "attention_mask": inputs.attention_mask.tolist(),
-                "max_new_tokens": max_tokens,
+                "attention_mask":   inputs.attention_mask.tolist(),
+                "max_new_tokens":   max_tokens,
             }
 
-            resp = requests.post(
-                f"{server_url}/generate",
-                json=payload,
-                timeout=300,
-            )
+            resp = requests.post(f"{server_url}/generate", json=payload, timeout=300)
 
             if not resp.ok:
-                raise RuntimeError(
-                    f"Server {resp.status_code}: {resp.text[:600]}"
-                )
+                raise RuntimeError(f"Server {resp.status_code}: {resp.text[:600]}")
 
             body = resp.json()
             if "error" in body:
                 raise RuntimeError(f"Server error: {body['error']}\n{body.get('traceback','')}")
 
-            # --- Step 5: decode ---
-            # No decryption needed on the output — the EE model's lm_head was
-            # also permuted so output logits correctly map to the real vocabulary.
-            # We skip special tokens and strip the prompt echo if present.
+            # --- OUTPUT DECODING ---
+            # The EE model's lm_head rows are permuted with sigma_inv, so output
+            # logits correctly index the real vocabulary — decode normally.
             gen_ids = body["generated_ids"]
             result = tokenizer.decode(gen_ids, skip_special_tokens=True)