Update app.py

app.py

@@ -1,55 +1,80 @@
 from flask import Flask, render_template, request
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer
 import numpy as np
-import requests, json
+import requests
+import json
 from huggingface_hub import hf_hub_download
 
 app = Flask(__name__)
 
+def get_sigma(hidden_size: int, seed: int):
+    """Client-side encryption key from secret seed"""
+    rng = np.random.default_rng(seed)
+    sigma = rng.permutation(hidden_size)
+    return sigma
+
 @app.route("/", methods=["GET", "POST"])
 def index():
     result = None
+    error = None
+
     if request.method == "POST":
         server_url = request.form["server_url"].rstrip("/")
-        prompt = request.form["prompt"]
+        ee_model_name = request.form["ee_model_name"].strip()
         ee_seed = int(request.form["ee_seed"])
-        ee_model_name = request.form["ee_model_name"]
-        max_tokens = int(request.form["max_tokens"])
+        prompt = request.form["prompt"].strip()
+        max_tokens = int(request.form.get("max_tokens", 256))
 
         try:
-            # Load tokenizer + ee_config from the EE model
-            tokenizer = AutoTokenizer.from_pretrained(ee_model_name, trust_remote_code=True)
+            # 1. Load config to know hidden_size + original model
             config_path = hf_hub_download(ee_model_name, "ee_config.json")
             with open(config_path) as f:
                 ee_config = json.load(f)
+            hidden_size = ee_config["hidden_size"]
+            original_model_name = ee_config["original_model"]
+
+            # 2. Generate encryption permutation (this is your secret key in action)
+            sigma = get_sigma(hidden_size, ee_seed)
+
+            # 3. Load tokenizer
+            tokenizer = AutoTokenizer.from_pretrained(ee_model_name, trust_remote_code=True)
 
-            # Load only embedding layer from EE model (transformed!)
-            embed_layer = AutoModelForCausalLM.from_pretrained(
-                ee_model_name, trust_remote_code=True, device_map="cpu"
-            ).model.embed_tokens
+            # 4. Load ORIGINAL (clean) embedding layer
+            embed_model = AutoModelForCausalLM.from_pretrained(
+                original_model_name,
+                torch_dtype=torch.float16,
+                device_map="cpu",
+                trust_remote_code=True
+            )
+            embed_layer = embed_model.model.embed_tokens
 
-            # Tokenize + get encrypted embeddings
+            # 5. Tokenize + compute normal embeddings
             inputs = tokenizer(prompt, return_tensors="pt")
             with torch.no_grad():
-                embeds = embed_layer(inputs.input_ids)  # already "encrypted" because we loaded transformed embed
+                normal_embeds = embed_layer(inputs.input_ids)   # shape: (1, seq_len, hidden_size)
+
+            # 6. === EXPLICIT ENCRYPTION (this is the key step you asked for) ===
+            # Permute the hidden dimension according to the secret sigma
+            encrypted_embeds = normal_embeds[..., sigma]   # now scrambled — provider sees nothing
 
-            # Send to server
+            # 7. Send ONLY encrypted embeddings to server
             payload = {
-                "encrypted_embeds": embeds.tolist(),
+                "encrypted_embeds": encrypted_embeds.tolist(),
                 "attention_mask": inputs.attention_mask.tolist(),
                 "max_new_tokens": max_tokens
             }
-            resp = requests.post(f"{server_url}/generate", json=payload, timeout=180)
+
+            resp = requests.post(f"{server_url}/generate", json=payload, timeout=300)
             resp.raise_for_status()
-            gen_ids = resp.json()["generated_ids"]
 
+            gen_ids = resp.json()["generated_ids"]
             result = tokenizer.decode(gen_ids, skip_special_tokens=True)
 
         except Exception as e:
-            result = f"Error: {str(e)}"
+            error = str(e)
 
-    return render_template("client.html", result=result)
+    return render_template("client.html", result=result, error=error)
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)