Простой пример переноса приложения машинного обучения в интерфейс командной строки

Когда мы создаем приложение любого типа, мы должны предоставить интерфейс, с помощью которого другие люди смогут его использовать. Я хочу показать вам простой пример переноса приложения ML в интерфейс командной строки.

Я думаю, можно с уверенностью сказать, что мы хотим получить наилучшие результаты при минимальных усилиях. В этом случае лучшим выбором будет fire библиотека. fire - это библиотека, разработанная и широко используемая Google. Давайте рассмотрим его функциональность на конкретных примерах:

main.py
import fire


def add(a: int, b: int):
    """
    Returns sum of a and b
    :param a: first argument
    :param b: second argument
    :return: sum of a and b
    """
    return a+b


if __name__ == "__main__":
    fire.Fire({
        "sum": add
    })

Здесь fire создает команду CLI sum из функции add. Из документа python он создает описание команды.

Вывод python main.py sum --help:

NAME
    main.py sum - Returns sum of a and b
SYNOPSIS
    main.py sum A B
DESCRIPTION
    Returns sum of a and b
POSITIONAL ARGUMENTS
    A
        first argument
    B
        second argument
NOTES
    You can also use flags syntax for POSITIONAL ARGUMENTS

Итак, теперь вы можете использовать эту команду несколькими способами:

python main.py 1 2

python main.py --a 1 --b 2

Кроме того, вы можете определить необязательные аргументы команды, указав значение по умолчанию в определении функции:

def add(a: int, b: int = 2):
    """
    Returns sum of a and b
    :param a: first argument
    :param b: second argument (default: 2)
    :return: sum of a and b
    """
    return a+b

Теперь результат python main.py sum --help:

NAME
    main.py sum - Returns sum of a and b
SYNOPSIS
    main.py sum A <flags>
DESCRIPTION
    Returns sum of a and b
POSITIONAL ARGUMENTS
    A
        first argument
FLAGS
    --b=B
        second argument (default: 2)
NOTES
    You can also use flags syntax for POSITIONAL ARGUMENTS

Итак, теперь вы можете использовать эту команду несколькими способами:

python main.py 1

python main.py 1 2

python main.py 1 --b 2

python main.py --a 1 --b 2

Больше не нужно слов, вы можете увидеть, насколько простой и быстрой может быть настройка приложения с интерфейсом командной строки.

Теперь давайте создадим скелет программы для типичного варианта использования приложения машинного обучения. Предположим, вам нужно создать приложение CLI с двумя функциями train и predict. Входными данными функции поезда являются данные поезда и некоторые параметры модели, выходными данными является файл обученной модели. Predict принимает входные данные для прогнозов, обученную модель и сохраняет прогнозы в выходной файл.

За исключением сопоставления функций, fire может создать интерфейс CLI из экземпляра класса python, что в данном случае может быть лучшим вариантом. В нашем случае код может быть таким:

import fire
import pickle
import logging

import pandas as pd
from sklearn.neighbors import KNeighborsClassifier

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

def load_model(model_path: str):
    """Loads model from `model_path`"""
    with open(model_path, 'rb') as file:
        saved_model = pickle.load(file)
        return saved_model

def save_model(model, model_path: str):
    """Saves `model` to `model_path`"""
    with open(model_path, 'wb') as file:
        pickle.dump(model, file)

class Classifier:
    """
    Some classifier, that makes some random classifications.
    """

    def train(self, train_data_path: str, model_path: str, k: int = 5):
        """
        Trains model on `train_data_path` data and saves trained model to `model_path`.
        Additionaly you can set KNN classifier `k` parameter.
        :param train_data_path: path to train data in csv format
        :param model_path: path to save model to.
        :param k: k-neighbors parameter of model.
        """
        logger.info(f"Loading train data from {train_data_path} ...")
        df = pd.read_csv(train_data_path)
        X = df.drop(columns=['y'])
        y = df['y']

        logger.info("Running model training...")
        model = KNeighborsClassifier(n_neighbors=k)
        model.fit(X, y)

        logger.info(f"Saving model to {model_path} ...")
        save_model(model, model_path)

        logger.info("Successfully trained model.")

    def predict(self, predict_data_path: str, model_path: str, output_path: str):
        """
        Predicts `predict_data_path` data using `model_path` model and saves predictions to `output_path`
        :param predict_data_path: path to data for predictions
        :param model_path: path to trained model
        :param output_path: path to save predictions
        """
        logger.info(f"Loading data for predictions from {predict_data_path} ...")
        X = pd.read_csv(predict_data_path)

        logger.info(f"Loading model from {model_path} ...")
        model = load_model(model_path)

        logger.info("Running model predictions...")
        y_pred = model.predict(X)

        logger.info(f"Saving predictions to {output_path} ...")
        pd.DataFrame(y_pred).to_csv(output_path)

        logger.info("Successfully predicted.")

if __name__ == "__main__":
    fire.Fire(Classifier)

Результат python main.py:

NAME
    main.py - Some classifier, that makes some random classifications.
SYNOPSIS
    main.py COMMAND
DESCRIPTION
    Some classifier, that makes some random classifications.
COMMANDS
    COMMAND is one of the following:
predict
       Predicts `predict_data_path` data using `model_path` model and saves predictions to `output_path`
train
       Trains model on `train_data_path` data and saves trained model to `model_path`. Additionaly you can set KNN classifier `k` parameter.

И соответствующая командная документация:

python main.py train --help :

NAME
    main.py train - Trains model on `train_data_path` data and saves trained model to `model_path`. Additionaly you can set KNN classifier `k` parameter.
SYNOPSIS
    main.py train TRAIN_DATA_PATH MODEL_PATH <flags>
DESCRIPTION
    Trains model on `train_data_path` data and saves trained model to `model_path`. Additionaly you can set KNN classifier `k` parameter.
POSITIONAL ARGUMENTS
    TRAIN_DATA_PATH
        path to train data in csv format
    MODEL_PATH
        path to save model to.
FLAGS
    --k=K
        k-neighbors parameter of model.
NOTES
    You can also use flags syntax for POSITIONAL ARGUMENTS

python main.py predict --help :

NAME
    main.py predict - Predicts `predict_data_path` data using `model_path` model and saves predictions to `output_path`
SYNOPSIS
    main.py predict PREDICT_DATA_PATH MODEL_PATH OUTPUT_PATH
DESCRIPTION
    Predicts `predict_data_path` data using `model_path` model and saves predictions to `output_path`
POSITIONAL ARGUMENTS
    PREDICT_DATA_PATH
        path to data for predictions
    MODEL_PATH
        path to trained model
    OUTPUT_PATH
        path to save predictions
NOTES
    You can also use flags syntax for POSITIONAL ARGUMENTS

Пример использования:

python main.py train train.csv knn.sav --k 7
python main.py predict test.csv knn.sav out.csv

Теперь вы можете использовать этот фрагмент для создания собственного приложения CLI. Несколько советов по дальнейшему улучшению:

  1. Использовать ведение журнала вместо печати
  2. Создание валидации аргументов
  3. Создайте шаблон README с кратким описанием использования приложения для конечного пользователя.
  4. Узнайте больше о других интересных функциях на https://github.com/google/python-fire

Надеюсь, эта статья будет полезной и информативной. Жду ваших отзывов!