wehub-resource-sync
27 lines
707 B
Python
27 lines
707 B
Python
import sklearn
|
|
|
|
import mlflow
|
|
|
|
# Use explicit model logging to control the conda environment and pip requirements
|
|
mlflow.sklearn.autolog(log_models=False)
|
|
|
|
# Load data
|
|
X, y = sklearn.datasets.load_diabetes(return_X_y=True)
|
|
X_train, X_test, y_train, y_test = sklearn.model_selection.train_test_split(
|
|
X, y, test_size=0.2, random_state=0
|
|
)
|
|
|
|
# Train model
|
|
with mlflow.start_run() as run:
|
|
print(f"MLflow run ID: {run.info.run_id}")
|
|
|
|
model = sklearn.linear_model.Ridge(alpha=0.03)
|
|
model.fit(X_train, y_train)
|
|
|
|
mlflow.sklearn.log_model(
|
|
model,
|
|
name="model",
|
|
signature=mlflow.models.infer_signature(X_train[:10], y_train[:10]),
|
|
input_example=X_train[:10],
|
|
)
|