File size: 3,699 Bytes
17c5137 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
from src.entity.config_entity import TRANSFORMER_OBJECT_FILE_NAME
from src.entity.config_entity import MODEL_FILE_NAME
from src.entity.config_entity import TARGET_ENCODER_OBJECT_FILE_NAME
from src.exception import CropException
from src.logger import logging
import os
import sys
from glob import glob
from typing import Optional
class ModelResolver:
def __init__(
self,
model_registry: str = "saved_models",
transformer_dir_name="transformer",
target_encoder_dir_name="target_encoder",
model_dir_name="model",
):
self.model_registry = model_registry
os.makedirs(self.model_registry, exist_ok=True)
self.transformer_dir_name = transformer_dir_name
self.target_encoder_dir_name = target_encoder_dir_name
self.model_dir_name = model_dir_name
def get_latest_dir_path(self) -> Optional[str]:
try:
dir_names = os.listdir(self.model_registry)
if len(dir_names) == 0:
return None
dir_names = list(map(int, dir_names))
latest_dir_name = max(dir_names)
return os.path.join(self.model_registry, f"{latest_dir_name}")
except Exception as e:
raise CropException(e, sys)
def get_latest_model_path(self):
try:
latest_dir = self.get_latest_dir_path()
if latest_dir is None:
raise Exception(f"Model is not available")
return os.path.join(latest_dir, self.model_dir_name, MODEL_FILE_NAME)
except Exception as e:
raise CropException(e, sys)
def get_latest_transformer_path(self):
try:
latest_dir = self.get_latest_dir_path()
if latest_dir is None:
raise Exception(f"Transformer is not availabel")
return os.path.join(latest_dir, self.transformer_dir_name, TRANSFORMER_OBJECT_FILE_NAME)
except Exception as e:
raise CropException(e, sys)
def get_latest_target_encoder_path(self):
try:
latest_dir = self.get_latest_dir_path()
if latest_dir is None:
raise Exception(f"Target encoder is not available")
return os.path.join(latest_dir, self.target_encoder_dir_name, TARGET_ENCODER_OBJECT_FILE_NAME)
except Exception as e:
raise CropException(e, sys)
def get_latest_save_dir_path(self):
try:
latest_dir = self.get_latest_dir_path()
if latest_dir == None:
return os.path.join(self.model_registry, f"{0}")
latest_dir_num = int(os.path.basename(self.get_latest_dir_path()))
return os.path.join(self.model_registry, f"{latest_dir_num + 1}")
except Exception as e:
raise CropException(e, sys)
def get_latest_save_model_path(self):
try:
latest_dir = self.get_latest_save_dir_path()
return os.path.join(latest_dir, self.model_dir_name, MODEL_FILE_NAME)
except Exception as e:
raise CropException(e, sys)
def get_latest_save_transformer_path(self):
try:
latest_dir = self.get_latest_save_dir_path()
return os.path.join(latest_dir, self.transformer_dir_name, TRANSFORMER_OBJECT_FILE_NAME)
except Exception as e:
raise CropException(e, sys)
def get_latest_save_target_encoder_path(self):
try:
latest_dir = self.get_latest_save_dir_path()
return os.path.join(latest_dir, self.target_encoder_dir_name, TARGET_ENCODER_OBJECT_FILE_NAME)
except Exception as e:
raise CropException(e, sys)
|