learnware/tests/test_hetero_market/example_learnwares/model0.py

from learnware.model import BaseModel
import numpy as np
import joblib
import os


class MyModel(BaseModel):
    def __init__(self):
        super(MyModel, self).__init__(input_shape=(20,), output_shape=(1,))
        dir_path = os.path.dirname(os.path.abspath(__file__))
        model_path = os.path.join(dir_path, "ridge.pkl")
        model = joblib.load(model_path)
        self.model = model

    def predict(self, X: np.ndarray) -> np.ndarray:
        return self.model.predict(X)