ludwig-ai--ludwig/tests/integration_tests/test_sequence_features.py

import contextlib
import copy
from io import StringIO

import pandas as pd
import pytest
import torch

from ludwig.api import LudwigModel
from ludwig.constants import DECODER, ENCODER_OUTPUT_STATE, LOGITS
from ludwig.data.dataset_synthesizer import build_synthetic_dataset
from ludwig.data.preprocessing import preprocess_for_training
from ludwig.features.feature_registries import update_config_with_metadata
from tests.integration_tests.utils import generate_data, run_experiment, sequence_feature

#
# this test is focused on testing input sequence features with all encoders
# and output sequence feature with Generator decoder.  Except for specified
# configuration parameters all other parameters assume default values.
#

TEST_VOCAB_SIZE = 132
TEST_HIDDEN_SIZE = 32
TEST_STATE_SIZE = 8
TEST_EMBEDDING_SIZE = 64
TEST_NUM_FILTERS = 24


# generates dataset that can be used for rest of test
@pytest.fixture(scope="module")
def generate_sequence_training_data():
    input_features = [
        sequence_feature(
            encoder={
                "vocab_size": TEST_VOCAB_SIZE,
                "embedding_size": TEST_EMBEDDING_SIZE,
                "state_size": TEST_STATE_SIZE,
                "hidden_size": TEST_HIDDEN_SIZE,
                "num_filters": TEST_NUM_FILTERS,
                "min_len": 5,
                "max_len": 10,
                "type": "rnn",
                "cell_type": "lstm",
            }
        )
    ]

    output_features = [
        sequence_feature(
            decoder={"type": "generator", "min_len": 5, "max_len": 10, "cell_type": "lstm", "attention": "bahdanau"}
        )
    ]

    # generate synthetic data set testing
    dataset = build_synthetic_dataset(150, copy.deepcopy(input_features) + copy.deepcopy(output_features))
    raw_data = "\n".join([r[0] + "," + r[1] for r in dataset])
    df = pd.read_csv(StringIO(raw_data))

    return df, input_features, output_features


# setups up minimal number of data structures required to support initialized
# input and output features.  The function returns initialized LudwigModel
# and batcher for training dataset
@contextlib.contextmanager
def setup_model_scaffolding(raw_df, input_features, output_features):
    # setup input feature for testing
    config = {"input_features": input_features, "output_features": output_features}

    # setup model scaffolding to for testing
    model = LudwigModel(config)
    training_set, _, _, training_set_metadata = preprocess_for_training(
        model.config, training_set=raw_df, skip_save_processed_input=True
    )
    model.training_set_metadata = training_set_metadata
    update_config_with_metadata(model.config_obj, training_set_metadata)
    model.model = model.create_model(model.config_obj)

    # setup batcher to go through synthetic data
    with training_set.initialize_batcher() as batcher:
        yield model, batcher


# TODO(#1333): Refactor this test once torch sequence generator work is complete.
# - Tests may be covered by other smaller scoped unit tests.
#
# tests output feature sequence with `Generator` decoder
# pytest parameters
#   dec_cell_type: decoder cell type
#   combiner_output_shapes: is a 2-tuple specifies the possible types of
#     tensors that the combiner may generate for sequences.
#     combiner_output_shapes[0]: specifies shape for hidden key
#     combiner_output_shapes[1]: is either None or 1 or 2-tuple representing
#       the encoder_output_state key. None: no encoder_output_state key,
#       1-tuple: generate tf.Tensor, 2-tuple: generate list with 2 tf.Tensors
# TODO(Justin): Move these to test_sequence_generator unit tests, and reintroduce decoder attention, beam_width, and
# num_layers when these are reimplemented.
@pytest.mark.parametrize(
    "dec_cell_type,combiner_output_shapes",
    [
        ("lstm", ((128, 10, TEST_STATE_SIZE), None)),
        ("rnn", ((128, 10, TEST_STATE_SIZE), ((128, TEST_STATE_SIZE), (128, TEST_STATE_SIZE)))),
        ("gru", ((128, 10, TEST_STATE_SIZE), ((128, TEST_STATE_SIZE),))),
    ],
    ids=["lstm_no_state", "rnn_dual_state", "gru_single_state"],
)
def test_sequence_decoders(
    dec_cell_type,
    combiner_output_shapes,
    generate_sequence_training_data,
):
    # retrieve pre-computed dataset and features
    raw_df = generate_sequence_training_data[0]
    input_features = generate_sequence_training_data[1]
    output_features = generate_sequence_training_data[2]
    output_feature_name = output_features[0]["name"]
    output_features[0][DECODER]["cell_type"] = dec_cell_type

    with setup_model_scaffolding(raw_df, input_features, output_features) as (model, _):
        # generate synthetic encoder_output tensors and make it look like
        # it came out of the combiner
        encoder_output = torch.randn(combiner_output_shapes[0])
        combiner_outputs = {"hidden": encoder_output}

        if combiner_output_shapes[1] is not None:
            if len(combiner_output_shapes[1]) > 1:
                encoder_output_state = (
                    torch.randn(combiner_output_shapes[1][0]),
                    torch.randn(combiner_output_shapes[1][1]),
                )
            else:
                encoder_output_state = torch.randn(combiner_output_shapes[1][0])

            combiner_outputs[ENCODER_OUTPUT_STATE] = encoder_output_state

        decoder = model.model.output_features.get(output_feature_name).decoder_obj
        decoder_out = decoder(combiner_outputs)

        # gather expected components of the shape
        batch_size = combiner_outputs["hidden"].shape[0]
        seq_size = output_features[0][DECODER]["max_len"] + 2  # For start and stop symbols.
        vocab_size = model.config_obj.output_features.to_list()[0][DECODER]["vocab_size"]

        # confirm shape and format of decoder output
        assert list(decoder_out[LOGITS].size()) == [batch_size, seq_size, vocab_size]


# final sanity test.  Checks a subset of sequence parameters
@pytest.mark.parametrize(
    "enc_encoder,enc_cell_type,dec_cell_type",
    [
        ("embed", "lstm", "lstm"),
        ("rnn", "rnn", "gru"),
        ("rnn", "gru", "rnn"),
    ],
    ids=["embed_lstm", "rnn_gru", "gru_rnn"],
)
def test_sequence_generator(enc_encoder, enc_cell_type, dec_cell_type, csv_filename):
    # Define input and output features
    input_features = [
        sequence_feature(encoder={"type": enc_encoder, "min_len": 5, "max_len": 10, "cell_type": enc_cell_type})
    ]
    output_features = [
        sequence_feature(decoder={"type": "generator", "min_len": 5, "max_len": 10, "cell_type": dec_cell_type})
    ]

    # Generate test data
    rel_path = generate_data(input_features, output_features, csv_filename)

    # run the experiment
    run_experiment(input_features, output_features, dataset=rel_path)