Feat/onnx export

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"

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"]
 

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_PATH> \
+    --model_arch <MODEL_ARCH> \
+    --output <OUTPUT_ONNX> \
+    [--component full|decoder] \
+    [--constants <CONSTANTS_YAML>] \
+    [--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.

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!(),
     }
 }

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;