Files
ray-project--ray/python/ray/tests/ludwig/ludwig_test_utils.py
T
2026-07-13 13:17:40 +08:00

644 lines
19 KiB
Python

# -*- coding: utf-8 -*-
# Copyright (c) 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.
# ==============================================================================
# This file is copied and adapted from
# https://github.com/ludwig-ai/ludwig/blob/master/tests/integration_tests/utils.py
import multiprocessing
import os
import random
import shutil
import sys
import traceback
import unittest
import uuid
from distutils.util import strtobool
from typing import Any, Dict, List, Optional
import cloudpickle
import numpy as np
import pandas as pd
from ludwig.api import LudwigModel
from ludwig.backend import LocalBackend
from ludwig.constants import COLUMN, NAME, PROC_COLUMN, VECTOR
from ludwig.data.dataset_synthesizer import DATETIME_FORMATS, build_synthetic_dataset
from ludwig.experiment import experiment_cli
from ludwig.features.feature_utils import compute_feature_hash
from ludwig.utils.data_utils import read_csv, replace_file_extension
ENCODERS = [
"embed",
"rnn",
"parallel_cnn",
"cnnrnn",
"stacked_parallel_cnn",
"stacked_cnn",
"transformer",
]
HF_ENCODERS_SHORT = ["distilbert"]
HF_ENCODERS = [
"bert",
"gpt",
"gpt2",
# "transformer_xl",
"xlnet",
"xlm",
"roberta",
"distilbert",
"ctrl",
"camembert",
"albert",
"t5",
"xlmroberta",
"longformer",
"flaubert",
"electra",
"mt5",
]
class LocalTestBackend(LocalBackend):
@property
def supports_multiprocessing(self):
return False
def parse_flag_from_env(key, default=False):
try:
value = os.environ[key]
except KeyError:
# KEY isn't set, default to `default`.
_value = default
else:
# KEY is set, convert it to True or False.
try:
_value = strtobool(value)
except ValueError:
# More values are supported, but let's keep the message simple.
raise ValueError("If set, {} must be yes or no.".format(key))
return _value
_run_slow_tests = parse_flag_from_env("RUN_SLOW", default=False)
def slow(test_case):
"""
Decorator marking a test as slow.
Slow tests are skipped by default. Set the RUN_SLOW environment variable
to a truth value to run them.
"""
if not _run_slow_tests:
test_case = unittest.skip("Skipping: this test is too slow")(test_case)
return test_case
def generate_data(
input_features: List[Dict[str, Any]],
output_features: List[Dict[str, Any]],
filename: str = "test_csv.csv",
num_examples: int = 25,
) -> str:
"""Generate synthetic data based on input/output feature specs.
Args:
input_features: Input feature schema.
output_features: Output feature schema.
filename: Path to the file where data is stored.
num_examples: Number of examples to generate.
Returns:
The path to the file where the generated data was written.
"""
features = input_features + output_features
df = build_synthetic_dataset(num_examples, features)
data = [next(df) for _ in range(num_examples)]
dataframe = pd.DataFrame(data[1:], columns=data[0])
dataframe.to_csv(filename, index=False)
return filename
def random_string(length=5):
return uuid.uuid4().hex[:length].upper()
def numerical_feature(normalization=None, **kwargs):
feature = {
"name": "num_" + random_string(),
"type": "number",
"preprocessing": {"normalization": normalization},
}
feature.update(kwargs)
feature[COLUMN] = feature[NAME]
feature[PROC_COLUMN] = compute_feature_hash(feature)
return feature
def category_feature(**kwargs):
feature = {
"type": "category",
"name": "category_" + random_string(),
"vocab_size": 10,
"embedding_size": 5,
}
feature.update(kwargs)
feature[COLUMN] = feature[NAME]
feature[PROC_COLUMN] = compute_feature_hash(feature)
return feature
def text_feature(**kwargs):
feature = {
"name": "text_" + random_string(),
"type": "text",
"reduce_input": None,
"vocab_size": 5,
"min_len": 7,
"max_len": 7,
"embedding_size": 8,
"state_size": 8,
}
feature.update(kwargs)
feature[COLUMN] = feature[NAME]
feature[PROC_COLUMN] = compute_feature_hash(feature)
return feature
def set_feature(**kwargs):
feature = {
"type": "set",
"name": "set_" + random_string(),
"vocab_size": 10,
"max_len": 5,
"embedding_size": 5,
}
feature.update(kwargs)
feature[COLUMN] = feature[NAME]
feature[PROC_COLUMN] = compute_feature_hash(feature)
return feature
def sequence_feature(**kwargs):
feature = {
"type": "sequence",
"name": "sequence_" + random_string(),
"vocab_size": 10,
"max_len": 7,
"encoder": "embed",
"embedding_size": 8,
"fc_size": 8,
"state_size": 8,
"num_filters": 8,
"hidden_size": 8,
}
feature.update(kwargs)
feature[COLUMN] = feature[NAME]
feature[PROC_COLUMN] = compute_feature_hash(feature)
return feature
def image_feature(folder, **kwargs):
feature = {
"type": "image",
"name": "image_" + random_string(),
"encoder": "resnet",
"preprocessing": {
"in_memory": True,
"height": 12,
"width": 12,
"num_channels": 3,
},
"resnet_size": 8,
"destination_folder": folder,
"fc_size": 8,
"num_filters": 8,
}
feature.update(kwargs)
feature[COLUMN] = feature[NAME]
feature[PROC_COLUMN] = compute_feature_hash(feature)
return feature
def audio_feature(folder, **kwargs):
feature = {
"name": "audio_" + random_string(),
"type": "audio",
"preprocessing": {
"audio_feature": {
"type": "fbank",
"window_length_in_s": 0.04,
"window_shift_in_s": 0.02,
"num_filter_bands": 80,
},
"audio_file_length_limit_in_s": 3.0,
},
"encoder": "stacked_cnn",
"should_embed": False,
"conv_layers": [
{
"filter_size": 400,
"pool_size": 16,
"num_filters": 32,
"regularize": "false",
},
{
"filter_size": 40,
"pool_size": 10,
"num_filters": 64,
"regularize": "false",
},
],
"fc_size": 256,
"destination_folder": folder,
}
feature.update(kwargs)
feature[COLUMN] = feature[NAME]
feature[PROC_COLUMN] = compute_feature_hash(feature)
return feature
def timeseries_feature(**kwargs):
feature = {
"name": "timeseries_" + random_string(),
"type": "timeseries",
"max_len": 7,
}
feature.update(kwargs)
feature[COLUMN] = feature[NAME]
feature[PROC_COLUMN] = compute_feature_hash(feature)
return feature
def binary_feature(**kwargs):
feature = {"name": "binary_" + random_string(), "type": "binary"}
feature.update(kwargs)
feature[COLUMN] = feature[NAME]
feature[PROC_COLUMN] = compute_feature_hash(feature)
return feature
def bag_feature(**kwargs):
feature = {
"name": "bag_" + random_string(),
"type": "bag",
"max_len": 5,
"vocab_size": 10,
"embedding_size": 5,
}
feature.update(kwargs)
feature[COLUMN] = feature[NAME]
feature[PROC_COLUMN] = compute_feature_hash(feature)
return feature
def date_feature(**kwargs):
feature = {
"name": "date_" + random_string(),
"type": "date",
"preprocessing": {
"datetime_format": random.choice(list(DATETIME_FORMATS.keys()))
},
}
feature.update(kwargs)
feature[COLUMN] = feature[NAME]
feature[PROC_COLUMN] = compute_feature_hash(feature)
return feature
def h3_feature(**kwargs):
feature = {"name": "h3_" + random_string(), "type": "h3"}
feature.update(kwargs)
feature[COLUMN] = feature[NAME]
feature[PROC_COLUMN] = compute_feature_hash(feature)
return feature
def vector_feature(**kwargs):
feature = {"type": VECTOR, "vector_size": 5, "name": "vector_" + random_string()}
feature.update(kwargs)
feature[COLUMN] = feature[NAME]
feature[PROC_COLUMN] = compute_feature_hash(feature)
return feature
def run_experiment(
input_features: Optional[List[Dict[str, Any]]],
output_features: Optional[List[Dict[str, Any]]],
skip_save_processed_input: bool = True,
config: Optional[Dict[str, Any]] = None,
backend: Optional[LocalBackend] = None,
**kwargs,
) -> None:
"""Run an experiment and clean up artifacts saved to disk.
Args:
input_features: List of input feature dictionaries.
output_features: List of output feature dictionaries.
skip_save_processed_input: Whether to skip persisting processed input
to disk.
config: Optional Ludwig configuration dictionary. If unset, a default
config is constructed from ``input_features`` and
``output_features``.
backend: Optional Ludwig backend to use. Defaults to
``LocalTestBackend()``.
**kwargs: Extra keyword arguments forwarded to the underlying
``experiment_cli`` call.
"""
if input_features is not None and output_features is not None:
# This if is necessary so that the caller can call with
# config_file (and not config)
config = {
"input_features": input_features,
"output_features": output_features,
"combiner": {"type": "concat", "fc_size": 14},
"training": {"epochs": 2},
}
args = {
"config": config,
"backend": backend or LocalTestBackend(),
"skip_save_training_description": True,
"skip_save_training_statistics": True,
"skip_save_processed_input": skip_save_processed_input,
"skip_save_progress": True,
"skip_save_unprocessed_output": True,
"skip_save_model": True,
"skip_save_predictions": True,
"skip_save_eval_stats": True,
"skip_collect_predictions": True,
"skip_collect_overall_stats": True,
"skip_save_log": True,
}
args.update(kwargs)
_, _, _, _, exp_dir_name = experiment_cli(**args)
shutil.rmtree(exp_dir_name, ignore_errors=True)
def generate_output_features_with_dependencies(main_feature, dependencies):
# helper function to generate multiple output features specifications
# with dependencies, support for 'test_experiment_multiple_seq_seq` unit
# test
# Parameters:
# main_feature: feature identifier, valid values 'feat1', 'feat2', 'feat3'
# dependencies: list of dependencies for 'main_feature', do not li
# Example:
# generate_output_features_with_dependencies('feat2', ['feat1', 'feat3'])
output_features = [
category_feature(vocab_size=2, reduce_input="sum"),
sequence_feature(vocab_size=10, max_len=5),
numerical_feature(),
]
# value portion of dictionary is a tuple: (position, feature_name)
# position: location of output feature in the above output_features list
# feature_name: Ludwig generated feature name
feature_names = {
"feat1": (0, output_features[0]["name"]),
"feat2": (1, output_features[1]["name"]),
"feat3": (2, output_features[2]["name"]),
}
# generate list of dependencies with real feature names
generated_dependencies = [feature_names[feat_name][1] for feat_name in dependencies]
# specify dependencies for the main_feature
output_features[feature_names[main_feature][0]][
"dependencies"
] = generated_dependencies
return output_features
def _subproc_wrapper(fn, queue, *args, **kwargs):
fn = cloudpickle.loads(fn)
try:
results = fn(*args, **kwargs)
except Exception as e:
traceback.print_exc(file=sys.stderr)
results = e
queue.put(results)
def spawn(fn):
def wrapped_fn(*args, **kwargs):
ctx = multiprocessing.get_context("spawn")
queue = ctx.Queue()
p = ctx.Process(
target=_subproc_wrapper,
args=(cloudpickle.dumps(fn), queue, *args),
kwargs=kwargs,
)
p.start()
p.join()
results = queue.get()
if isinstance(results, Exception):
raise RuntimeError(
f"Spawned subprocess raised {type(results).__name__}, "
f"check log output above for stack trace."
)
return results
return wrapped_fn
def run_api_experiment(
input_features: List[Dict[str, Any]],
output_features: List[Dict[str, Any]],
data_csv: str,
) -> None:
"""Run an experiment through Ludwig's Python API.
Args:
input_features: Input schema.
output_features: Output schema.
data_csv: Path to data.
"""
config = {
"input_features": input_features,
"output_features": output_features,
"combiner": {"type": "concat", "fc_size": 14},
"training": {"epochs": 2},
}
model = LudwigModel(config)
output_dir = None
try:
# Training with csv
_, _, output_dir = model.train(
dataset=data_csv,
skip_save_processed_input=True,
skip_save_progress=True,
skip_save_unprocessed_output=True,
)
model.predict(dataset=data_csv)
model_dir = os.path.join(output_dir, "model")
loaded_model = LudwigModel.load(model_dir)
# Necessary before call to get_weights() to materialize the weights
loaded_model.predict(dataset=data_csv)
model_weights = model.model.get_weights()
loaded_weights = loaded_model.model.get_weights()
for model_weight, loaded_weight in zip(model_weights, loaded_weights):
assert np.allclose(model_weight, loaded_weight)
finally:
# Remove results/intermediate data saved to disk
shutil.rmtree(output_dir, ignore_errors=True)
try:
# Training with dataframe
data_df = read_csv(data_csv)
_, _, output_dir = model.train(
dataset=data_df,
skip_save_processed_input=True,
skip_save_progress=True,
skip_save_unprocessed_output=True,
)
model.predict(dataset=data_df)
finally:
shutil.rmtree(output_dir, ignore_errors=True)
def create_data_set_to_use(data_format, raw_data):
# helper function for generating training and test data with specified
# format handles all data formats except for hdf5
# assumes raw_data is a csv dataset generated by
# tests.integration_tests.utils.generate_data() function
# support for writing to a fwf dataset based on this stackoverflow posting:
# https://stackoverflow.com/questions/16490261/python-pandas-write-dataframe-to-fixed-width-file-to-fwf
from ray._private.thirdparty.tabulate.tabulate import tabulate
def to_fwf(df, fname):
content = tabulate(df.values.tolist(), list(df.columns), tablefmt="plain")
open(fname, "w").write(content)
pd.DataFrame.to_fwf = to_fwf
dataset_to_use = None
if data_format == "csv":
dataset_to_use = raw_data
elif data_format in {"df", "dict"}:
dataset_to_use = pd.read_csv(raw_data)
if data_format == "dict":
dataset_to_use = dataset_to_use.to_dict(orient="list")
elif data_format == "excel":
dataset_to_use = replace_file_extension(raw_data, "xlsx")
pd.read_csv(raw_data).to_excel(dataset_to_use, index=False)
elif data_format == "excel_xls":
dataset_to_use = replace_file_extension(raw_data, "xls")
pd.read_csv(raw_data).to_excel(dataset_to_use, index=False)
elif data_format == "feather":
dataset_to_use = replace_file_extension(raw_data, "feather")
pd.read_csv(raw_data).to_feather(dataset_to_use)
elif data_format == "fwf":
dataset_to_use = replace_file_extension(raw_data, "fwf")
pd.read_csv(raw_data).to_fwf(dataset_to_use)
elif data_format == "html":
dataset_to_use = replace_file_extension(raw_data, "html")
pd.read_csv(raw_data).to_html(dataset_to_use, index=False)
elif data_format == "json":
dataset_to_use = replace_file_extension(raw_data, "json")
pd.read_csv(raw_data).to_json(dataset_to_use, orient="records")
elif data_format == "jsonl":
dataset_to_use = replace_file_extension(raw_data, "jsonl")
pd.read_csv(raw_data).to_json(dataset_to_use, orient="records", lines=True)
elif data_format == "parquet":
dataset_to_use = replace_file_extension(raw_data, "parquet")
pd.read_csv(raw_data).to_parquet(dataset_to_use, index=False)
elif data_format == "pickle":
dataset_to_use = replace_file_extension(raw_data, "pickle")
pd.read_csv(raw_data).to_pickle(dataset_to_use)
elif data_format == "stata":
dataset_to_use = replace_file_extension(raw_data, "stata")
pd.read_csv(raw_data).to_stata(dataset_to_use)
elif data_format == "tsv":
dataset_to_use = replace_file_extension(raw_data, "tsv")
pd.read_csv(raw_data).to_csv(dataset_to_use, sep="\t", index=False)
else:
ValueError("'{}' is an unrecognized data format".format(data_format))
return dataset_to_use
def train_with_backend(
backend,
config,
dataset=None,
training_set=None,
validation_set=None,
test_set=None,
predict=True,
evaluate=True,
):
model = LudwigModel(config, backend=backend)
output_dir = None
ret = False
try:
_, _, output_dir = model.train(
dataset=dataset,
training_set=training_set,
validation_set=validation_set,
test_set=test_set,
skip_save_processed_input=True,
skip_save_progress=True,
skip_save_unprocessed_output=True,
)
if dataset is None:
dataset = training_set
if predict:
preds, _ = model.predict(dataset=dataset)
assert backend.df_engine.compute(preds) is not None
if evaluate:
_, eval_preds, _ = model.evaluate(dataset=dataset)
assert backend.df_engine.compute(eval_preds) is not None
ret = True
finally:
# Remove results/intermediate data saved to disk
shutil.rmtree(output_dir, ignore_errors=True)
return ret