Exponential increases in microbial and viral genomic data demand transformational advances in scalable, generalizable frameworks for their interpretation. Standard homology-based functional analyses are hindered by the rapid divergence of microbial and especially viral genomes and proteins that significantly decreases the volume of usable data. Here, we present Protein Set Transformer (PST), a protein-based genome language model that models genomes as sets of proteins without considering sparsely available functional labels. Trained on >100k viruses, PST outperformed other homology- and language model-based approaches for relating viral genomes based on shared protein content. Further, PST demonstrated protein structural and functional awareness by clustering capsid-fold-containing proteins with known capsid proteins and uniquely clustering late gene proteins within related viruses. Our data establish PST as a valuable method for diverse viral genomics, ecology, and evolutionary applications. We posit that the PST framework can be a foundation model for microbial genomics when trained on suitable data.

Genomes used to create this dataset are publicly available, and all data within this dataset were generated by the study. See the manuscript for details.

Changes after Sep 19, 2024:

• Added

  • foldseek_databases.tar.gz
    • Precomputed foldseek 3Di databases for each test dataset
  • PST-TL-P__small.ckpt.gz
    • New pretrained model checkpoint for model trained with triplet loss and tuned with protein diversity groups
  • PST-TL-P__large.ckpt.gz
    • New pretrained model checkpoint for model trained with triplet loss and tuned with protein diversity groups
  • PST-MLM.tar.gz
    • New pretrained model checkpoints for models trained with masked language modeling loss
  • PST_training_set_PST-TL-P__large_protein_embeddings.h5
  • IMGVRv4_test_set_PST-TL-P__large_protein_embeddings.h5
  • MGnify_set_PST-TL-P__large_protein_embeddings.h5
  • PST_training_set_PST-TL-P__small_protein_embeddings.h5
  • IMGVRv4_test_set_PST-TL-P__small_protein_embeddings.h5
  • MGnify_set_PST-TL-P__small_protein_embeddings.h5

• Changed

  • esm-large_protein_embeddings.tar.gz
    • Now part of esm_embeddings.tar.gz
    • Includes MGnify test set
  • esm-small_protein_embeddings.tar.gz
    • Now part of esm_embeddings.tar.gz
    • Includes MGnify test set
  • fasta.tar.gz
    • Includes MGnify test set
  • genome_clusters.tar.gz
    • Includes MGnify test set
  • genome_embeddings.tar.gz
    • Split into different files:
      • PST-TL_genome_embeddings.tar.gz contain all PST-TL genome embeddings for each dataset
      • other_genome_embeddings.tar.gz contain all others
  • genslm_protein_embeddings.tar.gz
    • Converted into a single .h5 file called genslm_ORF_embeddings.h5
    • Includes MGnify test set
  • host_prediction.tar.gz
    • Knowledge graphs were reconstructed using a different vector similarity search method
    • Retrained models using new genome embeddings
  • protein_clusters.tar.gz
    • Includes MGnify test set
  • pst-large_protein_embeddings.tar.gz
    • Split into dataset specific files for easier access for each dataset:
      • PST_training_set_PST-TL-T__large_protein_embeddings.h5
      • IMGVRv4_test_set_PST-TL-T__large_protein_embeddings.h5
      • MGnify_set_PST-TL-T__large_protein_embeddings.h5
  • pst-small_protein_embeddings.tar.gz
    • Split into dataset specific files for easier access for each dataset:
      • PST_training_set_PST-TL-T__small_protein_embeddings.h5
      • IMGVRv4_test_set_PST-TL-T__small_protein_embeddings.h5
      • MGnify_set_PST-TL-T__small_protein_embeddings.h5
  • supplementary_data.tar.gz
    • Most figures were modified, so all supplementary datasets changed to reflect changes in manuscript
  • supplementary_tables.zip
    • Added 3 new supplementary tables and included the MGnify test dataset in the existing tables
  • trained_models.tar.gz
    • Split into separate files for easier access of individual models:
      • PST-TL-T__small.ckpt.gz
      • PST-TL-T__large.ckpt.gz

• Removed

  • aai.tar.gz
    • These were originally raw protein-protein alignments for the IMGVRv4 dataset
    • These have been summarized in supplementary_data.tar.gz
    • But the raw alignments had to be removed to make more storage

Changes after Jun 18, 2025:

• Added

  • supplementary_data.tar.gz now refers to Supplementary Data files as described in the manuscript. These are tables that are too large to be submitted at tables.

• Changed

  • source_data.tar.gz now refers to the source data used for the figures.

    • Each file has been renamed to refer to the specific figure panel.

    Old file	New file
    supplementary_tables/supplementary_table_1.tsv	supplementary_data/supplementary_data_1.tsv
    supplementary_tables/supplementary_table_2.tsv	supplementary_data/supplementary_data_2.tsv
    supplementary_tables/supplementary_table_3.tsv	supplementary_tables/supplementary_table_1.tsv
    supplementary_tables/supplementary_table_4.tsv	supplementary_tables/supplementary_table_2.tsv
    supplementary_tables/supplementary_table_5.tsv	supplementary_data/supplementary_data_3.tsv
    supplementary_tables/supplementary_table_6.tsv	supplementary_data/supplementary_data_4.tsv
    supplementary_tables/supplementary_table_7.tsv	supplementary_tables/supplementary_table_5.tsv
    supplementary_tables/supplementary_table_8.tsv	supplementary_data/supplementary_data_5.tsv
    supplementary_tables/supplementary_table_9.tsv	supplementary_tables/supplementary_table_6.tsv
    supplementary_tables/supplementary_table_10.tsv	supplementary_tables/supplementary_table_3.tsv
    supplementary_tables/supplementary_table_11.tsv	supplementary_tables/supplementary_table_4.tsv

Changes after Sep 30, 2025:

• Changed
  • The following files were reuploaded since there was an issue with the ordering of the embeddings:
    • IMGVRv4_test_set_PST-TL-T__large_protein_embeddings.h5
    • IMGVRv4_test_set_PST-TL-T__small_protein_embeddings.h5
    • IMGVRv4_test_set_PST-TL-P__large_protein_embeddings.h5
    • IMGVRv4_test_set_PST-TL-P__small_protein_embeddings.h5
    • MGnify_test_set_PST-TL-T__large_protein_embeddings.h5
    • MGnify_test_set_PST-TL-T__small_protein_embeddings.h5
    • MGnify_test_set_PST-TL-P__large_protein_embeddings.h5
    • MGnify_test_set_PST-TL-P__small_protein_embeddings.h5
    • esm_embeddings.tar.gz
    • PST-TL_genome_embeddings.tar.gz
    • genome_clusters.tar.gz
    • protein_clusters.tar.gz
    • other_genome_embeddings.tar.gz
    • host_prediction.tar.gz
  • The format of all data has not changed, just the actual content to reflect the order of genomes/proteins in Supplementary Tables 1 & 2.