Update app.py

app.py

@@ -8,22 +8,27 @@ from huggingface_hub import hf_hub_download
 
 app = Flask(__name__)
 
-# Cache tokenizer/embed layer so repeated requests don't reload from scratch
+# Cache tokenizer + embed layer so repeated requests don't reload
 _cache = {}
 
 
 def get_sigma(hidden_size: int, seed: int):
-    """Derive client-side encryption permutation from secret seed."""
+    """Derive encryption permutation from secret seed."""
     rng = np.random.default_rng(seed)
     return rng.permutation(hidden_size)
 
 
 def load_client_components(ee_model_name: str):
-    """Load (and cache) tokenizer + original embed layer for a given EE model."""
+    """
+    Load (and cache) only what the client needs:
+      - tokenizer from the EE model
+      - embedding layer from the ORIGINAL model (just embed_tokens, not the full LLM)
+      - hidden_size from ee_config
+    """
     if ee_model_name in _cache:
         return _cache[ee_model_name]
 
-    # 1. Fetch EE config to discover hidden_size + original model name
+    # 1. Read EE config to get hidden_size + original model name
     config_path = hf_hub_download(ee_model_name, "ee_config.json")
     with open(config_path) as f:
         ee_config = json.load(f)
@@ -31,17 +36,21 @@ def load_client_components(ee_model_name: str):
     hidden_size = ee_config["hidden_size"]
     original_model_name = ee_config["original_model"]
 
-    # 2. Load tokenizer (from the EE model)
+    # 2. Tokenizer from the EE model
     tokenizer = AutoTokenizer.from_pretrained(ee_model_name, trust_remote_code=True)
 
-    # 3. Load ONLY the original embedding layer (CPU is fine — no forward pass needed)
-    embed_model = AutoModelForCausalLM.from_pretrained(
+    # 3. Load ONLY the original model's embed_tokens layer — we don't need the full LLM,
+    #    but HF doesn't support partial loading so we load it fully then discard the rest.
+    #    float32 on CPU is fine — we're only doing one embedding lookup, no generation.
+    original_model = AutoModelForCausalLM.from_pretrained(
         original_model_name,
-        torch_dtype=torch.float16,
+        torch_dtype=torch.float32,
         device_map="cpu",
         trust_remote_code=True,
     )
-    embed_layer = embed_model.model.embed_tokens
+    embed_layer = original_model.model.embed_tokens
+    embed_layer.eval()
+    del original_model  # free memory — we only need the embed layer
 
     _cache[ee_model_name] = (tokenizer, embed_layer, hidden_size)
     return tokenizer, embed_layer, hidden_size
@@ -55,29 +64,32 @@ def index():
 
     if request.method == "POST":
         form_data = request.form.to_dict()
-        server_url   = request.form["server_url"].rstrip("/")
+        server_url    = request.form["server_url"].rstrip("/")
         ee_model_name = request.form["ee_model_name"].strip()
-        ee_seed      = int(request.form["ee_seed"])
-        prompt       = request.form["prompt"].strip()
-        max_tokens   = int(request.form.get("max_tokens", 256))
+        ee_seed       = int(request.form["ee_seed"])
+        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)
 
-            # Derive encryption key
+            # Derive encryption permutation from secret seed
             sigma = get_sigma(hidden_size, ee_seed)
 
             # Tokenize
             inputs = tokenizer(prompt, return_tensors="pt")
 
-            # Compute plain embeddings
+            # Compute plain embeddings from original model's embed layer
             with torch.no_grad():
                 normal_embeds = embed_layer(inputs.input_ids)  # (1, seq_len, hidden)
 
-            # Encrypt: permute hidden dimension — server sees only scrambled vectors
-            encrypted_embeds = normal_embeds[..., sigma]
+            # Encrypt: permute hidden dimension with secret key
+            # Server sees only scrambled vectors — can't recover original prompt
+            encrypted_embeds = normal_embeds[..., sigma]  # (1, seq_len, hidden)
+
+            # Cast to float16 to match server model dtype
+            encrypted_embeds = encrypted_embeds.to(torch.float16)
 
-            # Send to server
             payload = {
                 "encrypted_embeds": encrypted_embeds.tolist(),
                 "attention_mask": inputs.attention_mask.tolist(),
@@ -89,17 +101,22 @@ def index():
                 json=payload,
                 timeout=300,
             )
-            resp.raise_for_status()
+
+            # Surface the server's error body if it returns non-2xx
+            if not resp.ok:
+                raise RuntimeError(
+                    f"Server returned {resp.status_code}: {resp.text[:500]}"
+                )
 
             gen_ids = resp.json()["generated_ids"]
             result = tokenizer.decode(gen_ids, skip_special_tokens=True)
 
+        except RuntimeError as e:
+            error = str(e)
         except requests.exceptions.ConnectionError:
-            error = f"Could not connect to server at {server_url}. Is it running?"
-        except requests.exceptions.HTTPError as e:
-            error = f"Server returned an error: {e.response.status_code} — {e.response.text}"
+            error = f"Could not connect to {server_url} — is the server Space running?"
         except Exception as e:
-            error = str(e)
+            error = f"{type(e).__name__}: {e}"
 
     return render_template("client.html", result=result, error=error, form=form_data)