<commit_before># pep8: disable=E501

from __future__ import print_function

import collections
import shutil
import pytest
import tempfile

import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.python.keras import layers

import mlflow
import mlflow.tensorflow
import mlflow.keras

import os

SavedModelInfo = collections.namedtuple(
    "SavedModelInfo",
    ["path", "meta_graph_tags", "signature_def_key", "inference_df", "expected_results_df"])


@pytest.fixture
def random_train_data():
    return np.random.random((1000, 32))


@pytest.fixture
def random_one_hot_labels():
    n, n_class = (1000, 10)
    classes = np.random.randint(0, n_class, n)
    labels = np.zeros((n, n_class))
    labels[np.arange(n), classes] = 1
    return labels


@pytest.fixture(params=[True, False])
@pytest.mark.usefixtures("tracking_uri_mock")
def manual_run(request):
    if request.param:
        mlflow.start_run()
    yield
    mlflow.end_run()


def create_tf_keras_model():
    model = tf.keras.Sequential()

    model.add(layers.Dense(64, activation='relu', input_shape=(32,)))
    model.add(layers.Dense(64, activation='relu'))
    model.add(layers.Dense(10, activation='softmax'))

    model.compile(optimizer=tf.train.AdamOptimizer(0.001),
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])
    return model


@pytest.mark.large
@pytest.mark.parametrize('fit_variant', ['fit', 'fit_generator'])
def test_tf_keras_autolog_ends_auto_created_run(random_train_data, random_one_hot_labels,
                                                fit_variant):
    mlflow.tensorflow.autolog()

    data = random_train_data
    labels = random_one_hot_labels

    model = create_tf_keras_model()

    model.fit(data, labels, epochs=10)

    assert mlflow.active_run() is None


@pytest.mark.large
@pytest.mark.parametrize('fit_variant', ['fit', 'fit_generator'])
def test_tf_keras_autolog_persists_manually_created_run(random_train_data, random_one_hot_labels,
                                                        fit_variant):
    mlflow.tensorflow.autolog()
    with mlflow.start_run() as run:
        data = random_train_data
        labels = random_one_hot_labels

        model = create_tf_keras_model()

        model.fit(data, labels, epochs=10)

        assert mlflow.active_run()
        assert mlflow.active_run().info.run_id == run.info.run_id


@pytest.fixture
def tf_keras_random_data_run(random_train_data, random_one_hot_labels, manual_run, fit_variant):
    mlflow.tensorflow.autolog(every_n_iter=5)

    data = random_train_data
    labels = random_one_hot_labels

    model = create_tf_keras_model()

    if fit_variant == 'fit_generator':
        def generator():
            while True:
                yield data, labels
        model.fit_generator(generator(), epochs=10, steps_per_epoch=1)
    else:
        model.fit(data, labels, epochs=10, steps_per_epoch=1)

    client = mlflow.tracking.MlflowClient()
    return client.get_run(client.list_run_infos(experiment_id='0')[0].run_id)


@pytest.mark.large
@pytest.mark.parametrize('fit_variant', ['fit', 'fit_generator'])
def test_tf_keras_autolog_logs_expected_data(tf_keras_random_data_run):
    data = tf_keras_random_data_run.data
    assert 'epoch_acc' in data.metrics
    assert 'epoch_loss' in data.metrics
    # Testing explicitly passed parameters are logged correctly
    assert 'epochs' in data.params
    assert data.params['epochs'] == '10'
    assert 'steps_per_epoch' in data.params
    assert data.params['steps_per_epoch'] == '1'
    # Testing default parameters are logged correctly
    assert 'initial_epoch' in data.params
    assert data.params['initial_epoch'] == '0'
    # Testing unwanted parameters are not logged
    assert 'callbacks' not in data.params
    assert 'validation_data' not in data.params
    # Testing optimizer parameters are logged
    assert 'optimizer_name' in data.params
    assert data.params['optimizer_name'] == 'AdamOptimizer'
    assert 'model_summary' in data.tags
    assert 'Total params: 6,922' in data.tags['model_summary']
    client = mlflow.tracking.MlflowClient()
    all_epoch_acc = client.get_metric_history(tf_keras_random_data_run.info.run_id, 'epoch_acc')
    assert all((x.step - 1) % 5 == 0 for x in all_epoch_acc)
    artifacts = client.list_artifacts(tf_keras_random_data_run.info.run_id)
    artifacts = map(lambda x: x.path, artifacts)
    assert 'model_summary.txt' in artifacts


@pytest.mark.large
@pytest.mark.parametrize('fit_variant', ['fit', 'fit_generator'])
def test_tf_keras_autolog_model_can_load_from_artifact(tf_keras_random_data_run, random_train_data):
    client = mlflow.tracking.MlflowClient()
    artifacts = client.list_artifacts(tf_keras_random_data_run.info.run_id)
    artifacts = map(lambda x: x.path, artifacts)
    assert 'model' in artifacts
    assert 'tensorboard_logs' in artifacts
    model = mlflow.keras.load_model("runs:/" + tf_keras_random_data_run.info.run_id +
                                    "/model")
    model.predict(random_train_data)


@pytest.fixture
def tf_keras_random_data_run_with_callback(random_train_data, random_one_hot_labels, manual_run,
                                           callback, restore_weights, patience):
    mlflow.tensorflow.autolog(every_n_iter=1)

    data = random_train_data
    labels = random_one_hot_labels

    model = create_tf_keras_model()
    if callback == 'early':
        # min_delta is set as such to guarantee early stopping
        callback = tf.keras.callbacks.EarlyStopping(monitor='loss', patience=patience,
                                                    min_delta=99999999,
                                                    restore_best_weights=restore_weights)
    else:
        callback = tf.keras.callbacks.ProgbarLogger(count_mode='samples')

    history = model.fit(data, labels, epochs=10, callbacks=[callback])

    client = mlflow.tracking.MlflowClient()
    return client.get_run(client.list_run_infos(experiment_id='0')[0].run_id), history, callback


@pytest.mark.large<commit_msg>tensorflow test needs the mock...<commit_after># pep8: disable=E501

from __future__ import print_function

import collections
import shutil
import pytest
import tempfile

import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.python.keras import layers

import mlflow
import mlflow.tensorflow
import mlflow.keras

import os

SavedModelInfo = collections.namedtuple(
    "SavedModelInfo",
    ["path", "meta_graph_tags", "signature_def_key", "inference_df", "expected_results_df"])


@pytest.fixture
def random_train_data():
    return np.random.random((1000, 32))


@pytest.fixture
def random_one_hot_labels():
    n, n_class = (1000, 10)
    classes = np.random.randint(0, n_class, n)
    labels = np.zeros((n, n_class))
    labels[np.arange(n), classes] = 1
    return labels


@pytest.fixture(params=[True, False])
@pytest.mark.usefixtures("tracking_uri_mock")
def manual_run(request):
    if request.param:
        mlflow.start_run()
    yield
    mlflow.end_run()


def create_tf_keras_model():
    model = tf.keras.Sequential()

    model.add(layers.Dense(64, activation='relu', input_shape=(32,)))
    model.add(layers.Dense(64, activation='relu'))
    model.add(layers.Dense(10, activation='softmax'))

    model.compile(optimizer=tf.train.AdamOptimizer(0.001),
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])
    return model


@pytest.mark.large
@pytest.mark.parametrize('fit_variant', ['fit', 'fit_generator'])
def test_tf_keras_autolog_ends_auto_created_run(random_train_data, random_one_hot_labels,
                                                fit_variant):
    mlflow.tensorflow.autolog()

    data = random_train_data
    labels = random_one_hot_labels

    model = create_tf_keras_model()

    model.fit(data, labels, epochs=10)

    assert mlflow.active_run() is None


@pytest.mark.large
@pytest.mark.parametrize('fit_variant', ['fit', 'fit_generator'])
def test_tf_keras_autolog_persists_manually_created_run(random_train_data, random_one_hot_labels,
                                                        fit_variant):
    mlflow.tensorflow.autolog()
    with mlflow.start_run() as run:
        data = random_train_data
        labels = random_one_hot_labels

        model = create_tf_keras_model()

        model.fit(data, labels, epochs=10)

        assert mlflow.active_run()
        assert mlflow.active_run().info.run_id == run.info.run_id


@pytest.fixture
def tf_keras_random_data_run(random_train_data, random_one_hot_labels, manual_run, fit_variant):
    mlflow.tensorflow.autolog(every_n_iter=5)

    data = random_train_data
    labels = random_one_hot_labels

    model = create_tf_keras_model()

    if fit_variant == 'fit_generator':
        def generator():
            while True:
                yield data, labels
        model.fit_generator(generator(), epochs=10, steps_per_epoch=1)
    else:
        model.fit(data, labels, epochs=10, steps_per_epoch=1)

    client = mlflow.tracking.MlflowClient()
    return client.get_run(client.list_run_infos(experiment_id='0')[0].run_id)


@pytest.mark.large
@pytest.mark.parametrize('fit_variant', ['fit', 'fit_generator'])
def test_tf_keras_autolog_logs_expected_data(tf_keras_random_data_run):
    data = tf_keras_random_data_run.data
    assert 'epoch_acc' in data.metrics
    assert 'epoch_loss' in data.metrics
    # Testing explicitly passed parameters are logged correctly
    assert 'epochs' in data.params
    assert data.params['epochs'] == '10'
    assert 'steps_per_epoch' in data.params
    assert data.params['steps_per_epoch'] == '1'
    # Testing default parameters are logged correctly
    assert 'initial_epoch' in data.params
    assert data.params['initial_epoch'] == '0'
    # Testing unwanted parameters are not logged
    assert 'callbacks' not in data.params
    assert 'validation_data' not in data.params
    # Testing optimizer parameters are logged
    assert 'optimizer_name' in data.params
    assert data.params['optimizer_name'] == 'AdamOptimizer'
    assert 'model_summary' in data.tags
    assert 'Total params: 6,922' in data.tags['model_summary']
    client = mlflow.tracking.MlflowClient()
    all_epoch_acc = client.get_metric_history(tf_keras_random_data_run.info.run_id, 'epoch_acc')
    assert all((x.step - 1) % 5 == 0 for x in all_epoch_acc)
    artifacts = client.list_artifacts(tf_keras_random_data_run.info.run_id)
    artifacts = map(lambda x: x.path, artifacts)
    assert 'model_summary.txt' in artifacts


@pytest.mark.large
@pytest.mark.parametrize('fit_variant', ['fit', 'fit_generator'])
def test_tf_keras_autolog_model_can_load_from_artifact(tf_keras_random_data_run, random_train_data):
    client = mlflow.tracking.MlflowClient()
    artifacts = client.list_artifacts(tf_keras_random_data_run.info.run_id)
    artifacts = map(lambda x: x.path, artifacts)
    assert 'model' in artifacts
    assert 'tensorboard_logs' in artifacts
    model = mlflow.keras.load_model("runs:/" + tf_keras_random_data_run.info.run_id +
                                    "/model")
    model.predict(random_train_data)


@pytest.fixture
def tf_keras_random_data_run_with_callback(random_train_data, random_one_hot_labels, manual_run,
                                           callback, restore_weights, patience):
    mlflow.tensorflow.autolog(every_n_iter=1)

    data = random_train_data
    labels = random_one_hot_labels

    model = create_tf_keras_model()
    if callback == 'early':
        # min_delta is set as such to guarantee early stopping
        callback = tf.keras.callbacks.EarlyStopping(monitor='loss', patience=patience,
                                                    min_delta=99999999,
                                                    restore_best_weights=restore_weights)
    else:
        callback = tf.keras.callbacks.ProgbarLogger(count_mode='samples')

    history = model.fit(data, labels, epochs=10, callbacks=[callback])

    client = mlflow.tracking.MlflowClient()
    return client.get_run(client.list_run_infos(experiment_id='0')[0].run_id), history, callback


@pytest.mark.large
@pytest.mark.usefixtures("tracking_uri_mock")
{'fraction_non_alphanumeric': 0.09960809928151534, 'fraction_numerical': 0.010940561724363161, 'mean_word_length': 3.5417130144605116, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 22, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3396138', 'n_tokens_mistral': 4053, 'n_tokens_neox': 3923, 'n_words': 746}