Update app.py

app.py

@@ -12,27 +12,11 @@ _cache = {}
 
 
 def get_sigma(hidden_size: int, seed: int) -> np.ndarray:
-    """
-    Derive the encryption permutation from the secret seed.
-    This is the CLIENT'S secret key — it never leaves this Space.
-    The server only ever sees embeddings already scrambled by sigma.
-    """
     rng = np.random.default_rng(seed)
     return rng.permutation(hidden_size)
 
 
 def load_client_components(ee_model_name: str):
-    """
-    Load and cache:
-      - ee_config  → hidden_size + original model name
-      - tokenizer  → from EE model
-      - embed_layer → from the ORIGINAL (untransformed) model
-
-    The original embed_layer is used to produce plain vectors from token IDs.
-    The client then applies sigma to those plain vectors before sending.
-    The server's EE model has weights permuted with sigma_inv, so:
-        EE_model(sigma(plain_embed(tokens))) == original_model(plain_embed(tokens))
-    """
     if ee_model_name in _cache:
         return _cache[ee_model_name]
 
@@ -43,9 +27,10 @@ def load_client_components(ee_model_name: str):
     hidden_size = ee_config["hidden_size"]
     original_model_name = ee_config["original_model"]
 
+    # Tokenizer from EE model (same vocab as original)
     tokenizer = AutoTokenizer.from_pretrained(ee_model_name, trust_remote_code=True)
 
-    # Load ORIGINAL model just for its embed layer — discard everything else
+    # Load ORIGINAL model just to extract embed_tokens, then discard
     original_model = AutoModelForCausalLM.from_pretrained(
         original_model_name,
         torch_dtype=torch.float32,
@@ -70,34 +55,44 @@ def index():
         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"])   # SECRET — client only
+        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)
 
-            # --- CLIENT-SIDE ENCRYPTION ---
-
-            # Step 1: tokenize
-            inputs = tokenizer(prompt, return_tensors="pt")
-
-            # Step 2: embed with ORIGINAL model embed layer → plain vectors
+            # --- Step 1: Apply chat template ---
+            # Qwen3 (and most instruct models) require the prompt wrapped in the
+            # chat template before tokenization, otherwise the model sees raw text
+            # with no special tokens and produces garbage.
+            messages = [{"role": "user", "content": prompt}]
+            formatted = tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,  # appends <|im_start|>assistant\n
+            )
+
+            # --- Step 2: Tokenize the formatted prompt ---
+            inputs = tokenizer(formatted, return_tensors="pt")
+            input_ids = inputs.input_ids  # (1, seq_len)
+            input_len = input_ids.shape[1]
+
+            # --- Step 3: Embed with ORIGINAL model's embed layer ---
             with torch.no_grad():
-                plain_embeds = embed_layer(inputs.input_ids)  # (1, seq_len, hidden)
+                plain_embeds = embed_layer(input_ids)  # (1, seq_len, hidden)
 
-            # Step 3: apply sigma permutation — this is the encryption
-            # The server NEVER sees plain_embeds, only the scrambled version.
-            # Without knowing the seed, the server cannot recover the original.
+            # --- Step 4: Encrypt — permute hidden dim with secret sigma ---
             sigma = get_sigma(hidden_size, ee_seed)
-            encrypted_embeds = plain_embeds[..., sigma]        # (1, seq_len, hidden)
+            encrypted_embeds = plain_embeds[..., sigma]       # (1, seq_len, hidden)
             encrypted_embeds = encrypted_embeds.to(torch.float16)
 
-            # --- SEND TO SERVER ---
+            # --- Step 5: Send to server ---
             payload = {
                 "encrypted_embeds": encrypted_embeds.tolist(),
                 "attention_mask":   inputs.attention_mask.tolist(),
                 "max_new_tokens":   max_tokens,
+                "input_len":        input_len,  # so server can strip prompt tokens
             }
 
             resp = requests.post(f"{server_url}/generate", json=payload, timeout=300)
@@ -107,13 +102,11 @@ def index():
 
             body = resp.json()
             if "error" in body:
-                raise RuntimeError(f"Server error: {body['error']}\n{body.get('traceback','')}")
+                raise RuntimeError(f"Server error: {body['error']}\n{body.get('traceback', '')}")
 
-            # --- OUTPUT DECODING ---
-            # The EE model's lm_head rows are permuted with sigma_inv, so output
-            # logits correctly index the real vocabulary — decode normally.
+            # --- Step 6: Decode only the NEW tokens (strip echoed prompt) ---
             gen_ids = body["generated_ids"]
-            result = tokenizer.decode(gen_ids, skip_special_tokens=True)
+            result = tokenizer.decode(gen_ids, skip_special_tokens=True).strip()
 
         except RuntimeError as e:
             error = str(e)