Model training

pip install scikit-learn mlflow==3.5.1

# Original source code and more details can be found in:
# https://www.mlflow.org/docs/latest/tutorials-and-examples/tutorial.html

# The data set used in this example is from
# http://archive.ics.uci.edu/ml/datasets/Wine+Quality
# P. Cortez, A. Cerdeira, F. Almeida, T. Matos and J. Reis.
# Modeling wine preferences by data mining from physicochemical properties.
# In Decision Support Systems, Elsevier, 47(4):547-553, 2009.

import warnings
import sys

import pandas as pd
import numpy as np
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from sklearn.model_selection import train_test_split
from sklearn.linear_model import ElasticNet
from urllib.parse import urlparse
import mlflow
import mlflow.sklearn
from mlflow.models.signature import infer_signature

import logging

logging.basicConfig(level=logging.WARN)
logger = logging.getLogger(__name__)


def eval_metrics(actual, pred):
    rmse = np.sqrt(mean_squared_error(actual, pred))
    mae = mean_absolute_error(actual, pred)
    r2 = r2_score(actual, pred)
    return rmse, mae, r2


if __name__ == "__main__":
    warnings.filterwarnings("ignore")
    np.random.seed(40)

    # Read the wine-quality csv file from the URL
    csv_url = (
        "http://archive.ics.uci.edu/ml"
        "/machine-learning-databases/wine-quality/winequality-red.csv"
    )
    try:
        data = pd.read_csv(csv_url, sep=";")
    except Exception as e:
        logger.exception(
            "Unable to download training & test CSV, "
            "check your internet connection. Error: %s",
            e,
        )
        
    # Split the data into training and test sets. (0.75, 0.25) split.
    train, test = train_test_split(data)

    # The predicted column is "quality" which is a scalar from [3, 9]
    train_x = train.drop(["quality"], axis=1)
    test_x = test.drop(["quality"], axis=1)
    train_y = train[["quality"]]
    test_y = test[["quality"]]

    alpha = 0.5
    l1_ratio = 0.5

    experiment_name = "wine-classification"

    existing_experiment = mlflow.get_experiment_by_name(experiment_name)
    if existing_experiment is None:
        experiment_id = mlflow.create_experiment(
            name=experiment_name
        )
    else:
        experiment_id = existing_experiment.experiment_id

    mlflow.set_experiment(experiment_name)

    # Add or update tags to the created experiment.
    mlflow.set_experiment_tags({
        "project_name": "Fraud Prevention",
        "team": "Data Science Core",
        "priority": "High"
    })

    with mlflow.start_run(experiment_id=experiment_id):
        lr = ElasticNet(alpha=alpha, l1_ratio=l1_ratio, random_state=42)
        lr.fit(train_x, train_y)

        predicted_qualities = lr.predict(test_x)

        (rmse, mae, r2) = eval_metrics(test_y, predicted_qualities)

        print("Elasticnet model (alpha=%f, l1_ratio=%f):" % (alpha, l1_ratio))
        print("  RMSE: %s" % rmse)
        print("  MAE: %s" % mae)
        print("  R2: %s" % r2)

        mlflow.log_param("alpha", alpha)
        mlflow.log_param("l1_ratio", l1_ratio)
        mlflow.log_metric("rmse", rmse)
        mlflow.log_metric("r2", r2)
        mlflow.log_metric("mae", mae)
        # You can tag each run under an experiment independently.
        mlflow.set_tag("version", "1.0")

        tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
        model_signature = infer_signature(train_x, train_y)

        if tracking_url_type_store != "file":
            mlflow.sklearn.log_model(
                lr,
                "my-new-model",
                registered_model_name="ElasticnetWineModel",
                input_example=train_x.head(1),
                signature=model_signature,
            )
        else:
            mlflow.sklearn.log_model(lr, "model", signature=model_signature)
        
print("done.")

from torch.fft import Tensor
from mlflow.types import Schema, TensorSpec
import mlflow
import mlflow.pytorch
from mlflow.models import ModelSignature
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import numpy as np
import logging

logging.basicConfig(level=logging.WARN)
logger = logging.getLogger(__name__)

# Transformations for MNIST images
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# Load MNIST dataset
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('./data', train=False, download=True, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

# Simple neural network
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(28*28, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(-1, 28*28)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Define input schema
inputs_schema = Schema([TensorSpec(type=np.dtype(np.float32), shape=(-1, 1, 28, 28))])

# Define output schema
outputs_schema = Schema([TensorSpec(type=np.dtype(np.float32), shape=(-1, 10))])

# Create the signature
model_signature = ModelSignature(inputs=inputs_schema, outputs=outputs_schema)

model = Net()
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

experiment_name = "Testing"

existing_experiment = mlflow.get_experiment_by_name(experiment_name)
if existing_experiment is None:
    experiment_id = mlflow.create_experiment(
        name=experiment_name,
        artifact_location="mlflow-artifacts:/pytorch-artifacts"
    )
else:
    experiment_id = existing_experiment.experiment_id

mlflow.set_experiment(experiment_name)

# Add or update tags to the created experiment.
mlflow.set_experiment_tags({
    "project_name": "Fraud Prevention",
    "team": "Data Science Core",
    "priority": "High"
})

input_example = False

with mlflow.start_run():
    mlflow.log_param("lr", 0.001)
    mlflow.log_param("batch_size", 64)

    for epoch in range(5):  # train for 5 epochs
        model.train()
        train_loss = 0
        correct = 0
        total = 0

        for data, target in train_loader:
            input_example = data.numpy()
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            _, predicted = torch.max(output.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()

        epoch_loss = train_loss / len(train_loader)
        epoch_acc = correct / total

        # Log metrics per epoch
        mlflow.log_metric("train_loss", epoch_loss, step=epoch)
        mlflow.log_metric("train_accuracy", epoch_acc, step=epoch)
        # You can tag each run under an experiment independently.
        mlflow.set_tag("version", "1.0")

        print(f"Epoch {epoch+1}, Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc:.4f}")

    # Save trained model
    mlflow.pytorch.log_model(model, "mnist_model", signature=model_signature, input_example=input_example)