#!/usr/bin/env python

# # Simple Model Training Example
#
# This example is the API example for this Ludwig command line example
# (https://ludwig-ai.github.io/ludwig-docs/latest/examples/mnist/).
import logging
import shutil

import yaml

from ludwig.api import LudwigModel
from ludwig.datasets import mnist
import pandas as pd

with open("./config.yaml") as f:
    config = yaml.safe_load(f.read())

# modify config
config['trainer']['epochs'] = 5
config['trainer']['early_stop'] = -1

# use RAY backend
model = LudwigModel(config, backend = {"type":"ray"}, logging_level=logging.INFO)

# load data and select 20 rows
df = mnist.load(split=False)
df = df.drop('split',axis=1)
df1 = df[0:10]
df2 = df[-10:]
df = pd.concat([df1, df2],axis=0)

# clean out prior results
shutil.rmtree("./results", ignore_errors=True)

(train_stats, _, output_directory) = model.train(
    dataset=df,
    experiment_name="simple_image_experiment",
    model_name="single_model",
    skip_save_processed_input=True,
)