Package: prokbert
The ProkBERT model familyProkBERT is an advanced genomic language model specifically designed for microbiome analysis. This repository contains the ProkBERT package and utilities, as well as the LCA tokenizer and model definitions.Table of ContentsIntroductionFeaturesInstallationInstalling with pipInstalling with condaUsing DockerUsing Singularity (Apptainer)ApplicationsQuick startTutorials and examplesTokenization and segmentationVisualizing sequence representationsFinetuning example for promoter sequencesEvaluation and inference examplePretraining exampleProkBERT's resultsCiting this workIntroductionThe ProkBERT model family is a transformer-based, encoder-only architecture based onBERT. Built on transfer learning and self-supervised methodologies, ProkBERT models capitalize on the abundant available data, demonstrating adaptability across diverse scenarios. The models’ learned representations align with established biological understanding, shedding light on phylogenetic relationships. With the novel Local Context-Aware (LCA) tokenization, the ProkBERT family overcomes the context size limitations of traditional transformer models without sacrificing performance or the information-rich local context. In bioinformatics tasks like promoter prediction and phage identification, ProkBERT models excel. For promoter predictions, the best-performing model achieved an MCC of 0.74 for E. coli and 0.62 in mixed-species contexts. In phage identification, they all consistently outperformed tools like VirSorter2 and DeepVirFinder, registering an MCC of 0.85. Compact yet powerful, the ProkBERT models are efficient, generalizable, and swift.FeaturesTailored to microbes.Local Context-Aware (LCA) tokenization for better genomic sequence understanding.Pre-trained models available for immediate use and fine-tuning.High performance in various bioinformatics tasks.Facilitation of both supervised and unsupervised learning.InstallationInstalling with pipThe recommended way to install ProkBERT is through pip, which will handle most dependencies automatically:pipinstallprokbertInstalling with condaProkBERT is also available as a conda package from the Bioconda channel. To install it using conda, run:condainstallprokbert-cbiocondaUsing DockerBefore using the ProkBERT container with GPU support, make sure you have the following installed on your system:Python (3.10 or later)Docker(required if you plan to use the Docker image)NVIDIA Docker(required if you intend to use Docker with GPU support)To pull and run the ProkBERT Docker image, use:dockerpullobalasz/prokbertTo run the container with GPU support, use:dockerrun--gpusall-it--rm-v$(pwd):/appobalasz/prokbertpython/app/finetuning.py--helpThis command runs the ProkBERT Docker container with GPU support, mounts the current directory to the container (allowing access to your local files), and displays the help message for thefinetuning.pyscript.Using Singularity (Apptainer)For users who prefer Singularity (now also known as Apptainer), a Singularity (Apptainer) container definition is provided. Please refer to the official documentation for instructions on building and running Singularity containers. A prebuilt container is available on Zenodo:https://zenodo.org/records/10659030.Building the Singularity container:apptainerbuildprokbert.sifprokbert.defTo pull directly from Docker Hub and convert to a Singularity image file:singularitypullprokbert.sifdocker://obalasz/prokbertOnce you have your.siffile, you can run ProkBERT with the following command:singularityrun--nvprokbert.sifpython/opt/prokbert/examples/finetuning.py--helpThe--nvflag enables NVIDIA GPU support, assuming your system has NVIDIA drivers and CUDA installed. Remove this flag if you're running on a CPU-only system.You can also shell into the container or execute specific commands as follows:
Shell into the container:singularityshell--nvprokbert.sifNote: If you encounter any problems, please do not hesitate to contact us or open an issue. My email address isobalasz@gmail.com.ApplicationsProkBERT has been validated in several key genomic tasks, including:Learning meaningful repreresentation for seqeuences (zero-shot capibility)Accurate bacterial promoter prediction.Detailed phage sequence analysis within complex microbiome datasets.Quick StartOur models and datasets are avaialble on thehugginface page.
The models are easy to use with thetransformerspackage.
We provide examples and descriptions as notebooks in the next chapter and some example scsripts regarging how to preprocess your sequence data and how to finetune the available models. The examples are available in theexamplefolder of this repository.TLDR exampleTo load the model from Hugging Face:importtorchfromtransformersimportAutoTokenizer,AutoModelfromprokbert.prokbert_tokenizerimportProkBERTTokenizertokenizer=ProkBERTTokenizer(tokenization_params={'kmer':6,'shift':1})model=AutoModel.from_pretrained("nerualbioinfo/prokbert-mini",trust_remote_code=True)segment="TATGTAACATAATGCGACCAATAATCGTAATGAATATGAGAAGTGTGATATTATAACATTTCATGACTACTGCAAGACTAA"inputs=tokenizer(segment)['input_ids']encoded_sequence=tokenizer.batch_encode_plus([segment])# Converting the ids to tensors of int64encoded_sequence={key:torch.tensor(val,dtype=torch.int64)forkey,valinencoded_sequence.items()}model(**encoded_sequence)Tutorials and examples:Tokenization and segmentationFor examples of how to preprocess the raw sequence data, which are frequently stored in fasta format:
Examples:Segmentation (sequence preprocessing):colab linkTokenizationcolab linkPreprocessing for pretrainingVisualizing sequence representationsAn example of how to visualize the genomic features of ESKAPE pathogens. More description about the dataset is available on Hugging Face
Example:ESKAPE pathogen genomic features:colab linkFinetuning example for promoter sequencesHere we provide an example of a practical transfer learning task. It is formulated as a binary classification. We provide a notebook for presenting the basic concepts and a command line script as a template.
Examples:Finetuning for promoter identification task:colab linkPython script for the finetuning:linkUsage example:gitclonehttps://github.com/nbrg-ppcu/prokbertcdexamples
pythonfinetuning.py\--model_nameneuralbioinfo/prokbert-mini\--ftmodelmini_promoter\--model_outputpathfinetuning_outputs\--num_train_epochs1\--per_device_train_batch_size128For practical applications or for larger training tasks we recommend using theDistributed DataParallel.Evaluation and inference exampleIn this section, we provide a practical example for evaluating finetuned ProkBERT models. This is crucial for understanding the model's performance on specific tasks, such as promoter prediction or phage identification.We have prepared a detailed example that guides you through the process of evaluating our finetuned models. This includes loading the model, preparing the data, running the inference, and interpreting the results.Example:Evaluation of the Finetuned ModelsPretraining exampleHere you can find an example of pretraining ProkBERT from scratch. Pretraining is an essential step, allowing the model to learn the underlying patterns before being fine-tuned for downstream tasks. All of the pretrained models are available on Hugging Face.Pretrained Models:Modelk-merShiftHugging Face URLProkBERT-mini61LinkProkBERT-mini-c11LinkProkBERT-mini-long62LinkPreprocessing the example data:It is a good practice to preprocess larger sequence sets (>100MB) in advance. Below is an example script for achieving this by running on multiple fasta files:Clone the ProkBERT repository and navigate to the examples directory. Then, preprocess the example data into a format suitable for training. This involves converting sequences into k-mer representations. Run the following commands:gitclonehttps://github.com/nbrg-ppcu/prokbertcdprokbert/examples
pythonprokbert_seqpreprocess.py\--kmer6\--shift1\--fasta_file_dir../src/prokbert/data/pretraining\--out../src/prokbert/data/preprocessed/pretraining_k6s1.h5Parameters:--kmer: The size of the k-mer (number of bases) to use for sequence encoding.--shift: The shift size for sliding the k-mer window across sequences.--fasta_file_dir: Directory containing your FASTA files for pretraining.--out: Output file path for the preprocessed data in HDF5 format.Running the pretraining from scratch:Use the preprocessed HDF file as input for pretraining. Execute the commands below:pythonprokbert_pretrain.py\--kmer6\--shift1\--dataset_path../src/prokbert/data/preprocessed/pretraining_k6s1.h5\--model_nameprokbert_k6s1\--output_dir./tmppretraining\--model_outputpath./tmppretrainingParameters:--model_name: Name for the model configuration to be used or saved.--output_dir: Directory where the training logs and temporary files will be saved.--model_outputpath: Path where the final trained model should be saved.ProkBERT’s resultsProkBERT is a novel genomic language model family designed for microbiome studies, pretrained on extensive genomic datasets from the NCBI RefSeq database. The pretraining covered a diverse range of organisms, utilizing over 206.65 billion tokens from various genomes, encompassing bacteria, viruses, archaea, and fungi. The final dataset included 976,878 unique contigs from 17,178 assemblies, representing 3,882 distinct genera. Detailed methodology and tokenization strategies are elaborated in ourpaper.Genomic featuresProkBERT's embeddings provide insights into genomic structures and relationships, distinguishing between different genomic elements. The models showcase strong generalization abilities, aligning with established biological principles.Figure 1: UMAP embeddings of genomic segment representations, highlighting the distinction between genomic features.Promoter identificationProkBERT models have been applied to identify bacterial promoters, demonstrating superior accuracy and robustness compared to existing tools.Figure 2: Performance metrics of ProkBERT in promoter prediction.ProkBERT models demonstrated high performance with an impressive accuracy and MCC (Matthews Correlation Coefficient) of 0.87 and 0.74, outperforming other established tools such as CNNProm and iPro70-FMWin. This highlights ProkBERT's effectiveness in correctly identifying both promoters and non-promoters, with consistent results across various model variants. The evaluation also included a comparative analysis with newer tools like Promotech and iPromoter-BnCNN.Phage predictionIn phage sequence analysis, ProkBERT outperforms traditional methods, proving its efficacy in identifying phage sequences within complex genomic data.Figure 3: Comparative analysis showcasing ProkBERT's performance in phage prediction.Our evaluations demonstrate the performance of ProkBERT in classifying phage sequences. It achieves high sensitivity and specificity even in challenging cases where available sequence information is limited. However, this exercise also highlights an inherent limitation of ProkBERT, the restricted context window size.
In comparative benchmarks with varying short sequence lengths, ProkBERT consistently surpassed established tools like VirSorter2 and DeepVirFinderAvailable models and datasetsPretrained modelsModel Namek-merShiftHugging Face URLneuralbioinfo/prokbert-mini61Linkneuralbioinfo/prokbert-mini-long62Linkneuralbioinfo/prokbert-mini-c11LinkFinetuned models for promoter predictionModel Namek-merShiftHugging Face URLneuralbioinfo/prokbert-mini-promoter61Linkneuralbioinfo/prokbert-mini-long-promoter62Linkneuralbioinfo/prokbert-mini-c-promoter11LinkFinetuned models for phage identificationModel Namek-merShiftHugging Face URLneuralbioinfo/prokbert-mini-phage61Linkneuralbioinfo/prokbert-mini-long-phage62Linkneuralbioinfo/prokbert-mini-c-phage11LinkDatasetsDataset NameHugging Face URLneuralbioinfo/ESKAPE-genomic-featuresLinkneuralbioinfo/phage-test-10kLinkneuralbioinfo/bacterial_promotersLinkneuralbioinfo/ESKAPE-maskingLinkCiting this workIf you use the code or data in this package, please cite:@Article{ProkBERT2024,author={Ligeti, Balázs and Szepesi-Nagy, István and Bodnár, Babett and Ligeti-Nagy, Noémi and Juhász, János},journal={Frontiers in Microbiology},title={{ProkBERT} family: genomic language models for microbiome applications},year={2024},volume={14},URL={https://www.frontiersin.org/articles/10.3389/fmicb.2023.1331233},DOI={10.3389/fmicb.2023.1331233},ISSN={1664-302X}}