Create app.py

app.py

@@ -0,0 +1,55 @@
+from flask import Flask, render_template, request
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import numpy as np
+import requests, json
+from huggingface_hub import hf_hub_download
+
+app = Flask(__name__)
+
+@app.route("/", methods=["GET", "POST"])
+def index():
+    result = None
+    if request.method == "POST":
+        server_url = request.form["server_url"].rstrip("/")
+        prompt = request.form["prompt"]
+        ee_seed = int(request.form["ee_seed"])
+        ee_model_name = request.form["ee_model_name"]
+        max_tokens = int(request.form["max_tokens"])
+
+        try:
+            # Load tokenizer + ee_config from the EE model
+            tokenizer = AutoTokenizer.from_pretrained(ee_model_name, trust_remote_code=True)
+            config_path = hf_hub_download(ee_model_name, "ee_config.json")
+            with open(config_path) as f:
+                ee_config = json.load(f)
+
+            # 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
+
+            # Tokenize + get encrypted embeddings
+            inputs = tokenizer(prompt, return_tensors="pt")
+            with torch.no_grad():
+                embeds = embed_layer(inputs.input_ids)  # already "encrypted" because we loaded transformed embed
+
+            # Send to server
+            payload = {
+                "encrypted_embeds": 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.raise_for_status()
+            gen_ids = resp.json()["generated_ids"]
+
+            result = tokenizer.decode(gen_ids, skip_special_tokens=True)
+
+        except Exception as e:
+            result = f"Error: {str(e)}"
+
+    return render_template("client.html", result=result)
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=7860)