diff --git a/.gitmodules b/.gitmodules index 97bbd46..3ca2a80 100644 --- a/.gitmodules +++ b/.gitmodules @@ -20,3 +20,6 @@ [submodule "thirdparty/googletest"] path = thirdparty/googletest url = https://github.com/google/googletest.git +[submodule "thirdparty/benchmark"] + path = thirdparty/benchmark + url = https://github.com/google/benchmark.git diff --git a/CMakeLists.txt b/CMakeLists.txt index 7e48787..16f2966 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -38,6 +38,19 @@ if(NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/thirdparty/miniaudio/miniaudio.h") endif() +# Benchmark is a CLI-only static PIC dependency, even in shared core builds. +# Function scope leaves the caller's dependency/test/install settings untouched. +function(clap_add_benchmark) + set(BUILD_SHARED_LIBS OFF) + set(CMAKE_POSITION_INDEPENDENT_CODE ON) + set(BENCHMARK_ENABLE_TESTING OFF) + set(BENCHMARK_ENABLE_GTEST_TESTS OFF) + set(BENCHMARK_ENABLE_INSTALL OFF) + set(BENCHMARK_INSTALL_DOCS OFF) + set(BENCHMARK_ENABLE_WERROR OFF) + add_subdirectory(thirdparty/benchmark EXCLUDE_FROM_ALL) +endfunction() + # Scope BUILD_SHARED_LIBS to dependency creation: the core still honors the # caller's setting, without forcing cache values or duplicating its build graph. function(clap_add_dependencies) @@ -59,11 +72,16 @@ include(cmake/static-development.cmake) clap_static_development() if(CLAP_BUILD_CLI) + if(NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/thirdparty/benchmark/CMakeLists.txt") + message(FATAL_ERROR "Google Benchmark submodule is missing. Run: git submodule update --init --recursive") + endif() + # Keep CLI-only archives out of the core's static development SDK collector. + clap_add_benchmark() # The CLI target differs from the core library target, while its executable # keeps the existing "clap" name and command-line behavior. - add_executable(clap-cli src/cli.cpp src/embedding.cpp src/server.cpp src/quantize.cpp) + add_executable(clap-cli src/cli.cpp src/embedding.cpp src/bench.cpp src/server.cpp src/quantize.cpp) target_compile_features(clap-cli PRIVATE cxx_std_20) - target_link_libraries(clap-cli PRIVATE clap CLI11::CLI11 httplib::httplib nlohmann_json::nlohmann_json) + target_link_libraries(clap-cli PRIVATE clap CLI11::CLI11 httplib::httplib nlohmann_json::nlohmann_json benchmark::benchmark) set_target_properties(clap-cli PROPERTIES OUTPUT_NAME clap CXX_EXTENSIONS OFF) if(MSVC) @@ -423,6 +441,13 @@ if(BUILD_TESTING) "${MERT_REFERENCE_SOURCE}" "${MERT2_REFERENCE_SOURCE}" "${MUQ_REFERENCE_SOURCE}" "${MUQ_MULAN_REFERENCE_SOURCE}") set_tests_properties(test-encoder-quantized-parity PROPERTIES TIMEOUT 600) endif() + add_test(NAME test-clap-bench-packaging + COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/tests/bench_build.py" + "${CMAKE_COMMAND}" "${CMAKE_CURRENT_BINARY_DIR}" "$") + set_tests_properties(test-clap-bench-packaging PROPERTIES TIMEOUT 90) + add_test(NAME test-clap-bench-cli + COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/tests/bench_cli.py" "$") + set_tests_properties(test-clap-bench-cli PROPERTIES TIMEOUT 180) add_test(NAME test-clap-embedding-cli COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/tests/embedding_cli.py" "$") set_tests_properties(test-clap-embedding-cli PROPERTIES TIMEOUT 30) @@ -448,6 +473,8 @@ endforeach() if(CLAP_BUILD_CLI) install(TARGETS clap-cli RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}) + install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/thirdparty/benchmark/LICENSE" + DESTINATION ${CMAKE_INSTALL_LIBDIR}/licenses RENAME benchmark.txt) endif() # The runtime component is also usable by standalone C-library consumers. diff --git a/src/bench.cpp b/src/bench.cpp new file mode 100644 index 0000000..9cff973 --- /dev/null +++ b/src/bench.cpp @@ -0,0 +1,198 @@ +#include "audio.h" +#include "bench.h" +#include "clap.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { +void check(clap_status status) { + if (status != CLAP_OK) + throw std::runtime_error(clap_last_error()); +} + +// The native file reporter owns an internal stream which would otherwise be +// destroyed without checking late write/flush errors (for example /dev/full). +class CheckedReporter final : public benchmark::BenchmarkReporter { + public: + CheckedReporter(std::unique_ptr delegate, bool &failed) + : delegate_(std::move(delegate)), failed_(failed) {} + + bool ReportContext(const Context &context) override { + delegate_->SetOutputStream(&GetOutputStream()); + delegate_->SetErrorStream(&GetErrorStream()); + return delegate_->ReportContext(context); + } + void ReportRunsConfig(double min_time, bool explicit_iters, benchmark::IterationCount iters) override { + delegate_->ReportRunsConfig(min_time, explicit_iters, iters); + } + void ReportRuns(const std::vector &runs) override { delegate_->ReportRuns(runs); } + void Finalize() override { + delegate_->Finalize(); + GetOutputStream().flush(); + GetErrorStream().flush(); + failed_ = failed_ || !GetOutputStream() || !GetErrorStream(); + } + + private: + std::unique_ptr delegate_; + bool &failed_; +}; + +std::string native_option(const BenchOptions &options, const std::string &name, const char *environment, + const char *fallback, bool boolean = false) { + const char *value = std::getenv(environment); + std::string result = value ? value : fallback; + const auto flag = "--" + name; + const auto prefix = flag + "="; + for (const auto &argument : options.benchmark_args) { + if (argument.starts_with(prefix)) + result = argument.substr(prefix.size()); + else if (boolean && argument == flag) + result = "true"; + } + return result; +} + +bool native_tabular(const BenchOptions &options) { + auto value = native_option(options, "benchmark_counters_tabular", "BENCHMARK_COUNTERS_TABULAR", "false", true); + // Match the pinned library's boolean syntax, including its one-character aliases. + if (value.size() == 1) + return std::isalnum(static_cast(value[0])) && + std::string("0fFnN").find(value[0]) == std::string::npos; + std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) { return std::tolower(c); }); + return value != "false" && value != "no" && value != "off"; +} + +struct BenchmarkSession { + ~BenchmarkSession() { + benchmark::ClearRegisteredBenchmarks(); + benchmark::Shutdown(); + } +}; +} // namespace + +int cmd_bench(const BenchOptions &options) { + // Defaults are native flags, not registration overrides: later user flags + // retain all of Google Benchmark's parsing and override semantics. + std::vector arguments{"clap bench", "--benchmark_min_time=1s", "--benchmark_repetitions=3"}; + arguments.insert(arguments.end(), options.benchmark_args.begin(), options.benchmark_args.end()); + std::vector argv; + for (auto &argument : arguments) + argv.push_back(argument.data()); + int argc = static_cast(argv.size()); + argv.push_back(nullptr); + benchmark::Initialize(&argc, argv.data()); + const BenchmarkSession session; + if (benchmark::ReportUnrecognizedArguments(argc, argv.data())) + return 1; + + bool failed = false; + bool prepared = false; + std::string failure; + std::unique_ptr model(nullptr, clap_destroy); + clap::Audio decoded; + auto infer = [&] { + if (options.frames) { + clap_audio_frames sequence{}; + const std::unique_ptr owned(&sequence, + clap_free_audio_frames); + check(clap_get_audio_frames(model.get(), decoded.samples.data(), decoded.samples.size(), + decoded.sample_rate, &sequence)); + benchmark::DoNotOptimize(sequence.values); + benchmark::DoNotOptimize(sequence.frames); + benchmark::ClobberMemory(); + } else { + float *values = nullptr; + std::size_t count = 0; + const auto status = options.audio_path.empty() + ? clap_encode_text(model.get(), options.text.c_str(), &values, &count) + : clap_encode_audio(model.get(), decoded.samples.data(), decoded.samples.size(), + decoded.sample_rate, options.crop_offset, &values, &count); + const std::unique_ptr owned(values, clap_free); + check(status); + benchmark::DoNotOptimize(values); + benchmark::DoNotOptimize(count); + benchmark::ClobberMemory(); + } + }; + const char *name = options.frames ? "audio_frames" : options.audio_path.empty() ? "text" : "audio"; + benchmark::RegisterBenchmark( + name, + [&](benchmark::State &state) { + try { + // Callback setup precedes State's timer. It is deliberately lazy: + // list-tests and unmatched filters never load/decode/infer anything. + if (failed) { + state.SkipWithError(failure.c_str()); + return; + } + if (!prepared) { + if (!options.audio_path.empty()) + decoded = clap::decode_audio(options.audio_path); + auto runtime = clap_default_options(); + runtime.threads = options.threads; + runtime.device = options.device.c_str(); + CLAP *handle = nullptr; + const auto status = clap_create(options.model_path.c_str(), &runtime, &handle); + model.reset(handle); + check(status); + if (options.crop_requested && clap_audio_sample_rate(model.get()) != 48000) + throw std::runtime_error("--crop-offset requires a CLAP model"); + infer(); // Exactly one untimed warmup, shared across calibration/repetitions. + prepared = true; + } + for (auto _ : state) { + (void)_; + infer(); // Full public API, including preprocessing, allocation and free. + } + state.SetItemsProcessed(state.iterations()); + } catch (const std::exception &error) { + failed = true; + failure = error.what(); + state.SkipWithError(failure.c_str()); + } + }) + ->UseRealTime() + ->Unit(benchmark::kMillisecond); + benchmark::AddCustomContext("model", options.model_path); + benchmark::AddCustomContext("requested_threads", std::to_string(options.threads)); + const char *environment_device = std::getenv("CLAP_DEVICE"); + benchmark::AddCustomContext("requested_device", + options.device.empty() + ? (environment_device && *environment_device ? environment_device : "auto") + : options.device); + benchmark::AddCustomContext("input_mode", name); + + bool output_failed = false; + std::unique_ptr file_reporter; + if (!native_option(options, "benchmark_out", "BENCHMARK_OUT", "").empty()) { + const auto format = native_option(options, "benchmark_out_format", "BENCHMARK_OUT_FORMAT", "json"); + std::unique_ptr delegate; + if (format == "json") + delegate = std::make_unique(); + else if (format == "csv") { + // Preserve the pinned library's native CSV flag despite deprecation. + BENCHMARK_DISABLE_DEPRECATED_WARNING + delegate = std::make_unique(); + BENCHMARK_RESTORE_DEPRECATED_WARNING + } else + delegate = std::make_unique( + native_tabular(options) ? benchmark::ConsoleReporter::OO_Tabular : benchmark::ConsoleReporter::OO_None); + file_reporter = std::make_unique(std::move(delegate), output_failed); + } + const auto matches = benchmark::RunSpecifiedBenchmarks(nullptr, file_reporter.get()); + std::cout.flush(); + std::cerr.flush(); + if (output_failed) + std::cerr << "Failed to write benchmark report\n"; + return matches && !failed && !output_failed && std::cout && std::cerr ? 0 : 1; +} diff --git a/src/bench.h b/src/bench.h new file mode 100644 index 0000000..c0ead68 --- /dev/null +++ b/src/bench.h @@ -0,0 +1,19 @@ +#pragma once + +#include +#include +#include + +struct BenchOptions { + std::string model_path; + std::string text; + std::string audio_path; + std::string device; + int threads = 0; + std::size_t crop_offset = 0; + bool crop_requested = false; + bool frames = false; + std::vector benchmark_args; +}; + +int cmd_bench(const BenchOptions &options); diff --git a/src/cli.cpp b/src/cli.cpp index f03d2ca..eafd7c8 100644 --- a/src/cli.cpp +++ b/src/cli.cpp @@ -13,6 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ +#include "bench.h" #include "clap.h" #include "embedding.h" #include "quantize.h" @@ -28,6 +29,40 @@ int main(int argc, char **argv) { CLI::App app{"clap.cpp"}; argv = app.ensure_utf8(argv); + BenchOptions bench_options; + // Only bench's literal separator belongs to Google Benchmark. Leave other + // subcommands' parsing (and errors) unchanged, including their -- separator. + if (argc > 1 && std::string(argv[1]) == "bench") { + for (int i = 2; i < argc; ++i) { + if (std::string(argv[i]) == "--") { + for (int j = i + 1; j < argc; ++j) + bench_options.benchmark_args.emplace_back(argv[j]); + argc = i; + break; + } + } + } + auto *bench = app.add_subcommand("bench", "Benchmark native embedding latency"); + bench->footer("Pass native Google Benchmark flags after -- (for example --benchmark_format=json). " + "Defaults: --benchmark_min_time=1s --benchmark_repetitions=3."); + bench->add_option("-m,--model", bench_options.model_path, "Converted model GGUF") + ->required() + ->check(CLI::ExistingFile); + bench->add_option("-t,--threads", bench_options.threads, "ggml CPU threads (0: hardware concurrency)") + ->check(CLI::NonNegativeNumber) + ->capture_default_str(); + bench->add_option("--device", bench_options.device, "GGML device name, auto or cpu (default: CLAP_DEVICE or auto)"); + auto *bench_input = bench->add_option_group("Input", "Choose one text or audio file"); + bench_input->require_option(1, 1); + bench_input->add_option("--text", bench_options.text, "One UTF-8 text"); + auto *bench_audio = + bench_input->add_option("-a,--audio", bench_options.audio_path, "WAV, MP3 or FLAC")->check(CLI::ExistingFile); + auto *bench_crop = bench + ->add_option("--crop-offset", bench_options.crop_offset, + "Sample offset at 48 kHz into CLAP audio longer than ten seconds") + ->check(CLI::NonNegativeNumber); + bench->add_flag("--frames", bench_options.frames, + "Benchmark final-layer waveform-encoder frames instead of pooling"); std::string host = "127.0.0.1"; int port = 8080; std::string server_model; @@ -83,6 +118,14 @@ int main(int argc, char **argv) { } try { + if (*bench) { + bench_options.crop_requested = bench_crop->count() != 0; + if (bench_options.crop_requested && !*bench_audio) + throw std::runtime_error("--crop-offset requires --audio"); + if (bench_options.frames && (!*bench_audio || bench_options.crop_requested)) + throw std::runtime_error("--frames requires --audio and excludes --crop-offset"); + return cmd_bench(bench_options); + } if (*devices) { clap_device *list = nullptr; std::size_t count = 0; diff --git a/tests/bench_build.py b/tests/bench_build.py new file mode 100644 index 0000000..bbda493 --- /dev/null +++ b/tests/bench_build.py @@ -0,0 +1,22 @@ +"""Verify that only the CLI's license, not Benchmark SDKs, is installed.""" +from pathlib import Path +import subprocess +import sys +import tempfile + + +with tempfile.TemporaryDirectory(prefix="clap-bench-install-") as directory: + prefix = Path(directory) + command = [sys.argv[1], "--install", sys.argv[2], "--prefix", str(prefix)] + if len(sys.argv) > 3 and sys.argv[3]: + command += ["--config", sys.argv[3]] + result = subprocess.run(command, capture_output=True, text=True, timeout=60) + assert result.returncode == 0, result.stdout + result.stderr + # CMAKE_INSTALL_BINDIR is configurable (the libraries preset uses clap_cpp). + assert any(path.is_file() and path.name in {"clap", "clap.exe"} + for path in prefix.rglob("*")), "CLI not installed" + assert len(list(prefix.rglob("benchmark.txt"))) == 1, "Benchmark license missing" + unexpected = [str(path.relative_to(prefix)) for path in prefix.rglob("*") + if "benchmark" in path.name.lower() and path.name != "benchmark.txt"] + assert not unexpected, f"Benchmark leaked into installed SDK: {unexpected}" +print("CLI and Benchmark license installed; no Benchmark archives, headers or CMake packages") diff --git a/tests/bench_cli.py b/tests/bench_cli.py new file mode 100644 index 0000000..0ffaa26 --- /dev/null +++ b/tests/bench_cli.py @@ -0,0 +1,178 @@ +"""Hermetic Google Benchmark CLI integration tests (stdlib only).""" +import csv +import io +import json +import math +from pathlib import Path +import struct +import subprocess +import sys +import tempfile +import unittest +import wave + +from tiny_model import write_model, write_mert_model + +CLI = sys.argv[1] + + +class BenchTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.directory = tempfile.TemporaryDirectory(prefix="clap-bench-") + cls.addClassCleanup(cls.directory.cleanup) + cls.root = Path(cls.directory.name) + cls.model = cls.root / "tiny.gguf" + write_model(cls.model, waveform=True) + cls.invalid = cls.root / "invalid.gguf" + cls.invalid.write_bytes(b"not a model") + cls.audio = cls.root / "audio.wav" + with wave.open(str(cls.audio), "wb") as stream: + stream.setparams((1, 2, 24000, 0, "NONE", "not compressed")) + stream.writeframes(struct.pack("<2400h", *[ + int(4000 * math.sin(i * 0.071)) for i in range(2400)])) + + def run_cli(self, *args, ok=True): + result = subprocess.run([CLI, *map(str, args)], capture_output=True, + text=True, timeout=60) + if ok: + self.assertEqual(result.returncode, 0, result.stderr) + else: + self.assertNotEqual(result.returncode, 0, result.stdout + result.stderr) + return result + + def bench(self, *args, native=()): + return self.run_cli("bench", "-m", self.model, "-t", "1", *args, "--", + "--benchmark_min_time=0.001s", "--benchmark_repetitions=1", + "--benchmark_format=json", *native) + + def report(self, result, name, repetitions=1): + document = json.loads(result.stdout) + rows = [r for r in document["benchmarks"] if r.get("run_type") == "iteration"] + self.assertEqual(len(rows), repetitions, document) + for row in rows: + self.assertTrue(row["name"].startswith(name + "/"), row) + self.assertFalse(row.get("error_occurred"), row) + self.assertGreater(row["iterations"], 0) + self.assertGreater(row["real_time"], 0) + self.assertEqual(row["time_unit"], "ms") + self.assertGreater(row["items_per_second"], 0) + return document + + def test_frames_real_inference(self): + model = self.root / "mert.gguf" + write_mert_model(model) + result = self.run_cli("bench", "-m", model, "-t", "1", "-a", self.audio, + "--frames", "--", "--benchmark_min_time=0.001s", + "--benchmark_repetitions=1", "--benchmark_format=json") + self.report(result, "audio_frames") + self.run_cli("bench", "-m", model, "-t", "1", "-a", self.audio, + "--crop-offset", "0", "--", "--benchmark_min_time=0.001s", ok=False) + + def test_help(self): + result = self.run_cli("bench", "--help") + for option in ("--model", "--text", "--audio", "--threads", "--device", + "--frames", "--crop-offset", "--benchmark_"): + self.assertIn(option, result.stdout) + + def test_text_native_json_and_repetitions(self): + doc = self.report(self.bench("--text", "hello", "--device", "cpu", + native=("--benchmark_repetitions=2",)), "text", 2) + self.assertEqual(doc["context"]["requested_device"], "cpu") + + def test_defaults_and_filter(self): + result = self.run_cli("bench", "-m", self.model, "--text", "hello", "-t", "1", "--", + "--benchmark_format=json", "--benchmark_filter=^text/") + self.report(result, "text", 3) + # Native calibration is approximate and reuses iteration counts across + # repetitions; do not impose a wall-duration threshold on shared runners. + result = self.run_cli("bench", "-m", self.invalid, "--text", "", "--threads", "0", "--", + "--benchmark_list_tests=true") + self.assertIn("text", result.stdout) + result = self.run_cli("bench", "-m", self.invalid, "--text", "a", "--", "--help") + self.assertIn("--benchmark_filter", result.stdout + result.stderr) + + def test_audio_resampling_and_crop(self): + self.report(self.bench("-a", self.audio, "--crop-offset", "0"), "audio") + + def test_list_does_not_load_weights_or_decode(self): + for args, name in [(["--text", "hello"], "text"), + (["--audio", self.invalid], "audio"), + (["--audio", self.invalid, "--frames"], "audio_frames")]: + result = self.run_cli("bench", "-m", self.invalid, *args, "--", + "--benchmark_list_tests=true") + self.assertIn(name, result.stdout) + self.assertNotIn("GGUF", result.stderr) + + def test_validation(self): + cases = [[], ["--text", "a", "-a", self.audio], ["--text", "a", "--frames"], + ["--text", "a", "--crop-offset", "0"], + ["-a", self.audio, "--frames", "--crop-offset", "0"], + ["--text", "a", "-t", "-1"], ["-a", self.audio, "--crop-offset", "-1"], + ["--text", "a", "--unknown"], + ["--text", "a", "--benchmark_min_time=0.001s"], + ["--text", "a", "--", "--unknown"]] + for args in cases: + with self.subTest(args=args): + self.run_cli("bench", "-m", self.model, *args, ok=False) + self.run_cli("bench", "--text", "a", ok=False) + self.run_cli("embedding", "-m", self.model, "--text", "a", "--", + "--benchmark_list_tests", ok=False) + + def test_unknown_native_flags_precede_model_load(self): + result = self.run_cli("bench", "-m", self.invalid, "--text", "a", "--", + "--nonsense", ok=False) + self.assertIn("unrecognized", result.stderr.lower()) + self.assertNotIn("GGUF", result.stderr) + + def test_no_match_and_output_failure(self): + for flag in ("--benchmark_filter=does_not_match", "--benchmark_filter=[", + f"--benchmark_out={self.root / 'missing' / 'out.json'}"): + with self.subTest(flag=flag): + self.run_cli("bench", "-m", self.model, "--text", "a", "--", flag, + "--benchmark_min_time=0.001s", ok=False) + if Path("/dev/full").exists(): + self.run_cli("bench", "-m", self.model, "-t", "1", "--text", "a", "--", + "--benchmark_out=/dev/full", "--benchmark_min_time=0.001s", + "--benchmark_repetitions=1", ok=False) + + def test_native_output_file(self): + target = self.root / "report.json" + self.bench("--text", "a", native=(f"--benchmark_out={target}", + "--benchmark_out_format=json")) + self.report(subprocess.CompletedProcess([], 0, target.read_text()), "text") + + def test_native_csv_and_console(self): + target = self.root / "report.csv" + result = self.bench("--text", "a", native=("--benchmark_format=csv", + f"--benchmark_out={target}", "--benchmark_out_format=csv")) + for contents in (result.stdout, target.read_text()): + # Native CSV files include diagnostic context before their header. + contents = contents[contents.index("name,iterations,"):] + rows = list(csv.DictReader(io.StringIO(contents))) + self.assertTrue(rows[0]["name"].startswith("text/"), contents) + self.assertGreater(int(rows[0]["iterations"]), 0) + result = self.bench("--text", "a", native=("--benchmark_format=console",)) + self.assertIn("text/real_time", result.stdout) + + def test_native_tabular_file(self): + target = self.root / "table.txt" + for flag, enabled in [("--benchmark_counters_tabular=true", True), + ("--benchmark_counters_tabular", True), + ("--benchmark_counters_tabular=false", False)]: + with self.subTest(flag=flag): + self.bench("--text", "a", native=(f"--benchmark_out={target}", + "--benchmark_out_format=console", flag)) + header = next(line for line in target.read_text().splitlines() + if "Benchmark" in line and "Time" in line) + self.assertEqual("items_per_second" in header, enabled) + + def test_inference_errors_nonzero(self): + result = self.run_cli("bench", "-m", self.model, "-a", self.audio, "--frames", + "-t", "1", "--", "--benchmark_min_time=0.001s", + "--benchmark_repetitions=1", "--benchmark_format=json", ok=False) + self.assertTrue(json.loads(result.stdout)["benchmarks"][0]["error_occurred"]) + + +if __name__ == "__main__": + unittest.main(argv=[sys.argv[0], *sys.argv[2:]]) diff --git a/tests/tiny_model.py b/tests/tiny_model.py index 27db3b3..41ff38a 100644 --- a/tests/tiny_model.py +++ b/tests/tiny_model.py @@ -175,3 +175,61 @@ def reference(prompt): return dict(model=path.name, dimensions=5, cases=[reference(prompt) for prompt in ("hello", "world", "hello world!", "caf\u00e9", "a", "b")]) + + +def write_mert_model(path): + """Small valid waveform/frame encoder, without optional ML packages.""" + config = dict(model_type="mert_model", sample_rate=24000, + feat_extract_norm="group", feat_proj_layer_norm=True, + hidden_act="gelu", feat_extract_activation="gelu", hidden_size=8, + intermediate_size=16, num_hidden_layers=1, num_attention_heads=2, + num_conv_pos_embeddings=4, num_conv_pos_embedding_groups=2, + conv_dim=[4, 4], conv_kernel=[10, 3], conv_stride=[5, 2], + layer_norm_eps=1e-5, conv_bias=False) + tensors = {} + + def add(name, shape, norm=False): + serial = len(tensors) + 1 + values = [1.0 if norm else 0.1 * math.sin((i + 17 * serial) * 0.173) + for i in range(math.prod(shape))] + tensors[name] = (shape, struct.pack("<" + "f" * len(values), *values)) + + def norm(name, dim): + add(name + ".weight", [dim], True) + add(name + ".bias", [dim]) + + def linear(name, inputs, outputs): + add(name + ".weight", [inputs, outputs]) + add(name + ".bias", [outputs]) + + add("frontend.0.conv.weight", [10, 1, 4]) + norm("frontend.0.layer_norm", 4) + add("frontend.1.conv.weight", [3, 4, 4]) + norm("projection.layer_norm", 4) + linear("projection.projection", 4, 8) + add("position.weight", [4, 4, 8]) + add("position.bias", [8]) + norm("encoder.layer_norm", 8) + for part in ("q_proj", "k_proj", "v_proj", "out_proj"): + linear("layers.0.attention." + part, 8, 8) + norm("layers.0.layer_norm", 8) + norm("layers.0.final_layer_norm", 8) + linear("layers.0.ffn.in", 8, 16) + linear("layers.0.ffn.out", 16, 8) + + def string(value): + raw = value.encode("utf-8") + return struct.pack("