Skip to content

Commit

Permalink
preload enrichment model on start
Browse files Browse the repository at this point in the history
  • Loading branch information
cdzombak committed Jul 31, 2024
1 parent c6ca70d commit 85fedda
Showing 1 changed file with 20 additions and 1 deletion.
21 changes: 20 additions & 1 deletion ntfy.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ class NtfyRecord:
class EnrichmentConfig:
enable: bool = False
endpoint: str = ""
keep_alive: str = "240m"
keep_alive: str = "1440m"
model: str = "llava"
prompt_files: Dict[str, str] = dataclasses.field(default_factory=lambda: {})
timeout_s: float = 5.0
Expand Down Expand Up @@ -346,11 +346,30 @@ def _enrich(self, logger, n: ObjectNotification) -> ObjectNotification:
enriched_class=model_desc,
)

def _load_enrichment_model(self, logger):
if not self._config.enrichment.enable:
return
try:
requests.post(
self._config.enrichment.endpoint,
json={
"model": self._config.enrichment.model,
"keep_alive": self._config.enrichment.keep_alive,
},
)
logger.info(f"enrichment model {self._config.enrichment.model} preloaded")
except requests.RequestException as e:
logger.error(
f"enrichment model {self._config.enrichment.model} preload failed: {e}"
)

def _run(self):
logger = logging.getLogger(__name__)
logging.basicConfig(level=self._config.log_level, format=LOG_DEFAULT_FMT)
logger.info("starting notifier")

self._load_enrichment_model(logger)

while True:
n: Notification = self._input_queue.get()
logger.debug("received notification: " + n.message())
Expand Down

0 comments on commit 85fedda

Please sign in to comment.