Move example model training to tasks

Author: naglepuff

bats_ai/core/management/commands/examplelog.py

@@ -1,13 +1,20 @@
-from django.conf import settings
 import djclick as click
-from mlflow import MlflowClient
-from sklearn.datasets import load_diabetes
-from sklearn.ensemble import RandomForestRegressor
-from sklearn.model_selection import train_test_split
+import mlflow
 
+from bats_ai.core.tasks import example_train
 
-@click.command()
-def command():
-    click.echo("Running Mlflow experiment")
 
-    client = MlflowClient(tracking_uri=settings.MLFLOW_ENDPOINT)
+@click.command()
+@click.option('--experiment-name', type=click.STRING, required=False, default='Default')
+def command(experiment_name):
+    click.echo('Finding experiment')
+    experiment = mlflow.get_experiment_by_name(experiment_name)
+    if experiment:
+        click.echo(f'Creating a log for experiment {experiment_name}')
+        example_train.delay(experiment_name)
+        # train_body(experiment_name)
+    else:
+        click.echo(
+            f'Could not find experiment {experiment_name}.'
+            ' Use the create experiment command to create a new experiement.'
+        )

bats_ai/core/tasks.py

@@ -4,6 +4,7 @@
 from PIL import Image
 from celery import shared_task
 import cv2
+from django.conf import settings
 from django.core.files import File
 import numpy as np
 import scipy
@@ -243,3 +244,56 @@ def predict(compressed_spectrogram_id: int):
     recording_annotation.species.set(species)
     recording_annotation.save()
     return label, score, confs
+
+
+def train_body(experiment_name: str):
+    import mlflow
+    from mlflow.models import infer_signature
+    from sklearn import datasets
+    from sklearn.linear_model import LogisticRegression
+    from sklearn.metrics import accuracy_score
+    from sklearn.model_selection import train_test_split
+
+    X, y = datasets.load_iris(return_X_y=True)
+
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+
+    params = {
+        'solver': 'lbfgs',
+        'max_iter': 1000,
+        'multi_class': 'auto',
+        'random_state': 8888,
+    }
+
+    lr = LogisticRegression(**params)
+    lr.fit(X_train, y_train)
+
+    y_pred = lr.predict(X_test)
+
+    accuracy = accuracy_score(y_test, y_pred)
+
+    mlflow.set_tracking_uri(settings.MLFLOW_ENDPOINT)
+    mlflow.set_experiment(experiment_name)
+
+    print(mlflow.get_tracking_uri())
+    print(mlflow.get_artifact_uri())
+
+    mlflow.end_run()
+    with mlflow.start_run():
+        mlflow.log_params(params)
+        mlflow.log_metric('accuracy', accuracy)
+        mlflow.set_tag('Training Info', 'Basic LR model for iris data')
+
+        signature = infer_signature(X_train, lr.predict(X_train))
+        _ = mlflow.sklearn.log_model(
+            sk_model=lr,
+            artifact_path='iris_model',
+            signature=signature,
+            input_example=X_train,
+            registered_model_name='tracking-quickstart',
+        )
+
+
+@shared_task
+def example_train(experiment_name: str):
+    train_body(experiment_name)