clankers-cli 0.1.6

Command-line interface for clankeRS
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

use std::fs;
use std::process::Command;

use anyhow::{Context, Result};

// Embedded so `clankers import-pytorch` works from a plain `cargo install`
// with no repo checkout; written to a temp file at run time for python3.
const EXPORT_SCRIPT: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/assets/export_pytorch_to_onnx.py"
));

pub fn execute(model: &str, checkpoint: &str, output: &str, opset: Option<u32>) -> Result<()> {
    let script = std::env::temp_dir().join(format!(
        "clankers_export_pytorch_to_onnx_{}.py",
        std::process::id()
    ));
    fs::write(&script, EXPORT_SCRIPT)
        .with_context(|| format!("write export script to {}", script.display()))?;

    let mut cmd = Command::new("python3");
    cmd.arg(&script)
        .arg("--model")
        .arg(model)
        .arg("--checkpoint")
        .arg(checkpoint)
        .arg("--output")
        .arg(output);

    if let Some(opset) = opset {
        cmd.arg("--opset").arg(opset.to_string());
    }

    let status = cmd
        .status()
        .context("failed to run python3 (is Python 3 installed and on PATH?)");
    let _ = fs::remove_file(&script);

    if !status?.success() {
        anyhow::bail!("PyTorch export failed");
    }

    println!("clankeRS: exported ONNX model to {output}");
    Ok(())
}