From 92fbf3b396085bab114d6d4c25d096de3570ff3b Mon Sep 17 00:00:00 2001
From: zachcp <zachcp@users.noreply.github.com>
Date: Thu, 2 Jan 2025 18:16:38 -0500
Subject: [PATCH] 20250102 amplify (#97)

* update main.rs

* export ESM2 output

* fix log_softmax --> softmax
---
 ferritin-examples/examples/esm2-onnx/main.rs | 3 ++-
 ferritin-onnx-models/src/models/esm2/mod.rs  | 5 ++++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/ferritin-examples/examples/esm2-onnx/main.rs b/ferritin-examples/examples/esm2-onnx/main.rs
index fdb29a3..7d264a1 100644
--- a/ferritin-examples/examples/esm2-onnx/main.rs
+++ b/ferritin-examples/examples/esm2-onnx/main.rs
@@ -37,9 +37,10 @@ fn main() -> Result<()> {
     let esm2 = ESM2::new(esm_model)?;
     let protein = args.protein_string.as_ref().unwrap().as_str();
     let logits = esm2.run_model(protein)?;
+    println!("Outputs: {:?}", logits);
 
     let normed = esm2.extract_logits(&logits)?;
-    println!("{:?}", normed);
+    // println!("Normalized: {:?}", normed);
 
     Ok(())
 }
diff --git a/ferritin-onnx-models/src/models/esm2/mod.rs b/ferritin-onnx-models/src/models/esm2/mod.rs
index 6040175..c109732 100644
--- a/ferritin-onnx-models/src/models/esm2/mod.rs
+++ b/ferritin-onnx-models/src/models/esm2/mod.rs
@@ -101,14 +101,17 @@ impl ESM2 {
         )?;
         let outputs =
             model.run(ort::inputs!["input_ids" => tokens_array,"attention_mask" => mask_array]?)?;
+
         let logits = outputs["logits"].try_extract_tensor::<f32>()?.to_owned();
+
         Ok(ndarray_to_tensor_f32(logits)?)
     }
 
     // Softmax and simplify
     pub fn extract_logits(&self, tensor: &Tensor) -> Result<Vec<LogitPosition>> {
-        let tensor = ops::log_softmax(tensor, D::Minus1)?;
+        let tensor = ops::softmax(tensor, D::Minus1)?;
         let data = tensor.to_vec3::<f32>()?;
+        println!("Data: {:?}", data);
         let shape = tensor.dims();
         let mut logit_positions = Vec::new();
         for seq_pos in 0..shape[1] {