add support for printing input/output to NNs (#46)

Found this sitting on my local disk; guess it was useful at some point.

crates/cervo-core/src/batcher.rs

@@ -117,12 +117,7 @@ impl Batcher {
             }
         }
 
-        Ok(self
-            .scratch
-            .ids
-            .drain(..)
-            .zip(outputs.into_iter())
-            .collect::<_>())
+        Ok(self.scratch.ids.drain(..).zip(outputs).collect::<_>())
     }
 
     /// Check if there is any data to run on here.

crates/cervo-runtime/src/runtime.rs

@@ -129,7 +129,9 @@ impl Runtime {
         let mut unselected_jobs = Vec::new();
 
         while let Some(ticket) = self.queue.pop() {
-            let Some(model) = self.models.iter().find(|m| m.id == ticket.1) else {continue};
+            let Some(model) = self.models.iter().find(|m| m.id == ticket.1) else {
+                continue;
+            };
 
             if model.needs_to_execute()
                 && (selected_jobs.is_empty() || model.can_run_in_time(available_cpu_time))

crates/cervo/src/commands/run.rs

@@ -1,6 +1,6 @@
 use anyhow::{bail, Result};
 use cervo_asset::AssetData;
-use cervo_core::prelude::{Inferer, InfererExt, State};
+use cervo_core::prelude::{Inferer, InfererExt, Response, State};
 use clap::Parser;
 
 use std::{collections::HashMap, fs::File, path::PathBuf, time::Instant};
@@ -19,6 +19,12 @@ pub(crate) struct Args {
     /// An epsilon key to randomize noise.
     #[clap(short, long)]
     with_epsilon: Option<String>,
+
+    #[clap(long)]
+    print_output: bool,
+
+    #[clap(long)]
+    print_input: bool,
 }
 
 fn build_inputs_from_desc(count: u64, inputs: &[(String, Vec<usize>)]) -> HashMap<u64, State<'_>> {
@@ -39,6 +45,28 @@ fn build_inputs_from_desc(count: u64, inputs: &[(String, Vec<usize>)]) -> HashMa
         .collect()
 }
 
+fn indent_by(target: String, prefix_len: usize) -> String {
+    let prefix = " ".repeat(prefix_len);
+
+    target
+        .lines()
+        .map(|line| format!("{}{}", prefix, line))
+        .collect::<Vec<_>>()
+        .join("\n")
+}
+
+fn print_input(obs: &HashMap<u64, State<'_>>) {
+    let formatted = format!("{:#?}", obs);
+    let indented = indent_by(formatted, 4);
+    println!("Inputs:\n{}", indented);
+}
+
+fn print_output(obs: &HashMap<u64, Response<'_>>) {
+    let formatted = format!("{:#?}", obs);
+    let indented = indent_by(formatted, 4);
+    println!("Outputs:\n{}", indented);
+}
+
 pub(super) fn run(config: Args) -> Result<()> {
     let mut reader = File::open(&config.file)?;
     let inferer = if cervo_nnef::is_nnef_tar(&config.file) {
@@ -65,19 +93,36 @@ pub(super) fn run(config: Args) -> Result<()> {
             .collect::<Vec<_>>();
 
         let observations = build_inputs_from_desc(config.batch_size as u64, &shapes);
+
+        if config.print_input {
+            print_input(&observations);
+        }
+
         inferer.infer_batch(observations.clone())?;
 
         let start = Instant::now();
-        inferer.infer_batch(observations)?;
-        start.elapsed()
+        let res = inferer.infer_batch(observations)?;
+
+        let dur = start.elapsed();
+        if config.print_output {
+            print_output(&res);
+        }
+
+        dur
     } else {
         let shapes = inferer.input_shapes().to_vec();
         let observations = build_inputs_from_desc(config.batch_size as u64, &shapes);
         inferer.infer_batch(observations.clone())?;
 
         let start = Instant::now();
-        inferer.infer_batch(observations)?;
-        start.elapsed()
+        let res = inferer.infer_batch(observations)?;
+
+        let dur = start.elapsed();
+        if config.print_output {
+            print_output(&res);
+        }
+
+        dur
     };
 
     println!(