!pip install mlflow

import mlflow
import mlflow.sklearn
import pandas as pd
from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression, Ridge
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error

data = load_diabetes()
X = pd.DataFrame(data.data, columns=data.feature_names)
y = pd.Series(data.target)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

mlflow.set_experiment("diabetes-regression")

models_to_try = {
    "LinearRegression": LinearRegression(),
    "Ridge(alpha=1.0)": Ridge(alpha=1.0),
    "RandomForest(n_estimators=100)": RandomForestRegressor(n_estimators=100, random_state=42),
}

for name, mdl in models_to_try.items():
    with mlflow.start_run(run_name=name):
        mdl.fit(X_train, y_train)
        y_pred = mdl.predict(X_test)
        mse = mean_squared_error(y_test, y_pred)

        mlflow.log_param("model_type", name)
        mlflow.log_metric("mse", mse)
        
        # Use skops_trusted_types to allow saving complex estimators like RandomForest
        mlflow.sklearn.log_model(
            sk_model=mdl,
            artifact_path="model",
            skops_trusted_types=["sklearn.tree._tree.Tree"]
        )

        print(f"{name}: MSE = {mse:.2f}")

!pip install pyngrok

!pip install -q pyngrok
import subprocess, time
from pyngrok import ngrok

NGROK_TOKEN = "PASTE_YOUR_NGROK_AUTHTOKEN_HERE"
ngrok.set_auth_token(NGROK_TOKEN)

mlflow_process = subprocess.Popen(["mlflow", "ui", "--port", "5001"])
time.sleep(5)

public_url = ngrok.connect(5001)
print("MLflow UI available at:", public_url)

runs = mlflow.search_runs(experiment_names=["diabetes-regression"], order_by=["metrics.mse ASC"])
print(runs[["tags.mlflow.runName", "params.model_type", "metrics.mse"]])