Update app.py

app.py

@@ -14,7 +14,6 @@ ee_tokenizer = None
 ee_config = None
 loaded_model_name = None
 
-# Detect HF Space URL automatically
 SPACE_HOST = os.environ.get("SPACE_HOST", "")
 SPACE_URL = f"https://{SPACE_HOST}" if SPACE_HOST else "http://localhost:7860"
 
@@ -63,7 +62,6 @@ def index():
     )
 
 
-# === INFERENCE ENDPOINT ===
 @app.route("/generate", methods=["POST"])
 def generate():
     if ee_model is None:
@@ -74,17 +72,18 @@ def generate():
         if data is None:
             return jsonify({"error": "Request body must be JSON"}), 400
 
-        # Determine the model's actual dtype so we always match it
         model_dtype = next(ee_model.parameters()).dtype
 
-        # Build tensors, cast to model dtype + move to device in one step
+        # Cast incoming embeddings to model dtype + move to device
         encrypted_embeds = torch.tensor(data["encrypted_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] * encrypted_embeds.shape[1]])
-        ).to(device=ee_model.device)  # stays int64, that's correct
+            data.get("attention_mask", [[1] * input_seq_len])
+        ).to(device=ee_model.device)
 
         max_new = int(data.get("max_new_tokens", 256))
 
@@ -99,7 +98,11 @@ def generate():
                 pad_token_id=ee_tokenizer.eos_token_id,
             )
 
-        return jsonify({"generated_ids": output_ids[0].tolist()})
+        # 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 jsonify({"generated_ids": new_ids})
 
     except Exception as e:
         return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500