Update app.py

app.py

@@ -86,32 +86,19 @@ def generate():
                 for layer in data["past_key_values"]
             )
 
-        # Ensure model config has caching enabled
-        ee_model.config.use_cache = True
-
         with torch.no_grad():
             out = ee_model(
                 inputs_embeds=inputs_embeds,
                 attention_mask=attention_mask,
-                past_key_values=past_key_values,
-                use_cache=True,
+                use_cache=False,
                 output_hidden_states=True,
             )
 
-        # Final hidden state (sigma-space) — client decrypts + runs lm_head
+        # Return final hidden state in sigma-space — client applies sigma_inv + lm_head
         last_hidden = out.hidden_states[-1]  # (1, seq_len, hidden)
 
-        # Serialize KV cache — guard against None (some models/configs don't return it)
-        new_past = None
-        if out.past_key_values is not None:
-            new_past = [
-                [t.cpu().tolist() for t in layer]
-                for layer in out.past_key_values
-            ]
-
         return jsonify({
             "last_hidden": last_hidden.cpu().tolist(),
-            "past_key_values": new_past,
         })
 
     except Exception as e: