Update app.py

app.py

@@ -105,7 +105,7 @@ def generate_tokens(
 
         # Decrypt: apply sigma_inv to get plain hidden state
         last_hidden = torch.tensor(body["last_hidden"], dtype=torch.float32)  # (1, seq, hidden)
-        past_key_values = body["past_key_values"]
+        past_key_values = body.get("past_key_values")  # may be None
 
         # Take only the last position
         last_pos = last_hidden[:, -1:, :]              # (1, 1, hidden) sigma-space