Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
286 lines
12 KiB
C++
286 lines
12 KiB
C++
#include "engine/framework/assets/tensor_source.h"
|
|
#include "engine/framework/core/backend.h"
|
|
#include "engine/framework/core/backend_weight_store.h"
|
|
#include "engine/framework/io/json.h"
|
|
#include "engine/framework/modules/speech_encoders/whisper_embedding.h"
|
|
|
|
#include <ggml-backend.h>
|
|
#include <ggml.h>
|
|
|
|
#include <algorithm>
|
|
#include <cstddef>
|
|
#include <cstring>
|
|
#include <filesystem>
|
|
#include <fstream>
|
|
#include <iostream>
|
|
#include <memory>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
namespace {
|
|
|
|
constexpr size_t kWeightContextBytes = 256ull * 1024ull * 1024ull;
|
|
constexpr size_t kGraphContextBytes = 512ull * 1024ull * 1024ull;
|
|
|
|
struct GgmlContextDeleter {
|
|
void operator()(ggml_context * ctx) const noexcept {
|
|
if (ctx != nullptr) {
|
|
ggml_free(ctx);
|
|
}
|
|
}
|
|
};
|
|
|
|
std::string arg_value(int argc, char ** argv, const std::string & name, const std::string & fallback) {
|
|
for (int i = 1; i + 1 < argc; ++i) {
|
|
if (argv[i] == name) {
|
|
return argv[i + 1];
|
|
}
|
|
}
|
|
return fallback;
|
|
}
|
|
|
|
int int_arg(int argc, char ** argv, const std::string & name, int fallback) {
|
|
return std::stoi(arg_value(argc, argv, name, std::to_string(fallback)));
|
|
}
|
|
|
|
engine::core::BackendType parse_backend(const std::string & value) {
|
|
if (value == "cpu") {
|
|
return engine::core::BackendType::Cpu;
|
|
}
|
|
if (value == "cuda") {
|
|
return engine::core::BackendType::Cuda;
|
|
}
|
|
if (value == "vulkan") {
|
|
return engine::core::BackendType::Vulkan;
|
|
}
|
|
if (value == "best") {
|
|
return engine::core::BackendType::BestAvailable;
|
|
}
|
|
throw std::runtime_error("unsupported Whisper embedding backend: " + value);
|
|
}
|
|
|
|
std::vector<float> read_f32_file(const std::filesystem::path & path, size_t expected_count) {
|
|
std::ifstream input(path, std::ios::binary);
|
|
if (!input) {
|
|
throw std::runtime_error("failed to open input file: " + path.string());
|
|
}
|
|
input.seekg(0, std::ios::end);
|
|
const std::streamoff bytes = input.tellg();
|
|
input.seekg(0, std::ios::beg);
|
|
if (bytes < 0 || static_cast<size_t>(bytes) != expected_count * sizeof(float)) {
|
|
throw std::runtime_error("input file byte count does not match expected Whisper mel shape");
|
|
}
|
|
std::vector<float> values(expected_count);
|
|
input.read(reinterpret_cast<char *>(values.data()), static_cast<std::streamsize>(bytes));
|
|
if (!input) {
|
|
throw std::runtime_error("failed to read input file: " + path.string());
|
|
}
|
|
return values;
|
|
}
|
|
|
|
void write_f32_file(const std::filesystem::path & path, const std::vector<float> & values) {
|
|
std::ofstream output(path, std::ios::binary);
|
|
if (!output) {
|
|
throw std::runtime_error("failed to open output file: " + path.string());
|
|
}
|
|
output.write(
|
|
reinterpret_cast<const char *>(values.data()),
|
|
static_cast<std::streamsize>(values.size() * sizeof(float)));
|
|
if (!output) {
|
|
throw std::runtime_error("failed to write output file: " + path.string());
|
|
}
|
|
}
|
|
|
|
engine::modules::WhisperEmbeddingConfig load_config(const std::filesystem::path & model_dir) {
|
|
const auto root = engine::io::json::parse_file(model_dir / "config.json");
|
|
const auto & audio = root.require("audio_encoder");
|
|
engine::modules::WhisperEmbeddingConfig config;
|
|
config.n_mels = audio.require("n_mels").as_i64();
|
|
config.n_audio_ctx = audio.require("n_audio_ctx").as_i64();
|
|
config.n_audio_state = audio.require("n_audio_state").as_i64();
|
|
config.n_audio_head = audio.require("n_audio_head").as_i64();
|
|
config.n_audio_layer = audio.require("n_audio_layer").as_i64();
|
|
return config;
|
|
}
|
|
|
|
engine::modules::LinearWeights load_linear(
|
|
engine::core::BackendWeightStore & store,
|
|
const engine::assets::TensorSource & source,
|
|
const std::string & prefix,
|
|
int64_t out_features,
|
|
int64_t in_features,
|
|
bool use_bias) {
|
|
engine::modules::LinearWeights weights;
|
|
weights.weight = store.load_tensor(source, prefix + ".weight", engine::assets::TensorStorageType::F32, {out_features, in_features});
|
|
if (use_bias) {
|
|
weights.bias = store.load_f32_tensor(source, prefix + ".bias", {out_features});
|
|
}
|
|
return weights;
|
|
}
|
|
|
|
engine::modules::NormWeights load_norm(
|
|
engine::core::BackendWeightStore & store,
|
|
const engine::assets::TensorSource & source,
|
|
const std::string & prefix,
|
|
int64_t hidden_size) {
|
|
return {
|
|
store.load_f32_tensor(source, prefix + ".weight", {hidden_size}),
|
|
store.load_f32_tensor(source, prefix + ".bias", {hidden_size}),
|
|
};
|
|
}
|
|
|
|
engine::modules::WhisperEmbeddingWeights load_weights(
|
|
ggml_backend_t backend,
|
|
const std::filesystem::path & model_dir,
|
|
const engine::modules::WhisperEmbeddingConfig & config,
|
|
std::shared_ptr<engine::core::BackendWeightStore> & store_out) {
|
|
const auto source = engine::assets::open_tensor_source(model_dir / "model.safetensors");
|
|
auto store = std::make_shared<engine::core::BackendWeightStore>(
|
|
backend,
|
|
"whisper_embedding.weights",
|
|
kWeightContextBytes);
|
|
engine::modules::WhisperEmbeddingWeights weights;
|
|
weights.conv1 = {
|
|
store->load_tensor(*source, "encoder.conv1.weight", engine::assets::TensorStorageType::F32, {config.n_audio_state, config.n_mels, 3}),
|
|
store->load_f32_tensor(*source, "encoder.conv1.bias", {config.n_audio_state}),
|
|
};
|
|
weights.conv2 = {
|
|
store->load_tensor(*source, "encoder.conv2.weight", engine::assets::TensorStorageType::F32, {config.n_audio_state, config.n_audio_state, 3}),
|
|
store->load_f32_tensor(*source, "encoder.conv2.bias", {config.n_audio_state}),
|
|
};
|
|
weights.positional_embedding = store->load_f32_tensor(
|
|
*source,
|
|
"encoder.positional_embedding",
|
|
{config.n_audio_ctx, config.n_audio_state});
|
|
weights.layers.reserve(static_cast<size_t>(config.n_audio_layer));
|
|
for (int64_t layer = 0; layer < config.n_audio_layer; ++layer) {
|
|
const std::string prefix = "encoder.blocks." + std::to_string(layer);
|
|
engine::modules::WhisperEncoderLayerWeights layer_weights;
|
|
layer_weights.attention_norm = load_norm(*store, *source, prefix + ".attn_ln", config.n_audio_state);
|
|
layer_weights.attention.query = load_linear(*store, *source, prefix + ".attn.query", config.n_audio_state, config.n_audio_state, true);
|
|
layer_weights.attention.key = load_linear(*store, *source, prefix + ".attn.key", config.n_audio_state, config.n_audio_state, false);
|
|
layer_weights.attention.value = load_linear(*store, *source, prefix + ".attn.value", config.n_audio_state, config.n_audio_state, true);
|
|
layer_weights.attention.out = load_linear(*store, *source, prefix + ".attn.out", config.n_audio_state, config.n_audio_state, true);
|
|
layer_weights.mlp_norm = load_norm(*store, *source, prefix + ".mlp_ln", config.n_audio_state);
|
|
layer_weights.mlp.fc1_weight = store->load_tensor(
|
|
*source,
|
|
prefix + ".mlp.0.weight",
|
|
engine::assets::TensorStorageType::F32,
|
|
{config.n_audio_state * 4, config.n_audio_state});
|
|
layer_weights.mlp.fc1_bias = store->load_f32_tensor(*source, prefix + ".mlp.0.bias", {config.n_audio_state * 4});
|
|
layer_weights.mlp.fc2_weight = store->load_tensor(
|
|
*source,
|
|
prefix + ".mlp.2.weight",
|
|
engine::assets::TensorStorageType::F32,
|
|
{config.n_audio_state, config.n_audio_state * 4});
|
|
layer_weights.mlp.fc2_bias = store->load_f32_tensor(*source, prefix + ".mlp.2.bias", {config.n_audio_state});
|
|
weights.layers.push_back(std::move(layer_weights));
|
|
}
|
|
weights.final_norm = load_norm(*store, *source, "encoder.ln_post", config.n_audio_state);
|
|
store->upload();
|
|
store_out = std::move(store);
|
|
return weights;
|
|
}
|
|
|
|
class WhisperEmbeddingGraph {
|
|
public:
|
|
WhisperEmbeddingGraph(
|
|
ggml_backend_t backend,
|
|
engine::core::BackendType backend_type,
|
|
const engine::modules::WhisperEmbeddingConfig & config,
|
|
const engine::modules::WhisperEmbeddingWeights & weights)
|
|
: backend_(backend),
|
|
config_(config) {
|
|
ggml_init_params params{kGraphContextBytes, nullptr, true};
|
|
ctx_.reset(ggml_init(params));
|
|
if (ctx_ == nullptr) {
|
|
throw std::runtime_error("failed to initialize Whisper embedding graph context");
|
|
}
|
|
engine::core::ModuleBuildContext ctx{ctx_.get(), "whisper.embedding", backend_type};
|
|
auto input = engine::core::make_tensor(
|
|
ctx,
|
|
GGML_TYPE_F32,
|
|
engine::core::TensorShape::from_dims({1, config_.n_mels, config_.n_audio_ctx * 2}));
|
|
input_ = input.tensor;
|
|
auto output = engine::modules::WhisperEmbeddingModule(config_).build(ctx, input, weights);
|
|
output_ = output.tensor;
|
|
ggml_set_output(output_);
|
|
graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false);
|
|
ggml_build_forward_expand(graph_, output_);
|
|
gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));
|
|
if (gallocr_ == nullptr || !ggml_gallocr_reserve(gallocr_, graph_) || !ggml_gallocr_alloc_graph(gallocr_, graph_)) {
|
|
throw std::runtime_error("failed to allocate Whisper embedding graph");
|
|
}
|
|
}
|
|
|
|
~WhisperEmbeddingGraph() {
|
|
if (gallocr_ != nullptr) {
|
|
ggml_gallocr_free(gallocr_);
|
|
}
|
|
}
|
|
|
|
std::vector<float> run(const std::vector<float> & mel) {
|
|
if (mel.size() != static_cast<size_t>(config_.n_mels * config_.n_audio_ctx * 2)) {
|
|
throw std::runtime_error("Whisper embedding input shape mismatch");
|
|
}
|
|
ggml_backend_tensor_set(input_, mel.data(), 0, mel.size() * sizeof(float));
|
|
const ggml_status status = engine::core::compute_backend_graph(backend_, graph_);
|
|
ggml_backend_synchronize(backend_);
|
|
if (status != GGML_STATUS_SUCCESS) {
|
|
throw std::runtime_error("Whisper embedding graph compute failed");
|
|
}
|
|
std::vector<float> output(static_cast<size_t>(config_.n_audio_ctx * config_.n_audio_state));
|
|
ggml_backend_tensor_get(output_, output.data(), 0, output.size() * sizeof(float));
|
|
return output;
|
|
}
|
|
|
|
private:
|
|
ggml_backend_t backend_ = nullptr;
|
|
engine::modules::WhisperEmbeddingConfig config_;
|
|
std::unique_ptr<ggml_context, GgmlContextDeleter> ctx_;
|
|
ggml_tensor * input_ = nullptr;
|
|
ggml_tensor * output_ = nullptr;
|
|
ggml_cgraph * graph_ = nullptr;
|
|
ggml_gallocr_t gallocr_ = nullptr;
|
|
};
|
|
|
|
} // namespace
|
|
|
|
int main(int argc, char ** argv) {
|
|
try {
|
|
const std::filesystem::path model_dir = arg_value(argc, argv, "--model", "models/whisper-medium");
|
|
const std::filesystem::path input_path = arg_value(argc, argv, "--input", "build/whisper_embedding_input.f32");
|
|
const std::filesystem::path output_path = arg_value(argc, argv, "--output", "build/whisper_embedding_cpp.f32");
|
|
const std::string backend_name = arg_value(argc, argv, "--backend", "cpu");
|
|
const int device = int_arg(argc, argv, "--device", 0);
|
|
const int threads = int_arg(argc, argv, "--threads", 8);
|
|
|
|
engine::core::BackendConfig backend_config;
|
|
backend_config.type = parse_backend(backend_name);
|
|
backend_config.device = device;
|
|
backend_config.threads = threads;
|
|
ggml_backend_t backend = engine::core::init_backend(backend_config);
|
|
if (backend_config.type == engine::core::BackendType::Cpu) {
|
|
engine::core::set_backend_threads(backend, threads);
|
|
}
|
|
|
|
auto config = load_config(model_dir);
|
|
const auto mel = read_f32_file(input_path, static_cast<size_t>(config.n_mels * config.n_audio_ctx * 2));
|
|
{
|
|
std::shared_ptr<engine::core::BackendWeightStore> weight_store;
|
|
const auto weights = load_weights(backend, model_dir, config, weight_store);
|
|
WhisperEmbeddingGraph graph(backend, backend_config.type, config, weights);
|
|
const auto output = graph.run(mel);
|
|
write_f32_file(output_path, output);
|
|
}
|
|
ggml_backend_free(backend);
|
|
std::cout << "output=" << output_path << "\n";
|
|
std::cout << "shape=1," << config.n_audio_ctx << "," << config.n_audio_state << "\n";
|
|
return 0;
|
|
} catch (const std::exception & error) {
|
|
std::cerr << "whisper_embedding_parity failed: " << error.what() << "\n";
|
|
return 1;
|
|
}
|
|
}
|