bi/serve_embed
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add cache

Browse Source
This commit is contained in:
SFirouzi 2026-03-14 10:46:48 +03:30
parent 8eebc192e0
commit 3a38099749
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/serve_embed.py
View File

@ -2,6 +2,7 @@ from fastapi import FastAPI
from fastapi import HTTPException from fastapi import HTTPException
import uvicorn import uvicorn
from pydantic import BaseModel from pydantic import BaseModel
import torch
from models.embedder_gemma import TextEmbedderGemma from models.embedder_gemma import TextEmbedderGemma
from models.embedder_gemma_train import TextEmbedderGemmaTrain from models.embedder_gemma_train import TextEmbedderGemmaTrain
@ -52,4 +53,6 @@ def embed_gemma(request: EmbedRequest):
else: else:
raise HTTPException(status_code=400, detail="Invalid model") raise HTTPException(status_code=400, detail="Invalid model")
torch.cuda.empty_cache()
return {"data": [{"embedding": emb.tolist()} for emb in embeddings]} return {"data": [{"embedding": emb.tolist()} for emb in embeddings]}