Update app.py

app.py

@@ -2,17 +2,14 @@ from flask import Flask, render_template, request, flash, jsonify
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import login
-import traceback
 import os, json
 
 app = Flask(__name__)
 app.secret_key = os.urandom(24)
 
-# Globals
 ee_model = None
 ee_tokenizer = None
 ee_config = None
-loaded_model_name = None
 
 SPACE_HOST = os.environ.get("SPACE_HOST", "")
 SPACE_URL = f"https://{SPACE_HOST}" if SPACE_HOST else "http://localhost:7860"
@@ -20,91 +17,99 @@ SPACE_URL = f"https://{SPACE_HOST}" if SPACE_HOST else "http://localhost:7860"
 
 @app.route("/", methods=["GET", "POST"])
 def index():
-    global ee_model, ee_tokenizer, ee_config, loaded_model_name
+    global ee_model, ee_tokenizer, ee_config
 
     if request.method == "POST":
-        action = request.form.get("action")
+        ee_model_name = request.form["ee_model_name"].strip()
+        hf_token      = request.form["hf_token"].strip()
 
-        if action == "start_server":
-            ee_model_name = request.form["ee_model_name"].strip()
-            hf_token = request.form["hf_token"].strip()
+        try:
+            login(token=hf_token)
 
-            try:
-                login(token=hf_token)
-
-                ee_model = AutoModelForCausalLM.from_pretrained(
-                    ee_model_name,
-                    torch_dtype=torch.float16,
-                    device_map="auto",
-                    trust_remote_code=True,
-                )
-                ee_tokenizer = AutoTokenizer.from_pretrained(
-                    ee_model_name, trust_remote_code=True
-                )
+            ee_model = AutoModelForCausalLM.from_pretrained(
+                ee_model_name, torch_dtype=torch.float16,
+                device_map="auto", trust_remote_code=True
+            )
+            ee_tokenizer = AutoTokenizer.from_pretrained(
+                ee_model_name, trust_remote_code=True
+            )
 
-                from huggingface_hub import hf_hub_download
-                config_path = hf_hub_download(ee_model_name, "ee_config.json")
-                with open(config_path) as f:
-                    ee_config = json.load(f)
+            from huggingface_hub import hf_hub_download
+            config_path = hf_hub_download(ee_model_name, "ee_config.json")
+            with open(config_path) as f:
+                ee_config = json.load(f)
 
-                loaded_model_name = ee_model_name
-                flash(f"Model loaded successfully: {ee_model_name}", "success")
-                flash("Point your Client Space to this Space's URL below.", "info")
+            flash(f"✅ Model loaded: {ee_model_name}", "success")
+            flash("Point your Client Space to this Space's URL below.", "info")
 
-            except Exception as e:
-                flash(f"Error: {str(e)}", "danger")
+        except Exception as e:
+            flash(f"Error: {str(e)}", "danger")
 
     return render_template(
         "index.html",
         server_ready=(ee_model is not None),
-        model_name=loaded_model_name,
+        model_name=ee_config["original_model"] if ee_config else None,
         space_url=SPACE_URL,
     )
 
 
 @app.route("/generate", methods=["POST"])
 def generate():
+    """
+    Receives sigma-encrypted embeddings + optional past_key_values.
+    Returns last hidden state (still in sigma-space) + new KV cache.
+    Does NOT run lm_head — that stays on the client.
+    Server never sees token IDs, logits, or plaintext.
+    """
     if ee_model is None:
-        return jsonify({"error": "Server not started yet — load a model first"}), 400
+        return jsonify({"error": "Server not started yet"}), 400
 
     try:
         data = request.json
-        if data is None:
-            return jsonify({"error": "Request body must be JSON"}), 400
-
         model_dtype = next(ee_model.parameters()).dtype
 
-        # Cast incoming embeddings to model dtype + move to device
-        encrypted_embeds = torch.tensor(data["encrypted_embeds"]).to(
+        inputs_embeds = torch.tensor(data["inputs_embeds"]).to(
             dtype=model_dtype, device=ee_model.device
-        )  # (1, seq_len, hidden)
-
-        input_seq_len = encrypted_embeds.shape[1]
+        )
 
         attention_mask = torch.tensor(
-            data.get("attention_mask", [[1] * input_seq_len])
+            data.get("attention_mask", [[1] * inputs_embeds.shape[1]])
         ).to(device=ee_model.device)
 
-        max_new = int(data.get("max_new_tokens", 256))
+        past_key_values = None
+        if data.get("past_key_values"):
+            past_key_values = tuple(
+                tuple(
+                    torch.tensor(t).to(dtype=model_dtype, device=ee_model.device)
+                    for t in layer
+                )
+                for layer in data["past_key_values"]
+            )
 
         with torch.no_grad():
-            output_ids = ee_model.generate(
-                inputs_embeds=encrypted_embeds,
+            out = ee_model(
+                inputs_embeds=inputs_embeds,
                 attention_mask=attention_mask,
-                max_new_tokens=max_new,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.9,
-                pad_token_id=ee_tokenizer.eos_token_id,
+                past_key_values=past_key_values,
+                use_cache=True,
+                output_hidden_states=True,
             )
 
-        # output_ids includes the full sequence; return only the newly generated tokens
-        # (the client sent embeddings, not IDs, so output starts at position 0)
-        new_ids = output_ids[0].tolist()
+        # Return final hidden state in sigma-space — client decrypts + runs lm_head
+        last_hidden = out.hidden_states[-1]  # (1, seq_len, hidden)
+
+        new_past = [
+            [t.cpu().tolist() for t in layer]
+            for layer in out.past_key_values
+        ]
 
-        return jsonify({"generated_ids": new_ids})
+        return jsonify({
+            "last_hidden": last_hidden.cpu().tolist(),
+            "past_key_values": new_past,
+        })
 
     except Exception as e:
+        import traceback
         return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500