S(m)AMBA

Authors

• Vladyslav Humennyy
• Volodymyr Kuzma
• Roman Naumenko

Description

Custom implementation of Mamba state space model pretrained for NLP tasks.

S(m)AMBA reflects our vision of how S4 and Mamba should look like. We trained this model on four different datasets:

1. Shakespeare - collection of Shakespeare's writings.
2. Harry Tinny - "Boy Who Lived" chapter of the first book in the famous series.
3. Harry 1 - entire first book.
4. Harry Full - first 4 books in the series.

Demo

What we have achieved with this network:

Usage

Training

python train.py --config configs/<model>_config.yaml

For config file you can chose one of already available configs (in the configs directory), or create or own.

Inference

For inference you can also choose from predifined configs with pretrained models.

python inference.py \
    --config configs/<model>_config.yaml \
    --checkpoint weights/<model>.pth \
    --string "Example prompt." \
    --length 256

License

Unlicensed