Spaces:
Sleeping
Sleeping
File size: 1,184 Bytes
a8b81f3 |
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 |
from abc import ABC, abstractmethod
import joblib
import logging
import pickle
from pathlib import Path
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class BaseModel(ABC):
def __init__(self, model_path):
self.model_path = model_path
self.model = None
self.scaler = None
self.X_train = None
self.y_train = None
@abstractmethod
def train(self, X, y):
pass
@abstractmethod
def predict(self, X):
pass
def save_model(self):
model_data = {
'model': self.model,
'scaler': self.scaler,
'X_train': self.X_train,
'y_train': self.y_train
}
with open(self.model_path, 'wb') as f:
pickle.dump(model_data, f)
@classmethod
def load_model(cls):
instance = cls()
with open(instance.model_path, 'rb') as f:
model_data = pickle.load(f)
instance.model = model_data['model']
instance.scaler = model_data['scaler']
instance.X_train = model_data['X_train']
instance.y_train = model_data['y_train']
return instance |