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

add batch_size

Browse Source
This commit is contained in:
SFirouzi 2026-03-14 11:55:52 +03:30
parent cc60fce612
commit b36c680b3f
2 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
config/base.py
View File

@ -8,4 +8,4 @@ GEMMA_LORA_PATH = os.getenv("GEMMA_LORA_PATH")
BGE_MODEL_PATH = os.getenv("BGE_MODEL_PATH") BGE_MODEL_PATH = os.getenv("BGE_MODEL_PATH")
BGE_LORA_PATH = os.getenv("BGE_LORA_PATH") BGE_LORA_PATH = os.getenv("BGE_LORA_PATH")
BATCH_SIZE = 250 BATCH_SIZE = 100

1
src/serve_embed.py
View File

@ -55,6 +55,7 @@ def embed_gemma(request: EmbedRequest):
raise HTTPException(status_code=400, detail="Invalid model") raise HTTPException(status_code=400, detail="Invalid model")
except Exception as e: except Exception as e:
torch.cuda.empty_cache()
raise HTTPException(status_code=500, detail=str(e)) raise HTTPException(status_code=500, detail=str(e))
torch.cuda.empty_cache() torch.cuda.empty_cache()