shivendrra/BiosaicTokenizer

Biosaic Tokenizer

Overview

Biosaic(Bio-Mosaic) is a tokenizer library built for Enigma2. It contains: Tokenizer, Embedder for DNA & Amino Acid Protein Sequences. Has a VQ-VAE & Evoformer architecture based encoders that could convert sequences into embeddings and vice-versa for model training use-case.

Features

• Tokenization: converts the sequences into K-Mers. (for DNA only)
• Encoding: converts sequences into embeddings for classification, training purposes.
• Easy use: it's very basic and easy to use library.
• SoTA encoder: Evoformer & VQ-VAE model are inspired from the AlphaFold-2

Models

It has two different Models,

• for DNA tokenization & encoding: VQ-VAE
• for Protein Encodings: EvoFormer

VQ-VAE is around 160M parameter big(for now it's just around 40M just to test run). EvoFormer is around 136M parameter big (still in training).

Config:

class ModelConfig:
  d_model: int= 768
  in_dim: int= 4
  beta: float= 0.15
  dropout: float= 0.25
  n_heads: int= 16
  n_layers: int= 12

class ModelConfig:
  DEVICE       = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  A            = 4        # DNA alphabet
  C            = 21       # 21 letter for amino acid & 4 for dna
  d_msa        = 768
  d_pair       = 256
  n_heads      = 32
  n_blocks     = 28

Training:

For training the VQ-VAE & Evo-Former model, batch training is preferred, with it's own sepearte Dateset class that takes input of the strings and then Hot-encodes the DNA Sequences first and then fill them into batches according to train & val splits which is around 20% of the full dataset.

For VQ-VAE:

class TrainConfig:
  device        = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  learning_rate = 1e-4         # bumped from 1e-5
  weight_decay  = 1e-4
  amsgrad       = True
  warmup_epochs = 50           # linear warm‑up
  epochs        = 2000
  eval_interval = 100
  eval_iters    = 30
  batch_size    = 6
  block_size    = 256

For EvoFormer:

class TrainConfig:
  DEVICE       = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  LR           = 1e-4
  WD           = 1e-4
  AMS          = True
  WARMUP       = 50
  EPOCHS       = 500
  BATCH        = 8
  MSA_SEQ      = 32       # number of sequences in each MSA
  L_SEQ        = 256      # length of each sequence
  EVAL_ITERS   = 5
  EVAL_INTV    = 50