Signal: Fix async signal handling

Run unload async functions before exiting the program.

Signed-off-by: kingbri <bdashore3@proton.me>
This commit is contained in:
kingbri
2024-07-30 11:11:05 -04:00
parent fbf1455db1
commit 01c7702859
2 changed files with 12 additions and 6 deletions

View File

@@ -1,6 +1,7 @@
import gc
import pathlib
import torch
from loguru import logger
from typing import List, Optional
from common.utils import unwrap
@@ -50,6 +51,8 @@ class InfinityContainer:
gc.collect()
torch.cuda.empty_cache()
logger.info("Embedding model unloaded.")
async def generate(self, sentence_input: List[str]):
result_embeddings, usage = await self.engine.embed(sentence_input)