Create app.py

app.py

@@ -0,0 +1,133 @@
+# 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
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+app = FastAPI()
+
+# -----------------------------
+# Load model once on startup
+# -----------------------------
+MODEL_NAME = "jinaai/jina-embeddings-v4"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    MODEL_NAME, trust_remote_code=True, torch_dtype=torch.float16
+).to(device)
+model.eval()
+
+
+# -----------------------------
+# Request / Response Models
+# -----------------------------
+class EmbedRequest(BaseModel):
+    text: str
+    task: str = "retrieval"          # "retrieval", "text-matching", "code", etc.
+    prompt_name: Optional[str] = None
+    return_token_embeddings: bool = True   # False → for queries (pooled embedding)
+
+
+class EmbedResponse(BaseModel):
+    embeddings: List[List[float]]  # (num_tokens, hidden_dim) if token-level
+                                   # (1, hidden_dim) if pooled query
+
+
+class TokenizeRequest(BaseModel):
+    text: str
+
+
+class TokenizeResponse(BaseModel):
+    input_ids: List[int]
+
+
+class DecodeRequest(BaseModel):
+    input_ids: List[int]
+
+
+class DecodeResponse(BaseModel):
+    text: str
+
+
+# -----------------------------
+# Embedding Endpoint
+# -----------------------------
+@app.post("/embed", response_model=EmbedResponse)
+def embed(req: EmbedRequest):
+    text = req.text
+
+    # -----------------------------
+    # Case 1: Query → directly pooled embedding
+    # -----------------------------
+    if not req.return_token_embeddings:
+        with torch.no_grad():
+            emb = model.encode_text(
+                texts=[text],
+                task=req.task,
+                prompt_name=req.prompt_name or "query",
+                return_multivector=False
+            )
+        return {"embeddings": emb.tolist()}  # shape: (1, hidden_dim)
+
+    # -----------------------------
+    # Case 2: Long passages → sliding window token embeddings
+    # -----------------------------
+    enc = tokenizer(text, add_special_tokens=False, return_tensors="pt")
+    input_ids = enc["input_ids"].squeeze(0).to(device)  # (total_tokens,)
+    total_tokens = input_ids.size(0)
+
+    max_len = model.config.max_position_embeddings  # e.g., 32k for v4
+    stride = 50  # overlap for sliding window
+    embeddings = []
+    position = 0
+
+    while position < total_tokens:
+        end = min(position + max_len, total_tokens)
+        window_ids = input_ids[position:end].unsqueeze(0).to(device)
+
+        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,
+                )
+
+        window_embeds = outputs.squeeze(0).cpu()  # (window_len, hidden_dim)
+
+        # 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)
+    return {"embeddings": full_embeddings.tolist()}
+
+
+# -----------------------------
+# Tokenize Endpoint
+# -----------------------------
+@app.post("/tokenize", response_model=TokenizeResponse)
+def tokenize(req: TokenizeRequest):
+    enc = tokenizer(req.text, add_special_tokens=False)
+    return {"input_ids": enc["input_ids"]}
+
+
+# -----------------------------
+# Decode Endpoint
+# -----------------------------
+@app.post("/decode", response_model=DecodeResponse)
+def decode(req: DecodeRequest):
+    decoded = tokenizer.decode(req.input_ids)
+    return {"text": decoded}