Update app.py

app.py

@@ -1,9 +1,9 @@
 # app.py (FastAPI server to host the Jina Embedding model)
-# Must be set before importing Hugging Face libs
 import os
 os.environ["HF_HOME"] = "/tmp/huggingface"
 os.environ["HF_HUB_CACHE"] = "/tmp/huggingface/hub"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface/transformers"
+
 from fastapi import FastAPI
 from pydantic import BaseModel
 from typing import List, Optional
@@ -30,14 +30,25 @@ model.eval()
 # -----------------------------
 class EmbedRequest(BaseModel):
     text: str
-    task: str = "retrieval"          # "retrieval", "text-matching", "code", etc.
+    task: str = "retrieval"
     prompt_name: Optional[str] = None
-    return_token_embeddings: bool = True   # False → for queries (pooled embedding)
+    return_token_embeddings: bool = True
+    truncate_dim: Optional[int] = None  # for matryoshka embeddings
 
 
 class EmbedResponse(BaseModel):
-    embeddings: List[List[float]]  # (num_tokens, hidden_dim) if token-level
-                                   # (1, hidden_dim) if pooled query
+    embeddings: List[List[float]]
+
+
+class EmbedImageRequest(BaseModel):
+    image: str
+    task: str = "retrieval"
+    return_multivector: bool = True
+    truncate_dim: Optional[int] = None
+
+
+class EmbedImageResponse(BaseModel):
+    embeddings: List[List[float]]
 
 
 class TokenizeRequest(BaseModel):
@@ -57,34 +68,33 @@ class DecodeResponse(BaseModel):
 
 
 # -----------------------------
-# Embedding Endpoint
+# Embedding Endpoint (text)
 # -----------------------------
 @app.post("/embed", response_model=EmbedResponse)
 def embed(req: EmbedRequest):
     text = req.text
 
-    # -----------------------------
-    # Case 1: Query → directly pooled embedding
-    # -----------------------------
+    # Case 1: Query → pooled mean of multivectors
     if not req.return_token_embeddings:
         with torch.no_grad():
-            emb = model.encode_text(
+            outputs = model.encode_text(
                 texts=[text],
                 task=req.task,
                 prompt_name=req.prompt_name or "query",
-                return_multivector=False
+                return_multivector=True,
+                truncate_dim=req.truncate_dim,
             )
-        return {"embeddings": emb}  # shape: (1, hidden_dim)
+        # outputs[0] = (num_vectors, hidden_dim)
+        pooled = outputs[0].mean(dim=0).cpu().tolist()
+        return {"embeddings": [pooled]}
 
-    # -----------------------------
-    # Case 2: Long passages → sliding window token embeddings
-    # -----------------------------
+    # 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 = enc["input_ids"].squeeze(0).to(device)
     total_tokens = input_ids.size(0)
 
-    max_len = model.config.max_position_embeddings  # e.g., 32k for v4
-    stride = 50  # overlap for sliding window
+    max_len = model.config.max_position_embeddings  # ~32k
+    stride = 50
     embeddings = []
     position = 0
 
@@ -94,27 +104,41 @@ def embed(req: EmbedRequest):
 
         with torch.no_grad():
             outputs = model.encode_text(
-                    texts=[tokenizer.decode(window_ids[0])],
-                    task=req.task,
-                    prompt_name=req.prompt_name or "passage",
-                    return_multivector=True,
-                )
+                texts=[tokenizer.decode(window_ids[0])],
+                task=req.task,
+                prompt_name=req.prompt_name or "passage",
+                return_multivector=True,
+                truncate_dim=req.truncate_dim,
+            )
 
-        window_embeds = outputs[0].cpu()  # (window_len, hidden_dim)
+        window_embeds = outputs[0].cpu()
 
-        # Drop overlapping tokens except in first window
         if position > 0:
             window_embeds = window_embeds[stride:]
 
         embeddings.append(window_embeds)
-
-        # Advance window
         position += max_len - stride
 
-    full_embeddings = torch.cat(embeddings, dim=0)  # (total_tokens, hidden_dim)
+    full_embeddings = torch.cat(embeddings, dim=0)
     return {"embeddings": full_embeddings}
 
 
+# -----------------------------
+# Embedding Endpoint (image)
+# -----------------------------
+@app.post("/embed_image", response_model=EmbedImageResponse)
+def embed_image(req: EmbedImageRequest):
+    with torch.no_grad():
+        outputs = model.encode_image(
+            images=[req.image],
+            task=req.task,
+            return_multivector=req.return_multivector,
+            truncate_dim=req.truncate_dim,
+        )
+    pooled = outputs[0].mean(dim=0).cpu()
+    return {"embeddings": [pooled]}
+
+
 # -----------------------------
 # Tokenize Endpoint
 # -----------------------------
@@ -130,4 +154,4 @@ def tokenize(req: TokenizeRequest):
 @app.post("/decode", response_model=DecodeResponse)
 def decode(req: DecodeRequest):
     decoded = tokenizer.decode(req.input_ids)
-    return {"text": decoded}
+    return {"text": decoded}