import mlflow
import mlflow.sklearn
from sklearn.datasets import load_diabetes
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score
# Load data
data = load_diabetes(as_frame=True)
df = data.frame
X = df.drop("target", axis=1)
y = df["target"]
# Split data
X_train, X_test, y_train, y_test = train_test_split(
X,
y,
test_size=0.2,
random_state=42
)
# MLflow
mlflow.set_tracking_uri("sqlite:///mlflow.db")
mlflow.set_experiment("MLOps Demo 1")
with mlflow.start_run():
# Train
model = LinearRegression()
model.fit(X_train, y_train)
# Predict
predictions = model.predict(X_test)
# Metrics
mse = mean_squared_error(
y_test,
predictions
)
rmse = mse ** 0.5
r2 = r2_score(
y_test,
predictions
)
# Log parameters
mlflow.log_param(
"model",
"LinearRegression"
)
mlflow.log_param(
"test_size",
0.2
)
# Log metrics
mlflow.log_metric("mse", mse)
mlflow.log_metric("rmse", rmse)
mlflow.log_metric("r2", r2)
# Register model
mlflow.sklearn.log_model(
model,
name="diabetes_model",
registered_model_name="DiabetesRegression"
)
print("Training completed")
print(f"MSE: {mse:.4f}")
print(f"RMSE: {rmse:.4f}")
print(f"R2: {r2:.4f}")