enzymetk 0.0.1__py3-none-any.whl

enzymetk/__init__.py

@@ -0,0 +1,56 @@
+###############################################################################
+#                                                                             #
+#    This program is free software: you can redistribute it and/or modify     #
+#    it under the terms of the GNU General Public License as published by     #
+#    the Free Software Foundation, either version 3 of the License, or        #
+#    (at your option) any later version.                                      #
+#                                                                             #
+#    This program is distributed in the hope that it will be useful,          #
+#    but WITHOUT ANY WARRANTY; without even the implied warranty of           #
+#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the            #
+#    GNU General Public License for more details.                             #
+#                                                                             #
+#    You should have received a copy of the GNU General Public License        #
+#    along with this program. If not, see <http://www.gnu.org/licenses/>.     #
+#                                                                             #
+###############################################################################
+
+"""
+Author: Ariane Mora
+Date: March 2025
+"""
+__title__ = 'enzymetk'
+__description__ = 'Toolkit for enzymes and what not'
+__url__ = 'https://github.com/arianemora/enzyme-tk/'
+__version__ = '0.0.1'
+__author__ = 'Ariane Mora'
+__author_email__ = 'ariane.n.mora@gmail.com'
+__license__ = 'GPL3'
+
+# from enzymetk.step import *
+# from enzymetk.generate_msa_step import ClustalOmega
+# from enzymetk.annotateEC_CLEAN_step import CLEAN
+# from enzymetk.annotateEC_proteinfer_step import ProteInfer
+# from enzymetk.dock_chai_step import Chai
+# from enzymetk.dock_vina_step import Vina
+# from enzymetk.embedchem_chemberta_step import ChemBERT
+# from enzymetk.embedchem_rxnfp_step import RxnFP
+# from enzymetk.embedchem_selformer_step import SelFormer
+# from enzymetk.embedchem_unimol_step import UniMol
+# from enzymetk.embedprotein_esm_step import EmbedESM
+# from enzymetk.generate_tree_step import FastTree
+# from enzymetk.inpaint_ligandMPNN_step import LigandMPNN
+# from enzymetk.metagenomics_porechop_trim_reads_step import PoreChop
+# from enzymetk.metagenomics_prokka_annotate_genes import Prokka
+# #from enzymetk.predict_activity_step import 
+# from enzymetk.predict_catalyticsite_step import ActiveSitePred
+# from enzymetk.sequence_search_blast import BLAST
+# from enzymetk.similarity_foldseek_step import FoldSeek
+# from enzymetk.similarity_mmseqs_step import MMseqs
+# from enzymetk.similarity_reaction_step import ReactionDist
+# from enzymetk.similarity_substrate_step import SubstrateDist
+
+
+
+
+

enzymetk/annotateEC_CLEAN_step.py

@@ -0,0 +1,122 @@
+"""
+Install clean and then you need to activate the environment and install and run via that. 
+
+Honestly it's a bit hacky the way they do it, not bothered to change things so have to save the data to their
+repo and then copy it out of it.
+"""
+from enzymetk.step import Step
+import pandas as pd
+import numpy as np
+from multiprocessing.dummy import Pool as ThreadPool
+from tempfile import TemporaryDirectory
+import os
+import subprocess
+import random
+import string
+import os
+from tqdm import tqdm
+
+
+class CLEAN(Step):
+    
+    def __init__(self, id_col: str, seq_col: str, clean_dir: str, num_threads: int = 1, 
+                 ec1_filter: list = None, ec2_filter: list = None, ec3_filter: list = None, ec4_filter: list = None, 
+                 env_name: str = 'clean', args: list = None):
+        self.env_name = env_name
+        self.args = args
+        self.id_col = id_col
+        self.clean_dir = clean_dir
+        self.seq_col = seq_col # This is the column which has the sequence in it 
+        self.num_threads = num_threads  
+        self.ec1_filter = ec1_filter        
+        self.ec2_filter = ec2_filter
+        self.ec3_filter = ec3_filter
+        self.ec4_filter = ec4_filter
+
+    def __filter_df(self, df: pd.DataFrame) -> pd.DataFrame:
+        # ------------- Separate out ECs ------------------
+        df['id'] = [c.split(',')[0] for c in df[0].values]
+        df['ec'] = [c.split(',')[1:] for c in df[0].values]
+        df = df.drop(columns=0)
+        df = df.explode('ec')
+        df['score'] = [float(ec.split('/')[1]) for ec in df['ec'].values]
+        df['ec'] = [str(ec.split('/')[0]) for ec in df['ec'].values]
+        df['predicted_ecs'] = [ec.split(':')[1] for ec in df['ec'].values]
+        df['EC1'] = [r.split('.')[0] for r in df['predicted_ecs'].values]
+        df['EC2'] = [r.split('.')[1] for r in df['predicted_ecs'].values]
+        df['EC3'] = [r.split('.')[2] for r in df['predicted_ecs'].values]
+        df['EC4'] = [r.split('.')[3] for r in df['predicted_ecs'].values]
+        
+        if self.ec1_filter is not None:
+            df = df[df['EC1'].isin(self.ec1_filter)]
+        if self.ec2_filter is not None:
+            df = df[df['EC2'].isin(self.ec2_filter)]
+        if self.ec3_filter is not None:
+            df = df[df['EC3'].isin(self.ec3_filter)]
+        if self.ec4_filter is not None:
+            df = df[df['EC4'].isin(self.ec4_filter)]
+            
+        df = df.sort_values(by='score', ascending=False)
+        # Drop duplicates based on id only keeping the highest score
+        df.drop_duplicates(subset='id', keep='first', inplace=True)
+        return df
+    
+    def __execute(self, data: list) -> np.array:
+        df, tmp_dir = data
+        # Make sure in the directory of proteinfer
+        # Create the fasta file based on the id and the sequence value columns
+        tmp_label = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
+        input_filename = f'{tmp_dir}CLEAN_{tmp_label}.fasta'
+        
+        # write fasta file which is the input for proteinfer
+        with open(input_filename, 'w+') as fout:
+            for entry, seq in df[[self.id_col, self.seq_col]].values:
+                fout.write(f'>{entry.strip()}\n{seq.strip()}\n')
+        # Run it multi threaded
+        os.chdir(self.clean_dir)
+
+        tmp_label = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
+        # Since clean is GPU hungry, we only run CLEAN on the ones that proteInfer has predicted to be class 3.
+        # Need to first copy the data to the CLEAN folder because it's stupid
+        os.chdir(f'{self.clean_dir}')
+        cmd = ['cp',  input_filename, f'{self.clean_dir}data/inputs/{tmp_label}.fasta']
+        self.run(cmd)
+        # Run clean with clean environment
+        cmd = ['conda', 'run', '-n', self.env_name, 'python3', f'{self.clean_dir}CLEAN_infer_fasta.py', 
+                        '--fasta_data', tmp_label]
+        if self.args is not None:
+            # Add the args to the command
+            cmd.extend(self.args)
+        self.run(cmd)
+        # Copy across the results file
+        df = pd.read_csv(f'{self.clean_dir}results/inputs/{tmp_label}_maxsep.csv', header=None, sep='\t')
+        cmd = ['rm', f'{self.clean_dir}data/inputs/{tmp_label}.fasta']
+        self.run(cmd)
+        cmd = ['rm', f'{self.clean_dir}results/inputs/{tmp_label}_maxsep.csv']
+        self.run(cmd)   
+        
+        # Change back to the current folder     
+        dir_path = os.path.dirname(os.path.realpath(__file__))
+        os.chdir(dir_path)
+        
+        return df
+    
+    def execute(self, df: pd.DataFrame) -> pd.DataFrame:
+        with TemporaryDirectory() as tmp_dir:
+            if self.num_threads > 1:
+                output_filenames = []
+                df_list = np.array_split(df, self.num_threads)
+                for df_chunk in tqdm(df_list):
+                    try:
+                        output_filenames.append(self.__execute([df_chunk, tmp_dir]))
+                    except Exception as e:
+                         print(f"Error in executing ESM2 model: {e}")
+                         continue
+                df = pd.DataFrame()
+                print(output_filenames)
+                for sub_df in output_filenames:
+                    df = pd.concat([df, sub_df])
+                return df
+            else:
+                return self.__execute([df, tmp_dir])
+                return df

enzymetk/annotateEC_CREEP_step.py

@@ -0,0 +1,82 @@
+from enzymetk.step import Step
+import pandas as pd
+from tempfile import TemporaryDirectory
+import subprocess
+import logging
+import numpy as np
+import os
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+"""
+import os
+
+Example file as input:
+Reaction,EC number,Reaction Text,EC3,EC2,EC1
+O=C(OCC(CC)CCCC)C1=CC=CC=C1C(OCC(CC)CCCC)=O>>OC(C2=CC=CC=C2C(O)=O)=O,3.1.1.60,DEHP->PA,3.1.1,3.1,3
+CCCCC(CC)COC(=O)c1ccccc1C(=O)OCC(CC)CCCC.O>>CCCCC(CC)CO.CCCCC(CC)COC(=O)c1ccccc1C(=O)O,3.1.1.60,DEHP-MEHP,3.1.1,3.1,3
+
+os.system(f'
+python step_02_extract_CREEP.py --pretrained_folder=/disk1/share/software/CREEP/data/bioremediation_split --dataset=/disk1/share/software/CREEP/output/DEHP/bioremediation_reaction_test.csv --modality=reaction
+')
+
+os.system(f'python downstream_retrieval.py --pretrained_folder=CREEP/$OUTPUT_DIR --query_dataset=$TEST_SET --reference_dataset=all_ECs --query_modality=reaction --reference_modality=protein')
+"""
+    
+class CREEP(Step):
+    
+    def __init__(self, id_col: str, value_col: str, CREEP_dir: str, CREEP_cache_dir: str, modality: str, reference_modality: str, 
+                 env_name: str = 'CREEP', args_extract: list = None, args_retrieval: list = None):
+        self.env_name = env_name
+        self.id_col = id_col
+        self.value_col = value_col  
+        self.modality = modality
+        self.reference_modality = reference_modality
+        self.CREEP_dir = CREEP_dir
+        self.CREEP_cache_dir = CREEP_cache_dir
+        self.args_extract = args_extract
+        self.args_retrieval = args_retrieval
+            
+    def __execute(self, df: pd.DataFrame, tmp_dir: str) -> pd.DataFrame:
+        tmp_dir = '/disk1/ariane/vscode/degradeo/pipeline/tmp/'
+        input_filename = f'{tmp_dir}/creepasjkdkajshdkja.csv'
+        df.to_csv(input_filename, index=False)
+        cmd = ['conda', 'run', '-n', self.env_name, 'python', f'{self.CREEP_dir}scripts/step_02_extract_CREEP.py', '--pretrained_folder', 
+                                 f'{self.CREEP_cache_dir}output/easy_split', 
+                                  '--dataset', input_filename,
+                                  '--cache_dir', self.CREEP_dir, 
+                                  '--modality', self.modality.strip(), 
+                                  '--output_dir', f'{tmp_dir}']
+        if self.args_extract is not None:
+            cmd.extend(self.args_extract)
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        cmd = ['conda', 'run', '-n', self.env_name, 'python', f'{self.CREEP_dir}scripts/downstream_retrieval.py', '--pretrained_folder',
+                                 f'{self.CREEP_cache_dir}output/easy_split', 
+                                 '--query_dataset', input_filename, 
+                                 '--reference_dataset', 'all_ECs',
+                                 '--query_modality', self.modality.strip(),
+                                 '--cache_dir', self.CREEP_cache_dir, 
+                                 '--output_dir', f'{tmp_dir}',
+                                 '--reference_modality', self.reference_modality]
+        if self.args_retrieval is not None:
+            cmd.extend(self.args_retrieval)
+        self.run(cmd)
+        output_filename = f'{tmp_dir}/creep_reaction2protein_retrieval_similarities.npy'
+        return output_filename
+    
+    def execute(self, df: pd.DataFrame) -> pd.DataFrame:
+        with TemporaryDirectory() as tmp_dir:
+            output_filename = self.__execute(df, tmp_dir)
+            df = pd.read_csv(f"{self.CREEP_dir}/data/processed_data/EC_list.txt", header=None)
+            data = np.load(output_filename)
+            all_ecs = np.load(f"{self.CREEP_dir}/data/output/easy_split/representations/all_ECs_cluster_centers.npy", allow_pickle=True)
+            rxn_data = np.load(f"{self.CREEP_dir}/data/output/easy_split/representations/easy_reaction_test_representations.npy", allow_pickle=True)
+            data_dict = rxn_data.item()
+            print(data_dict)
+            data_dict = data_dict['reaction_repr_array']
+            data_rxn = all_ecs.item()
+            data_rxn = data_rxn['protein_repr_array']
+            for i, d in enumerate(data):
+                df[f'sim_{i}'] = d
+            return df

enzymetk/annotateEC_proteinfer_step.py

@@ -0,0 +1,136 @@
+from enzymetk.step import Step
+import pandas as pd
+import numpy as np
+from multiprocessing.dummy import Pool as ThreadPool
+from tempfile import TemporaryDirectory
+import os
+import subprocess
+
+
+class ProteInfer(Step):
+    
+    def __init__(self, id_col: str, seq_col: str, proteinfer_dir: str, num_threads: int = 1, 
+                 ec1_filter: list = None, ec2_filter: list = None, ec3_filter: list = None, ec4_filter: list = None, 
+                 env_name: str = 'proteinfer', args: list = None):
+        """Initialize the CLEAN step for enzyme classification.
+        
+        Filters are lists of strings which are the EC values to keep. If None then keep all EC values.
+        
+        Parameters
+        ----------
+        id_col : str
+            Name of the column containing sequence identifiers in the input DataFrame
+        seq_col : str
+            Name of the column containing protein sequences in the input DataFrame
+        clean_dir : str
+            Path to the CLEAN software directory containing the CLEAN_infer_fasta.py script
+        num_threads : int, optional
+            Number of parallel threads to use for processing (default=1)
+        ec1_filter : list, optional
+            List of EC1 values to filter by (default=None) if None then keep all EC1 values also use '-' to keep missing values
+        ec2_filter : list, optional
+            List of EC2 values to filter by (default=None) if None then keep all EC2 values also use '-' to keep missing values 
+        ec3_filter : list, optional
+            List of EC3 values to filter by (default=None) if None then keep all EC3 values also use '-' to keep missing values e.g. ['3', '-']
+        ec4_filter : list, optional
+            List of EC4 values to filter by (default=None) if None then keep all EC4 values also use '-' to keep missing values e.g. ['1', '-']
+            
+        Notes
+        -----
+        CLEAN requires a GPU and the 'clean' conda environment to be installed.
+        The CLEAN software directory should contain the following structure:
+        - data/inputs/ : Directory for temporary fasta files
+        - results/inputs/ : Directory where CLEAN outputs results
+        """
+        self.env_name = env_name
+        self.args = args
+        self.id_col = id_col
+        self.proteinfer_dir = proteinfer_dir
+        self.seq_col = seq_col # This is the column which has the sequence in it 
+        self.num_threads = num_threads
+        self.ec1_filter = ec1_filter
+        self.ec2_filter = ec2_filter
+        self.ec3_filter = ec3_filter
+        self.ec4_filter = ec4_filter
+        
+    def __execute(self, data: list) -> np.array:
+        df, tmp_dir = data
+        # Make sure in the directory of proteinfer
+        # Create the fasta file based on the id and the sequence value columns
+        input_filename = f'{tmp_dir}proteinfer.fasta'
+        output_filename = f'{tmp_dir}proteinfer.txt'
+        
+        # write fasta file which is the input for proteinfer
+        with open(input_filename, 'w+') as fout:
+            for entry, seq in df[[self.id_col, self.seq_col]].values:
+                fout.write(f'>{entry.strip()}\n{seq.strip()}\n')
+        
+        os.chdir(self.proteinfer_dir)
+        cmd = ['conda', 'run', '-n', self.env_name, 'python3', 
+                os.path.join(self.proteinfer_dir, f'proteinfer.py'),
+                '-i', input_filename,
+                '-o', output_filename]
+        if self.args is not None:
+            # Add the args to the command
+            cmd.extend(self.args)
+        self.run(cmd)
+        df = pd.read_csv(output_filename, sep='\t')
+        
+        # Change back to the current folder     
+        dir_path = os.path.dirname(os.path.realpath(__file__))
+        os.chdir(dir_path)
+        
+        return df
+    
+    def __clean_df(self, results: pd.DataFrame) -> pd.DataFrame:
+        """ 
+        Clean the proteinfer formatted file
+        """ 
+        results['predicted_ecs'] = [ec.split(':')[1] if 'EC:' in ec else 'None' for ec in results['predicted_label'].values]
+        # Remobe missing ECs
+        results = results[results['predicted_ecs'] != 'None']
+        
+        # ------------- Separate out ECs ------------------
+        results['EC1'] = [r.split('.')[0] for r in results['predicted_ecs'].values]
+        results['EC2'] = [r.split('.')[1] for r in results['predicted_ecs'].values]
+        results['EC3'] = [r.split('.')[2] for r in results['predicted_ecs'].values]
+        results['EC4'] = [r.split('.')[3] for r in results['predicted_ecs'].values]
+        # Filter to only have one EC per seqeunce 
+        # ------------- Group ------------------
+        # Now we want to group by the sequence_name and keep only the highest confidence level assignment
+        df = results.groupby('sequence_name')
+        rows = []
+        for grp in df:
+            top_row = grp[1].sort_values(by='predicted_label', ascending=False).values[0]
+            rows.append(top_row)
+        df = pd.DataFrame(rows, columns=results.columns)
+                
+        # ------------- Filter to EC XXXX ------------------
+        if self.ec1_filter is not None:
+            df = df[df['EC1'].isin(self.ec1_filter)]
+        if self.ec2_filter is not None:
+            df = df[df['EC2'].isin(self.ec2_filter)]
+        if self.ec3_filter is not None:
+            df = df[df['EC3'].isin(self.ec3_filter)]
+        if self.ec4_filter is not None:
+            df = df[df['EC4'].isin(self.ec4_filter)]
+        return df
+        
+    def execute(self, df: pd.DataFrame) -> pd.DataFrame:
+        with TemporaryDirectory() as tmp_dir:
+            if self.num_threads > 1:
+                data = []
+                df_list = np.array_split(df, self.num_threads)
+                pool = ThreadPool(self.num_threads)
+                for df_chunk in df_list:
+                    data.append([df_chunk, tmp_dir])
+                results = pool.map(self.__execute, data)
+                df = pd.DataFrame()
+                for dfs in results:
+                    df = pd.concat([df, dfs])
+                #df = self.__clean_df(df)
+                return df
+            else:
+                df = self.__execute([df, tmp_dir])
+                #df = self.__clean_df(df)
+                return df

enzymetk/dock_chai_step.py

@@ -0,0 +1,51 @@
+from enzymetk.step import Step
+import pandas as pd
+from docko.chai import run_chai
+import logging
+import numpy as np
+
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+    
+class Chai(Step):
+    
+    def __init__(self, id_col: str, seq_col: str, substrate_col: str, output_dir: str, num_threads: int):
+        self.id_col = id_col
+        self.seq_col = seq_col
+        self.substrate_col = substrate_col
+        self.output_dir = output_dir or None
+        self.num_threads = num_threads or 1
+
+    def __execute(self, df: pd.DataFrame, tmp_dir: str) -> pd.DataFrame:
+        output_filenames = []
+        for run_id, seq, substrate in df[[self.id_col, self.seq_col, self.substrate_col]].values:
+            # Might have an issue if the things are not correctly installed in the same dicrectory 
+            if not isinstance(substrate, str):
+                substrate = ''
+            print(run_id, seq, substrate)
+            run_chai(run_id, # name
+                    seq, # sequence
+                    substrate, # ligand as smiles
+                    tmp_dir)
+            output_filenames.append(f'{tmp_dir}/{run_id}/')
+        return output_filenames
+    
+    def execute(self, df: pd.DataFrame) -> pd.DataFrame:
+        if self.output_dir:
+            if self.num_threads > 1:
+                output_filenames = []
+                df_list = np.array_split(df, self.num_threads)
+                for df_chunk in df_list:
+                    output_filenames += self.__execute(df_chunk, self.output_dir)
+                    
+                df['output_dir'] = output_filenames
+                return df
+            
+            else:
+                output_filenames = self.__execute(df, self.output_dir)
+                df['output_dir'] = output_filenames
+                return df
+        else:
+            print('No output directory provided')

enzymetk/dock_vina_step.py

@@ -0,0 +1,63 @@
+from enzymetk.step import Step
+import pandas as pd
+from docko.docko import *
+import logging
+import numpy as np
+import os
+from multiprocessing.dummy import Pool as ThreadPool
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+    
+class Vina(Step):
+    
+    def __init__(self, id_col: str, structure_col: str, sequence_col: str, 
+                 substrate_col: str, substrate_name_col: str, active_site_col: str, output_dir: str, num_threads: int):
+        print('Expects active site residues as a string separated by |. Zero indexed.')
+        self.id_col = id_col
+        self.structure_col = structure_col
+        self.sequence_col = sequence_col
+        self.substrate_col = substrate_col
+        self.substrate_name_col = substrate_name_col
+        self.active_site_col = active_site_col  # Expects active site residues as a string separated by |
+        self.output_dir = output_dir or None
+        self.num_threads = num_threads or 1
+
+    def __execute(self, df: pd.DataFrame) -> pd.DataFrame:
+        output_filenames = []
+        # ToDo: update to create from sequence if the path doesn't exist.
+        for label, structure_path, seq, substrate_smiles, substrate_name, residues in df[[self.id_col, self.structure_col, self.sequence_col, self.substrate_col, self.substrate_name_col, self.active_site_col]].values:
+            os.system(f'mkdir {self.output_dir}{label}')
+            try:
+                residues = str(residues)
+                residues = [int(r) + 1 for r in residues.split('|')]
+                if not os.path.exists(f'{structure_path}'):
+                    # Try get the AF2 structure we expect the label to be the uniprot id
+                    get_alphafold_structure(label, f'{self.output_dir}{label}/{label}_AF2.pdb')
+                    structure_path = f'{self.output_dir}{label}/{label}_AF2.pdb'
+                clean_one_pdb(f'{structure_path}', f'{self.output_dir}{label}/{label}.pdb')
+                pdb_to_pdbqt_protein(f'{self.output_dir}{label}/{label}.pdb', f'{self.output_dir}{label}/{label}.pdbqt')
+                score = dock(sequence='', protein_name=label, smiles=substrate_smiles, ligand_name=substrate_name, residues=residues, 
+                            protein_dir=f'{self.output_dir}', ligand_dir=f'{self.output_dir}', output_dir=f'{self.output_dir}{label}/', pH=7.4,
+                            method='vina', size_x=10.0, size_y=10.0, size_z=10.0)
+                output_filename = f'{self.output_dir}{label}/{label}.pdb'
+                output_filenames.append(output_filename)
+            except Exception as e:
+                print(f'Error docking {label}: {e}')
+                output_filenames.append(None)
+        return output_filenames
+        
+    
+    def execute(self, df: pd.DataFrame) -> pd.DataFrame:
+        if self.output_dir:
+            if self.num_threads > 1:
+                pool = ThreadPool(self.num_threads)
+                df_list = np.array_split(df, self.num_threads)
+                results = pool.map(self.__execute, df_list)
+            else:
+                results = self.__execute(df)
+            df['output_dir'] = results
+            return df
+        else:
+            print('No output directory provided')

enzymetk/embedchem_chemberta_step.py

@@ -0,0 +1,61 @@
+from enzymetk.step import Step
+import pandas as pd
+import numpy as np
+from multiprocessing.dummy import Pool as ThreadPool
+from transformers import AutoModel, AutoTokenizer
+
+class ChemBERT(Step):
+    
+    def __init__(self, id_col: str, value_col: str, num_threads: int):
+        self.id_col = id_col
+        self.value_col = value_col
+        self.num_threads = num_threads
+        model_version = 'seyonec/PubChem10M_SMILES_BPE_450k'
+        self.model = AutoModel.from_pretrained(model_version, output_attentions=True)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_version)
+        self.seq_len_limit = 500
+        self.embedding_len = 768
+        
+        
+    def __execute(self, data: list) -> np.array:
+        results = []
+        for v in data:
+            i, smiles = v[0], v[1]
+            print(smiles)
+            encoded_input = self.tokenizer(
+                smiles,
+                truncation=True,
+                max_length=self.seq_len_limit,
+                padding='max_length',
+                return_tensors='pt')
+            output = self.model(**encoded_input)
+            results.append((i, output['last_hidden_state'][:, 0][0].detach().numpy())) 
+        return results
+
+    
+    def execute(self, df: pd.DataFrame) -> pd.DataFrame:
+        if self.num_threads > 1:
+            data = []
+            df_list = np.array_split(df, self.num_threads)
+            pool = ThreadPool(self.num_threads)
+            for df_chunk in df_list:
+                data.append([(i, v) for i, v in df_chunk[[self.id_col, self.value_col]].values])
+            results = pool.map(self.__execute, data)
+            all_results_map = {}
+            for r in results:
+                for j in r:
+                    all_results_map[j[0]] = j[1]
+            encodings = []
+            for uid in df[self.id_col].values:
+                if all_results_map.get(uid) is None:
+                    encodings.append(np.zeros(self.embedding_len))
+                else:
+                    encodings.append(all_results_map.get(uid))
+            df['chemberta'] = encodings
+            return df
+        
+        else:
+            data = [(i, v) for i, v in df[[self.id_col, self.value_col]].values]
+            results = self.__execute(data)
+            df['chemberta'] = [r[1] for r in results]
+            return df

enzymetk/embedchem_rxnfp_run.py

@@ -0,0 +1,28 @@
+from rxnfp.transformer_fingerprints import RXNBERTFingerprintGenerator, get_default_model_and_tokenizer
+import pandas as pd
+import pickle
+import argparse
+
+
+def run_rxnfp(output_filename, input_filename, label):
+    df = pd.read_csv(input_filename)
+    rxns = df[label].values
+    model, tokenizer = get_default_model_and_tokenizer()
+    rxnfp_generator = RXNBERTFingerprintGenerator(model, tokenizer)
+    fps = rxnfp_generator.convert_batch(rxns)
+    df['rxnfp'] = fps
+    with open(output_filename, 'wb') as file:
+        pickle.dump(df, file)
+        
+def parse_args():
+    parser = argparse.ArgumentParser(description="Run rxnfp on a dataset")
+    parser.add_argument('-out', '--out', required=True, help='Path to the output directory')
+    parser.add_argument('-input', '--input', type=str, required=True, help='path to the dataframe')
+    parser.add_argument('-label', '--label', type=str, required=True, help='label of the column')
+    return parser.parse_args()
+
+def main():
+    args = parse_args()
+    run_rxnfp(args.out, args.input, args.label)
+
+main()

enzymetk/embedchem_rxnfp_step.py

@@ -0,0 +1,55 @@
+from enzymetk.step import Step
+import pandas as pd
+from tempfile import TemporaryDirectory
+import pickle
+import subprocess
+from pathlib import Path
+import logging
+import numpy as np
+from tqdm import tqdm
+import random
+import string
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+    
+class RxnFP(Step):
+    
+    def __init__(self, smiles_col: str, num_threads: int, env_name: str = 'rxnfp'):
+        self.value_col = smiles_col
+        self.num_threads = num_threads or 1
+        self.env_name = env_name
+
+    def __execute(self, df: pd.DataFrame, tmp_dir: str) -> pd.DataFrame:
+        tmp_label = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
+
+        output_filename = f'{tmp_dir}/rxnfp_{tmp_label}.pkl'
+        input_filename = f'{tmp_dir}/input_{tmp_label}.csv'
+        df.to_csv(input_filename, index=False)
+        cmd = ['conda', 'run', '-n', self.env_name, 'python', Path(__file__).parent/'embedchem_rxnfp_run.py', '--out', output_filename, 
+                                '--input', input_filename, '--label', self.value_col]
+        self.run(cmd)
+        # Might have an issue if the things are not correctly installed in the same dicrectory 
+        
+        return output_filename
+    
+    def execute(self, df: pd.DataFrame) -> pd.DataFrame:
+        with TemporaryDirectory() as tmp_dir:
+            if self.num_threads > 1:
+                output_filenames = []
+                df_list = np.array_split(df, self.num_threads)
+                for df_chunk in tqdm(df_list, total=len(df_list)):
+                    output_filenames.append(self.__execute(df_chunk, tmp_dir))
+                    
+                df = pd.DataFrame()
+                for p in output_filenames:
+                    with open(f'{p}', 'rb') as file:
+                        tmp_df = pickle.load(file)
+                    df = pd.concat([df, tmp_df])
+                return df
+            
+            else:
+                output_filename = self.__execute(df, tmp_dir)
+                with open(f'{output_filename}', 'rb') as file:
+                    return pickle.load(file)