Clean and Lint Pass (#99)

* ruct lint pass * a bit more * new safetensors messe sup conversion
ferritin-bio · Jan 3, 2025 · 68190a2 · 68190a2
1 parent b45f47a
commit 68190a2
Show file tree

Hide file tree

Showing 12 changed files with 42 additions and 49 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -9,7 +9,6 @@ members = [
     "ferritin-plms",
     "ferritin-pymol",
     "ferritin-test-data",
-    # "ferritin-tui",
 ]
 resolver = "2"
 
@@ -29,7 +28,7 @@ candle-hf-hub = "0.3.3"
 candle-metal-kernels = "0.8.1"
 candle-nn = "0.8.1"
 candle-transformers = "0.8.1"
-itertools = "0.13.0"
+itertools = "0.14.0"
 once_cell = "1.20.2"
 pdbtbx = "0.12.0"
 rand = "0.8.5"

diff --git a/ferritin-core/src/featurize/structure_features.rs b/ferritin-core/src/featurize/structure_features.rs
@@ -1,7 +1,7 @@
 //!  Protein->Tensor utiilities useful for Machine Learning
-use super::utilities::{aa1to_int, aa3to1, get_nearest_neighbours, int_to_aa1, AAAtom};
+use super::utilities::{aa1to_int, aa3to1, get_nearest_neighbours, AAAtom};
 use crate::AtomCollection;
-use candle_core::{DType, Device, Error as CandleError, IndexOp, Result, Tensor, D};
+use candle_core::{DType, Device, IndexOp, Result, Tensor, D};
 use itertools::MultiUnzip;
 use pdbtbx::Element;
 use strum::IntoEnumIterator;

diff --git a/ferritin-core/src/featurize/utilities.rs b/ferritin-core/src/featurize/utilities.rs
@@ -1,4 +1,4 @@
-use candle_core::{DType, Device, IndexOp, Result, Tensor, D};
+use candle_core::{IndexOp, Result, Tensor, D};
 use strum::{Display, EnumIter, EnumString};
 
 #[rustfmt::skip]
@@ -329,11 +329,7 @@ pub fn get_nearest_neighbours(
 #[cfg(test)]
 mod tests {
     use super::*;
-    // use crate::ligandmpnn::proteinfeatures::LMPNNFeatures;
-    use crate::AtomCollection;
-    use ferritin_test_data::TestFile;
     use pdbtbx;
-    use pdbtbx::Element;
 
     #[test]
     fn test_residue_codes() {

diff --git a/ferritin-examples/Cargo.toml b/ferritin-examples/Cargo.toml
@@ -21,13 +21,10 @@ ferritin-bevy = { path = "../ferritin-bevy" }
 ferritin-core = { path = "../ferritin-core" }
 ferritin-plms = { path = "../ferritin-plms" }
 ferritin-onnx-models = { path = "../ferritin-onnx-models" }
-ndarray-safetensors = "0.2"
 ndarray = "0.16"
 ort = "=2.0.0-rc.9"
 pdbtbx.workspace = true
-safetensors.workspace = true
 serde_json.workspace = true
-# wonnx = "0.5.1"
 
 
 [target.'cfg(target_os = "macos")'.features]

diff --git a/ferritin-examples/examples/esm2-onnx/main.rs b/ferritin-examples/examples/esm2-onnx/main.rs
@@ -1,12 +1,6 @@
-use anyhow::{Error as E, Result};
+use anyhow::Result;
 use clap::Parser;
 use ferritin_onnx_models::{ESM2Models, ESM2};
-use ndarray::Array2;
-use ort::{
-    execution_providers::CUDAExecutionProvider,
-    session::{builder::GraphOptimizationLevel, Session},
-};
-use std::env;
 
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]

diff --git a/ferritin-examples/examples/ligandmpnn-onnx/main.rs b/ferritin-examples/examples/ligandmpnn-onnx/main.rs
@@ -8,8 +8,8 @@ fn main() -> Result<()> {
     let (protfile, _handle) = TestFile::protein_01().create_temp()?;
     let (pdb, _) = pdbtbx::open(protfile).expect("PDB/CIF");
     let ac = AtomCollection::from(&pdb);
-    let model = LigandMPNN::new().unwrap();
-    let logits = model.run_model(ac, 10, 0.1).unwrap();
+    let model = LigandMPNN::new()?;
+    let logits = model.run_model(ac, 10, 0.1)?;
     println!("{:?}", logits);
 
     Ok(())

diff --git a/ferritin-onnx-models/src/models/esm2/mod.rs b/ferritin-onnx-models/src/models/esm2/mod.rs
@@ -13,15 +13,14 @@ use anyhow::{anyhow, Result};
 use candle_core::{Tensor, D};
 use candle_hf_hub::api::sync::Api;
 use candle_nn::ops;
-use ndarray::s;
 use ndarray::Array2;
 use ort::{
     execution_providers::CUDAExecutionProvider,
     session::{
-        builder::{GraphOptimizationLevel, SessionBuilder},
-        output, Session,
-    },
-    value::Value,
+        builder::{GraphOptimizationLevel, SessionBuilder}
+        , Session,
+    }
+    ,
 };
 use serde::{Deserialize, Serialize};
 use std::path::PathBuf;
@@ -60,7 +59,7 @@ impl ESM2 {
         })
     }
     pub fn load_model_path(model: ESM2Models) -> Result<PathBuf> {
-        let api = Api::new().unwrap();
+        let api = Api::new()?;
         let repo_id = match model {
             ESM2Models::ESM2_T6_8M => "zcpbx/esm2-t6-8m-UR50D-onnx",
             ESM2Models::ESM2_T12_35M => "zcpbx/esm2-t12-35M-UR50D-onnx",
@@ -69,7 +68,7 @@ impl ESM2 {
         }
         .to_string();
 
-        let model_path = api.model(repo_id).get("model.onnx").unwrap();
+        let model_path = api.model(repo_id).get("model.onnx")?;
         Ok(model_path)
     }
     pub fn load_tokenizer() -> Result<Tokenizer> {

diff --git a/ferritin-plms/Cargo.toml b/ferritin-plms/Cargo.toml
@@ -21,7 +21,7 @@ candle-core.workspace = true
 candle-metal-kernels = { workspace = true, optional = true }
 candle-nn.workspace = true
 candle-transformers.workspace = true
-clap = "4.5.23"
+clap = { version = "4.5", features = ["derive"] }
 ferritin-core = { path = "../ferritin-core" }
 ferritin-test-data = { path = "../ferritin-test-data" }
 itertools.workspace = true
@@ -41,4 +41,4 @@ candle-metal-kernels.workspace = true
 [dev-dependencies]
 assert_cmd = "2.0.16"
 ferritin-test-data = { path = "../ferritin-test-data" }
-tempfile = "3.14.0"
+tempfile = "3.15.0"
diff --git a/ferritin-plms/src/esm/layers/blocks.rs b/ferritin-plms/src/esm/layers/blocks.rs
@@ -1,6 +1,6 @@
 use super::attention::MultiHeadAttention;
 use super::geom_attention::GeometricReasoningOriginalImpl;
-use crate::esm::models::esmc::{ESMCConfig, Ffn_Type};
+use crate::esm::models::esmc::{ESMCConfig, FfnType};
 // use crate::esm::utils::structure::affine3d::Affine3D;
 use candle_core::{Module, Result, Tensor, D};
 use candle_nn::{self as nn, VarBuilder};
@@ -155,8 +155,8 @@ impl UnifiedTransformerBlock {
         let geom_attn = None;
 
         let ffn = match ffn_type {
-            Ffn_Type::SWIGLU => SwiGLU::load(vb.pp("ffn"), config)?,
-            _ => unimplemented!(), // Ffn_Type::GLU => unimplemented!(),
+            FfnType::SWIGLU => SwiGLU::load(vb.pp("ffn"), config)?,
+            _ => unimplemented!(), // FfnType::GLU => unimplemented!(),
         };
 
         Ok(Self {

diff --git a/ferritin-plms/src/esm/models/esmc.rs b/ferritin-plms/src/esm/models/esmc.rs
@@ -38,7 +38,7 @@ impl ESMTokenizer {
 }
 
 #[derive(Clone, Copy)]
-pub enum Ffn_Type {
+pub enum FfnType {
     SWIGLU,
     GLU,
 }
@@ -49,7 +49,7 @@ pub struct ESMCConfig {
     pub n_heads: usize,
     pub n_layers: usize,
     pub v_head_transformer: Option<usize>,
-    pub ffn_type: Ffn_Type,
+    pub ffn_type: FfnType,
     pub tokenizer: ESMTokenizer,
     // oringal above.
     pub use_plain_attn: bool,
@@ -80,7 +80,7 @@ impl ESMCConfig {
             n_heads: 15,
             n_layers: 30,
             v_head_transformer: None,
-            ffn_type: Ffn_Type::SWIGLU,
+            ffn_type: FfnType::SWIGLU,
             tokenizer: ESMTokenizer::Esm3OpenSmall,
             use_plain_attn: true,
             n_layers_geom: 1,

diff --git a/ferritin-test-data/Cargo.toml b/ferritin-test-data/Cargo.toml
@@ -7,4 +7,4 @@ license.workspace = true
 description.workspace = true
 
 [dependencies]
-tempfile = "3.14.0"
+tempfile = "3.15.0"