Update app.py

app.py

@@ -8,73 +8,101 @@ from huggingface_hub import hf_hub_download
 
 app = Flask(__name__)
 
+# Cache tokenizer/embed layer so repeated requests don't reload from scratch
+_cache = {}
+
+
 def get_sigma(hidden_size: int, seed: int):
-    """Client-side encryption key from secret seed"""
+    """Derive client-side encryption permutation from secret seed."""
     rng = np.random.default_rng(seed)
-    sigma = rng.permutation(hidden_size)
-    return sigma
+    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."""
+    if ee_model_name in _cache:
+        return _cache[ee_model_name]
+
+    # 1. Fetch EE config to discover 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)
+
+    hidden_size = ee_config["hidden_size"]
+    original_model_name = ee_config["original_model"]
+
+    # 2. Load 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(
+        original_model_name,
+        torch_dtype=torch.float16,
+        device_map="cpu",
+        trust_remote_code=True,
+    )
+    embed_layer = embed_model.model.embed_tokens
+
+    _cache[ee_model_name] = (tokenizer, embed_layer, hidden_size)
+    return tokenizer, embed_layer, hidden_size
+
 
 @app.route("/", methods=["GET", "POST"])
 def index():
     result = None
     error = None
+    form_data = {}
 
     if request.method == "POST":
-        server_url = request.form["server_url"].rstrip("/")
+        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"])
-        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:
-            # 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)
+            tokenizer, embed_layer, hidden_size = load_client_components(ee_model_name)
 
-            # 3. Load tokenizer
-            tokenizer = AutoTokenizer.from_pretrained(ee_model_name, trust_remote_code=True)
-
-            # 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
+            # Derive encryption key
+            sigma = get_sigma(hidden_size, ee_seed)
 
-            # 5. Tokenize + compute normal embeddings
+            # Tokenize
             inputs = tokenizer(prompt, return_tensors="pt")
+
+            # Compute plain embeddings
             with torch.no_grad():
-                normal_embeds = embed_layer(inputs.input_ids)   # shape: (1, seq_len, hidden_size)
+                normal_embeds = embed_layer(inputs.input_ids)  # (1, seq_len, hidden)
 
-            # 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
+            # Encrypt: permute hidden dimension — server sees only scrambled vectors
+            encrypted_embeds = normal_embeds[..., sigma]
 
-            # 7. Send ONLY encrypted embeddings to server
+            # Send to server
             payload = {
                 "encrypted_embeds": encrypted_embeds.tolist(),
                 "attention_mask": inputs.attention_mask.tolist(),
-                "max_new_tokens": max_tokens
+                "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,
+            )
             resp.raise_for_status()
 
             gen_ids = resp.json()["generated_ids"]
             result = tokenizer.decode(gen_ids, skip_special_tokens=True)
 
+        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}"
         except Exception as e:
             error = str(e)
 
-    return render_template("client.html", result=result, error=error)
+    return render_template("client.html", result=result, error=error, form=form_data)
+
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)