src/inference/predict.py

import mlflow
import mlflow.sklearn


def load_model_from_run(run_id: str):
    mlflow.set_tracking_uri("file:./mlruns")  # FIXED

    model_uri = f"runs:/{run_id}/model"
    model = mlflow.sklearn.load_model(model_uri)

    return model

def predict(model, X):
    return model.predict(X)