examples/hamilton_ui/components/models.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

"""This module contains specific incarnations of models."""

from sklearn import base

from hamilton.function_modifiers import source, subdag

try:
    import model_fitting
except ImportError:
    from . import model_fitting


@subdag(
    model_fitting,
    inputs={
        "data_set": source("data_set_v1"),
    },
    config={"clf": "svm", "shuffle_train_test_split": True, "test_size_fraction": 0.2},
)
def svm_model(
    fit_clf: base.ClassifierMixin, training_accuracy: float, testing_accuracy: float
) -> dict:
    return {
        "svm": fit_clf,
        "training_accuracy": training_accuracy,
        "testing_accuracy": testing_accuracy,
    }


@subdag(
    model_fitting,
    inputs={
        "data_set": source("data_set_v1"),
    },
    config={
        "clf": "logistic",
        "shuffle_train_test_split": True,
        "test_size_fraction": 0.2,
        "penalty": "l2",
    },
)
def lr_model(
    fit_clf: base.ClassifierMixin, training_accuracy: float, testing_accuracy: float
) -> dict:
    return {
        "logistic": fit_clf,
        "training_accuracy": training_accuracy,
        "testing_accuracy": testing_accuracy,
    }


def best_model(svm_model: dict, lr_model: dict) -> dict:
    """Returns the best model based on the testing accuracy."""
    if svm_model["testing_accuracy"] > lr_model["testing_accuracy"]:
        return svm_model
    else:
        return lr_model