Update app.py

app.py

@@ -8,6 +8,7 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from typing import List, Optional
 import torch
+import gc
 from transformers import AutoModel, AutoTokenizer
 
 app = FastAPI()
@@ -72,74 +73,88 @@ class DecodeResponse(BaseModel):
 # -----------------------------
 @app.post("/embed", response_model=EmbedResponse)
 def embed(req: EmbedRequest):
-    text = req.text
-
-    # -----------------------------
-    # 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="query",
-                return_multivector=True,   # always token-level
-                truncate_dim=req.truncate_dim,
-            )
-        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)
-
-    max_len = min(15_000, model.config.max_position_embeddings)  # ~32k
-    stride = 50
-    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="passage",
-                return_multivector=True,   # always token-level
-                truncate_dim=req.truncate_dim,
-            )
-
-        window_embeds = outputs[0].cpu()
-
-        if position > 0:
-            window_embeds = window_embeds[stride:]
-
-        embeddings.append(window_embeds)
-        position += max_len - stride
-
-    full_embeddings = torch.cat(embeddings, dim=0).tolist()
-    return {"embeddings": full_embeddings}
+    try:
+        text = req.text
+    
+        # -----------------------------
+        # Case 1: Query → mean pool across token embeddings
+        # -----------------------------
+        if (req.prompt_name or "").lower() == "query":
+            with torch.inference_mode():
+                outputs = model.encode_text(
+                    texts=[text],
+                    task=req.task,
+                    prompt_name="query",
+                    return_multivector=True,   # always token-level
+                    truncate_dim=req.truncate_dim,
+                )
+            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)
+    
+        max_len = min(15_000, model.config.max_position_embeddings)  # ~32k
+        stride = 50
+        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.inference_mode():
+                outputs = model.encode_text(
+                    texts=[tokenizer.decode(window_ids[0])],
+                    task=req.task,
+                    prompt_name="passage",
+                    return_multivector=True,   # always token-level
+                    truncate_dim=req.truncate_dim,
+                )
+    
+            window_embeds = outputs[0].cpu()
+    
+            if position > 0:
+                window_embeds = window_embeds[stride:]
+    
+            embeddings.append(window_embeds)
+            position += max_len - stride
+    
+        full_embeddings = torch.cat(embeddings, dim=0).tolist()
+        return {"embeddings": full_embeddings}
+    finally:
+        # --- Cleanup CUDA memory ---
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
 
 # -----------------------------
 # 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]}
-
+    try:
+        with torch.inference_mode():
+            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]}
+
+    finally:
+        # --- Cleanup CUDA memory ---
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
 
 # -----------------------------
 # Tokenize Endpoint