refactor(ml): modularization and styling (#2835)

* basic refactor and styling

* removed batching

* module entrypoint

* removed unused imports

* model superclass,  model cache now in app state

* fixed cache dir and enforced abstract method

---------

Co-authored-by: Alex Tran <alex.tran1502@gmail.com>
This commit is contained in:
Mert
2023-06-24 23:18:09 -04:00
committed by GitHub
parent 837ad24f58
commit a2f5674bbb
12 changed files with 281 additions and 182 deletions

View File

@@ -1,5 +1,10 @@
from pathlib import Path
from pydantic import BaseSettings
from .schemas import ModelType
class Settings(BaseSettings):
cache_folder: str = "/cache"
classification_model: str = "microsoft/resnet-50"
@@ -15,8 +20,12 @@ class Settings(BaseSettings):
min_face_score: float = 0.7
class Config(BaseSettings.Config):
env_prefix = 'MACHINE_LEARNING_'
env_prefix = "MACHINE_LEARNING_"
case_sensitive = False
def get_cache_dir(model_name: str, model_type: ModelType) -> Path:
return Path(settings.cache_folder, model_type.value, model_name)
settings = Settings()