Update app.py

app.py

@@ -74,21 +74,24 @@ class DecodeResponse(BaseModel):
 def embed(req: EmbedRequest):
     text = req.text
 
-    # Case 1: Query → pooled mean of multivectors
-    if not req.return_token_embeddings:
+    # -----------------------------
+    # Case 1: Query → mean pool across token embeddings
+    # -----------------------------
+    if (req.prompt_name or "").lower() == "query":
         with torch.no_grad():
             outputs = model.encode_text(
                 texts=[text],
                 task=req.task,
-                prompt_name=req.prompt_name or "query",
-                return_multivector=req.return_token_embeddings,
+                prompt_name="query",
+                return_multivector=True,   # always token-level
                 truncate_dim=req.truncate_dim,
             )
-        # outputs[0] = (num_vectors, hidden_dim)
-        pooled = outputs[0].mean(dim=0).cpu()
-        return {"embeddings": [pooled]}
+        pooled = outputs[0].mean(dim=0).cpu().tolist()
+        return {"embeddings": [pooled]}   # wrap in batch dimension
 
+    # -----------------------------
     # Case 2: Passage → sliding window, token-level embeddings
+    # -----------------------------
     enc = tokenizer(text, add_special_tokens=False, return_tensors="pt")
     input_ids = enc["input_ids"].squeeze(0).to(device)
     total_tokens = input_ids.size(0)
@@ -106,8 +109,8 @@ def embed(req: EmbedRequest):
             outputs = model.encode_text(
                 texts=[tokenizer.decode(window_ids[0])],
                 task=req.task,
-                prompt_name=req.prompt_name or "passage",
-                return_multivector=req.return_token_embeddings,
+                prompt_name="passage",
+                return_multivector=True,   # always token-level
                 truncate_dim=req.truncate_dim,
             )
 
@@ -119,10 +122,9 @@ def embed(req: EmbedRequest):
         embeddings.append(window_embeds)
         position += max_len - stride
 
-    full_embeddings = torch.cat(embeddings, dim=0)
+    full_embeddings = torch.cat(embeddings, dim=0).tolist()
     return {"embeddings": full_embeddings}
 
-
 # -----------------------------
 # Embedding Endpoint (image)
 # -----------------------------