Repository: bioinf-mcb/Metagenomic-DeepFRI (language: python)

Issue:
`ruff` linting
Docs: https://github.com/charliermarsh/ruff


Discussion / hints:

Fix patch:
diff --git a/.gitmodules b/.gitmodules
deleted file mode 100644
index 10b7ae1..0000000
--- a/.gitmodules
+++ /dev/null
@@ -1,3 +0,0 @@
-[submodule "meta_deepFRI/DeepFRI"]
-	path = meta_deepFRI/DeepFRI
-	url = https://github.com/bioinf-mcb/DeepFRI
diff --git a/CMakeLists.txt b/CMakeLists.txt
index e1fb2cd..09587b4 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -4,10 +4,12 @@ project(GenomicLibrary)
 set(CMAKE_CXX_STANDARD 17)
 
 add_library(AtomDistanceIO SHARED
-        meta_deepFRI/CPP_lib/library_definition.cpp
-        meta_deepFRI/CPP_lib/atoms_file_io.h
-        meta_deepFRI/CPP_lib/python_utils.h
-        meta_deepFRI/CPP_lib/load_contact_maps.h)
+        mDeepFRI/CPP_lib/library_definition.cpp
+        mDeepFRI/CPP_lib/atoms_file_io.h
+        mDeepFRI/CPP_lib/atoms_file_io.cpp
+        mDeepFRI/CPP_lib/python_utils.h
+        mDeepFRI/CPP_lib/load_contact_maps.h
+        mDeepFRI/CPP_lib/load_contact_maps.cpp)
 
 
 target_include_directories(AtomDistanceIO PUBLIC ${PY_INCLUDE_PATH})
diff --git a/README.md b/README.md
index 7842046..a19de08 100644
--- a/README.md
+++ b/README.md
@@ -5,9 +5,9 @@ Do you have **thousands of protein sequences** with **unknown structures**, but
 molecular function, biological process, cellular component and enzyme commission **predicted by DeepFRI Graph Convolutional Network?**
 
 This is the right project for this task! Pipeline in a nutshell:
-1. Search for similar target protein sequences using MMseqs2
-2. Align target protein contact map to fit your query protein with unknown structure
-3. Run predictions on query sequence combined with aligned target contact map or sequence alone if no alignment was found
+1. Search for similar target protein sequences using MMseqs2.
+2. Align target protein contact map to fit your query protein with unknown structure.
+3. Run predictions on query sequence combined with aligned target contact map or sequence alone if no alignment was found.
 
 ### Built With
 
@@ -22,7 +22,7 @@ This is the right project for this task! Pipeline in a nutshell:
 
 1. Clone repo locally
 ```{code-block} bash
-git clone https://github.com/bioinf-mcb/Metagenomic-DeepFRI --recursive
+git clone https://github.com/bioinf-mcb/Metagenomic-DeepFRI
 cd Metagenomic-DeepFRI
 ```
 2. Setup conda environment
@@ -30,40 +30,31 @@ cd Metagenomic-DeepFRI
 conda env create --name deepfri --file environment.yml
 conda activate deepfri
 ```
-3. Install `meta-DeepFRI`
+3. Install `mDeepFRI`
 ```{code-block} bash
 pip install .
 ```
 4. Verify installation
 ```{code-block} bash
 pytest
-deepfri --help
+mDeepFRI --help
 ```
 
-## 2. Retrieve `DeepFRI` model weights
-### CPU weights
-```bash
-wget https://users.flatironinstitute.org/~renfrew/DeepFRI_data/newest_trained_models.tar.gz
-tar -xf newest_trained_models.tar.gz
+# Usage
+## 1. Download models
+Run command:
 ```
-
-### GPU weights
-```bash
-wget https://users.flatironinstitute.org/~renfrew/DeepFRI_data/trained_models.tar.gz
-tar -xf trained_models.tar.gz
+mDeepFRI get-models --output path/to/weights/folder
 ```
 
-**Attention:** Using DeepFRI with GPU requires extra installation steps, described in section 4. GPU Setup of [tensorfow installation guide](https://www.tensorflow.org/install/pip).
-
-# Usage
-## 1. Prepare database
+## 2. Prepare database
 
 1. Upload structure (`.pdb` or `.mmcif`) files to a folder in your system.
 2. Run command:
-   ```
-   deepfri_db_build --input path/to/folder/with/strucures --output path/to/database
-   ```
-**Tip:** building a database from AF20Swissprot (~550k predicted structures) on 32 CPU cores took ~30 min.
+```
+mDeepFRI build-db --input path/to/folder/with/strucures --output path/to/database -t threads
+```
+**Tip:** building a database from AF2Swissprot (~550k predicted structures) on 32 CPU cores took ~30 min.
 
 Use parameter `-max_len` to define maximal length of the protein. Due to initial DeepFRI training set default value is set to `1000`.
 
@@ -82,12 +73,13 @@ Accepted formats are: `.pdb`, `.cif`, `.ent` both raw and compressed with `.gz`
 Protein ID is used as a filename. A new protein whose ID already exists in the database will be skipped.
 Use `--overwrite` flag to overwrite existing sequences and atoms positions.
 
-## 2. Predict protein function
-   ```
-   deepfri -i /path/to/protein/sequences -db /path/to/database/folder/from/previous/step -w /path/to/deepfri/weights/folder -o /output_path
-   ```
+## 3. Predict protein function
+```
+mDeepFRI predict-function -i /path/to/protein/sequences -d /path/to/database/folder/from/previous/step -w /path/to/deepfri/weights/folder -o /output_path
+```
 **Attention:** Single instance of DeepFRI on GPU requires 10GB VRAM.
-Other available parameters can be found upon command `deepfri --help`.
+
+Other available parameters can be found upon command `mDeepFRI --help`.
 
 ## Results
 Finished folder will contain:
@@ -95,7 +87,7 @@ Finished folder will contain:
 2. `mmseqs2_search_results.m8`
 3. `alignments.json` - results of alignment search implemented in `utils.search_alignments.py`
 4. `metadata*` - files with some useful info
-5. `results*` - multiple files from DeepFRI. Organized by model type ['GCN' / 'CNN'] and its mode ['mf', 'bp', 'cc', 'ec'] for the total of 8 files.
+5. `results*` - multiple files from DeepFRI. Organized by model type (`GCN` or `CNN`) and its mode (`mf`, `bp`, `cc`, `ec`) for the total of 8 files.
 Sometimes results from one model can be missing which means that all query proteins sequences were aligned correctly or none of them were aligned.
    ```
    mf = molecular_function
@@ -103,16 +95,20 @@ Sometimes results from one model can be missing which means that all query prote
    cc = cellular_component
    ec = enzyme_commission
    ```
+
+## GPU / CPU utilization
+If CUDA is installed on your machine, `metaDeepFRI` will automatically use it for prediction, no additional installations are needed. If not, the model will use CPUs. If argument `threads` is provided, the prediction will run on multiple CPU cores.
+
 ## Citations
 If you use this software please cite:
-    - Gligorijević et al. "Structure-based protein function prediction using graph convolutional networks" Nat. Comms. (2021). https://doi.org/10.1038/s41467-021-23303-9
-    - Steinegger & Söding "MMseqs2 enables sensitive protein sequence searching for the analysis of massive data sets" Nat. Biotechnol. https://doi.org/10.1038/nbt.3988
-    - Maranga et al. "Comprehensive Functional Annotation of Metagenomes and Microbial Genomes Using a Deep Learning-Based Method" mSystems (2023) https://doi.org/10.1128/msystems.01178-22
-    - Daily "Parasail: SIMD C library for global, semi-global, and local pairwise sequence alignments" BMC Bioinform. (2016) https://doi.org/10.1186/s12859-016-0930-z
+- Gligorijević et al. "Structure-based protein function prediction using graph convolutional networks" Nat. Comms. (2021). https://doi.org/10.1038/s41467-021-23303-9
+- Steinegger & Söding "MMseqs2 enables sensitive protein sequence searching for the analysis of massive data sets" Nat. Biotechnol. (2017) https://doi.org/10.1038/nbt.3988
+- Maranga et al. "Comprehensive Functional Annotation of Metagenomes and Microbial Genomes Using a Deep Learning-Based Method" mSystems (2023) https://doi.org/10.1128/msystems.01178-22
+- Daily "Parasail: SIMD C library for global, semi-global, and local pairwise sequence alignments" BMC Bioinform. (2016) https://doi.org/10.1186/s12859-016-0930-z
 
 ## Contributing
 
-If you have a suggestion that would make this project better, email me or fork the repo and create a pull request.
+If you have a suggestion that would make this project better, please send an e-mail or fork the repo and create a pull request.
 
 ### Contact
 
diff --git a/environment.yml b/environment.yml
index 3152f86..c093434 100644
--- a/environment.yml
+++ b/environment.yml
@@ -3,16 +3,16 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - tensorflow=2.8.0
-  - biopython=1.79
   - parasail-python=1.3.3
-  - boost=1.80.0
-  - boost-cpp=1.80.0
-  - numpy=1.24.2
-  - pandas==1.3.5
+  - click
+  - boost>=1.80.0
+  - boost-cpp>=1.80.0
+  - numpy>=1.19
+  - pandas>=1.2
   - mmseqs2=14.7e284
-  - pip=23.0.1
-  - pysam=0.21.0
-  - pytest=7.2.2
-  - python=3.10
-  - cmake=3.22.1
+  - pip
+  - pysam>=0.16.0
+  - pytest
+  - python
+  - cmake>=3.19
+  - onnxruntime>=1.10
diff --git a/meta_deepFRI/CPP_lib/README.md b/mDeepFRI/CPP_lib/README.md
similarity index 100%
rename from meta_deepFRI/CPP_lib/README.md
rename to mDeepFRI/CPP_lib/README.md
diff --git a/meta_deepFRI/CPP_lib/__init__.py b/mDeepFRI/CPP_lib/__init__.py
similarity index 100%
rename from meta_deepFRI/CPP_lib/__init__.py
rename to mDeepFRI/CPP_lib/__init__.py
diff --git a/mDeepFRI/CPP_lib/atoms_file_io.cpp b/mDeepFRI/CPP_lib/atoms_file_io.cpp
new file mode 100644
index 0000000..1750a67
--- /dev/null
+++ b/mDeepFRI/CPP_lib/atoms_file_io.cpp
@@ -0,0 +1,45 @@
+//
+// Created by soliareofastora on 24.12.2021.
+//
+#include <boost/python.hpp>
+#include <boost/python/numpy.hpp>
+#include <fstream>
+
+namespace py = boost::python;
+namespace np = py::numpy;
+
+static void SaveAtomsFile(const np::ndarray &position_array,
+                          const np::ndarray &groups_array,
+                          const std::string &save_path) {
+  int chain_length = (int)groups_array.shape(0);
+  int *group_indexes = reinterpret_cast<int *>(groups_array.get_data());
+  int atom_count = group_indexes[chain_length - 1];
+  float *atom_positions = reinterpret_cast<float *>(position_array.get_data());
+
+  std::ofstream writer(save_path, std::ios::out | std::ios::binary);
+  writer.write(reinterpret_cast<const char *>(&chain_length), 4);
+  writer.write(reinterpret_cast<const char *>(group_indexes), 4 * chain_length);
+  writer.write(reinterpret_cast<const char *>(atom_positions),
+               4 * atom_count * 3);
+  writer.close();
+};
+
+static std::tuple<int, int *, float *>
+LoadAtomsFile(const std::string &file_path) {
+  std::ifstream reader(file_path, std::ios::in | std::ios::binary);
+
+  int chain_length;
+  reader.read(reinterpret_cast<char *>(&chain_length), 4);
+
+  int *group_indexes = new int[chain_length];
+  reader.read(reinterpret_cast<char *>(group_indexes), 4 * chain_length);
+
+  chain_length--;
+  int atom_count = group_indexes[chain_length];
+
+  float *atoms_positions = new float[atom_count * 3];
+  reader.read(reinterpret_cast<char *>(atoms_positions), 4 * atom_count * 3);
+  reader.close();
+
+  return std::make_tuple(chain_length, group_indexes, atoms_positions);
+};
diff --git a/mDeepFRI/CPP_lib/atoms_file_io.h b/mDeepFRI/CPP_lib/atoms_file_io.h
new file mode 100644
index 0000000..fdbfb519
--- /dev/null
+++ b/mDeepFRI/CPP_lib/atoms_file_io.h
@@ -0,0 +1,18 @@
+#ifndef ATOMS_FILE_IO
+#define ATOMS_FILE_IO
+
+#include <boost/python.hpp>
+#include <boost/python/numpy.hpp>
+#include <fstream>
+
+namespace py = boost::python;
+namespace np = py::numpy;
+
+static void SaveAtomsFile(const np::ndarray &position_array,
+                          const np::ndarray &groups_array,
+                          const std::string &save_path);
+
+static std::tuple<int, int *, float *>
+LoadAtomsFile(const std::string &file_path);
+
+#endif
diff --git a/meta_deepFRI/CPP_lib/library_definition.cpp b/mDeepFRI/CPP_lib/library_definition.cpp
similarity index 90%
rename from meta_deepFRI/CPP_lib/library_definition.cpp
rename to mDeepFRI/CPP_lib/library_definition.cpp
index 031dc65..536a2cf 100644
--- a/meta_deepFRI/CPP_lib/library_definition.cpp
+++ b/mDeepFRI/CPP_lib/library_definition.cpp
@@ -11,7 +11,7 @@
 namespace py = boost::python;
 namespace np = py::numpy;
 
-BOOST_PYTHON_MODULE (libAtomDistanceIO) {
+BOOST_PYTHON_MODULE(libAtomDistanceIO) {
   py::def("initialize", Initialize);
   py::def("save_atoms", SaveAtomsFile);
   py::def("load_contact_map", LoadContactMap);
diff --git a/meta_deepFRI/CPP_lib/load_contact_maps.h b/mDeepFRI/CPP_lib/load_contact_maps.cpp
similarity index 50%
rename from meta_deepFRI/CPP_lib/load_contact_maps.h
rename to mDeepFRI/CPP_lib/load_contact_maps.cpp
index 379f648..f33e178 100644
--- a/meta_deepFRI/CPP_lib/load_contact_maps.h
+++ b/mDeepFRI/CPP_lib/load_contact_maps.cpp
@@ -1,10 +1,3 @@
-//
-// Created by soliareofastora on 04.01.2022.
-//
-
-#ifndef LOAD_CONTACT_MAPS
-#define LOAD_CONTACT_MAPS
-
 #include <cmath>
 #include <iostream>
 #include <queue>
@@ -12,27 +5,44 @@
 #include <sys/stat.h>
 
 #include "atoms_file_io.h"
+#include "load_contact_maps.h"
 #include "python_utils.h"
 
-static float Distance(float* array, int i, int j) {
-  return sqrtf(powf(array[i * 3] - array[j * 3], 2) + powf(array[i * 3 + 1] - array[j * 3 + 1], 2) + powf(array[i * 3 + 2] - array[j * 3 + 2], 2));
+static float Distance(std::unique_ptr<float[]> &array, int i, int j) {
+  int xIndex = i * 3;
+  int yIndex = xIndex + 1;
+  int zIndex = xIndex + 2;
+
+  float xDist = array.get()[xIndex] - array.get()[xIndex + 3 * j];
+  float yDist = array.get()[yIndex] - array.get()[yIndex + 3 * j];
+  float zDist = array.get()[zIndex] - array.get()[zIndex + 3 * j];
+
+  float xSquared = xDist * xDist;
+  float ySquared = yDist * yDist;
+  float zSquared = zDist * zDist;
+
+  float distance = std::hypot(xSquared, ySquared, zSquared);
+
+  return distance;
 }
 
-inline bool exists (const std::string& name) {
+inline bool exists(const std::string &name) {
   struct stat buffer;
-  return (stat (name.c_str(), &buffer) == 0);
+  return (stat(name.c_str(), &buffer) == 0);
 }
 
-
-static std::pair<bool*, int> LoadDenseContactMap(const std::string& file_path, const float angstrom_contact_threshold){
+static std::pair<bool *, int>
+LoadDenseContactMap(const std::string &file_path,
+                    const float angstrom_contact_threshold) {
   int chain_length;
-  int* group_indexes;
-  float* atoms_positions;
-  std::tie(chain_length, group_indexes, atoms_positions) = LoadAtomsFile(file_path);
+  int *group_indexes;
+  float *atoms_positions;
+  std::tie(chain_length, group_indexes, atoms_positions) =
+      LoadAtomsFile(file_path);
 
   // allocate output array
-  bool* const output_data = new bool[(int) pow(chain_length, 2)];
-  std::memset(output_data, 0, (int) pow(chain_length, 2));
+  bool *const output_data = new bool[(int)pow(chain_length, 2)];
+  std::memset(output_data, 0, (int)pow(chain_length, 2));
 
   // fill up output array with atom contacts
   for (int group_a = 0; group_a < chain_length; ++group_a) {
@@ -41,9 +51,12 @@ static std::pair<bool*, int> LoadDenseContactMap(const std::string& file_path, c
     for (int group_b = group_a + 1; group_b < chain_length; ++group_b) {
       bool group_connected = false;
 
-      for (int atom_a = group_indexes[group_a]; atom_a < group_indexes[group_a + 1]; ++atom_a) {
-        for (int atom_b = group_indexes[group_b]; atom_b < group_indexes[group_b + 1]; ++atom_b) {
-          if (Distance(atoms_positions, atom_a, atom_b) <= angstrom_contact_threshold) {
+      for (int atom_a = group_indexes[group_a];
+           atom_a < group_indexes[group_a + 1]; ++atom_a) {
+        for (int atom_b = group_indexes[group_b];
+             atom_b < group_indexes[group_b + 1]; ++atom_b) {
+          if (Distance(atoms_positions, atom_a, atom_b) <=
+              angstrom_contact_threshold) {
             group_connected = true;
             output_data[group_a * chain_length + group_b] = true;
             output_data[group_a + group_b * chain_length] = true;
@@ -62,11 +75,12 @@ static std::pair<bool*, int> LoadDenseContactMap(const std::string& file_path, c
   return std::make_pair(output_data, chain_length);
 }
 
-
-static std::vector<std::pair<int, int>>* LoadSparseContactMap(const std::string& file_path, const float angstrom_contact_threshold){
+static std::vector<std::pair<int, int>> *
+LoadSparseContactMap(const std::string &file_path,
+                     const float angstrom_contact_threshold) {
   int chain_length;
-  int* group_indexes;
-  float* atoms_positions;
+  int *group_indexes;
+  float *atoms_positions;
 
   // verbose error if file_path doesn't exist
   if (!exists(file_path)) {
@@ -74,19 +88,23 @@ static std::vector<std::pair<int, int>>* LoadSparseContactMap(const std::string&
     return nullptr;
   }
 
-  std::tie(chain_length, group_indexes, atoms_positions) = LoadAtomsFile(file_path);
+  std::tie(chain_length, group_indexes, atoms_positions) =
+      LoadAtomsFile(file_path);
 
   // fill up vector with sparse atom contacts
-  auto* sparse_contacts = new std::vector<std::pair<int, int>>();
+  auto *sparse_contacts = new std::vector<std::pair<int, int>>();
   sparse_contacts->reserve(chain_length * 10);
 
   for (int group_a = 0; group_a < chain_length; ++group_a) {
     for (int group_b = group_a + 1; group_b < chain_length; ++group_b) {
       bool group_connected = false;
 
-      for (int atom_a = group_indexes[group_a]; atom_a < group_indexes[group_a + 1]; ++atom_a) {
-        for (int atom_b = group_indexes[group_b]; atom_b < group_indexes[group_b + 1]; ++atom_b) {
-          if (Distance(atoms_positions, atom_a, atom_b) <= angstrom_contact_threshold) {
+      for (int atom_a = group_indexes[group_a];
+           atom_a < group_indexes[group_a + 1]; ++atom_a) {
+        for (int atom_b = group_indexes[group_b];
+             atom_b < group_indexes[group_b + 1]; ++atom_b) {
+          if (Distance(atoms_positions, atom_a, atom_b) <=
+              angstrom_contact_threshold) {
             group_connected = true;
             sparse_contacts->emplace_back(group_a, group_b);
             break;
@@ -104,17 +122,22 @@ static std::vector<std::pair<int, int>>* LoadSparseContactMap(const std::string&
   return sparse_contacts;
 }
 
-
-static np::ndarray LoadContactMap(const std::string& file_path, const float angstrom_contact_threshold) {
-  bool* contact_map;
+static np::ndarray LoadContactMap(const std::string &file_path,
+                                  const float angstrom_contact_threshold) {
+  bool *contact_map;
   int chain_length;
-  std::tie(contact_map, chain_length) = LoadDenseContactMap(file_path, angstrom_contact_threshold);
+  std::tie(contact_map, chain_length) =
+      LoadDenseContactMap(file_path, angstrom_contact_threshold);
   return CreateNumpyArray(contact_map, chain_length);
 }
 
-
-static np::ndarray LoadAlignedContactMap(const std::string& file_path, float angstrom_contact_threshold, const std::string& query_alignment, const std::string& target_alignment, const int generated_contacts) {
-  std::vector<std::pair<int, int>>* sparse_target_contacts = LoadSparseContactMap(file_path, angstrom_contact_threshold);
+static np::ndarray LoadAlignedContactMap(const std::string &file_path,
+                                         float angstrom_contact_threshold,
+                                         const std::string &query_alignment,
+                                         const std::string &target_alignment,
+                                         const int generated_contacts) {
+  std::vector<std::pair<int, int>> *sparse_target_contacts =
+      LoadSparseContactMap(file_path, angstrom_contact_threshold);
   std::vector<std::pair<int, int>> sparse_query_contacts;
   sparse_query_contacts.reserve(sparse_target_contacts->size());
 
@@ -141,19 +164,21 @@ static np::ndarray LoadAlignedContactMap(const std::string& file_path, float ang
   }
 
   for (int i = 0; i < sparse_target_contacts->size(); ++i) {
-    int contact_x = target_to_query_indexes[sparse_target_contacts->operator[](i).first];
+    int contact_x =
+        target_to_query_indexes[sparse_target_contacts->operator[](i).first];
     if (contact_x < 0) {
       continue;
     }
-    int contact_y = target_to_query_indexes[sparse_target_contacts->operator[](i).second];
+    int contact_y =
+        target_to_query_indexes[sparse_target_contacts->operator[](i).second];
     if (contact_y < 0) {
       continue;
     }
     sparse_query_contacts.emplace_back(contact_x, contact_y);
   }
 
-  bool* const output_data = new bool[(int) pow(query_index, 2)];
-  std::memset(output_data, 0, (int) pow(query_index, 2));
+  bool *const output_data = new bool[(int)pow(query_index, 2)];
+  std::memset(output_data, 0, (int)pow(query_index, 2));
 
   for (int i = 0; i < query_index; ++i) {
     output_data[i * query_index + i] = true;
@@ -172,5 +197,3 @@ static np::ndarray LoadAlignedContactMap(const std::string& file_path, float ang
   delete sparse_target_contacts;
   return CreateNumpyArray(output_data, query_index);
 }
-
-#endif
diff --git a/mDeepFRI/CPP_lib/load_contact_maps.h b/mDeepFRI/CPP_lib/load_contact_maps.h
new file mode 100644
index 0000000..b8e8c5c
--- /dev/null
+++ b/mDeepFRI/CPP_lib/load_contact_maps.h
@@ -0,0 +1,34 @@
+//
+// Created by soliareofastora on 04.01.2022.
+//
+
+#ifndef LOAD_CONTACT_MAPS
+#define LOAD_CONTACT_MAPS
+
+#include <cmath>
+#include <iostream>
+#include <queue>
+#include <string>
+#include <sys/stat.h>
+
+#include "atoms_file_io.h"
+#include "python_utils.h"
+
+static std::pair<bool *, int>
+LoadDenseContactMap(const std::string &file_path,
+                    const float angstrom_contact_threshold);
+
+static std::vector<std::pair<int, int>> *
+LoadSparseContactMap(const std::string &file_path,
+                     const float angstrom_contact_threshold);
+
+static np::ndarray LoadContactMap(const std::string &file_path,
+                                  const float angstrom_contact_threshold);
+
+static np::ndarray LoadAlignedContactMap(const std::string &file_path,
+                                         float angstrom_contact_threshold,
+                                         const std::string &query_alignment,
+                                         const std::string &target_alignment,
+                                         const int generated_contacts);
+
+#endif
diff --git a/meta_deepFRI/CPP_lib/python_utils.h b/mDeepFRI/CPP_lib/python_utils.h
similarity index 100%
rename from meta_deepFRI/CPP_lib/python_utils.h
rename to mDeepFRI/CPP_lib/python_utils.h
diff --git a/mDeepFRI/__init__.py b/mDeepFRI/__init__.py
new file mode 100644
index 0000000..12c7e7c
--- /dev/null
+++ b/mDeepFRI/__init__.py
@@ -0,0 +1,45 @@
+import os
+
+__version__ = "2.0.0"
+
+repo_link = "https://huggingface.co/crusher083/mDeepFRI/resolve/main"
+
+model_names = [
+    "DeepCNN-MERGED_bp.onnx",
+    "DeepCNN-MERGED_cc.onnx",
+    "DeepCNN-MERGED_mf.onnx",
+    "DeepCNN-MERGED_ec.onnx",
+    "DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_bp.onnx",
+    "DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_cc.onnx",
+    "DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_ec.onnx",
+    "DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_mf.onnx",
+]
+
+model_links = [
+    os.path.join(repo_link, model_name) for model_name in model_names
+]
+
+config_names = [
+    "DeepCNN-MERGED_bp_model_params.json",
+    "DeepCNN-MERGED_cc_model_params.json",
+    "DeepCNN-MERGED_mf_model_params.json",
+    "DeepCNN-MERGED_ec_model_params.json",
+    "DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_bp_model_params.json",
+    "DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_cc_model_params.json",
+    "DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_ec_model_params.json",
+    "DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_mf_model_params.json",
+    "model_config.json"
+]
+
+config_links = [
+    os.path.join(repo_link, config_name) for config_name in config_names
+]
+
+TARGET_MMSEQS_DB_NAME = "targetDB"
+SEQUENCES = "seq"
+ATOMS = "atom"
+SEQ_ATOMS_DATASET_PATH = "seq_atom_db"
+
+ALIGNMENTS = "alignments.json"
+MERGED_SEQUENCES = 'merged_sequences.faa'
+MMSEQS_SEARCH_RESULTS = 'mmseqs2_search_results.m8'
diff --git a/mDeepFRI/cli.py b/mDeepFRI/cli.py
new file mode 100644
index 0000000..1d409e3
--- /dev/null
+++ b/mDeepFRI/cli.py
@@ -0,0 +1,241 @@
+import logging
+from pathlib import Path
+
+import click
+
+from mDeepFRI import __version__
+from mDeepFRI.database import build_database
+from mDeepFRI.pipeline import metagenomic_deepfri
+from mDeepFRI.utils.utils import download_model_weights
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+@click.group()
+@click.option("--debug/--no-debug", default=False)
+@click.version_option(version=__version__)
+@click.pass_context
+def cli(ctx, debug):
+    """mDeepFRI"""
+
+    ctx.ensure_object(dict)
+    ctx.obj["debug"] = debug
+
+    if debug:
+        logger.setLevel(logging.DEBUG)
+        logger.debug("Debug mode is on.")
+        logging.getLogger("requests").setLevel(logging.DEBUG)
+    else:
+        logger.setLevel(logging.INFO)
+        logger.info("Debug mode is off.")
+        logging.getLogger("requests").setLevel(logging.INFO)
+
+
+@cli.command
+@click.option(
+    "-o",
+    "--output",
+    required=True,
+    type=click.Path(exists=False),
+    help="Path to folder where the database will be created.",
+)
+@click.pass_context
+def get_models(ctx, output):
+    """Download model weights for mDeepFRI."""
+
+    if ctx.obj["debug"] is True:
+        logger.setLevel(logging.DEBUG)
+        logging.getLogger("requests").setLevel(logging.DEBUG)
+        logging.getLogger("urllib3").setLevel(logging.DEBUG)
+    else:
+        logger.setLevel(logging.INFO)
+        logging.getLogger("requests").setLevel(logging.INFO)
+        logging.getLogger("urllib3").setLevel(logging.INFO)
+
+    output_path = Path(output)
+    output_path.mkdir(parents=True, exist_ok=True)
+    download_model_weights(output_path)
+
+
+# add subcommand for building database
+@cli.command()
+@click.option(
+    "-i",
+    "--input",
+    required=True,
+    multiple=True,
+    type=click.Path(exists=True),
+    help="Path to structure files or folders containing structure files.",
+)
+@click.option(
+    "-o",
+    "--output",
+    required=True,
+    type=click.Path(exists=False),
+    help="Path to folder where the database will be created.",
+)
+@click.option(
+    "-t",
+    "--threads",
+    default=1,
+    type=int,
+    help="Number of threads to use. Default is 1.",
+)
+@click.option(
+    "-m",
+    "--max-protein-length",
+    default=1_000,
+    type=int,
+    help=
+    "If protein is longer than this value, it will be truncated. Default is 1000 aa.",
+)
+@click.pass_context
+def build_db(ctx, input, output, threads, max_protein_length):
+    """Build a database for meta_deepFRI."""
+
+    if ctx.obj["debug"] is True:
+        logger.setLevel(logging.DEBUG)
+
+    input_seqs = [Path(seqs) for seqs in input]
+    output_path = Path(output)
+    build_database(input_paths=input_seqs,
+                   output_path=output_path,
+                   overwrite=True,
+                   threads=threads,
+                   max_protein_length=max_protein_length)
+
+
+@cli.command()
+@click.option(
+    "-i",
+    "--input",
+    required=True,
+    type=click.Path(exists=True),
+    help="Path to an input protein sequences (FASTA file, may be gzipped).",
+)
+@click.option(
+    "-d",
+    "--db-path",
+    required=True,
+    type=click.Path(exists=True),
+    help="Path to a structures database folder.",
+)
+@click.option(
+    "-w",
+    "--weights",
+    required=True,
+    type=click.Path(exists=True),
+    help="Path to a folder containing model weights.",
+)
+@click.option(
+    "-o",
+    "--output",
+    required=True,
+    type=click.Path(exists=False),
+    help="Path to output file.",
+)
+@click.option(
+    "-f",
+    "--output-format",
+    default="tsv",
+    type=click.Choice(["tsv", "csv", "json"]),
+    help="Output format. Default is TSV.",
+)
+@click.option(
+    "-p",
+    "--processing-modes",
+    default=["bp", "cc", "ec", "mf"],
+    type=click.Choice(["bp", "cc", "ec", "mf"]),
+    multiple=True,
+    help="Processing modes. Default is all "
+    "(biological process, cellular component, enzyme comission, molecular function).",
+)
+@click.option(
+    "-a",
+    "--angstrom-contact-thresh",
+    default=6,
+    type=float,
+    help="Angstrom contact threshold. Default is 6.",
+)
+@click.option(
+    "--generate-contacts",
+    default=2,
+    type=int,
+    help="Gap fill threshold during contact map alignment.",
+)
+@click.option(
+    "--mmseqs-min-bit-score",
+    default=None,
+    type=float,
+    help="Minimum bit score for MMseqs2 alignment.",
+)
+@click.option(
+    "--mmseqs-max-evalue",
+    default=0.001,
+    type=float,
+    help="Maximum e-value for MMseqs2 alignment.",
+)
+@click.option(
+    "--mmseqs-min-identity",
+    default=0.5,
+    type=float,
+    help="Minimum identity for MMseqs2 alignment.",
+)
+@click.option(
+    "--alignment-matrix",
+    default="blosum62",
+    type=click.Choice(
+        ["blosum62", "pam250", "blosum80", "blosum45", "pam120", "pam160"]),
+    help="Alignment matrix for contact map alignment.",
+)
+@click.option(
+    "--alignment-gap-open",
+    default=10,
+    type=int,
+    help="Gap open penalty for contact map alignment.",
+)
+@click.option(
+    "--alignment-gap-extend",
+    default=1,
+    type=int,
+    help="Gap extend penalty for contact map alignment.",
+)
+@click.option(
+    "--alignment-min-identity",
+    default=0.5,
+    type=float,
+    help="Minimum identity for contact map alignment.",
+)
+@click.option(
+    "-t",
+    "--threads",
+    default=1,
+    type=int,
+    help="Number of threads to use. Default is 1.",
+)
+@click.pass_context
+def predict_function(ctx, input, db_path, weights, output, output_format,
+                     processing_modes, angstrom_contact_thresh,
+                     generate_contacts, mmseqs_min_bit_score,
+                     mmseqs_max_evalue, mmseqs_min_identity, alignment_matrix,
+                     alignment_gap_open, alignment_gap_extend,
+                     alignment_min_identity, threads):
+    """Predict protein function from sequence."""
+    if ctx.obj["debug"] is True:
+        logger.setLevel(logging.DEBUG)
+
+    output_path = Path(output)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    metagenomic_deepfri(Path(input), Path(db_path), Path(weights), output_path,
+                        output_format, processing_modes,
+                        angstrom_contact_thresh, generate_contacts,
+                        mmseqs_min_bit_score, mmseqs_max_evalue,
+                        mmseqs_min_identity, alignment_matrix,
+                        alignment_gap_open, alignment_gap_extend,
+                        alignment_min_identity, threads)
+
+
+if __name__ == "__main__":
+    cli(prog_name="mDeepFRI")
diff --git a/scripts/build_database.py b/mDeepFRI/database.py
similarity index 66%
rename from scripts/build_database.py
rename to mDeepFRI/database.py
index 4b5aa54..bb22b8b 100644
--- a/scripts/build_database.py
+++ b/mDeepFRI/database.py
@@ -1,4 +1,3 @@
-import argparse
 import json
 # Create logger
 import logging
@@ -9,13 +8,12 @@
 
 import numpy as np
 
-from meta_deepFRI.config.names import SEQ_ATOMS_DATASET_PATH
-from meta_deepFRI.CPP_lib import \
-    libAtomDistanceIO  # type: ignore[attr-defined]
-from meta_deepFRI.structure_files.parse_structure_file import (
+from mDeepFRI import SEQ_ATOMS_DATASET_PATH
+from mDeepFRI.CPP_lib import libAtomDistanceIO  # type: ignore[attr-defined]
+from mDeepFRI.structure_files.parse_structure_file import (
     process_structure_file, search_structure_files)
-from meta_deepFRI.utils.mmseqs import create_target_database
-from meta_deepFRI.utils.utils import shutdown
+from mDeepFRI.utils.mmseqs import create_target_database
+from mDeepFRI.utils.utils import shutdown
 
 logging.basicConfig(
     level=logging.DEBUG,
@@ -25,52 +23,6 @@
 logger = logging.getLogger(__name__)
 
 
-def parse_args():
-    parser = argparse.ArgumentParser(description="Read structure files from folders --input" \
-                                                 "to extract sequence and atom positions. "
-                                                 "Create and index new --output MMSEQS2 database",
-                                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-
-    # logic described here is implemented in parse_input_paths
-    parser.add_argument(
-        "-i",
-        "--input",
-        nargs='+',
-        required=True,
-        default=None,
-        help="List of folder or file paths containing structure files.")
-
-    parser.add_argument(
-        "-o",
-        "--output",
-        required=True,
-        default=None,
-        help="Path to folder where new MMSeqs2 database will be created.")
-
-    parser.add_argument(
-        "-t",
-        "--threads",
-        required=False,
-        default=1,
-        type=int,
-        help=
-        "Number of threads to use. If not provided, the program will use single core."
-    )
-    parser.add_argument(
-        "-max_len",
-        "--max_protein_length",
-        required=False,
-        default=1_000,
-        type=int,
-        help="If protein chain is longer than this value, it will be truncated"
-    )
-
-    # parser.add_argument("--overwrite", action="store_true",
-    #                     help="Flag to override existing sequences and atom positions")
-    return parser.parse_args()
-    # yapf: enable
-
-
 def build_database(
     input_paths: List[str],
     output_path: str,
@@ -171,19 +123,3 @@ def build_database(
               indent=4)
 
     create_target_database(seq_atoms_path, output_path)
-
-
-def main() -> None:
-    args = parse_args()
-    input_seqs = [pathlib.Path(seqs) for seqs in args.input]
-    output_path = pathlib.Path(args.output)
-
-    build_database(input_paths=input_seqs,
-                   output_path=output_path,
-                   overwrite=True,
-                   threads=args.threads,
-                   max_protein_length=args.max_protein_length)
-
-
-if __name__ == '__main__':
-    main()
diff --git a/meta_deepFRI/metagenomic_deepfri.py b/mDeepFRI/pipeline.py
similarity index 91%
rename from meta_deepFRI/metagenomic_deepfri.py
rename to mDeepFRI/pipeline.py
index a52dab3..a48ff28 100644
--- a/meta_deepFRI/metagenomic_deepfri.py
+++ b/mDeepFRI/pipeline.py
@@ -6,15 +6,13 @@
 
 from pysam.libcfaidx import FastxFile
 
-from meta_deepFRI.config.names import (ATOMS, SEQ_ATOMS_DATASET_PATH,
-                                       TARGET_MMSEQS_DB_NAME)
-from meta_deepFRI.CPP_lib import \
-    libAtomDistanceIO  # type: ignore[attr-defined]
-from meta_deepFRI.DeepFRI.deepfrier import Predictor
-from meta_deepFRI.utils.fasta_file_io import SeqFileLoader
-from meta_deepFRI.utils.mmseqs import run_mmseqs_search
-from meta_deepFRI.utils.search_alignments import search_alignments
-from meta_deepFRI.utils.utils import load_deepfri_config
+from mDeepFRI import ATOMS, SEQ_ATOMS_DATASET_PATH, TARGET_MMSEQS_DB_NAME
+from mDeepFRI.CPP_lib import libAtomDistanceIO  # type: ignore[attr-defined]
+from mDeepFRI.predict import Predictor
+from mDeepFRI.utils.fasta_file_io import SeqFileLoader
+from mDeepFRI.utils.mmseqs import run_mmseqs_search
+from mDeepFRI.utils.search_alignments import search_alignments
+from mDeepFRI.utils.utils import load_deepfri_config
 
 logging.basicConfig(
     level=logging.DEBUG,
@@ -145,7 +143,7 @@ def check_deepfri_weights(weights: pathlib.Path) -> pathlib.Path:
 
     for net in ["cnn", "gcn"]:
         for model_type, model_path in models_config[net]["models"].items():
-            model_name = weights / (pathlib.Path(model_path).name + ".hdf5")
+            model_name = weights / (pathlib.Path(model_path).name + ".onnx")
             config_name = weights / (pathlib.Path(model_path).name +
                                      "_model_params.json")
             assert model_name.exists(
@@ -235,7 +233,7 @@ def metagenomic_deepfri(query_file: pathlib.Path, database: pathlib.Path,
             output_file_name = output_path / f"results_gcn_{mode}"
 
             gcn_params = deepfri_models_config["gcn"]["models"][mode]
-            gcn = Predictor.Predictor(gcn_params, gcn=True)
+            gcn = Predictor(gcn_params, gcn=True, threads=threads)
 
             for query_id, alignment in alignments.items():
                 logging.info("Predicting %s", query_id)
@@ -250,12 +248,10 @@ def metagenomic_deepfri(query_file: pathlib.Path, database: pathlib.Path,
                     alignment["target_sequence"],  # target alignment
                     generate_contacts)
 
-                # conversion of cmap to an adequate format
-                cmap = generated_query_contact_map.reshape(
-                    1, *generated_query_contact_map.shape) * 1
-
                 # running the actual prediction
-                gcn.predict_with_cmap(query_seq, cmap, query_id)
+                gcn.predict_function(seqres=query_seq,
+                                     cmap=generated_query_contact_map,
+                                     chain=query_id)
 
                 if "tsv" in output_format:
                     gcn.export_tsv(output_file_name.with_suffix('.tsv'))
@@ -272,10 +268,11 @@ def metagenomic_deepfri(query_file: pathlib.Path, database: pathlib.Path,
             output_file_name = output_path / f"results_cnn_{mode}"
 
             cnn_params = deepfri_models_config["cnn"]["models"][mode]
-            cnn = Predictor.Predictor(cnn_params, gcn=False)
+            cnn = Predictor(cnn_params, gcn=False, threads=threads)
             for query_id in unaligned_queries:
                 logging.info("Predicting %s", query_id)
-                cnn.predict_from_sequence(query_seqs[query_id], query_id)
+                cnn.predict_function(seqres=query_seqs[query_id],
+                                     chain=query_id)
 
             if "tsv" in output_format:
                 cnn.export_tsv(output_file_name.with_suffix('.tsv'))
diff --git a/mDeepFRI/predict.py b/mDeepFRI/predict.py
new file mode 100644
index 0000000..b324873
--- /dev/null
+++ b/mDeepFRI/predict.py
@@ -0,0 +1,157 @@
+import csv
+import json
+
+import numpy as np
+import onnxruntime as rt
+
+from mDeepFRI.utils.bio_utils import seq2onehot
+
+
+## TODO: do not accumulate predictions, write them out to csv asap to avoid unneccessary RAM usage
+class Predictor(object):
+    """
+    Class for loading trained models and computing GO/EC predictions and class activation maps (CAMs).
+    """
+    def __init__(self, model_prefix: str, gcn: bool = True, threads: int = 0):
+        self.model_prefix = model_prefix
+        self.gcn = gcn
+
+        # Not clear how parameter influences GPU exec
+        if rt.get_device() == 'CPU':
+            self.threads = threads
+        elif rt.get_device() == 'GPU':
+            self.threads = 0
+
+        self._load_model()
+        self.prot2goterms = {}
+        self.goidx2chains = {}
+
+    def _load_model(self):
+        session_options = rt.SessionOptions()
+        session_options.intra_op_num_threads = self.threads
+        self.session = rt.InferenceSession(
+            self.model_prefix + '.onnx',
+            providers=['CUDAExecutionProvider', 'CPUExecutionProvider'],
+            sess_options=session_options,
+        )
+
+        # load parameters
+        with open(self.model_prefix + "_model_params.json") as json_file:
+            metadata = json.load(json_file)
+
+        self.gonames = np.asarray(metadata['gonames'])
+        self.goterms = np.asarray(metadata['goterms'])
+        self.thresh = 0.1 * np.ones(len(self.goterms))
+
+    def predict_function(
+        self,
+        seqres: np.ndarray,
+        cmap: np.ndarray = None,
+        chain: str = "",
+    ):
+        """
+        Computes GO/EC predictions for a single protein chain from sequence and contact map.
+
+        Args:
+            seqres (str): protein sequence
+            cmap (np.array): contact map
+            chain (str): chain ID
+
+        Returns:
+            None
+        """
+        self.Y_hat = np.zeros((1, len(self.goterms)), dtype=float)
+        self.data = {}
+        self.test_prot_list = [chain]
+        S = seq2onehot(seqres)
+        S = S.reshape(1, *S.shape)
+        inputDetails = self.session.get_inputs()
+        if self.gcn:
+            A = cmap.reshape(1, *cmap.shape)
+            prediction = self.session.run(
+                None, {
+                    inputDetails[0].name: A.astype(np.float32),
+                    inputDetails[1].name: S.astype(np.float32)
+                })[0]
+            self.data[chain] = [[A, S], seqres]
+        else:
+            prediction = self.session.run(
+                None, {inputDetails[0].name: S.astype(np.float32)})[0]
+            self.data[chain] = [[S], seqres]
+
+        y = prediction[:, :, 0].reshape(-1)
+        self.Y_hat[0] = y
+        self.prot2goterms[chain] = []
+        go_idx = np.where(y >= self.thresh)[0]
+        for idx in go_idx:
+            if idx not in self.goidx2chains:
+                self.goidx2chains[idx] = set()
+            self.goidx2chains[idx].add(chain)
+            self.prot2goterms[chain].append(
+                (self.goterms[idx], self.gonames[idx], float(y[idx])))
+
+    def export_json(self, output_fn: str):
+        """
+        Exports predictions to .json format
+
+        Args:
+            output_fn (str): output file name
+
+        Returns:
+            None
+        """
+        data = []
+        for prot, goterms in self.prot2goterms.items():
+            sorted_rows = sorted(goterms, key=lambda x: x[2], reverse=True)
+            for row in sorted_rows:
+                data.append({
+                    'Protein': prot,
+                    'GO_term/EC_number': row[0],
+                    'Score': '{:.5f}'.format(row[2]),
+                    'GO_term/EC_number name': row[1]
+                })
+
+        with open(output_fn, 'w', encoding="utf-8") as f:
+            json.dump(data, f, indent=4)
+
+    def __save_file(self,
+                    output_fn: str,
+                    delimiter: str,
+                    quotechar: str = '"'):
+        """
+        Exports predictions to .csv or .tsv format
+
+        Args:
+            output_fn (str): output file name
+            delimiter (str): delimiter for .csv or .tsv file
+            quotechar (str): quotechar for .csv file
+
+        Returns:
+            None
+        """
+
+        with open(output_fn, 'w', encoding="utf-8") as f:
+            writer = csv.writer(f, delimiter=delimiter, quotechar=quotechar)
+            writer.writerow([
+                'Protein', 'GO_term/EC_number', 'Score',
+                'GO_term/EC_number name'
+            ])
+            for prot, goterms in self.prot2goterms.items():
+                sorted_rows = sorted(goterms, key=lambda x: x[2], reverse=True)
+                for row in sorted_rows:
+                    writer.writerow(
+                        [prot, row[0], '{:.5f}'.format(row[2]), row[1]])
+
+    def export_csv(self, output_fn: str):
+        """
+        Exports predictions to .csv format
+        """
+
+        self.__save_file(output_fn, ",")
+
+    def export_tsv(self, output_fn: str):
+        """
+        Exports predictions to .tsv format
+        """
+
+        self.__save_file(output_fn, "\t")
diff --git a/meta_deepFRI/__init__.py b/mDeepFRI/structure_files/__init__.py
similarity index 100%
rename from meta_deepFRI/__init__.py
rename to mDeepFRI/structure_files/__init__.py
diff --git a/meta_deepFRI/structure_files/parse_structure_file.py b/mDeepFRI/structure_files/parse_structure_file.py
similarity index 96%
rename from meta_deepFRI/structure_files/parse_structure_file.py
rename to mDeepFRI/structure_files/parse_structure_file.py
index 004b436..1607fc6 100644
--- a/meta_deepFRI/structure_files/parse_structure_file.py
+++ b/mDeepFRI/structure_files/parse_structure_file.py
@@ -6,11 +6,10 @@
 
 import numpy as np
 
-from meta_deepFRI.config.names import ATOMS, SEQUENCES
-from meta_deepFRI.CPP_lib import \
-    libAtomDistanceIO  # type: ignore[attr-defined]
-from meta_deepFRI.structure_files.parsers import parse_mmcif, parse_pdb
-from meta_deepFRI.utils import bio_utils
+from mDeepFRI import ATOMS, SEQUENCES
+from mDeepFRI.CPP_lib import libAtomDistanceIO  # type: ignore[attr-defined]
+from mDeepFRI.structure_files.parsers import parse_mmcif, parse_pdb
+from mDeepFRI.utils import bio_utils
 
 # need to parse different type of files? Add a file name pattern with a parser in this dict.
 # read structure_files_parsers/README.md for more information about how to create new parser.
diff --git a/meta_deepFRI/structure_files/parsers.py b/mDeepFRI/structure_files/parsers.py
similarity index 100%
rename from meta_deepFRI/structure_files/parsers.py
rename to mDeepFRI/structure_files/parsers.py
diff --git a/meta_deepFRI/structure_files/__init__.py b/mDeepFRI/utils/__init__.py
similarity index 100%
rename from meta_deepFRI/structure_files/__init__.py
rename to mDeepFRI/utils/__init__.py
diff --git a/meta_deepFRI/utils/bio_utils.py b/mDeepFRI/utils/bio_utils.py
similarity index 82%
rename from meta_deepFRI/utils/bio_utils.py
rename to mDeepFRI/utils/bio_utils.py
index 60ac179..b91e5fd 100644
--- a/meta_deepFRI/utils/bio_utils.py
+++ b/mDeepFRI/utils/bio_utils.py
@@ -1,6 +1,7 @@
-# copied from Biopython to remove dependency
+import numpy as np
 import parasail
 
+# copied from Biopython to remove dependency
 protein_letters_1to3 = {
     "A": "Ala",
     "C": "Cys",
@@ -45,6 +46,28 @@
     PROTEIN_LETTERS[str.upper(k)] = v
 PROTEIN_LETTERS["UNK"] = "X"
 
+
+def seq2onehot(seq):
+    """Create 26-dim embedding"""
+    chars = [
+        '-', 'D', 'G', 'U', 'L', 'N', 'T', 'K', 'H', 'Y', 'W', 'C', 'P', 'V',
+        'S', 'O', 'I', 'E', 'F', 'X', 'Q', 'A', 'B', 'Z', 'R', 'M'
+    ]
+    vocab_size = len(chars)
+    vocab_embed = dict(zip(chars, range(vocab_size)))
+
+    # Convert vocab to one-hot
+    vocab_one_hot = np.zeros((vocab_size, vocab_size), int)
+    for _, val in vocab_embed.items():
+        vocab_one_hot[val, val] = 1
+
+    embed_x = [vocab_embed[v] for v in seq]
+    seqs_x = np.array([vocab_one_hot[j, :] for j in embed_x])
+
+    return seqs_x
+
+
+# parasail martices
 substitution_matrices = {
     "blosum30": parasail.blosum30,
     "blosum35": parasail.blosum35,
diff --git a/meta_deepFRI/utils/fasta_file_io.py b/mDeepFRI/utils/fasta_file_io.py
similarity index 95%
rename from meta_deepFRI/utils/fasta_file_io.py
rename to mDeepFRI/utils/fasta_file_io.py
index d67a4a1..e004fc8 100644
--- a/meta_deepFRI/utils/fasta_file_io.py
+++ b/mDeepFRI/utils/fasta_file_io.py
@@ -2,7 +2,7 @@
 
 from pysam.libcfaidx import FastxFile
 
-from meta_deepFRI.config.names import ATOMS, SEQUENCES
+from mDeepFRI import ATOMS, SEQUENCES
 
 
 class SeqFileLoader:
diff --git a/meta_deepFRI/utils/mmseqs.py b/mDeepFRI/utils/mmseqs.py
similarity index 94%
rename from meta_deepFRI/utils/mmseqs.py
rename to mDeepFRI/utils/mmseqs.py
index 4faaa60..cd47c70 100644
--- a/meta_deepFRI/utils/mmseqs.py
+++ b/mDeepFRI/utils/mmseqs.py
@@ -3,9 +3,9 @@
 
 import pandas as pd
 
-from meta_deepFRI.config.names import (MERGED_SEQUENCES, MMSEQS_SEARCH_RESULTS,
-                                       TARGET_MMSEQS_DB_NAME)
-from meta_deepFRI.utils.utils import merge_files_binary, run_command
+from mDeepFRI import (MERGED_SEQUENCES, MMSEQS_SEARCH_RESULTS,
+                      TARGET_MMSEQS_DB_NAME)
+from mDeepFRI.utils.utils import merge_files_binary, run_command
 
 MMSEQS_COLUMN_NAMES = [
     "query", "target", "identity", "alignment_length", "mismatches",
diff --git a/meta_deepFRI/utils/search_alignments.py b/mDeepFRI/utils/search_alignments.py
similarity index 96%
rename from meta_deepFRI/utils/search_alignments.py
rename to mDeepFRI/utils/search_alignments.py
index 7d5c0d4..42187ba 100644
--- a/meta_deepFRI/utils/search_alignments.py
+++ b/mDeepFRI/utils/search_alignments.py
@@ -7,9 +7,9 @@
 import pandas as pd
 import parasail
 
-from meta_deepFRI.config.names import ALIGNMENTS
-from meta_deepFRI.utils.bio_utils import substitution_matrices
-from meta_deepFRI.utils.fasta_file_io import SeqFileLoader
+from mDeepFRI import ALIGNMENTS
+from mDeepFRI.utils.bio_utils import substitution_matrices
+from mDeepFRI.utils.fasta_file_io import SeqFileLoader
 
 
 def alignment_sequences_identity(query: str, target: str) -> float:
diff --git a/meta_deepFRI/utils/utils.py b/mDeepFRI/utils/utils.py
similarity index 84%
rename from meta_deepFRI/utils/utils.py
rename to mDeepFRI/utils/utils.py
index 1e2f268..fd22d47 100644
--- a/meta_deepFRI/utils/utils.py
+++ b/mDeepFRI/utils/utils.py
@@ -1,4 +1,5 @@
 import json
+import logging
 import pathlib
 import shlex
 import shutil
@@ -7,6 +8,8 @@
 
 import requests
 
+from mDeepFRI import config_links, model_links
+
 
 def run_command(command, timeout=-1):
     if isinstance(command, str):
@@ -42,6 +45,19 @@ def download_file(url, path):
             shutil.copyfileobj(r.raw, f)
 
 
+def download_model_weights(output_path: pathlib.Path):
+
+    total_len = len(model_links)
+    for i, model_link in enumerate(model_links):
+        download_file(model_link, output_path / model_link.split("/")[-1])
+        logging.debug("Downloading model weights... %s/%s", i + 1, total_len)
+
+    total_len = len(config_links)
+    for i, config in enumerate(config_links):
+        download_file(config, output_path / config.split("/")[-1])
+        logging.debug("Downloading model configs... %s/%s", i + 1, total_len)
+
+
 def merge_files_binary(file_paths: list, output_path: pathlib.Path) -> None:
     with open(output_path, 'wb') as writer:
         for input_file in file_paths:
diff --git a/meta_deepFRI/CPP_lib/atoms_file_io.h b/meta_deepFRI/CPP_lib/atoms_file_io.h
deleted file mode 100644
index faf738b..0000000
--- a/meta_deepFRI/CPP_lib/atoms_file_io.h
+++ /dev/null
@@ -1,49 +0,0 @@
-//
-// Created by soliareofastora on 24.12.2021.
-//
-
-#ifndef ATOMS_FILE_IO
-#define ATOMS_FILE_IO
-
-#include <boost/python.hpp>
-#include <boost/python/numpy.hpp>
-#include <fstream>
-
-namespace py = boost::python;
-namespace np = py::numpy;
-
-
-static void SaveAtomsFile(const np::ndarray &position_array, const np::ndarray &groups_array, const std::string &save_path) {
-  int chain_length = (int) groups_array.shape(0);
-  int* group_indexes = reinterpret_cast<int*>(groups_array.get_data());
-  int atom_count = group_indexes[chain_length - 1];
-  float* atom_positions = reinterpret_cast<float*>(position_array.get_data());
-
-  std::ofstream writer(save_path, std::ios::out | std::ios::binary);
-  writer.write(reinterpret_cast<const char*>(&chain_length), 4);
-  writer.write(reinterpret_cast<const char*>(group_indexes), 4 * chain_length);
-  writer.write(reinterpret_cast<const char*>(atom_positions), 4 * atom_count * 3);
-  writer.close();
-}
-
-
-static std::tuple<int, int*, float*> LoadAtomsFile(const std::string& file_path){
-  std::ifstream reader(file_path, std::ios::in | std::ios::binary);
-
-  int chain_length;
-  reader.read(reinterpret_cast<char*>(&chain_length), 4);
-
-  int* group_indexes = new int[chain_length];
-  reader.read(reinterpret_cast<char*>(group_indexes), 4 * chain_length);
-
-  chain_length--;
-  int atom_count = group_indexes[chain_length];
-
-  float* atoms_positions = new float[atom_count * 3];
-  reader.read(reinterpret_cast<char*>(atoms_positions), 4 * atom_count * 3);
-  reader.close();
-
-  return std::make_tuple(chain_length, group_indexes, atoms_positions);
-}
-
-#endif
diff --git a/meta_deepFRI/DeepFRI b/meta_deepFRI/DeepFRI
deleted file mode 160000
index 4ea2a78..0000000
--- a/meta_deepFRI/DeepFRI
+++ /dev/null
@@ -1,1 +0,0 @@
-Subproject commit 4ea2a78a574e387ac131f0f73786ad4eed8a2c7c
diff --git a/meta_deepFRI/config/__init__.py b/meta_deepFRI/config/__init__.py
deleted file mode 100644
index 7ef77eb..0000000
--- a/meta_deepFRI/config/__init__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-import multiprocessing
-
-CPU_COUNT = multiprocessing.cpu_count()
diff --git a/meta_deepFRI/config/names.py b/meta_deepFRI/config/names.py
deleted file mode 100644
index c8d17e0..0000000
--- a/meta_deepFRI/config/names.py
+++ /dev/null
@@ -1,17 +0,0 @@
-# folder names
-DEFAULT_NAME = "default"
-TARGET_MMSEQS_DB_NAME = "targetDB"
-SEQUENCES = "seq"
-ATOMS = "atom"
-SEQ_ATOMS_DATASET_PATH = "seq_atom_db"
-
-# file names
-PROJECT_CONFIG = "project_config.json"
-JOB_CONFIG = "job_config.json"
-TASK_CONFIG = "task_config.json"
-TARGET_DB_CONFIG = "target_db_config.json"
-
-ALIGNMENTS = "alignments.json"
-MERGED_SEQUENCES = 'merged_sequences.faa'
-TASK_SEQUENCES = "task_sequences.faa"
-MMSEQS_SEARCH_RESULTS = 'mmseqs2_search_results.m8'
diff --git a/meta_deepFRI/utils/__init__.py b/meta_deepFRI/utils/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/scripts/README.md b/scripts/README.md
deleted file mode 100644
index c55b1a5..0000000
--- a/scripts/README.md
+++ /dev/null
@@ -1,16 +0,0 @@
-
-# Build database:
-
-
-## structure_files.parse_structure_file.process_structure_file:
-1. reads the structure file extracting sequence and atom positions
-2. skip short sequences and truncate ones that size is over max_protein_length inside target_db_config.json
-3. save extracted data:
-- sequence - protein_id.faa file containing single sequence
-```
-f.write(f">{protein_id}\n{sequence}\n")
-```
-SEQ_ATOMS_DATASET_PATH / SEQUENCES / (protein_id + ".faa")
-- atom positions - binary file containing positions of all atoms and correlated amino acid chain index
-SEQ_ATOMS_DATASET_PATH / ATOMS / (protein_id + ".bin")
-For more information on how those binary are saved check out source code at CPP_lib/atoms_file_io.h
diff --git a/scripts/__init__.py b/scripts/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/scripts/default_config/README.md b/scripts/default_config/README.md
deleted file mode 100644
index 661b92c..0000000
--- a/scripts/default_config/README.md
+++ /dev/null
@@ -1,55 +0,0 @@
-# CONFIGS
-
-## Example project_config.json file
-```
-{
-    "DEEPFRI_PROCESSING_MODES": [
-        "mf",
-        "bp",
-        "cc",
-        "ec"
-    ],
-    "MAX_QUERY_CHAIN_LENGTH": 2500,
-    "ANGSTROM_CONTACT_THRESHOLD": 6,
-    "GENERATE_CONTACTS": 2,
-    "PAIRWISE_ALIGNMENT_MATCH": 2,
-    "PAIRWISE_ALIGNMENT_MISSMATCH": -1,
-    "PAIRWISE_ALIGNMENT_GAP_OPEN": -0.5,
-    "PAIRWISE_ALIGNMENT_GAP_CONTINUATION": -0.1,
-    "MMSEQS_MIN_BIT_SCORE": -99999,
-    "MMSEQS_MAX_EVAL": 99999,
-    "ALIGNMENT_MIN_SEQUENCE_IDENTITY": 0.3
-}
-```
-MAX_TARGET_CHAIN_LENGTH is used in update_target_mmseqs_database.py:67 to truncate longer sequences
-MAX_TARGET_CHAIN_LENGTH = 2500
-MAX_QUERY_CHAIN_LENGTH is used to filter out query sequences that are too long
-MAX_QUERY_CHAIN_LENGTH = 2500
-
-DeepFRI was trained with default value of 10 as seen in DeepFRI/train_DeepFRI.py
-ANGSTROM_CONTACT_THRESHOLD = 6
-GENERATE_CONTACTS are used to fill gaps in target sequence during contact map alignment
-GENERATE_CONTACTS = 2
-
-parameters used to filter mmseqs2 search results before aligning
-MMSEQS_MIN_BIT_SCORE = -99999
-MMSEQS_MAX_EVAL = 99999
-MMSEQS_MIN_IDENTITY = 0.5
-
-parameters used to calculate alignment score
-PAIRWISE_ALIGNMENT_MATCH = 2
-PAIRWISE_ALIGNMENT_MISSMATCH = -1
-PAIRWISE_ALIGNMENT_GAP_OPEN = -0.5
-PAIRWISE_ALIGNMENT_GAP_CONTINUATION = -0.1
-
-parameter to filter alignments based on sequence identity
-ALIGNMENT_MIN_SEQUENCE_IDENTITY value should be between 0 and 1
-ALIGNMENT_MIN_SEQUENCE_IDENTITY = 0.3
-
-this list determines types of outputs of metagenomic_deepfri
-mf = molecular_function
-bp = biological_process
-cc = cellular_component
-ec = enzyme_commission
-['mf', 'bp', 'cc', 'ec']
-DEEPFRI_PROCESSING_MODES = ['mf', 'bp', 'cc', 'ec']
diff --git a/scripts/main_pipeline.py b/scripts/main_pipeline.py
deleted file mode 100644
index e302171..0000000
--- a/scripts/main_pipeline.py
+++ /dev/null
@@ -1,132 +0,0 @@
-import argparse
-from pathlib import Path
-
-from meta_deepFRI.metagenomic_deepfri import metagenomic_deepfri
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(
-    description="This script for prediction of protein function from sequence using" \
-                "structure database, contact map alignment and DeepFRI model.",
-    formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-
-    # logic described here is implemented in parse_input_paths
-    parser.add_argument(
-        "-i",
-        "--input",
-        required=True,
-        default=None,
-        help="Filepath to a query protein FASTA files (.faa or .faa.gz)")
-    parser.add_argument("-db",
-                        "--db_path",
-                        required=True,
-                        default=None,
-                        help="Database path.")
-    parser.add_argument("-o",
-                        "--output_path",
-                        required=True,
-                        default=None,
-                        help="Output directory path.")
-    parser.add_argument("--output_format",
-                        nargs="*",
-                        choices=['tsv', 'json', 'csv'],
-                        default=['tsv'],
-                        help="Output format.")
-
-    # DeepFRI parameters
-    parser.add_argument(
-        "-w",
-        "--weights",
-        required=False,
-        default="trained_models",
-        help=
-        "Path to a folder containing downloaded deepfri weights and model_config.json."
-    )
-    parser.add_argument("--processing_modes", nargs="*", choices=['mf', 'bp', 'cc', 'ec'],
-                        default=['mf', 'bp', 'cc', 'ec'],
-                        help="DeepFRI processing modes. mf - molecular function, " \
-                            "bp - biological process, cc - cellular component, " \
-                            "ec - enzyme comission.")
-
-    # Contact map alignment parameters
-    parser.add_argument("--angstrom_contact_thresh",
-                        required=False,
-                        default=6,
-                        type=int,
-                        help="Angstrom contact threshold.")
-    parser.add_argument("--generate_contacts",
-                        required=False,
-                        default=2,
-                        type=int,
-                        help="Gap fill during contact map alignment.")
-
-    # MMSeqs2 parameters
-    parser.add_argument("--mmseqs_min_bit_score",
-                        required=False,
-                        default=None,
-                        type=float,
-                        help="Minimum bit score for MMSeqs2 search.")
-    parser.add_argument("--mmseqs_max_evalue",
-                        required=False,
-                        default=None,
-                        type=float,
-                        help="Maximum e-value for MMSeqs2 search.")
-    parser.add_argument("--mmseqs_min_identity",
-                        required=False,
-                        default=0.5,
-                        type=float,
-                        help="Minimum sequence identity for MMSeqs2 search.")
-
-    # Pairwise alignment parameters
-    parser.add_argument(
-        "--alignment_matrix",
-        required=False,
-        default="blosum62",
-        type=str,
-        help="Substitution matrix for pairwise sequence alignment.")
-    parser.add_argument(
-        "--alignment_gap_open",
-        required=False,
-        default=11,
-        type=int,
-        help="Gap open score penalty for pairwise sequence alignment.")
-    parser.add_argument(
-        "--alignment_gap_extend",
-        required=False,
-        default=1,
-        type=int,
-        help="Gap continuation score penalty for pairwise sequence alignment.")
-    parser.add_argument(
-        "--alignment_min_identity",
-        required=False,
-        default=0.5,
-        type=float,
-        help="Minimum sequence identity for pairwise sequence alignment.")
-    parser.add_argument(
-        "-t",
-        "--threads",
-        required=False,
-        default=1,
-        type=int,
-        help="Number of threads to use for parallel processing.")
-
-    return parser.parse_args()
-
-
-def main():
-
-    args = parse_args()
-    output_path = Path(args.output_path)
-    output_path.mkdir(parents=True, exist_ok=True)
-
-    metagenomic_deepfri(Path(args.input), Path(args.db_path), Path(
-        args.weights), output_path, args.output_format, args.processing_modes,
-                        args.angstrom_contact_thresh, args.generate_contacts,
-                        args.mmseqs_min_bit_score, args.mmseqs_max_evalue,
-                        args.mmseqs_min_identity, args.alignment_matrix,
-                        args.alignment_gap_open, args.alignment_gap_extend,
-                        args.alignment_min_identity, args.threads)
-
-
-if __name__ == '__main__':
-    main()
diff --git a/setup.py b/setup.py
index 612e398..747f6e4 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 from setuptools import Extension, find_namespace_packages, setup
 from setuptools.command.build_ext import build_ext as build_ext_orig
 
-# TODO: Add __init__.py file to meta_deepFRI/DeepFRI/
+from mDeepFRI import __version__
 
 
 class CMakeExtension(Extension):
@@ -57,8 +57,8 @@ def build_cmake(self, ext):
 
 
 setup(
-    name="meta_deepFRI",
-    version="0.2.0",
+    name="mDeepFRI",
+    version=__version__,
     description=
     "Pipeline for searching and aligning contact maps for proteins, then running DeepFri's GCN.",
     author="Piotr Kucharski, Valentyn Bezshapkin",
@@ -69,11 +69,10 @@ def build_cmake(self, ext):
     include_package_data=True,
     entry_points={
         "console_scripts": [
-            "deepfri = scripts.main_pipeline:main",
-            "deepfri_db_build = scripts.build_database:main",
+            "mDeepFRI = mDeepFRI.cli:main",
         ],
     },
-    ext_modules=[CMakeExtension("meta_deepFRI/CPP_lib/library_definition")],
+    ext_modules=[CMakeExtension("mDeepFRI/CPP_lib/library_definition")],
     cmdclass={
         'build_ext': build_ext,
     },
diff --git a/weight_convert/README.md b/weight_convert/README.md
new file mode 100644
index 0000000..2de9c4d
--- /dev/null
+++ b/weight_convert/README.md
@@ -0,0 +1,10 @@
+# TF2 to ONNX conversion
+
+This part contains scripts for converting original TF2 models to ONNX format, tests for reproducible output and measurements for performance.
+- Every original [DeepFRI CPU model](https://users.flatironinstitute.org/~renfrew/DeepFRI_data/newest_trained_models.tar.gz) was converted from Tensorflow format to ONNX using `tf2onnx` package.
+- Prediction were made on 100 random proteins and the numerical difference was not larger than `1e-5`.
+- The performance of the models was measured on a single CPU core.
+- The time was compared for both formats.
+
+The plot shows a relative speedup after conversion of models to ONNX:
+![](onnx_vs_tf2.png)
diff --git a/weight_convert/convert_models2onnx.py b/weight_convert/convert_models2onnx.py
new file mode 100644
index 0000000..82a276f
--- /dev/null
+++ b/weight_convert/convert_models2onnx.py
@@ -0,0 +1,80 @@
+import argparse
+import json
+import logging
+from pathlib import Path
+
+import tf2onnx
+# Is used to load weigths for the model
+# Can be replaced with tf.keras.models.load_model
+from DeepFRI.deepfrier.Predictor import Predictor
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='[%(asctime)s] %(module)s.%(funcName)s %(levelname)s: %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S')
+
+logger = logging.getLogger(__name__)
+
+
+# create argparse
+def parse_args():
+
+    parser = argparse.ArgumentParser(
+        description="Convert tensorflow models to ONNX",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+    parser.add_argument("-w",
+                        "--weights",
+                        required=True,
+                        default=None,
+                        help="Path to folder containing TF models.")
+
+    parser.add_argument("-o",
+                        "--output",
+                        required=False,
+                        default=None,
+                        help="Path to folder where ONNX model will be saved.")
+
+    return parser.parse_args()
+
+
+def convert_model(model_path, gcn, output_path):
+
+    pred = Predictor(model_prefix=model_path, gcn=gcn)
+    tf2onnx.convert.from_keras(pred.model,
+                               output_path=str(output_path),
+                               opset=15)
+
+
+# Newest CPU models
+
+if __name__ == "__main__":
+
+    args = parse_args()
+
+    model_path = args.weights
+    try:
+        with open(model_path + "/model_config.json") as f:
+            config = json.loads(f.read())
+    except FileNotFoundError:
+        raise FileNotFoundError(
+            "File %s/model_config.json not found. Download models from"
+            "https://users.flatironinstitute.org/~renfrew/DeepFRI_data/newest_trained_models.tar.gz",
+            model_path)
+
+    output_dir = Path(args.output)
+    output_dir.mkdir(exist_ok=True)
+
+    for net_type in ["cnn", "gcn"]:
+        for mode in ["ec", "mf", "bp", "cc"]:
+            model_prefix = str(
+                Path(model_path) /
+                config[net_type]["models"][mode].split("/")[-1])
+            model_name = Path(model_prefix).name
+            out_path = Path(args.output) / f"{model_name}.onnx"
+            logging.info("Converting model %s to %s", model_prefix + ".hdf5",
+                         out_path)
+            convert_model(model_prefix, config[net_type]["gcn"], out_path)
+            logging.info(
+                "Tensorflow to ONNX conversion of %s - %s is finished",
+                net_type, mode)
diff --git a/weight_convert/inference_times.csv.gz b/weight_convert/inference_times.csv.gz
new file mode 100644
index 0000000..c19fe81
Binary files /dev/null and b/weight_convert/inference_times.csv.gz differ
diff --git a/weight_convert/onnx_vs_tf2.png b/weight_convert/onnx_vs_tf2.png
new file mode 100644
index 0000000..b257243
Binary files /dev/null and b/weight_convert/onnx_vs_tf2.png differ
diff --git a/weight_convert/performance_plot.py b/weight_convert/performance_plot.py
new file mode 100644
index 0000000..cc43900
--- /dev/null
+++ b/weight_convert/performance_plot.py
@@ -0,0 +1,23 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import seaborn as sns
+
+df = pd.read_csv('inference_times.csv.gz')
+df = df.assign(speedup=lambda x: x['tf_time'] / x['onnx_time'])
+
+fig, axes = plt.subplots(1, 4, figsize=(10, 5), sharey=True, sharex=True)
+
+for mode, axis in zip(df["mode"].unique(), axes.flatten()):
+    data = df.query(f"mode == \"{mode}\"")
+    data = data.sort_values("prot_len")
+
+    # round protein length to highest 25
+    data = data.assign(prot_len=lambda x: np.ceil(x['prot_len'] / 50) * 50)
+    sns.lineplot(x="prot_len", y="speedup", hue="net_type", data=data, ax=axis)
+    axis.set_xlabel("Protein length")
+    axis.set_title(f"Mode {mode}")
+
+axes[0].set_ylabel("Speedup ONNX vs TF2 (times)")
+fig.tight_layout()
+fig.savefig("onnx_vs_tf2.png", dpi=150)
diff --git a/weight_convert/random_100_protein_prediction.ipynb b/weight_convert/random_100_protein_prediction.ipynb
new file mode 100644
index 0000000..f748aab
--- /dev/null
+++ b/weight_convert/random_100_protein_prediction.ipynb
@@ -0,0 +1,277 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2023-04-28 10:40:46.260998: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /nfs/nas22/fs2201/biol_micro_unix_modules/modules/software/binutils/2.37/lib:/nfs/nas22/fs2201/biol_micro_unix_modules/modules/software/Perl/5.32.0-GCCcore-10.2.0/lib:/nfs/nas22/fs2201/biol_micro_unix_modules/modules/software/DB/18.1.40-GCCcore-10.2.0/lib:/nfs/nas22/fs2201/biol_micro_unix_modules/modules/software/libreadline/8.0-GCCcore-10.2.0/lib:/nfs/nas22/fs2201/biol_micro_unix_modules/modules/software/gettext/0.21-GCCcore-10.2.0/lib:/nfs/nas22/fs2201/biol_micro_unix_modules/modules/software/ncurses/6.2-GCCcore-10.2.0/lib:/nfs/nas22/fs2201/biol_micro_unix_modules/modules/software/libxml2/2.9.10-GCCcore-10.2.0/lib:/nfs/nas22/fs2201/biol_micro_unix_modules/modules/software/XZ/5.2.5-GCCcore-10.2.0/lib:/nfs/nas22/fs2201/biol_micro_unix_modules/modules/software/expat/2.2.9-GCCcore-10.2.0/lib:/nfs/nas22/fs2201/biol_micro_unix_modules/modules/software/cURL/7.72.0-GCCcore-10.2.0/lib:/nfs/nas22/fs2201/biol_micro_unix_modules/modules/software/zlib/1.2.11-GCCcore-10.2.0/lib:/nfs/nas22/fs2201/biol_micro_unix_modules/modules/software/GCCcore/10.2.0/lib64:/nfs/nas22/fs2201/biol_micro_unix_modules/modules/software/GCCcore/10.2.0/lib\n",
+      "2023-04-28 10:40:46.261036: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.\n"
+     ]
+    }
+   ],
+   "source": [
+    "import time\n",
+    "import csv\n",
+    "import onnxruntime as rt\n",
+    "import numpy as np\n",
+    "from pathlib import Path\n",
+    "import json\n",
+    "from meta_deepFRI.DeepFRI.deepfrier.Predictor import Predictor\n",
+    "from meta_deepFRI.DeepFRI.deepfrier.utils import seq2onehot\n",
+    "from meta_deepFRI.utils.bio_utils import protein_letters_1to3\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "# write a wrapper to record time of function execution \n",
+    "def timeit(func):\n",
+    "    def wrapper(*args, **kwargs):\n",
+    "        start = time.time()\n",
+    "        result = func(*args, **kwargs)\n",
+    "        end = time.time() - start\n",
+    "        return result, end\n",
+    "    return wrapper\n",
+    "\n",
+    "@timeit\n",
+    "def predict_tf(pred: Predictor, seq=None, cmap=None):\n",
+    "    if cmap is None:\n",
+    "        result_tf = pred.model(seq, training=False)\n",
+    "    else:\n",
+    "        result_tf = pred.model([cmap, seq], training=False)\n",
+    "    return result_tf\n",
+    "\n",
+    "@timeit\n",
+    "def predict_onnx(session, seq, cmap=None):\n",
+    "    inputDetails = session.get_inputs()\n",
+    "    if cmap is None:\n",
+    "        result_onnx = session.run(None, {inputDetails[0].name: seq.astype(np.float32)})\n",
+    "    else:\n",
+    "        result_onnx = session.run(None, {inputDetails[0].name: cmap.astype(np.float32), inputDetails[1].name: seq.astype(np.float32)})\n",
+    "    \n",
+    "    return result_onnx\n",
+    "\n",
+    "def generate_random_protein(prot_length):\n",
+    "    aa_array = np.random.choice(list(protein_letters_1to3.keys()), size=(1, prot_length))\n",
+    "    return \"\".join(list(*aa_array))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 88,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "DeepCNN-MERGED_enzyme_commission\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 100/100 [00:07<00:00, 13.08it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "DeepCNN-MERGED_enzyme_commission passed the test!\n",
+      "DeepCNN-MERGED_molecular_function\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 100/100 [00:07<00:00, 14.10it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "DeepCNN-MERGED_molecular_function passed the test!\n",
+      "DeepCNN-MERGED_biological_process\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 100/100 [00:08<00:00, 11.25it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "DeepCNN-MERGED_biological_process passed the test!\n",
+      "DeepCNN-MERGED_cellular_component\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 100/100 [00:07<00:00, 14.25it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "DeepCNN-MERGED_cellular_component passed the test!\n",
+      "DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_ec\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 100/100 [01:56<00:00,  1.17s/it]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_ec passed the test!\n",
+      "DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_mf\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 100/100 [01:59<00:00,  1.19s/it]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_mf passed the test!\n",
+      "DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_bp\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 100/100 [02:00<00:00,  1.20s/it]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_bp passed the test!\n",
+      "DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_cc\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 100/100 [01:58<00:00,  1.18s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "DeepFRI-MERGED_GraphConv_gcd_512-512-512_fcd_1024_ca_10.0_cc passed the test!\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "model_path =  \"/nfs/nas22/fs2202/biol_micro_sunagawa/Projects/EAN/PROPHAGE_REFSEQ_EAN/scratch/databases/meta_deepfri_data/newest_models/trained_models\"\n",
+    "with open(model_path + \"/model_config.json\") as f:\n",
+    "    config = json.loads(f.read())\n",
+    "\n",
+    "inference_times = []\n",
+    "\n",
+    "for net_type in [\"cnn\", \"gcn\"]:\n",
+    "        for mode in [\"ec\", \"mf\", \"bp\", \"cc\"]:\n",
+    "           \n",
+    "            tf_model_prefix = str(Path(model_path) / config[net_type][\"models\"][mode].split(\"/\")[-1])\n",
+    "            model_name = Path(tf_model_prefix).name\n",
+    "            print(model_name)\n",
+    "            pred = Predictor(model_prefix = tf_model_prefix, gcn=config[net_type][\"gcn\"])\n",
+    "            session = rt.InferenceSession(f'../onnx_deepfri_models/{model_name}.onnx', \n",
+    "                                          providers=['CPUExecutionProvider'])\n",
+    "\n",
+    "            # Testing GCNs\n",
+    "            for seed in tqdm(range(1000), total=100):\n",
+    "                np.random.seed(seed)\n",
+    "                # zero out cmap\n",
+    "                cmap = None\n",
+    "                \n",
+    "                # generate random protein and contact map\n",
+    "                prot_len = np.random.randint(60, 1000, size=1)[0]\n",
+    "                seq = generate_random_protein(prot_length=prot_len)\n",
+    "                one_hot = seq2onehot(seq)\n",
+    "                one_hot = one_hot.reshape(1, *one_hot.shape)\n",
+    "                \n",
+    "                if net_type == \"gcn\":\n",
+    "                    cmap = np.random.randint(0, 2, size=(1, prot_len, prot_len), dtype=int)\n",
+    "                    \n",
+    "                tf_pred, tf_time = predict_tf(pred, seq=one_hot, cmap=cmap)\n",
+    "                onnx_pred, onnx_time = predict_onnx(session, seq=one_hot, cmap=cmap)\n",
+    "\n",
+    "                # Skip if predictions all nan\n",
+    "                if np.isnan(onnx_pred[0]).all():\n",
+    "                    continue\n",
+    "\n",
+    "                # Compare predictions from TF and ONNX\n",
+    "                assert np.isclose(tf_pred, onnx_pred, atol=10e-5).all()\n",
+    "\n",
+    "                # Record inference time\n",
+    "                inference_times.append([prot_len, net_type, mode, tf_time, onnx_time])\n",
+    "\n",
+    "            print(f\"{model_name} passed the test!\")\n",
+    "\n",
+    "# write inference to csv\n",
+    "with open(\"inference_times.csv\", \"w\") as f:\n",
+    "    writer = csv.writer(f)\n",
+    "    writer.writerow([\"prot_len\", \"net_type\", \"mode\", \"tf_time\", \"onnx_time\"])\n",
+    "    writer.writerows(inference_times)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "deepfri",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.8"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}


Test patch:
diff --git a/tests/data/1S3P-A.bin b/tests/data/1S3P-A.bin
new file mode 100644
index 0000000..d62b4d5
Binary files /dev/null and b/tests/data/1S3P-A.bin differ
diff --git a/tests/data/1S3P-A.bin.cmap.pkl b/tests/data/1S3P-A.bin.cmap.pkl
new file mode 100644
index 0000000..8bb98a7
Binary files /dev/null and b/tests/data/1S3P-A.bin.cmap.pkl differ
diff --git a/tests/test_alignment.py b/tests/test_alignment.py
index 93f0553..a194fa4 100644
--- a/tests/test_alignment.py
+++ b/tests/test_alignment.py
@@ -1,6 +1,6 @@
 import pytest
 
-from meta_deepFRI.utils.search_alignments import alignment_sequences_identity
+from mDeepFRI.utils.search_alignments import alignment_sequences_identity
 
 
 @pytest.mark.parametrize('query,target,expected', [('AASDS', 'AASDS', 1),
diff --git a/tests/test_cli.py b/tests/test_cli.py
new file mode 100644
index 0000000..2db87e3
--- /dev/null
+++ b/tests/test_cli.py
@@ -0,0 +1,43 @@
+import pytest
+from click.testing import CliRunner
+
+from mDeepFRI import cli
+
+
+@pytest.fixture
+def runner():
+    return CliRunner()
+
+
+# test click cli
+def test_cli_help(runner):
+    result = runner.invoke(cli.main, ["--help"])
+    assert result.exit_code == 0
+    assert "--help" in result.output
+
+
+def test_cli_build_db(runner):
+    result = runner.invoke(cli.main, ["build-db", "--help"])
+    assert result.exit_code == 0
+    assert "--help" in result.output
+    assert "-m" in result.output
+    assert "-t" in result.output
+
+
+def test_cli_get_models(runner):
+    result = runner.invoke(cli.main, ["get-models", "--help"])
+    assert result.exit_code == 0
+    assert "--help" in result.output
+    assert "-o" in result.output
+
+
+def test_cli_predict_function(runner):
+    result = runner.invoke(cli.main, ["predict-function", "--help"])
+    assert result.exit_code == 0
+    assert "--help" in result.output
+    assert "-i" in result.output
+    assert "-o" in result.output
+    assert "-m" in result.output
+    assert "-t" in result.output
+    assert "-d" in result.output
+    assert "-w" in result.output
diff --git a/tests/test_cpp_lib.py b/tests/test_cpp_lib.py
index 606cf20..240f877 100644
--- a/tests/test_cpp_lib.py
+++ b/tests/test_cpp_lib.py
@@ -1,9 +1,27 @@
+import filecmp
+import pickle
+import tempfile
 from typing import Callable
 
-from meta_deepFRI.CPP_lib.libAtomDistanceIO import (initialize,
-                                                    load_aligned_contact_map,
-                                                    load_contact_map,
-                                                    save_atoms)
+import numpy as np
+import pytest
+
+from mDeepFRI.CPP_lib.libAtomDistanceIO import (initialize,
+                                                load_aligned_contact_map,
+                                                load_contact_map, save_atoms)
+from mDeepFRI.structure_files.parsers import parse_pdb
+
+
+@pytest.fixture()
+def reference_binary():
+    return "tests/data/1S3P-A.bin"
+
+
+@pytest.fixture()
+def expected_contact_map():
+    with open("tests/data/1S3P-A.bin.cmap.pkl", "rb") as f:
+        expected_contact_map = pickle.load(f)
+    return expected_contact_map
 
 
 def error_wrapper(func: Callable, *args, error_contains: str = None):
@@ -24,3 +42,25 @@ def test_functions():
                   error_contains="did not match C++ signature")
     error_wrapper(load_aligned_contact_map,
                   error_contains="did not match C++ signature")
+
+
+def test_load_contact_map(reference_binary, expected_contact_map):
+    initialize()
+    contact_map_cpp = load_contact_map(reference_binary, 6)
+    assert np.array_equal(contact_map_cpp, expected_contact_map)
+
+
+def test_save_atoms(expected_contact_map):
+    with open("tests/data/structures/1S3P-A.pdb") as f:
+        _, positions, groups = parse_pdb(f)
+    _, group_index = np.unique(groups, return_index=True)
+    group_index = np.sort(group_index)
+    group_indexes = np.append(group_index, positions.shape[0]).astype(np.int32)
+    initialize()
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        save_atoms(positions, group_indexes, tmpdirname + "/1S3P-A.bin")
+        assert filecmp.cmp(tmpdirname + "/1S3P-A.bin", "tests/data/1S3P-A.bin")
+
+
+def test_load_aligned_contact_map():
+    pass
diff --git a/tests/test_deepfri.py b/tests/test_deepfri.py
index b82e4b0..66b2780 100644
--- a/tests/test_deepfri.py
+++ b/tests/test_deepfri.py
@@ -24,7 +24,7 @@ def execute_command(command):
 @pytest.fixture
 def deepfri_database():
     stdout, stderr = execute_command(
-        f'deepfri_db_build -i {INPUT_STRUCTURES} -o {OUTPUT_DATABASE}')
+        f'mDeepFRI build-db -i {INPUT_STRUCTURES} -o {OUTPUT_DATABASE}')
     print(stdout)
 
 
@@ -35,7 +35,7 @@ def test_database(deepfri_database):
 @pytest.fixture
 def deepri_results(deepfri_database):
     stdout, stderr = execute_command(
-        f"deepfri -i {QUERY_FILE} -db {OUTPUT_DATABASE} -o {RESULTS}")
+        f"mDeepFRI -i {QUERY_FILE} -db {OUTPUT_DATABASE} -o {RESULTS}")
     print(stdout)
 
 
diff --git a/tests/test_main_pipeline.py b/tests/test_main_pipeline.py
index b9fc997..03a2d60 100644
--- a/tests/test_main_pipeline.py
+++ b/tests/test_main_pipeline.py
@@ -1,8 +1,8 @@
 import pytest
 
-from meta_deepFRI.utils import utils
+from mDeepFRI.utils import utils
 
 
 def test_deepfri_errors():
     with pytest.raises(RuntimeError):
-        utils.run_command("deepfri").strip()
+        utils.run_command("mDeepFri").strip()
diff --git a/tests/test_parsers.py b/tests/test_parsers.py
index 94a71e2..94d170a 100644
--- a/tests/test_parsers.py
+++ b/tests/test_parsers.py
@@ -4,7 +4,7 @@
 import numpy as np
 import pytest
 
-from meta_deepFRI.structure_files.parsers import parse_mmcif, parse_pdb
+from mDeepFRI.structure_files.parsers import parse_mmcif, parse_pdb
 
 
 @pytest.fixture
diff --git a/tests/test_structure_parsers.py b/tests/test_structure_parsers.py
index 370b30b..f645bb0 100644
--- a/tests/test_structure_parsers.py
+++ b/tests/test_structure_parsers.py
@@ -3,8 +3,8 @@
 import numpy as np
 import pytest
 
-from meta_deepFRI.structure_files.parse_structure_file import (
-    SeqAtoms, read_structure_file)
+from mDeepFRI.structure_files.parse_structure_file import (SeqAtoms,
+                                                           read_structure_file)
 
 
 @pytest.mark.parametrize(
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 5184825..07a2ac8 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1,6 +1,6 @@
 import pytest
 
-from meta_deepFRI.utils import bio_utils, utils
+from mDeepFRI.utils import bio_utils, utils
 
 
 def test_protein_letters():
diff --git a/weight_convert/onnx_test.csv b/weight_convert/onnx_test.csv
new file mode 100644
index 0000000..e239203
--- /dev/null
+++ b/weight_convert/onnx_test.csv
@@ -0,0 +1,18 @@
+Protein,GO_term/EC_number,Score,GO_term/EC_number name
+test,GO:0044419,0.27301,biological process involved in interspecies interaction between organisms
+test,GO:0016043,0.20897,cellular component organization
+test,GO:0032501,0.17729,multicellular organismal process
+test,GO:0043412,0.16177,macromolecule modification
+test,GO:0051704,0.15849,multi-organism process
+test,GO:0036211,0.15131,protein modification process
+test,GO:0006464,0.15099,cellular protein modification process
+test,GO:0009056,0.14845,catabolic process
+test,GO:1901575,0.13204,organic substance catabolic process
+test,GO:0044403,0.12657,biological process involved in symbiotic interaction
+test,GO:0051716,0.12328,cellular response to stimulus
+test,GO:0006629,0.12220,lipid metabolic process
+test,GO:0032502,0.12044,developmental process
+test,GO:0006950,0.11660,response to stress
+test,GO:0016032,0.11410,viral process
+test,GO:0048856,0.11072,anatomical structure development
+test,GO:0016070,0.10155,RNA metabolic process
diff --git a/weight_convert/test.csv b/weight_convert/test.csv
new file mode 100644
index 0000000..e239203
--- /dev/null
+++ b/weight_convert/test.csv
@@ -0,0 +1,18 @@
+Protein,GO_term/EC_number,Score,GO_term/EC_number name
+test,GO:0044419,0.27301,biological process involved in interspecies interaction between organisms
+test,GO:0016043,0.20897,cellular component organization
+test,GO:0032501,0.17729,multicellular organismal process
+test,GO:0043412,0.16177,macromolecule modification
+test,GO:0051704,0.15849,multi-organism process
+test,GO:0036211,0.15131,protein modification process
+test,GO:0006464,0.15099,cellular protein modification process
+test,GO:0009056,0.14845,catabolic process
+test,GO:1901575,0.13204,organic substance catabolic process
+test,GO:0044403,0.12657,biological process involved in symbiotic interaction
+test,GO:0051716,0.12328,cellular response to stimulus
+test,GO:0006629,0.12220,lipid metabolic process
+test,GO:0032502,0.12044,developmental process
+test,GO:0006950,0.11660,response to stress
+test,GO:0016032,0.11410,viral process
+test,GO:0048856,0.11072,anatomical structure development
+test,GO:0016070,0.10155,RNA metabolic process