# Copyright (c) 2023 Predibase, Inc., 2019 Uber Technologies, Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== import json import os import shutil from unittest import mock import pytest import torch import yaml from ludwig.api import LudwigModel from ludwig.callbacks import Callback from ludwig.constants import BATCH_SIZE, ENCODER, TRAINER, TYPE from ludwig.globals import MODEL_FILE_NAME, MODEL_HYPERPARAMETERS_FILE_NAME from ludwig.utils.data_utils import read_csv from tests.integration_tests.utils import ( category_feature, generate_data, get_weights, image_feature, run_api_experiment, sequence_feature, text_feature, ) def run_api_experiment_separated_datasets(input_features, output_features, data_csv): """Helper method to avoid code repetition in running an experiment. :param input_features: input schema :param output_features: output schema :param data_csv: path to data :return: None """ config = { "input_features": input_features, "output_features": output_features, "combiner": {"type": "concat", "output_size": 14}, TRAINER: {"epochs": 2, BATCH_SIZE: 128}, } model = LudwigModel(config) # Training with dataframe data_df = read_csv(data_csv) train_df = data_df.sample(frac=0.8) test_df = data_df.drop(train_df.index).sample(frac=0.5) validation_df = data_df.drop(train_df.index).drop(test_df.index) basename, ext = os.path.splitext(data_csv) train_fname = basename + ".train" + ext val_fname = basename + ".validation" + ext test_fname = basename + ".test" + ext output_dirs = [] try: train_df.to_csv(train_fname) validation_df.to_csv(val_fname) test_df.to_csv(test_fname) # Training with csv _, _, output_dir = model.train( training_set=train_fname, skip_save_processed_input=True, skip_save_progress=True, skip_save_unprocessed_output=True, ) output_dirs.append(output_dir) _, _, output_dir = model.train( training_set=train_fname, validation_set=val_fname, skip_save_processed_input=True, skip_save_progress=True, skip_save_unprocessed_output=True, ) output_dirs.append(output_dir) _, _, output_dir = model.train( training_set=train_fname, validation_set=val_fname, test_set=test_fname, skip_save_processed_input=True, skip_save_progress=True, skip_save_unprocessed_output=True, ) output_dirs.append(output_dir) _, output_dir = model.predict(dataset=test_fname) output_dirs.append(output_dir) finally: # Remove results/intermediate data saved to disk os.remove(train_fname) os.remove(val_fname) os.remove(test_fname) for output_dir in output_dirs: shutil.rmtree(output_dir, ignore_errors=True) output_dirs = [] try: _, _, output_dir = model.train( training_set=train_df, skip_save_processed_input=True, skip_save_progress=True, skip_save_unprocessed_output=True, ) output_dirs.append(output_dir) _, _, output_dir = model.train( training_set=train_df, validation_set=validation_df, skip_save_processed_input=True, skip_save_progress=True, skip_save_unprocessed_output=True, ) output_dirs.append(output_dir) _, _, output_dir = model.train( training_set=train_df, validation_set=validation_df, test_set=test_df, skip_save_processed_input=True, skip_save_progress=True, skip_save_unprocessed_output=True, ) output_dirs.append(output_dir) _, output_dir = model.predict(dataset=data_df) output_dirs.append(output_dir) finally: for output_dir in output_dirs: shutil.rmtree(output_dir, ignore_errors=True) def test_api_intent_classification(csv_filename): # Single sequence input, single category output input_features = [sequence_feature(encoder={"reduce_output": "sum"})] output_features = [category_feature(decoder={"vocab_size": 5}, reduce_input="sum")] # Generate test data rel_path = generate_data(input_features, output_features, csv_filename) # Test representative encoders (embed=simple, rnn=recurrent, transformer=attention) for encoder in ["embed", "rnn", "transformer"]: input_features[0][ENCODER][TYPE] = encoder run_api_experiment(input_features, output_features, data_csv=rel_path) def test_api_intent_classification_separated(csv_filename): # Single sequence input, single category output input_features = [sequence_feature(encoder={"reduce_output": "sum"})] output_features = [category_feature(decoder={"vocab_size": 5}, reduce_input="sum")] # Generate test data rel_path = generate_data(input_features, output_features, csv_filename) # Test representative encoders (embed=simple, rnn=recurrent, transformer=attention) for encoder in ["embed", "rnn", "transformer"]: input_features[0][ENCODER][TYPE] = encoder run_api_experiment_separated_datasets(input_features, output_features, data_csv=rel_path) def test_api_train_online(csv_filename): input_features = [sequence_feature(encoder={"reduce_output": "sum"})] output_features = [category_feature(decoder={"vocab_size": 5}, reduce_input="sum")] data_csv = generate_data(input_features, output_features, csv_filename) config = { "input_features": input_features, "output_features": output_features, "combiner": {"type": "concat", "output_size": 14}, } model = LudwigModel(config) for _ in range(2): model.train_online(dataset=data_csv) model.predict(dataset=data_csv) def test_api_training_set(tmpdir): input_features = [sequence_feature(encoder={"reduce_output": "sum"})] output_features = [category_feature(decoder={"vocab_size": 5}, reduce_input="sum")] data_csv = generate_data(input_features, output_features, os.path.join(tmpdir, "dataset.csv")) val_csv = shutil.copyfile(data_csv, os.path.join(tmpdir, "validation.csv")) test_csv = shutil.copyfile(data_csv, os.path.join(tmpdir, "test.csv")) config = { "input_features": input_features, "output_features": output_features, "combiner": {"type": "concat", "output_size": 14}, } model = LudwigModel(config) model.train(training_set=data_csv, validation_set=val_csv, test_set=test_csv) model.predict(dataset=test_csv) # Train again, this time the HDF5 cache will be used model.train(training_set=data_csv, validation_set=val_csv, test_set=test_csv) def test_api_training_determinism(tmpdir): input_features = [sequence_feature(encoder={"reduce_output": "sum"})] output_features = [category_feature(decoder={"vocab_size": 5}, reduce_input="sum")] data_csv = generate_data(input_features, output_features, os.path.join(tmpdir, "dataset.csv")) config = { "input_features": input_features, "output_features": output_features, "combiner": {"type": "concat", "output_size": 14}, "trainer": {BATCH_SIZE: 128}, # batch size must be fixed for determinism } # Train the model 3 times: # # 1. seed x # 2. seed y # 3. seed x # # Check that models (1) and (3) produce the same weights, # but (1) and (2) do not rand_x = 42 rand_y = 24 model_1 = LudwigModel(config) model_1.train(dataset=data_csv, output_directory=tmpdir, random_seed=rand_x) model_2 = LudwigModel(config) model_2.train(dataset=data_csv, output_directory=tmpdir, random_seed=rand_y) model_3 = LudwigModel(config) model_3.train(dataset=data_csv, output_directory=tmpdir, random_seed=rand_x) model_weights_1 = get_weights(model_1.model) model_weights_2 = get_weights(model_2.model) model_weights_3 = get_weights(model_3.model) divergence = False for weight_1, weight_2 in zip(model_weights_1, model_weights_2): if not torch.allclose(weight_1, weight_2): divergence = True break assert divergence, "model_1 and model_2 have identical weights with different seeds!" for weight_1, weight_3 in zip(model_weights_1, model_weights_3): assert torch.allclose(weight_1, weight_3) def run_api_commands( input_features, output_features, data_csv, output_dir, skip_save_training_description=False, skip_save_training_statistics=False, skip_save_model=False, skip_save_progress=False, skip_save_log=False, skip_save_processed_input=False, skip_save_unprocessed_output=False, skip_save_predictions=False, skip_save_eval_stats=False, skip_collect_predictions=False, skip_collect_overall_stats=False, ): """Helper method to avoid code repetition in running an experiment. :param input_features: input schema :param output_features: output schema :param data_csv: path to data :return: None """ config = { "input_features": input_features, "output_features": output_features, "combiner": {"type": "concat", "output_size": 14}, TRAINER: {"epochs": 2, BATCH_SIZE: 128}, } model = LudwigModel(config) # Training with csv model.train( dataset=data_csv, skip_save_training_description=skip_save_training_description, skip_save_training_statistics=skip_save_training_statistics, skip_save_model=skip_save_model, skip_save_progress=skip_save_progress, skip_save_log=skip_save_log, skip_save_processed_input=skip_save_processed_input, output_directory=output_dir, ) model.predict( dataset=data_csv, skip_save_unprocessed_output=skip_save_unprocessed_output, skip_save_predictions=skip_save_predictions, output_directory=output_dir, ) model.evaluate( dataset=data_csv, skip_save_unprocessed_output=skip_save_unprocessed_output, skip_save_predictions=skip_save_predictions, skip_save_eval_stats=skip_save_eval_stats, collect_predictions=not skip_collect_predictions, collect_overall_stats=not skip_collect_overall_stats, output_directory=output_dir, ) model.experiment( dataset=data_csv, skip_save_training_description=skip_save_training_description, skip_save_training_statistics=skip_save_training_statistics, skip_save_model=skip_save_model, skip_save_progress=skip_save_progress, skip_save_log=skip_save_log, skip_save_processed_input=skip_save_processed_input, skip_save_unprocessed_output=skip_save_unprocessed_output, skip_save_predictions=skip_save_predictions, skip_save_eval_stats=skip_save_eval_stats, skip_collect_predictions=skip_collect_predictions, skip_collect_overall_stats=skip_collect_overall_stats, output_directory=output_dir, ) @pytest.mark.parametrize( "skip_save_training_description,skip_save_training_statistics,skip_save_model," "skip_save_progress,skip_save_log,skip_save_processed_input", [ (False, False, False, False, False, False), # all saving enabled (True, True, True, True, True, True), # all saving disabled (True, False, True, False, True, False), # alternating pattern ], ids=["all_save", "all_skip", "mixed"], ) def test_api_skip_parameters_train( tmpdir, csv_filename, skip_save_training_description, skip_save_training_statistics, skip_save_model, skip_save_progress, skip_save_log, skip_save_processed_input, ): # Single sequence input, single category output input_features = [category_feature(encoder={"vocab_size": 5})] output_features = [category_feature(decoder={"vocab_size": 5})] # Generate test data rel_path = generate_data(input_features, output_features, os.path.join(tmpdir, csv_filename)) run_api_commands( input_features, output_features, data_csv=rel_path, output_dir=tmpdir, skip_save_training_description=skip_save_training_description, skip_save_training_statistics=skip_save_training_statistics, skip_save_model=skip_save_model, skip_save_progress=skip_save_progress, skip_save_log=skip_save_log, skip_save_processed_input=skip_save_processed_input, ) @pytest.mark.parametrize("skip_save_unprocessed_output", [False, True]) @pytest.mark.parametrize("skip_save_predictions", [False, True]) def test_api_skip_parameters_predict( tmpdir, csv_filename, skip_save_unprocessed_output, skip_save_predictions, ): # Single sequence input, single category output input_features = [category_feature(encoder={"vocab_size": 5})] output_features = [category_feature(decoder={"vocab_size": 5})] # Generate test data rel_path = generate_data(input_features, output_features, os.path.join(tmpdir, csv_filename)) run_api_commands( input_features, output_features, data_csv=rel_path, output_dir=tmpdir, skip_save_unprocessed_output=skip_save_unprocessed_output, skip_save_predictions=skip_save_predictions, ) @pytest.mark.parametrize( "skip_save_unprocessed_output,skip_save_predictions,skip_save_eval_stats," "skip_collect_predictions,skip_collect_overall_stats", [ (False, False, False, False, False), # all saving enabled (True, True, True, True, True), # all saving disabled (True, False, True, False, True), # alternating pattern ], ids=["all_save", "all_skip", "mixed"], ) def test_api_skip_parameters_evaluate( tmpdir, csv_filename, skip_save_unprocessed_output, skip_save_predictions, skip_save_eval_stats, skip_collect_predictions, skip_collect_overall_stats, ): # Single sequence input, single category output input_features = [category_feature(encoder={"vocab_size": 5})] output_features = [category_feature(decoder={"vocab_size": 5})] # Generate test data rel_path = generate_data(input_features, output_features, os.path.join(tmpdir, csv_filename)) run_api_commands( input_features, output_features, data_csv=rel_path, output_dir=tmpdir, skip_save_unprocessed_output=skip_save_unprocessed_output, skip_save_predictions=skip_save_predictions, skip_save_eval_stats=skip_save_eval_stats, skip_collect_predictions=skip_collect_predictions, skip_collect_overall_stats=skip_collect_overall_stats, ) @pytest.mark.parametrize( "epochs,batch_size,num_examples,steps_per_checkpoint", [ (1, 8, 16, 1), (2, 4, 32, 2), (2, 8, 16, 2), ], ids=["small", "large", "mixed"], ) def test_api_callbacks(tmpdir, csv_filename, epochs, batch_size, num_examples, steps_per_checkpoint): mock_callback = mock.Mock(wraps=Callback()) steps_per_epoch = num_examples / batch_size total_checkpoints = (steps_per_epoch / steps_per_checkpoint) * epochs total_batches = epochs * (num_examples / batch_size) input_features = [sequence_feature(encoder={"reduce_output": "sum"})] output_features = [category_feature(decoder={"vocab_size": 5}, reduce_input="sum")] config = { "input_features": input_features, "output_features": output_features, "combiner": {"type": "concat", "output_size": 14}, TRAINER: { "epochs": epochs, "batch_size": batch_size, "steps_per_checkpoint": steps_per_checkpoint, "early_stop": 0, # Disable early stopping. }, } model = LudwigModel(config, callbacks=[mock_callback]) data_csv = generate_data( input_features, output_features, os.path.join(tmpdir, csv_filename), num_examples=num_examples ) val_csv = shutil.copyfile(data_csv, os.path.join(tmpdir, "validation.csv")) test_csv = shutil.copyfile(data_csv, os.path.join(tmpdir, "test.csv")) model.train(training_set=data_csv, validation_set=val_csv, test_set=test_csv) assert mock_callback.on_epoch_start.call_count == epochs assert mock_callback.on_epoch_end.call_count == epochs assert mock_callback.should_early_stop.call_count == total_checkpoints assert mock_callback.on_validation_start.call_count == total_checkpoints assert mock_callback.on_validation_end.call_count == total_checkpoints assert mock_callback.on_test_start.call_count == total_checkpoints assert mock_callback.on_test_end.call_count == total_checkpoints assert mock_callback.on_batch_start.call_count == total_batches assert mock_callback.on_batch_end.call_count == total_batches assert mock_callback.on_eval_end.call_count == total_checkpoints assert mock_callback.on_eval_start.call_count == total_checkpoints @pytest.mark.parametrize( "epochs,batch_size,num_examples,checkpoints_per_epoch", [ (1, 8, 32, 1), (2, 4, 64, 2), (2, 8, 32, 4), ], ids=["single_checkpoint", "multi_checkpoint", "frequent_checkpoint"], ) def test_api_callbacks_checkpoints_per_epoch( tmpdir, csv_filename, epochs, batch_size, num_examples, checkpoints_per_epoch ): mock_callback = mock.Mock(wraps=Callback()) total_checkpoints = epochs * checkpoints_per_epoch total_batches = epochs * (num_examples / batch_size) input_features = [sequence_feature(encoder={"reduce_output": "sum"})] output_features = [category_feature(decoder={"vocab_size": 5}, reduce_input="sum")] config = { "input_features": input_features, "output_features": output_features, "combiner": {"type": "concat", "output_size": 14}, TRAINER: { "epochs": epochs, "batch_size": batch_size, "checkpoints_per_epoch": checkpoints_per_epoch, "early_stop": 0, # Disable early stopping. }, } model = LudwigModel(config, callbacks=[mock_callback]) data_csv = generate_data( input_features, output_features, os.path.join(tmpdir, csv_filename), num_examples=num_examples ) val_csv = shutil.copyfile(data_csv, os.path.join(tmpdir, "validation.csv")) test_csv = shutil.copyfile(data_csv, os.path.join(tmpdir, "test.csv")) model.train(training_set=data_csv, validation_set=val_csv, test_set=test_csv) assert mock_callback.on_epoch_start.call_count == epochs assert mock_callback.on_epoch_end.call_count == epochs assert mock_callback.should_early_stop.call_count == total_checkpoints assert mock_callback.on_validation_start.call_count == total_checkpoints assert mock_callback.on_validation_end.call_count == total_checkpoints assert mock_callback.on_test_start.call_count == total_checkpoints assert mock_callback.on_test_end.call_count == total_checkpoints assert mock_callback.on_batch_start.call_count == total_batches assert mock_callback.on_batch_end.call_count == total_batches assert mock_callback.on_eval_end.call_count == total_checkpoints assert mock_callback.on_eval_start.call_count == total_checkpoints def test_api_callbacks_default_train_steps(tmpdir, csv_filename): # Default for train_steps is -1: use epochs. train_steps = None epochs = 3 batch_size = 8 num_examples = 20 mock_callback = mock.Mock(wraps=Callback()) input_features = [sequence_feature(encoder={"reduce_output": "sum"})] output_features = [category_feature(decoder={"vocab_size": 5}, reduce_input="sum")] config = { "input_features": input_features, "output_features": output_features, "combiner": {"type": "concat", "output_size": 14}, TRAINER: {"epochs": epochs, "train_steps": train_steps, "batch_size": batch_size}, } model = LudwigModel(config, callbacks=[mock_callback]) model.train( training_set=generate_data( input_features, output_features, os.path.join(tmpdir, csv_filename), num_examples=num_examples ) ) assert mock_callback.on_epoch_start.call_count == epochs def test_api_callbacks_fixed_train_steps(tmpdir, csv_filename): train_steps = 4 batch_size = 8 num_examples = 20 mock_callback = mock.Mock(wraps=Callback()) input_features = [sequence_feature(encoder={"reduce_output": "sum"})] output_features = [category_feature(decoder={"vocab_size": 5}, reduce_input="sum")] config = { "input_features": input_features, "output_features": output_features, "combiner": {"type": "concat", "output_size": 14}, TRAINER: {"train_steps": train_steps, "batch_size": batch_size}, } model = LudwigModel(config, callbacks=[mock_callback]) model.train( training_set=generate_data( input_features, output_features, os.path.join(tmpdir, csv_filename), num_examples=num_examples ) ) # With 20 examples (14 train at 70% split), batch_size=8, steps_per_epoch=2. # So 4 train steps => 2 epochs. assert mock_callback.on_epoch_start.call_count == 2 def test_api_callbacks_fixed_train_steps_partial_epochs(tmpdir, csv_filename): # If train_steps is set manually, epochs is ignored. train_steps = 3 epochs = 2 batch_size = 8 num_examples = 20 mock_callback = mock.Mock(wraps=Callback()) input_features = [sequence_feature(encoder={"reduce_output": "sum"})] output_features = [category_feature(decoder={"vocab_size": 5}, reduce_input="sum")] config = { "input_features": input_features, "output_features": output_features, "combiner": {"type": "concat", "output_size": 14}, TRAINER: {"epochs": epochs, "train_steps": train_steps, "batch_size": batch_size}, } model = LudwigModel(config, callbacks=[mock_callback]) model.train( training_set=generate_data( input_features, output_features, os.path.join(tmpdir, csv_filename), num_examples=num_examples ) ) # With 20 examples, batch_size=8, steps_per_epoch=2. 3 train steps => 1 full epoch. assert mock_callback.on_epoch_end.call_count == 1 def test_api_callbacks_batch_size_1(tmpdir, csv_filename): epochs = 1 batch_size = 1 num_examples = 16 mock_callback = mock.Mock(wraps=Callback()) input_features = [sequence_feature(encoder={"reduce_output": "sum"})] output_features = [category_feature(decoder={"vocab_size": 5}, reduce_input="sum")] config = { "input_features": input_features, "output_features": output_features, "combiner": {"type": "concat", "output_size": 14}, TRAINER: {"epochs": epochs, "batch_size": batch_size}, } model = LudwigModel(config, callbacks=[mock_callback]) model.train( training_set=generate_data( input_features, output_features, os.path.join(tmpdir, csv_filename), num_examples=num_examples ) ) # There are exactly 1 epoch start, even with batch_size = 1. assert mock_callback.on_epoch_start.call_count == 1 assert mock_callback.on_epoch_end.call_count == 1 assert mock_callback.on_batch_start.call_count == 16 assert mock_callback.on_batch_end.call_count == 16 def test_api_callbacks_fixed_train_steps_less_than_one_epoch(tmpdir, csv_filename): # If train_steps is set manually, epochs is ignored. # With 80 examples at 70% split = 56 train examples, batch_size=8 => 7 steps per epoch. # train_steps=6 < 7, so less than one full epoch. train_steps = total_batches = 6 steps_per_checkpoint = 2 batch_size = 8 num_examples = 80 mock_callback = mock.Mock(wraps=Callback()) input_features = [sequence_feature(encoder={"reduce_output": "sum"})] output_features = [category_feature(decoder={"vocab_size": 5}, reduce_input="sum")] config = { "input_features": input_features, "output_features": output_features, "combiner": {"type": "concat", "output_size": 14}, TRAINER: { "train_steps": train_steps, "steps_per_checkpoint": steps_per_checkpoint, "batch_size": batch_size, }, } model = LudwigModel(config, callbacks=[mock_callback]) model.train( training_set=generate_data( input_features, output_features, os.path.join(tmpdir, csv_filename), num_examples=num_examples ) ) assert mock_callback.on_epoch_start.call_count == 1 assert mock_callback.on_epoch_end.call_count == 0 # The total number of batches is the number of train_steps assert mock_callback.on_batch_end.call_count == total_batches # The total number of evals is the number of times checkpoints are made assert mock_callback.on_eval_end.call_count == train_steps // steps_per_checkpoint def test_saved_weights_in_checkpoint(tmpdir): image_dest_folder = os.path.join(tmpdir, "generated_images") input_features = [ text_feature(), image_feature(image_dest_folder), ] output_features = [category_feature(name="class", output_feature=True)] data_csv = generate_data(input_features, output_features, os.path.join(tmpdir, "dataset.csv")) val_csv = shutil.copyfile(data_csv, os.path.join(tmpdir, "validation.csv")) test_csv = shutil.copyfile(data_csv, os.path.join(tmpdir, "test.csv")) config = { "input_features": input_features, "output_features": output_features, TRAINER: {BATCH_SIZE: 128}, } model = LudwigModel(config) _, _, output_dir = model.train( training_set=data_csv, validation_set=val_csv, test_set=test_csv, output_directory=tmpdir ) config_save_path = os.path.join(output_dir, MODEL_FILE_NAME, MODEL_HYPERPARAMETERS_FILE_NAME) with open(config_save_path) as f: saved_config = json.load(f) saved_input_features = saved_config["input_features"] for saved_input_feature in saved_input_features: assert "encoder" in saved_input_feature input_feature_encoder = saved_input_feature["encoder"] assert "saved_weights_in_checkpoint" in input_feature_encoder assert input_feature_encoder["saved_weights_in_checkpoint"] def test_constant_metadata(tmpdir): input_features = [category_feature(encoder={"vocab_size": 5})] output_features = [category_feature(name="class", decoder={"vocab_size": 5}, output_feature=True)] data_csv1 = generate_data(input_features, output_features, os.path.join(tmpdir, "dataset1.csv")) val_csv1 = shutil.copyfile(data_csv1, os.path.join(tmpdir, "validation1.csv")) test_csv1 = shutil.copyfile(data_csv1, os.path.join(tmpdir, "test1.csv")) config = { "input_features": input_features, "output_features": output_features, } model = LudwigModel(config) model.train(training_set=data_csv1, validation_set=val_csv1, test_set=test_csv1, output_directory=tmpdir) metadata1 = model.training_set_metadata data_csv2 = generate_data(input_features, output_features, os.path.join(tmpdir, "dataset2.csv"), num_examples=10) val_csv2 = shutil.copyfile(data_csv2, os.path.join(tmpdir, "validation2.csv")) test_csv2 = shutil.copyfile(data_csv2, os.path.join(tmpdir, "test2.csv")) model.train(training_set=data_csv2, validation_set=val_csv2, test_set=test_csv2, output_directory=tmpdir) metadata2 = model.training_set_metadata assert metadata1 == metadata2 @pytest.mark.integration_tests_i @pytest.mark.parametrize( "input_max_sequence_length, global_max_sequence_length, expect_raise", [ (5, "null", True), ("null", 5, True), (5, 5, True), (100, 100, False), (100, "null", False), ("null", "null", False), ], ) def test_llm_template_too_long(tmpdir, input_max_sequence_length, global_max_sequence_length, expect_raise): zero_shot_config = yaml.safe_load(f""" model_type: llm base_model: hf-internal-testing/tiny-random-GPTJForCausalLM input_features: - name: instruction type: text preprocessing: max_sequence_length: {input_max_sequence_length} output_features: - name: output type: text preprocessing: global_max_sequence_length: {global_max_sequence_length} """) zero_shot_config["prompt"] = {} zero_shot_config["prompt"]["template"] = ( "This is a very long template that is longer than the max sequence length {instruction}" ) input_features = [text_feature(name="instruction")] output_features = [text_feature(name="output", output_feature=True)] data_csv1 = generate_data(input_features, output_features, os.path.join(tmpdir, "dataset1.csv")) model = LudwigModel(zero_shot_config) if expect_raise: with pytest.raises(ValueError): model.preprocess(dataset=data_csv1, output_directory=tmpdir) else: model.preprocess(dataset=data_csv1, output_directory=tmpdir)