Update app.py

app.py

@@ -7,7 +7,6 @@ import json
 from huggingface_hub import hf_hub_download
 
 app = Flask(__name__)
-
 _cache = {}
 
 
@@ -19,17 +18,6 @@ def get_sigma(hidden_size: int, seed: int):
 
 
 def load_client_components(ee_model_name: str):
-    """
-    Client holds:
-      - tokenizer      (from original model)
-      - embed_tokens   (original, unmodified)
-      - lm_head        (original, unmodified)
-      - hidden_size
-
-    embed_tokens and lm_head never leave the client.
-    The server only has the transformer body with permuted weights.
-    sigma is derived from the seed — also never leaves the client.
-    """
     if ee_model_name in _cache:
         return _cache[ee_model_name]
 
@@ -42,7 +30,6 @@ def load_client_components(ee_model_name: str):
 
     tokenizer = AutoTokenizer.from_pretrained(original_model_name, trust_remote_code=True)
 
-    # Load original model to extract embed + lm_head, then discard the rest
     original_model = AutoModelForCausalLM.from_pretrained(
         original_model_name,
         torch_dtype=torch.float32,
@@ -51,7 +38,7 @@ def load_client_components(ee_model_name: str):
     )
     embed_layer = original_model.model.embed_tokens
     lm_head     = original_model.lm_head
-    final_norm  = original_model.model.norm  # final RMSNorm before lm_head
+    final_norm  = original_model.model.norm
     embed_layer.eval()
     lm_head.eval()
     final_norm.eval()
@@ -61,39 +48,37 @@ def load_client_components(ee_model_name: str):
     return tokenizer, embed_layer, lm_head, final_norm, hidden_size
 
 
-def generate_tokens(
-    server_url, tokenizer, embed_layer, lm_head, final_norm,
-    sigma_t, sigma_inv_t, formatted_prompt, max_new_tokens
-):
+def generate_tokens(server_url, tokenizer, embed_layer, lm_head, final_norm,
+                    sigma_t, sigma_inv_t, formatted_prompt, max_new_tokens):
     """
-    Token-by-token generation loop:
-      1. Client embeds current tokens → applies sigma → sends to server
-      2. Server returns last hidden state (sigma-space) + KV cache
-      3. Client applies sigma_inv → runs final_norm + lm_head → next token
-      4. Repeat until eos or max_tokens
+    Token-by-token generation. No KV cache — client accumulates all embeddings
+    and sends the full growing sequence each step.
+
+    Each step:
+      1. Encrypt all token embeddings so far with sigma
+      2. Send to server → get back last hidden state (sigma-space)
+      3. Decrypt last position: apply sigma_inv
+      4. Run final_norm + lm_head locally → next token
     """
     inputs = tokenizer(formatted_prompt, return_tensors="pt")
-    input_ids = inputs.input_ids
-    attention_mask = inputs.attention_mask
-
-    generated_ids = []
-    past_key_values = None
+    input_ids = inputs.input_ids  # (1, seq_len)
 
-    # First forward: send full prompt embeddings
+    # Build initial encrypted embeddings for full prompt
     with torch.no_grad():
-        plain_embeds = embed_layer(input_ids)           # (1, seq, hidden)
-    encrypted_embeds = plain_embeds[..., sigma_t]       # encrypt
-    encrypted_embeds = encrypted_embeds.to(torch.float16)
+        all_plain_embeds = embed_layer(input_ids)  # (1, seq_len, hidden)
 
-    current_mask = attention_mask
+    generated_ids = []
 
     for step in range(max_new_tokens):
+        # Encrypt the full sequence so far
+        all_encrypted = all_plain_embeds[..., sigma_t].to(torch.float16)  # (1, seq, hidden)
+        seq_len = all_encrypted.shape[1]
+        attention_mask = torch.ones(1, seq_len, dtype=torch.long)
+
         payload = {
-            "inputs_embeds":  encrypted_embeds.tolist(),
-            "attention_mask": current_mask.tolist(),
+            "inputs_embeds":  all_encrypted.tolist(),
+            "attention_mask": attention_mask.tolist(),
         }
-        if past_key_values is not None:
-            payload["past_key_values"] = past_key_values
 
         resp = requests.post(f"{server_url}/generate", json=payload, timeout=120)
         if not resp.ok:
@@ -103,18 +88,15 @@ def generate_tokens(
         if "error" in body:
             raise RuntimeError(f"Server error: {body['error']}")
 
-        # Decrypt: apply sigma_inv to get plain hidden state
+        # Decrypt last position only
         last_hidden = torch.tensor(body["last_hidden"], dtype=torch.float32)  # (1, seq, hidden)
-        past_key_values = body.get("past_key_values")  # may be None
+        last_pos_sigma = last_hidden[:, -1:, :]           # (1, 1, hidden) sigma-space
+        last_pos_plain = last_pos_sigma[..., sigma_inv_t] # (1, 1, hidden) plain-space
 
-        # Take only the last position
-        last_pos = last_hidden[:, -1:, :]              # (1, 1, hidden) sigma-space
-        plain_hidden = last_pos[..., sigma_inv_t]      # (1, 1, hidden) plain-space
-
-        # Client-side: final norm + lm_head → logits
+        # Client-side: final norm + lm_head → next token
         with torch.no_grad():
-            normed = final_norm(plain_hidden)
-            logits = lm_head(normed)                   # (1, 1, vocab)
+            normed  = final_norm(last_pos_plain)
+            logits  = lm_head(normed)                     # (1, 1, vocab)
 
         next_token_id = logits[0, -1, :].argmax().item()
         generated_ids.append(next_token_id)
@@ -122,14 +104,11 @@ def generate_tokens(
         if next_token_id == tokenizer.eos_token_id:
             break
 
-        # Prepare next step: embed + encrypt the single new token
+        # Append new token's plain embedding to the growing sequence
         next_id_tensor = torch.tensor([[next_token_id]])
         with torch.no_grad():
-            next_plain_embed = embed_layer(next_id_tensor)   # (1, 1, hidden)
-        encrypted_embeds = next_plain_embed[..., sigma_t].to(torch.float16)
-
-        # Extend attention mask by 1
-        current_mask = torch.ones(1, 1, dtype=attention_mask.dtype)
+            next_embed = embed_layer(next_id_tensor)      # (1, 1, hidden)
+        all_plain_embeds = torch.cat([all_plain_embeds, next_embed], dim=1)
 
     return generated_ids
 
@@ -137,11 +116,11 @@ def generate_tokens(
 @app.route("/", methods=["GET", "POST"])
 def index():
     result = None
-    error = None
+    error  = None
     form_data = {}
 
     if request.method == "POST":
-        form_data = request.form.to_dict()
+        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"])
@@ -154,8 +133,7 @@ def index():
 
             sigma_t, sigma_inv_t = get_sigma(hidden_size, ee_seed)
 
-            # Apply chat template
-            messages = [{"role": "user", "content": prompt}]
+            messages  = [{"role": "user", "content": prompt}]
             formatted = tokenizer.apply_chat_template(
                 messages, tokenize=False, add_generation_prompt=True
             )