diff --git a/Cargo.toml b/Cargo.toml index 21ef8f3..e884808 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,6 +11,9 @@ resolver = "2" [patch.crates-io] redeem-properties = { path = "redeem-properties" } +candle-core = { git = "https://github.com/singjc/candle.git", branch = "add/candle-onnx-export" } +candle-nn = { git = "https://github.com/singjc/candle.git", branch = "add/candle-onnx-export" } +candle-transformers = { git = "https://github.com/singjc/candle.git", branch = "add/candle-onnx-export" } [profile.release] lto = "fat" diff --git a/redeem-cli/Cargo.toml b/redeem-cli/Cargo.toml index 53db0b1..6b8efad 100644 --- a/redeem-cli/Cargo.toml +++ b/redeem-cli/Cargo.toml @@ -10,8 +10,11 @@ name = "redeem" path = "src/main.rs" [dependencies] -redeem-properties = { path = "../redeem-properties" } +redeem-properties = { path = "../redeem-properties", features = ["onnx-export"] } redeem-classifiers = { path = "../redeem-classifiers" } +candle-core = { git = "https://github.com/singjc/candle.git", branch = "add/candle-onnx-export" } +candle-nn = { git = "https://github.com/singjc/candle.git", branch = "add/candle-onnx-export" } +candle-onnx-export = { git = "https://github.com/singjc/candle.git", branch = "add/candle-onnx-export" } env_logger = "0.11.8" log = "0.4" clap = { version="4.0", features = ["cargo", "unicode"] } @@ -24,19 +27,19 @@ maud = "0.27.0" plotly = "0.12.1" rand = "0.8" -[dependencies.candle-core] -version = "0.8.4" -default-features = false -features = [] +# [dependencies.candle-core] +# version = "0.11.0" +# default-features = false +# features = [] -[dependencies.candle-nn] -version = "0.8.4" -default-features = false -features = [] +# [dependencies.candle-nn] +# version = "0.11.0" +# default-features = false +# features = [] [features] default = [] -cuda = ["candle-core/cuda"] +cuda = ["candle-core/cuda", "candle-nn/cuda", "redeem-properties/cuda"] svm = ["redeem-classifiers/svm"] xgboost = ["redeem-classifiers/xgboost"] diff --git a/redeem-cli/README.md b/redeem-cli/README.md index 97d65c3..40d37f5 100644 --- a/redeem-cli/README.md +++ b/redeem-cli/README.md @@ -102,6 +102,60 @@ redeem properties inference inference_config.json \ -o predictions.tsv ``` +### Properties — ONNX Export + +Export a supported property model to ONNX. Full-model export is currently implemented for +`rt_cnn_tf` and `ccs_cnn_tf` through the `ToOnnx` implementations in `redeem-properties`. +The previous decoder/head-only export remains available with `--component decoder`. + +```bash +redeem properties to-onnx \ + --model \ + --model_arch \ + --output \ + [--component full|decoder] \ + [--constants ] \ + [--external-data] +``` + +| Argument | Description | +|----------|-------------| +| `-m`, `--model` | Path to the trained model file (`.safetensors`, `.pt`, `.pth`, or `.pkl`) | +| `-a`, `--model_arch` | Model architecture (`rt_cnn_tf`, `ccs_cnn_tf`, `rt_cnn_lstm`, `ccs_cnn_lstm`, `ms2_bert`) | +| `-o`, `--output` | Path to write the ONNX model | +| `--component` | Component to export. Defaults to `full`; use `decoder` for the legacy head-only export | +| `--constants` | Optional model constants YAML | +| `--external-data` | Store initializer bytes in a sidecar `.onnx.data` file | +| `--data-file` | Optional external data file name recorded in the ONNX model | + +**Full RT transformer example:** + +```bash +redeem properties to-onnx \ + --model redeem-properties/assets/pretrained_models/redeem/20251205_100_epochs_min_max_rt_cnn_tf.safetensors \ + --model_arch rt_cnn_tf \ + --output redeem-properties/assets/pretrained_models/redeem/redeem_rt_cnn_tf.onnx +``` + +**Full CCS transformer example:** + +```bash +redeem properties to-onnx \ + --model redeem-properties/assets/pretrained_models/redeem/20251205_500_epochs_early_stopped_100_min_max_ccs_cnn_tf.safetensors \ + --model_arch ccs_cnn_tf \ + --output redeem-properties/assets/pretrained_models/redeem/redeem_ccs_cnn_tf.onnx +``` + +**Decoder/head-only export:** + +```bash +redeem properties to-onnx \ + --model model.safetensors \ + --model_arch rt_cnn_tf \ + --component decoder \ + --output rt_decoder.onnx +``` + ### Classifiers — Score Score a Percolator `.pin` file using the semi-supervised classifier. diff --git a/redeem-cli/src/main.rs b/redeem-cli/src/main.rs index 46af3d5..b917808 100644 --- a/redeem-cli/src/main.rs +++ b/redeem-cli/src/main.rs @@ -11,6 +11,7 @@ use redeem_cli::classifiers::score::score::{ }; use redeem_cli::properties::inference::inference; use redeem_cli::properties::inference::input::PropertyInferenceConfig; +use redeem_cli::properties::onnx; use redeem_cli::properties::train::input::PropertyTrainConfig; use redeem_cli::properties::train::trainer; @@ -140,6 +141,91 @@ fn main() -> Result<()> { .value_parser(clap::value_parser!(PathBuf)) .value_hint(ValueHint::FilePath), ) + ) + .subcommand( + Command::new("to-onnx") + .visible_alias("to_onnx") + .about("Export a supported property model to ONNX") + .arg( + Arg::new("model_path") + .short('m') + .long("model") + .help("Path to the trained model file (*.safetensors, *.pt, *.pth, or *.pkl)") + .required(true) + .value_parser(clap::value_parser!(PathBuf)) + .value_hint(ValueHint::FilePath), + ) + .arg( + Arg::new("model_arch") + .short('a') + .long("model_arch") + .help("Model architecture to export") + .value_parser([ + "rt_cnn_lstm", + "rt_cnn_tf", + "ccs_cnn_lstm", + "ccs_cnn_tf", + "ms2_bert", + ]) + .required(true), + ) + .arg( + Arg::new("output_file") + .short('o') + .long("output") + .help("Path to the output ONNX model (*.onnx)") + .required(true) + .value_parser(clap::value_parser!(PathBuf)) + .value_hint(ValueHint::FilePath), + ) + .arg( + Arg::new("constants") + .long("constants") + .help("Optional PeptDeep/Redeem model constants YAML") + .value_parser(clap::value_parser!(PathBuf)) + .value_hint(ValueHint::FilePath), + ) + .arg( + Arg::new("component") + .long("component") + .help("Model component to export. Full export is supported for rt_cnn_tf and ccs_cnn_tf; decoder keeps the legacy head-only export.") + .value_parser(["full", "decoder"]) + .default_value("full"), + ) + .arg( + Arg::new("device") + .long("device") + .help("Device used to load model tensors") + .value_parser(clap::builder::NonEmptyStringValueParser::new()) + .default_value("cpu"), + ) + .arg( + Arg::new("opset") + .long("opset") + .help("ONNX opset version") + .value_parser(clap::value_parser!(u32)) + .default_value("18"), + ) + .arg( + Arg::new("external_data") + .long("external-data") + .help("Write initializer tensor bytes to a sidecar .onnx.data file") + .action(ArgAction::SetTrue), + ) + .arg( + Arg::new("external_data_threshold") + .long("external-data-threshold") + .help("Minimum initializer byte size to write to external data when --external-data is set") + .value_parser(clap::value_parser!(usize)) + .default_value("1024"), + ) + .arg( + Arg::new("data_file") + .long("data-file") + .help("External data file name to record in the ONNX model") + .value_parser(clap::value_parser!(PathBuf)) + .value_hint(ValueHint::FilePath), + ), ), ) .subcommand( @@ -294,6 +380,15 @@ fn handle_properties(matches: &ArgMatches) -> Result<()> { } } } + Some(("to-onnx", onnx_matches)) | Some(("to_onnx", onnx_matches)) => { + match onnx::run_to_onnx(onnx_matches) { + Ok(_) => Ok(()), + Err(e) => { + log::error!("ONNX export failed: {:#}", e); + std::process::exit(1) + } + } + } _ => unreachable!(), } } diff --git a/redeem-cli/src/properties/mod.rs b/redeem-cli/src/properties/mod.rs index beb9fad..8eedcf4 100644 --- a/redeem-cli/src/properties/mod.rs +++ b/redeem-cli/src/properties/mod.rs @@ -1,4 +1,5 @@ pub mod inference; pub mod load_data; +pub mod onnx; pub mod train; pub mod util; diff --git a/redeem-cli/src/properties/onnx.rs b/redeem-cli/src/properties/onnx.rs new file mode 100644 index 0000000..57ac8c7 --- /dev/null +++ b/redeem-cli/src/properties/onnx.rs @@ -0,0 +1,294 @@ +use anyhow::{bail, Context, Result}; +use candle_core::Device; +use candle_onnx_export::{ + Dim, ExportContext, ExportOptions, OnnxGraph, Shape, TensorElementType, ToOnnx, WeightFormat, +}; +use clap::ArgMatches; +use redeem_properties::models::{ + ccs_cnn_lstm_model::CCSCNNLSTMModel, ccs_cnn_tf_model::CCSCNNTFModel, + model_interface::ModelInterface, rt_cnn_lstm_model::RTCNNLSTMModel, + rt_cnn_transformer_model::RTCNNTFModel, +}; +use redeem_properties::onnx_export::export_decoder_head_from_varmap; +use std::path::{Path, PathBuf}; + +/// Keep small shape/control initializers embedded so ONNX shape inference can read them. +const DEFAULT_EXTERNAL_DATA_THRESHOLD: usize = 1024; + +/// Export a supported `redeem-properties` model component to ONNX. +/// +/// `--component decoder` exports only the decoder/head. `--component full` exports the complete +/// model graph for architectures that have a `ToOnnx` implementation in `redeem-properties`. +pub fn run_to_onnx(matches: &ArgMatches) -> Result<()> { + let model_path = matches + .get_one::("model_path") + .context("missing required --model")? + .clone(); + let output_path = matches + .get_one::("output_file") + .context("missing required --output")? + .clone(); + let model_arch = matches + .get_one::("model_arch") + .context("missing required --model_arch")? + .as_str(); + let constants_path = matches.get_one::("constants").cloned(); + let component = matches + .get_one::("component") + .map(String::as_str) + .unwrap_or("decoder"); + let device_name = matches + .get_one::("device") + .map(String::as_str) + .unwrap_or("cpu"); + let opset_version = matches + .get_one::("opset") + .copied() + .unwrap_or(18) as i64; + let external_data_threshold = matches + .get_one::("external_data_threshold") + .copied() + .unwrap_or(DEFAULT_EXTERNAL_DATA_THRESHOLD); + + let device = parse_device(device_name)?; + + let mut graph = OnnxGraph::new(format!("{model_arch}_{component}")); + let output = match component { + "decoder" => { + let (decoder_prefix, decoder_dim) = decoder_layout(model_arch)?; + let input = graph.add_input( + format!("{decoder_prefix}_input"), + TensorElementType::Float32, + Shape::from_dims([Dim::from("batch"), Dim::from(decoder_dim)]), + ); + load_and_export_decoder( + &mut graph, + model_arch, + &model_path, + constants_path.as_ref(), + device, + decoder_prefix, + input, + )? + } + "full" => load_and_export_full( + &mut graph, + model_arch, + &model_path, + constants_path.as_ref(), + device, + )?, + other => bail!("unsupported ONNX export component: {other}"), + }; + graph.add_output(output); + + let weight_format = if matches.get_flag("external_data") { + WeightFormat::External { + data_filename: external_data_filename(&output_path, matches.get_one::("data_file"))?, + size_threshold: external_data_threshold, + } + } else { + WeightFormat::Embedded + }; + + graph + .save( + &output_path, + ExportOptions { + opset_version, + weight_format, + ..ExportOptions::default() + }, + ) + .map_err(|err| anyhow::anyhow!("failed to save ONNX model: {err}"))?; + + eprintln!( + "[ReDeeM::Properties] Exported {model_arch} {component} component to {:?}", + output_path + ); + Ok(()) +} + +fn load_and_export_decoder( + graph: &mut OnnxGraph, + model_arch: &str, + model_path: &Path, + constants_path: Option<&PathBuf>, + device: Device, + decoder_prefix: &str, + input: candle_onnx_export::Value, +) -> Result { + match model_arch { + "rt_cnn_lstm" => { + let model = RTCNNLSTMModel::new( + model_path.to_path_buf(), + constants_path.cloned(), + 0, + 0, + 0, + true, + device, + )?; + export_decoder_head_from_varmap(graph, model.get_varmap(), decoder_prefix, input, decoder_prefix) + .map_err(|err| anyhow::anyhow!("failed to export decoder/head: {err}")) + } + "rt_cnn_tf" => { + let model = RTCNNTFModel::new( + model_path.to_path_buf(), + constants_path.cloned(), + 0, + 0, + 0, + true, + device, + )?; + export_decoder_head_from_varmap(graph, model.get_varmap(), decoder_prefix, input, decoder_prefix) + .map_err(|err| anyhow::anyhow!("failed to export decoder/head: {err}")) + } + "ccs_cnn_lstm" => { + let model = CCSCNNLSTMModel::new( + model_path.to_path_buf(), + constants_path.cloned(), + 0, + 0, + 0, + true, + device, + )?; + export_decoder_head_from_varmap(graph, model.get_varmap(), decoder_prefix, input, decoder_prefix) + .map_err(|err| anyhow::anyhow!("failed to export decoder/head: {err}")) + } + "ccs_cnn_tf" => { + let model = CCSCNNTFModel::new( + model_path.to_path_buf(), + constants_path.cloned(), + 0, + 0, + 0, + true, + device, + )?; + export_decoder_head_from_varmap(graph, model.get_varmap(), decoder_prefix, input, decoder_prefix) + .map_err(|err| anyhow::anyhow!("failed to export decoder/head: {err}")) + } + "ms2_bert" => bail!( + "ms2_bert decoder export is not wired yet because the MS2 output/modloss heads need \ + sequence-shaped output metadata" + ), + other => bail!("unsupported model architecture: {other}"), + } +} + +fn load_and_export_full( + graph: &mut OnnxGraph, + model_arch: &str, + model_path: &Path, + constants_path: Option<&PathBuf>, + device: Device, +) -> Result { + match model_arch { + "rt_cnn_tf" => { + let model = RTCNNTFModel::new( + model_path.to_path_buf(), + constants_path.cloned(), + 0, + 0, + 0, + true, + device, + )?; + let input = redeem_properties::onnx_export::add_rt_cnn_tf_input(graph, "input"); + export_full_model(graph, &model, input) + } + "ccs_cnn_tf" => { + let model = CCSCNNTFModel::new( + model_path.to_path_buf(), + constants_path.cloned(), + 0, + 0, + 0, + true, + device, + )?; + let input = redeem_properties::onnx_export::add_ccs_cnn_tf_input(graph, "input"); + export_full_model(graph, &model, input) + } + "rt_cnn_lstm" | "ccs_cnn_lstm" | "ms2_bert" => bail!( + "full ONNX export for {model_arch} is not implemented yet; supported full exports: \ + rt_cnn_tf, ccs_cnn_tf" + ), + other => bail!("unsupported model architecture: {other}"), + } +} + +fn export_full_model( + graph: &mut OnnxGraph, + model: &M, + input: candle_onnx_export::Value, +) -> Result { + let mut ctx = ExportContext::new(graph); + let outputs = model + .to_onnx(&mut ctx, &[input]) + .map_err(|err| anyhow::anyhow!("failed to export full model: {err}"))?; + outputs + .into_iter() + .next() + .context("full model ONNX export produced no outputs") +} + +fn decoder_layout(model_arch: &str) -> Result<(&'static str, usize)> { + match model_arch { + "rt_cnn_lstm" => Ok(("rt_decoder", 256)), + "rt_cnn_tf" => Ok(("rt_decoder", 192)), + "ccs_cnn_lstm" => Ok(("ccs_decoder", 257)), + "ccs_cnn_tf" => Ok(("ccs_decoder", 129)), + "ms2_bert" => bail!( + "ms2_bert decoder export is not wired yet because the MS2 output/modloss heads need \ + sequence-shaped output metadata" + ), + other => bail!("unsupported model architecture: {other}"), + } +} + +fn parse_device(device_name: &str) -> Result { + match device_name { + "cpu" => Ok(Device::Cpu), + #[cfg(feature = "cuda")] + "cuda" => Ok(Device::new_cuda(0)?), + #[cfg(feature = "cuda")] + name if name.starts_with("cuda:") => { + let index = name + .strip_prefix("cuda:") + .and_then(|value| value.parse::().ok()) + .context("invalid CUDA device; expected cuda or cuda:")?; + Ok(Device::new_cuda(index)?) + } + other => bail!("unsupported device '{other}'; use cpu{}", cuda_help_suffix()), + } +} + +fn cuda_help_suffix() -> &'static str { + #[cfg(feature = "cuda")] + { + ", cuda, or cuda:" + } + #[cfg(not(feature = "cuda"))] + { + " or build redeem-cli with the cuda feature" + } +} + +fn external_data_filename(output_path: &Path, data_file: Option<&PathBuf>) -> Result { + if let Some(data_file) = data_file { + return data_file + .file_name() + .map(|name| name.to_string_lossy().into_owned()) + .context("--data-file must include a file name"); + } + + let onnx_name = output_path + .file_name() + .map(|name| name.to_string_lossy()) + .context("--output must include a file name")?; + Ok(format!("{onnx_name}.data")) +} diff --git a/redeem-cli/tests/cli_smoke.rs b/redeem-cli/tests/cli_smoke.rs index 8cdb80a..a47999f 100644 --- a/redeem-cli/tests/cli_smoke.rs +++ b/redeem-cli/tests/cli_smoke.rs @@ -81,6 +81,34 @@ fn properties_inference_no_config_prints_template() { .stderr(predicate::str::contains("No config file provided")); } +#[test] +fn properties_help_lists_to_onnx() { + cmd() + .args(["properties", "--help"]) + .assert() + .success() + .stdout(predicate::str::contains("to-onnx")); +} + +#[test] +fn properties_to_onnx_help_mentions_full_and_decoder_components() { + cmd() + .args(["properties", "to-onnx", "--help"]) + .assert() + .success() + .stdout(predicate::str::contains("full")) + .stdout(predicate::str::contains("decoder")); +} + +#[test] +fn properties_to_onnx_missing_args_errors() { + cmd() + .args(["properties", "to-onnx"]) + .assert() + .failure() + .stderr(predicate::str::contains("required")); +} + // --------------------------------------------------------------------------- // Classifiers subcommand // --------------------------------------------------------------------------- diff --git a/redeem-openms-ffi/Cargo.toml b/redeem-openms-ffi/Cargo.toml index 13ea04b..e751da4 100644 --- a/redeem-openms-ffi/Cargo.toml +++ b/redeem-openms-ffi/Cargo.toml @@ -12,7 +12,8 @@ crate-type = ["staticlib", "rlib"] [dependencies] anyhow = "1.0" -candle-core = { version = "0.8.4", default-features = false } +# candle-core = { version = "0.8.4", default-features = false } +candle-core = { git = "https://github.com/singjc/candle.git", branch = "add/candle-onnx-export" } csv = "1.1" once_cell = "1.20" redeem-properties = { version = "0.1.2", default-features = false } diff --git a/redeem-properties-py/Cargo.toml b/redeem-properties-py/Cargo.toml index 34ffbba..e23a0c1 100644 --- a/redeem-properties-py/Cargo.toml +++ b/redeem-properties-py/Cargo.toml @@ -13,10 +13,11 @@ redeem-properties = { path = "../redeem-properties" } pyo3 = { version = "0.22", features = ["extension-module"] } numpy = "0.22" anyhow = "1.0" +candle-core = { git = "https://github.com/singjc/candle.git", branch = "add/candle-onnx-export" } -[dependencies.candle-core] -version = "0.8.4" -default-features = false +# [dependencies.candle-core] +# version = "0.8.4" +# default-features = false [features] default = [] diff --git a/redeem-properties/Cargo.toml b/redeem-properties/Cargo.toml index 324cf9b..2215417 100644 --- a/redeem-properties/Cargo.toml +++ b/redeem-properties/Cargo.toml @@ -36,32 +36,40 @@ tqdm = "0.8.0" rayon = "1.5" sysinfo = "0.33.1" rustyms = { version = "0.11", default-features = false } +candle-core = { git = "https://github.com/singjc/candle.git", branch = "add/candle-onnx-export" } +candle-nn = { git = "https://github.com/singjc/candle.git", branch = "add/candle-onnx-export" } +candle-transformers = { git = "https://github.com/singjc/candle.git", branch = "add/candle-onnx-export" } + +candle-onnx-export = {git = "https://github.com/singjc/candle.git", branch = "add/candle-onnx-export", optional = true} # Optional embedding of pretrained model files into the crate binary. [dependencies.include_dir] version = "0.7" optional = true -[dependencies.candle-core] -version = "0.8.4" -default-features = false -features = [] +# [dependencies.candle-core] +# version = "0.11.0" +# default-features = false +# features = [] + +# [dependencies.candle-nn] +# version = "0.11.0" +# default-features = false +# features = [] + +# [dependencies.candle-transformers] +# version = "0.11.0" +# default-features = false +# features = [] -[dependencies.candle-nn] -version = "0.8.4" -default-features = false -features = [] -[dependencies.candle-transformers] -version = "0.8.4" -default-features = false -features = [] [features] default = ["legacy-peptdeep-models", "pretrained-download"] cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"] embed-pretrained = ["include_dir"] legacy-peptdeep-models = [] +onnx-export = ["dep:candle-onnx-export"] pretrained-download = ["dep:reqwest", "dep:zip"] diff --git a/redeem-properties/assets/pretrained_models/redeem/redeem_ccs_cnn_tf.onnx b/redeem-properties/assets/pretrained_models/redeem/redeem_ccs_cnn_tf.onnx new file mode 100644 index 0000000..019fc2f Binary files /dev/null and b/redeem-properties/assets/pretrained_models/redeem/redeem_ccs_cnn_tf.onnx differ diff --git a/redeem-properties/assets/pretrained_models/redeem/redeem_rt_cnn_tf.onnx b/redeem-properties/assets/pretrained_models/redeem/redeem_rt_cnn_tf.onnx new file mode 100644 index 0000000..c616eb6 Binary files /dev/null and b/redeem-properties/assets/pretrained_models/redeem/redeem_rt_cnn_tf.onnx differ diff --git a/redeem-properties/src/lib.rs b/redeem-properties/src/lib.rs index 77ebf5a..0712684 100644 --- a/redeem-properties/src/lib.rs +++ b/redeem-properties/src/lib.rs @@ -1,4 +1,6 @@ pub mod building_blocks; pub mod models; +#[cfg(feature = "onnx-export")] +pub mod onnx_export; pub mod pretrained; pub mod utils; diff --git a/redeem-properties/src/models/ccs_cnn_lstm_model.rs b/redeem-properties/src/models/ccs_cnn_lstm_model.rs index a66ea55..5566cc7 100644 --- a/redeem-properties/src/models/ccs_cnn_lstm_model.rs +++ b/redeem-properties/src/models/ccs_cnn_lstm_model.rs @@ -219,6 +219,19 @@ impl ModelInterface for CCSCNNLSTMModel { } } +#[cfg(feature = "onnx-export")] +impl candle_onnx_export::ToOnnx for CCSCNNLSTMModel { + fn to_onnx( + &self, + _ctx: &mut candle_onnx_export::ExportContext<'_>, + _inputs: &[candle_onnx_export::Value], + ) -> candle_onnx_export::Result> { + Err(candle_onnx_export::Error::UnsupportedTensor( + "full ONNX export for ccs_cnn_lstm is not implemented yet".to_string(), + )) + } +} + // // Forward Module Trait Implementation // impl Module for CCSCNNLSTMModel { // fn forward(&self, input: &Tensor) -> Result { diff --git a/redeem-properties/src/models/ccs_cnn_tf_model.rs b/redeem-properties/src/models/ccs_cnn_tf_model.rs index 46ddc4e..05580fb 100644 --- a/redeem-properties/src/models/ccs_cnn_tf_model.rs +++ b/redeem-properties/src/models/ccs_cnn_tf_model.rs @@ -331,6 +331,22 @@ impl ModelInterface for CCSCNNTFModel { } } +#[cfg(feature = "onnx-export")] +impl candle_onnx_export::ToOnnx for CCSCNNTFModel { + fn to_onnx( + &self, + ctx: &mut candle_onnx_export::ExportContext<'_>, + inputs: &[candle_onnx_export::Value], + ) -> candle_onnx_export::Result> { + let input = inputs + .first() + .cloned() + .unwrap_or_else(|| crate::onnx_export::add_ccs_cnn_tf_input(ctx.graph, "input")); + let output = crate::onnx_export::export_ccs_cnn_tf(&self.varmap, ctx, input)?; + Ok(vec![output]) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/redeem-properties/src/models/ms2_bert_model.rs b/redeem-properties/src/models/ms2_bert_model.rs index d14b09b..b917902 100644 --- a/redeem-properties/src/models/ms2_bert_model.rs +++ b/redeem-properties/src/models/ms2_bert_model.rs @@ -410,6 +410,19 @@ impl ModelInterface for MS2BertModel { } } +#[cfg(feature = "onnx-export")] +impl candle_onnx_export::ToOnnx for MS2BertModel { + fn to_onnx( + &self, + _ctx: &mut candle_onnx_export::ExportContext<'_>, + _inputs: &[candle_onnx_export::Value], + ) -> candle_onnx_export::Result> { + Err(candle_onnx_export::Error::UnsupportedTensor( + "full ONNX export for ms2_bert is not implemented yet".to_string(), + )) + } +} + // // Module Trait Implementation // impl Module for MS2BertModel { // fn forward(&self, input: &Tensor) -> Result { diff --git a/redeem-properties/src/models/rt_cnn_lstm_model.rs b/redeem-properties/src/models/rt_cnn_lstm_model.rs index 27ad574..ff07e5f 100644 --- a/redeem-properties/src/models/rt_cnn_lstm_model.rs +++ b/redeem-properties/src/models/rt_cnn_lstm_model.rs @@ -325,6 +325,19 @@ impl ModelInterface for RTCNNLSTMModel { } } +#[cfg(feature = "onnx-export")] +impl candle_onnx_export::ToOnnx for RTCNNLSTMModel { + fn to_onnx( + &self, + _ctx: &mut candle_onnx_export::ExportContext<'_>, + _inputs: &[candle_onnx_export::Value], + ) -> candle_onnx_export::Result> { + Err(candle_onnx_export::Error::UnsupportedTensor( + "full ONNX export for rt_cnn_lstm is not implemented yet".to_string(), + )) + } +} + // Module Trait Implementation // impl Module for RTCNNLSTMModel { diff --git a/redeem-properties/src/models/rt_cnn_transformer_model.rs b/redeem-properties/src/models/rt_cnn_transformer_model.rs index d2b7e93..4df7d77 100644 --- a/redeem-properties/src/models/rt_cnn_transformer_model.rs +++ b/redeem-properties/src/models/rt_cnn_transformer_model.rs @@ -392,3 +392,19 @@ impl ModelInterface for RTCNNTFModel { todo!("Implement print_weights for RTCNNTFModel"); } } + +#[cfg(feature = "onnx-export")] +impl candle_onnx_export::ToOnnx for RTCNNTFModel { + fn to_onnx( + &self, + ctx: &mut candle_onnx_export::ExportContext<'_>, + inputs: &[candle_onnx_export::Value], + ) -> candle_onnx_export::Result> { + let input = inputs + .first() + .cloned() + .unwrap_or_else(|| crate::onnx_export::add_rt_cnn_tf_input(ctx.graph, "input")); + let output = crate::onnx_export::export_rt_cnn_tf(&self.varmap, ctx, input)?; + Ok(vec![output]) + } +} diff --git a/redeem-properties/src/onnx_export.rs b/redeem-properties/src/onnx_export.rs new file mode 100644 index 0000000..ade50a3 --- /dev/null +++ b/redeem-properties/src/onnx_export.rs @@ -0,0 +1,877 @@ +//! ONNX export helpers for `redeem-properties` models. +//! +//! These helpers intentionally build the ONNX graph from the model `VarMap` and the same checkpoint +//! names used by the Candle loaders. Candle layer structs do not expose every tensor directly, so +//! the VarMap is the most reliable source of exportable weights. + +use candle_nn::VarMap; +use candle_onnx_export::{ + candle::tensor_to_initializer, ops, Attribute, Dim, Error as OnnxError, ExportContext, Node, + OnnxGraph, Result as OnnxResult, Shape, TensorData, TensorElementType, Value, +}; + +use crate::building_blocks::building_blocks::{AA_EMBEDDING_SIZE, MOD_FEATURE_SIZE}; + +/// Maximum sequence length available in the checkpoint positional encoding. +pub const TF_MAX_LEN: i64 = 100; + +/// Input feature width for RT transformer models: AA index + modification features. +pub const RT_TF_INPUT_FEATURES: i64 = 1 + MOD_FEATURE_SIZE as i64; + +/// Input feature width for CCS transformer models: AA index + modification features + charge. +pub const CCS_TF_INPUT_FEATURES: i64 = 1 + MOD_FEATURE_SIZE as i64 + 1; + +/// Adds a complete `rt_cnn_tf` graph body. +pub fn export_rt_cnn_tf(varmap: &VarMap, ctx: &mut ExportContext<'_>, input: Value) -> OnnxResult { + let input = limit_transformer_input_seq(ctx.graph, input, "rt_input_max_len")?; + let aa_indices = slice_feature(ctx.graph, input.clone(), 0, 1, "rt_aa_index")?; + let aa_indices = ops::squeeze_axes(ctx.graph, aa_indices, &[2], "rt_aa_indices")?; + let mod_x = slice_feature( + ctx.graph, + input, + 1, + 1 + MOD_FEATURE_SIZE as i64, + "rt_mod_features", + )?; + + let encoded = export_mod_cnn_transformer_attn_sum( + ctx.graph, + varmap, + "rt_encoder", + aa_indices, + mod_x, + None, + TransformerSpec { + mod_hidden_dim: 8, + hidden_dim: 192, + ff_dim: 768, + num_heads: 4, + num_layers: 2, + max_len: TF_MAX_LEN, + use_padding_mask: true, + }, + "rt_encoder_out", + )?; + + let decoded = + export_decoder_head_from_varmap(ctx.graph, varmap, "rt_decoder", encoded, "rt_decoder")?; + let prediction = ops::squeeze_axes(ctx.graph, decoded, &[1], "rt_prediction")?; + Ok(batch_prediction_value(prediction)) +} + +/// Adds a complete `ccs_cnn_tf` graph body. +pub fn export_ccs_cnn_tf( + varmap: &VarMap, + ctx: &mut ExportContext<'_>, + input: Value, +) -> OnnxResult { + let input = limit_transformer_input_seq(ctx.graph, input, "ccs_input_max_len")?; + let aa_indices = slice_feature(ctx.graph, input.clone(), 0, 1, "ccs_aa_index")?; + let aa_indices = ops::squeeze_axes(ctx.graph, aa_indices, &[2], "ccs_aa_indices")?; + let mod_x = slice_feature( + ctx.graph, + input.clone(), + 1, + 1 + MOD_FEATURE_SIZE as i64, + "ccs_mod_features", + )?; + let charge = slice_feature( + ctx.graph, + input, + 1 + MOD_FEATURE_SIZE as i64, + 1 + MOD_FEATURE_SIZE as i64 + 1, + "ccs_charge_raw", + )?; + let charge = ops::slice_i64( + ctx.graph, + charge, + &[0], + &[1], + Some(&[1_i64][..]), + None, + "ccs_charge_first_position", + )?; + let charge = ops::squeeze_axes(ctx.graph, charge, &[2], "ccs_charge")?; + + let encoded = export_mod_cnn_transformer_attn_sum( + ctx.graph, + varmap, + "ccs_encoder", + aa_indices, + mod_x, + Some(charge.clone()), + TransformerSpec { + mod_hidden_dim: 8, + hidden_dim: 128, + ff_dim: 256, + num_heads: 4, + num_layers: 2, + max_len: TF_MAX_LEN, + use_padding_mask: false, + }, + "ccs_encoder_out", + )?; + + let decoder_input = ops::concat(ctx.graph, &[encoded, charge], 1, "ccs_decoder_input")?; + let decoded = export_decoder_head_from_varmap( + ctx.graph, + varmap, + "ccs_decoder", + decoder_input, + "ccs_decoder", + )?; + let prediction = ops::squeeze_axes(ctx.graph, decoded, &[1], "ccs_prediction")?; + Ok(batch_prediction_value(prediction)) +} + +fn batch_prediction_value(value: Value) -> Value { + Value::new( + value.name, + value.elem_type, + Shape::from_dims([Dim::from("batch")]), + ) +} + +/// Returns a standard model input value for the RT transformer checkpoint family. +pub fn add_rt_cnn_tf_input(graph: &mut OnnxGraph, name: &str) -> Value { + graph.add_input( + name, + TensorElementType::Float32, + Shape::from_dims([ + Dim::from("batch"), + Dim::from("seq"), + Dim::Fixed(RT_TF_INPUT_FEATURES), + ]), + ) +} + +/// Returns a standard model input value for the CCS transformer checkpoint family. +pub fn add_ccs_cnn_tf_input(graph: &mut OnnxGraph, name: &str) -> Value { + graph.add_input( + name, + TensorElementType::Float32, + Shape::from_dims([ + Dim::from("batch"), + Dim::from("seq"), + Dim::Fixed(CCS_TF_INPUT_FEATURES), + ]), + ) +} + +#[derive(Debug, Clone, Copy)] +struct TransformerSpec { + mod_hidden_dim: usize, + hidden_dim: usize, + ff_dim: usize, + num_heads: usize, + num_layers: usize, + max_len: i64, + use_padding_mask: bool, +} + +fn export_mod_cnn_transformer_attn_sum( + graph: &mut OnnxGraph, + varmap: &VarMap, + prefix: &str, + aa_indices: Value, + mod_x: Value, + charge: Option, + spec: TransformerSpec, + output_name: &str, +) -> OnnxResult { + let mod_x = export_mod_embedding_fix_first_k( + graph, + varmap, + &format!("{prefix}.mod_nn"), + mod_x, + spec.mod_hidden_dim, + &format!("{prefix}_mod_embedding"), + )?; + + let mut features = vec![one_hot_aa(graph, aa_indices.clone(), &format!("{prefix}_aa_one_hot"))?, mod_x.clone()]; + if let Some(charge) = charge { + features.push(expand_charge_to_sequence( + graph, + charge, + mod_x.clone(), + &format!("{prefix}_charge_seq"), + )?); + } + let x = ops::concat(graph, &features, 2, &format!("{prefix}_aa_mod_features"))?; + let padding_mask = if spec.use_padding_mask { + Some(padding_mask_from_aa_indices( + graph, + aa_indices, + &format!("{prefix}_padding_mask"), + )?) + } else { + None + }; + + let x = export_seq_cnn(graph, varmap, &format!("{prefix}.input_cnn"), x, &format!("{prefix}_cnn"))?; + let x = linear_last_dim( + graph, + varmap, + &format!("{prefix}.proj_cnn_to_transformer"), + x, + &format!("{prefix}_projected"), + )?; + let x = export_transformer( + graph, + varmap, + &format!("{prefix}.input_transformer"), + x, + padding_mask, + spec, + &format!("{prefix}_transformer"), + )?; + export_attention_sum( + graph, + varmap, + &format!("{prefix}.attn_sum.attn.0"), + x, + output_name, + ) +} + +fn export_mod_embedding_fix_first_k( + graph: &mut OnnxGraph, + varmap: &VarMap, + prefix: &str, + mod_x: Value, + mod_hidden_dim: usize, + output_prefix: &str, +) -> OnnxResult { + let k = 6_i64; + if mod_hidden_dim as i64 <= k { + return Err(OnnxError::InvalidGraph(format!( + "mod embedding hidden dimension {mod_hidden_dim} must be greater than fixed prefix width {k}" + ))); + } + let transformed_width = mod_hidden_dim as i64 - k; + let weight_name = format!("{prefix}.nn.weight"); + let weight = tensor_from_varmap(varmap, &weight_name)?; + let weight_dims = weight.shape().dims(); + if weight_dims.len() != 2 || weight_dims[0] as i64 != transformed_width { + return Err(OnnxError::InvalidGraph(format!( + "expected {weight_name} to have output width {transformed_width}, got shape {:?}", + weight_dims + ))); + } + let first_k = ops::slice_i64( + graph, + mod_x.clone(), + &[0], + &[k], + Some(&[-1_i64][..]), + None, + &format!("{output_prefix}_first_k"), + )?; + let rest = ops::slice_i64( + graph, + mod_x, + &[k], + &[MOD_FEATURE_SIZE as i64], + Some(&[-1_i64][..]), + None, + &format!("{output_prefix}_rest"), + )?; + let transformed = linear_last_dim_with_names( + graph, + varmap, + &weight_name, + None, + rest, + &format!("{output_prefix}_transformed"), + )?; + let output = ops::concat(graph, &[first_k, transformed], -1, output_prefix)?; + Ok(output) +} + +fn export_seq_cnn( + graph: &mut OnnxGraph, + varmap: &VarMap, + prefix: &str, + input: Value, + output_prefix: &str, +) -> OnnxResult { + let transposed = ops::transpose(graph, input, &[0, 2, 1], &format!("{output_prefix}_ncl"))?; + let short = conv1d_from_varmap( + graph, + varmap, + &format!("{prefix}.cnn_short"), + transposed.clone(), + [1, 1], + &format!("{output_prefix}_short"), + )?; + let medium = conv1d_from_varmap( + graph, + varmap, + &format!("{prefix}.cnn_medium"), + transposed.clone(), + [2, 2], + &format!("{output_prefix}_medium"), + )?; + let long = conv1d_from_varmap( + graph, + varmap, + &format!("{prefix}.cnn_long"), + transposed.clone(), + [3, 3], + &format!("{output_prefix}_long"), + )?; + let concat = ops::concat(graph, &[transposed, short, medium, long], 1, &format!("{output_prefix}_cat"))?; + ops::transpose(graph, concat, &[0, 2, 1], output_prefix) +} + +fn export_transformer( + graph: &mut OnnxGraph, + varmap: &VarMap, + prefix: &str, + input: Value, + padding_mask: Option, + spec: TransformerSpec, + output_prefix: &str, +) -> OnnxResult { + let seq_len = sequence_length_vector(graph, input.clone(), &format!("{output_prefix}_seq_len"))?; + let pos_encoding = sinusoidal_position_encoding(spec.max_len as usize, spec.hidden_dim); + let pe_name = graph.unique_name(&format!("{output_prefix}_pos_encoding")); + let pe = TensorData::from_f32(pe_name, &[1, spec.max_len, spec.hidden_dim as i64], &pos_encoding)?; + let pe = graph.add_initializer_value(pe); + let pe = slice_position_encoding_to_seq( + graph, + pe, + seq_len, + &format!("{output_prefix}_pos_encoding_sliced"), + )?; + let mut x = ops::add(graph, input, pe, &format!("{output_prefix}_pos_add"))?; + + for layer_idx in 0..spec.num_layers { + x = export_transformer_layer( + graph, + varmap, + &format!("{prefix}.layer_{layer_idx}"), + x, + padding_mask.clone(), + spec, + &format!("{output_prefix}_layer_{layer_idx}"), + )?; + } + + Ok(x) +} + +fn export_transformer_layer( + graph: &mut OnnxGraph, + varmap: &VarMap, + prefix: &str, + input: Value, + padding_mask: Option, + spec: TransformerSpec, + output_prefix: &str, +) -> OnnxResult { + let attn = export_multi_head_attention( + graph, + varmap, + prefix, + input.clone(), + padding_mask, + spec, + &format!("{output_prefix}_attn"), + )?; + let attn_residual = ops::add(graph, input, attn, &format!("{output_prefix}_attn_residual"))?; + let norm1 = layer_norm_from_varmap( + graph, + varmap, + &format!("{prefix}.norm1"), + attn_residual, + &format!("{output_prefix}_norm1"), + )?; + let ff = export_feed_forward( + graph, + varmap, + prefix, + norm1.clone(), + &format!("{output_prefix}_ff"), + )?; + let ff_residual = ops::add(graph, norm1, ff, &format!("{output_prefix}_ff_residual"))?; + layer_norm_from_varmap( + graph, + varmap, + &format!("{prefix}.norm2"), + ff_residual, + &format!("{output_prefix}_norm2"), + ) +} + +fn export_multi_head_attention( + graph: &mut OnnxGraph, + varmap: &VarMap, + prefix: &str, + input: Value, + padding_mask: Option, + spec: TransformerSpec, + output_prefix: &str, +) -> OnnxResult { + let head_dim = spec.hidden_dim / spec.num_heads; + let q = project_heads(graph, varmap, &format!("{prefix}.proj_q"), input.clone(), spec, &format!("{output_prefix}_q"))?; + let k = project_heads(graph, varmap, &format!("{prefix}.proj_k"), input.clone(), spec, &format!("{output_prefix}_k"))?; + let v = project_heads(graph, varmap, &format!("{prefix}.proj_v"), input, spec, &format!("{output_prefix}_v"))?; + + let k_t = ops::transpose(graph, k, &[0, 1, 3, 2], &format!("{output_prefix}_k_t"))?; + let mut scores = ops::matmul(graph, q, k_t, &format!("{output_prefix}_scores"))?; + let scale_name = graph.unique_name(&format!("{output_prefix}_scale")); + let scale = graph.add_initializer_value(TensorData::scalar_f32(scale_name, (head_dim as f32).sqrt())?); + scores = ops::div(graph, scores, scale, &format!("{output_prefix}_scores_scaled"))?; + + if let Some(mask) = padding_mask { + let mask = ops::unsqueeze_axes(graph, mask, &[1, 2], &format!("{output_prefix}_mask_unsqueeze"))?; + let neg_name = graph.unique_name(&format!("{output_prefix}_mask_neg")); + let neg = graph.add_initializer_value(TensorData::scalar_f32(neg_name, -1e9)?); + let mask = ops::mul(graph, mask, neg, &format!("{output_prefix}_mask_scaled"))?; + scores = ops::add(graph, scores, mask, &format!("{output_prefix}_scores_masked"))?; + } + + let attn = ops::softmax(graph, scores, -1, &format!("{output_prefix}_weights"))?; + let context = ops::matmul(graph, attn, v, &format!("{output_prefix}_context"))?; + let context = ops::transpose(graph, context, &[0, 2, 1, 3], &format!("{output_prefix}_context_bthd"))?; + let context = reshape_with_shape( + graph, + context, + &[0, 0, spec.hidden_dim as i64], + &format!("{output_prefix}_context_flat"), + )?; + linear_last_dim(graph, varmap, &format!("{prefix}.proj_out"), context, output_prefix) +} + +fn project_heads( + graph: &mut OnnxGraph, + varmap: &VarMap, + prefix: &str, + input: Value, + spec: TransformerSpec, + output_prefix: &str, +) -> OnnxResult { + let head_dim = spec.hidden_dim / spec.num_heads; + let projected = linear_last_dim(graph, varmap, prefix, input, &format!("{output_prefix}_linear"))?; + let reshaped = reshape_with_shape( + graph, + projected, + &[ + 0, + 0, + spec.num_heads as i64, + head_dim as i64, + ], + &format!("{output_prefix}_reshape"), + )?; + ops::transpose(graph, reshaped, &[0, 2, 1, 3], output_prefix) +} + +fn export_feed_forward( + graph: &mut OnnxGraph, + varmap: &VarMap, + prefix: &str, + input: Value, + output_prefix: &str, +) -> OnnxResult { + let hidden = linear_last_dim(graph, varmap, &format!("{prefix}.lin1"), input, &format!("{output_prefix}_lin1"))?; + let hidden = ops::relu(graph, hidden, &format!("{output_prefix}_relu"))?; + linear_last_dim(graph, varmap, &format!("{prefix}.lin2"), hidden, output_prefix) +} + +fn sequence_length_vector( + graph: &mut OnnxGraph, + input: Value, + output_name: &str, +) -> OnnxResult { + let shape = ops::shape(graph, input, &format!("{output_name}_shape"))?; + ops::slice_i64( + graph, + shape, + &[1], + &[2], + Some(&[0_i64][..]), + None, + output_name, + ) +} + +fn limit_transformer_input_seq( + graph: &mut OnnxGraph, + input: Value, + output_name: &str, +) -> OnnxResult { + ops::slice_i64( + graph, + input, + &[0], + &[TF_MAX_LEN], + Some(&[1_i64][..]), + None, + output_name, + ) +} + +fn slice_position_encoding_to_seq( + graph: &mut OnnxGraph, + pos_encoding: Value, + seq_len: Value, + output_name: &str, +) -> OnnxResult { + let starts_name = graph.unique_name(&format!("{output_name}_starts")); + graph.add_initializer(TensorData::vec_i64(starts_name.clone(), &[0])?); + let axes_name = graph.unique_name(&format!("{output_name}_axes")); + graph.add_initializer(TensorData::vec_i64(axes_name.clone(), &[1])?); + + ops::slice( + graph, + pos_encoding, + starts_name, + seq_len, + Some(axes_name), + Option::::None, + output_name, + ) +} + +fn reshape_with_shape( + graph: &mut OnnxGraph, + input: Value, + shape: &[i64], + output_name: &str, +) -> OnnxResult { + let shape_name = graph.unique_name(&format!("{output_name}_shape")); + graph.add_initializer(TensorData::vec_i64(shape_name.clone(), shape)?); + ops::reshape(graph, input, shape_name, output_name) +} + +fn export_attention_sum( + graph: &mut OnnxGraph, + varmap: &VarMap, + prefix: &str, + input: Value, + output_name: &str, +) -> OnnxResult { + let weights = linear_last_dim(graph, varmap, prefix, input.clone(), &format!("{output_name}_scores"))?; + let weights = ops::softmax(graph, weights, 1, &format!("{output_name}_weights"))?; + let weighted = ops::mul(graph, input, weights, &format!("{output_name}_weighted"))?; + ops::reduce_sum_axes(graph, weighted, &[1], false, output_name) +} + +pub fn export_decoder_head_from_varmap( + graph: &mut OnnxGraph, + varmap: &VarMap, + prefix: &str, + input: Value, + output_prefix: &str, +) -> OnnxResult { + let data = varmap + .data() + .lock() + .map_err(|_| OnnxError::InvalidGraph("failed to lock VarMap".into()))?; + let nn_prefix = format!("{prefix}.nn."); + let mut indices = data + .keys() + .filter_map(|key| { + key.strip_prefix(&nn_prefix) + .and_then(|rest| rest.strip_suffix(".weight")) + .and_then(|idx| idx.parse::().ok()) + }) + .collect::>(); + indices.sort_unstable(); + indices.dedup(); + drop(data); + + if indices.is_empty() { + return Err(OnnxError::MissingTensor(format!("{prefix}.nn.*.weight"))); + } + + let mut value = input; + for idx in indices { + let weight = format!("{prefix}.nn.{idx}.weight"); + let bias = format!("{prefix}.nn.{idx}.bias"); + let rank = tensor_rank(varmap, &weight)?; + value = match rank { + 1 => { + ensure_initializer(graph, varmap, &weight)?; + ops::prelu(graph, value, weight, &format!("{output_prefix}_prelu_{idx}"))? + } + 2 => linear_2d_with_names( + graph, + varmap, + &weight, + has_tensor(varmap, &bias).then_some(bias.as_str()), + value, + &format!("{output_prefix}_linear_{idx}"), + )?, + _ => { + return Err(OnnxError::UnsupportedTensor(format!( + "decoder tensor {weight} has unsupported rank {rank}" + ))) + } + }; + } + + let scale_weight = format!("{prefix}.scale.weight"); + if has_tensor(varmap, &scale_weight) { + let scale_bias = format!("{prefix}.scale.bias"); + value = linear_2d_with_names( + graph, + varmap, + &scale_weight, + has_tensor(varmap, &scale_bias).then_some(scale_bias.as_str()), + value, + &format!("{output_prefix}_scale"), + )?; + } + + Ok(value) +} + +fn linear_2d_with_names( + graph: &mut OnnxGraph, + varmap: &VarMap, + weight_name: &str, + bias_name: Option<&str>, + input: Value, + output_name: &str, +) -> OnnxResult { + ensure_initializer(graph, varmap, weight_name)?; + if let Some(bias) = bias_name { + ensure_initializer(graph, varmap, bias)?; + } + ops::linear(graph, input, weight_name, bias_name, output_name) +} + +fn linear_last_dim( + graph: &mut OnnxGraph, + varmap: &VarMap, + prefix: &str, + input: Value, + output_name: &str, +) -> OnnxResult { + let weight = format!("{prefix}.weight"); + let bias = format!("{prefix}.bias"); + linear_last_dim_with_names( + graph, + varmap, + &weight, + has_tensor(varmap, &bias).then_some(bias.as_str()), + input, + output_name, + ) +} + +fn linear_last_dim_with_names( + graph: &mut OnnxGraph, + varmap: &VarMap, + weight_name: &str, + bias_name: Option<&str>, + input: Value, + output_name: &str, +) -> OnnxResult { + let weight = initializer_value(graph, varmap, weight_name)?; + let weight_t_name = graph.unique_name(&format!("{}_weight_t", sanitize(output_name))); + let weight_t = ops::transpose(graph, weight, &[1, 0], &weight_t_name)?; + let mut output = ops::matmul(graph, input, weight_t, output_name)?; + if let Some(bias) = bias_name { + ensure_initializer(graph, varmap, bias)?; + output = ops::add(graph, output, bias, &format!("{output_name}_bias"))?; + } + Ok(output) +} + +fn layer_norm_from_varmap( + graph: &mut OnnxGraph, + varmap: &VarMap, + prefix: &str, + input: Value, + output_name: &str, +) -> OnnxResult { + let weight = format!("{prefix}.weight"); + let bias = format!("{prefix}.bias"); + ensure_initializer(graph, varmap, &weight)?; + ensure_initializer(graph, varmap, &bias)?; + ops::layer_normalization(graph, input, weight, Some(bias), -1, 1e-5, output_name) +} + +fn conv1d_from_varmap( + graph: &mut OnnxGraph, + varmap: &VarMap, + prefix: &str, + input: Value, + pads: [i64; 2], + output_name: &str, +) -> OnnxResult { + let weight = format!("{prefix}.weight"); + let bias = format!("{prefix}.bias"); + ensure_initializer(graph, varmap, &weight)?; + ensure_initializer(graph, varmap, &bias)?; + ops::conv1d(graph, input, &weight, Some(&bias), pads, [1], output_name) +} + +fn one_hot_aa(graph: &mut OnnxGraph, aa_indices: Value, output_name: &str) -> OnnxResult { + let indices = ops::cast( + graph, + aa_indices, + TensorElementType::Int64, + &format!("{output_name}_indices"), + )?; + let depth_name = graph.unique_name(&format!("{output_name}_depth")); + let depth = graph.add_initializer_value(TensorData::scalar_i64( + depth_name, + AA_EMBEDDING_SIZE as i64, + )?); + let values_name = graph.unique_name(&format!("{output_name}_values")); + let values = graph.add_initializer_value(TensorData::from_f32(values_name, &[2], &[0.0, 1.0])?); + let output = Value::new(output_name, TensorElementType::Float32, Shape::unknown()); + let node_name = graph.unique_name(&format!("{output_name}_onehot")); + graph.add_node( + Node::new( + node_name, + "OneHot", + vec![indices.name, depth.name, values.name], + vec![output.name.clone()], + ) + .with_attr(Attribute::int("axis", -1)), + ); + Ok(output) +} + +fn padding_mask_from_aa_indices( + graph: &mut OnnxGraph, + aa_indices: Value, + output_name: &str, +) -> OnnxResult { + let zero_name = graph.unique_name(&format!("{output_name}_zero")); + let zero = graph.add_initializer_value(TensorData::scalar_f32(zero_name, 0.0)?); + let equal = Value::new( + graph.unique_name(&format!("{output_name}_bool")), + TensorElementType::Bool, + aa_indices.shape.clone(), + ); + let node_name = graph.unique_name(&format!("{output_name}_equal")); + graph.add_node(Node::new( + node_name, + "Equal", + vec![aa_indices.name, zero.name], + vec![equal.name.clone()], + )); + ops::cast(graph, equal, TensorElementType::Float32, output_name) +} + +fn expand_charge_to_sequence( + graph: &mut OnnxGraph, + charge: Value, + sequence_like: Value, + output_name: &str, +) -> OnnxResult { + let shape = ops::shape(graph, sequence_like, &format!("{output_name}_shape"))?; + let first_two_dims = ops::slice_i64( + graph, + shape, + &[0], + &[2], + Some(&[0_i64][..]), + None, + &format!("{output_name}_first_two_dims"), + )?; + let one_name = graph.unique_name(&format!("{output_name}_one_dim")); + let one = graph.add_initializer_value(TensorData::vec_i64(one_name, &[1])?); + let target_shape = ops::concat(graph, &[first_two_dims, one], 0, &format!("{output_name}_target_shape"))?; + let charge = ops::unsqueeze_axes(graph, charge, &[1], &format!("{output_name}_unsqueeze"))?; + let output = Value::new(output_name, TensorElementType::Float32, Shape::unknown()); + let node_name = graph.unique_name(&format!("{output_name}_expand")); + graph.add_node(Node::new( + node_name, + "Expand", + vec![charge.name, target_shape.name], + vec![output.name.clone()], + )); + Ok(output) +} + +fn slice_feature( + graph: &mut OnnxGraph, + input: Value, + start: i64, + end: i64, + output_name: &str, +) -> OnnxResult { + ops::slice_i64( + graph, + input, + &[start], + &[end], + Some(&[2_i64][..]), + None, + output_name, + ) +} + +fn ensure_initializer(graph: &mut OnnxGraph, varmap: &VarMap, name: &str) -> OnnxResult<()> { + if graph.has_initializer(name) { + return Ok(()); + } + let tensor = tensor_from_varmap(varmap, name)?; + graph.add_initializer(tensor_to_initializer(name.to_string(), &tensor)?); + Ok(()) +} + +fn initializer_value(graph: &mut OnnxGraph, varmap: &VarMap, name: &str) -> OnnxResult { + let tensor = tensor_from_varmap(varmap, name)?; + graph.add_initializer_value(tensor_to_initializer(name.to_string(), &tensor)?); + Ok(Value::new( + name, + TensorElementType::Float32, + Shape::from_dims( + tensor + .shape() + .dims() + .iter() + .copied() + .map(|dim| Dim::Fixed(dim as i64)) + .collect::>(), + ), + )) +} + +fn tensor_from_varmap(varmap: &VarMap, name: &str) -> OnnxResult { + let data = varmap + .data() + .lock() + .map_err(|_| OnnxError::InvalidGraph("failed to lock VarMap".into()))?; + data.get(name) + .map(|var| var.as_tensor().clone()) + .ok_or_else(|| OnnxError::MissingTensor(name.to_string())) +} + +fn has_tensor(varmap: &VarMap, name: &str) -> bool { + varmap + .data() + .lock() + .map(|data| data.contains_key(name)) + .unwrap_or(false) +} + +fn tensor_rank(varmap: &VarMap, name: &str) -> OnnxResult { + Ok(tensor_from_varmap(varmap, name)?.shape().dims().len()) +} + +fn sinusoidal_position_encoding(seq_len: usize, model_dim: usize) -> Vec { + let mut pe = vec![0.0; seq_len * model_dim]; + for pos in 0..seq_len { + for i in 0..model_dim { + let angle = pos as f32 / 10000_f32.powf(2.0 * ((i / 2) as f32) / model_dim as f32); + pe[pos * model_dim + i] = if i % 2 == 0 { angle.sin() } else { angle.cos() }; + } + } + pe +} + +fn sanitize(name: &str) -> String { + name.chars() + .map(|ch| if ch.is_ascii_alphanumeric() { ch } else { '_' }) + .collect() +}