Skip to content

Commit

Permalink
add flag
Browse files Browse the repository at this point in the history
  • Loading branch information
tgolsson committed Nov 26, 2024
1 parent 815e02a commit d2d03d7
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 2 deletions.
5 changes: 3 additions & 2 deletions cli/src/dump.rs
Original file line number Diff line number Diff line change
Expand Up @@ -175,13 +175,14 @@ pub fn handle(
}

let compress_submodels = sub_matches.is_present("compress-submodels");
let deterministic = sub_matches.is_present("nnef-deterministic");
if let Some(path) = sub_matches.value_of("nnef") {
let nnef = super::nnef(matches);
if let Some(mut typed) = model.downcast_ref::<TypedModel>().cloned() {
rename_outputs(&mut typed, sub_matches)?;
let file = fs::File::create(path)?;
let encoder = flate2::write::GzEncoder::new(file, flate2::Compression::default());
nnef.write_to_tar_with_config(&typed, encoder, compress_submodels, false)
nnef.write_to_tar_with_config(&typed, encoder, compress_submodels, deterministic)
.context("Writing model to tgz")?;
} else {
bail!("Only typed model can be dumped")
Expand All @@ -193,7 +194,7 @@ pub fn handle(
if let Some(mut typed) = model.downcast_ref::<TypedModel>().cloned() {
rename_outputs(&mut typed, sub_matches)?;
let file = fs::File::create(path)?;
nnef.write_to_tar_with_config(&typed, file, compress_submodels, false)
nnef.write_to_tar_with_config(&typed, file, compress_submodels, deterministic)
.context("Writing model to tar")?;
} else {
bail!("Only typed model can be dumped")
Expand Down
5 changes: 5 additions & 0 deletions cli/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -383,6 +383,11 @@ fn dump_subcommand<'a>() -> clap::Command<'a> {
Arg::new("compress-submodels")
.long("compress-submodels")
.help("Compress submodels if any (as a .tgz file)"),
)
.arg(
Arg::new("nnef-deterministic")
.long("nnef-deterministic")
.help("If provided, will try to make output .nnef.tar files deterministic"),
)
.arg(
Arg::new("nnef-graph")
Expand Down

0 comments on commit d2d03d7

Please sign in to comment.