PyxLSTM

Introduction to PyxLSTM

PyxLSTM is a cutting-edge Python library designed to provide an efficient and scalable implementation of the Extended Long Short-Term Memory (xLSTM) architecture. This innovative framework builds on the foundation of the traditional LSTM by incorporating several enhancements, namely exponential gating, memory mixing, and a matrix memory structure. These enhancements result in better performance and scalability for various sequence modeling tasks, making PyxLSTM a powerful tool for developers and researchers alike.

Features

PyxLSTM comes packed with numerous features that make it ideal for a wide range of applications:

• Variants: The library implements both scalar LSTM (sLSTM) and matrix LSTM (mLSTM) variants of the xLSTM architecture.
• Flexible Architecture: It supports the use of pre and post up-projection block structures, allowing for the creation of flexible model architectures.
• Ease of Use: High-level model definition and training utilities are provided for user convenience.
• Comprehensive Tools: Scripts for training, evaluation, and text generation are included, alongside data processing utilities and customizable dataset classes.
• Lightweight Design: The library’s modular and lightweight design ensures seamless integration into existing projects.
• Reliability: PyxLSTM is extensively tested and well-documented, ensuring both reliability and usability.
• Versatility: It is suitable for diverse sequence modeling tasks, from language modeling to text generation and more.

Installation

Getting started with PyxLSTM is simple. For a regular installation, users can utilize pip:

pip install PyxLSTM

For development purposes, including testing dependencies:

pip install PyxLSTM[dev]

Alternatively, users can clone the repository directly from GitHub:

git clone https://github.com/muditbhargava66/PyxLSTM.git
cd PyxLSTM
pip install -r requirements.txt
pip install -e .

Usage

The library is designed to simplify the process of using xLSTM for language modeling and other tasks. Below is a basic example illustrating how to employ PyxLSTM for language modeling:

import torch
from xLSTM.model import xLSTM
from xLSTM.data import LanguageModelingDataset, Tokenizer
from xLSTM.utils import load_config, set_seed, get_device
from xLSTM.training import train

# Load configuration
config = load_config("path/to/config.yaml")
set_seed(config.seed)
device = get_device()

# Initialize tokenizer and dataset 
tokenizer = Tokenizer(config.vocab_file)
train_dataset = LanguageModelingDataset(config.train_data, tokenizer, config.max_length)

# Create xLSTM model
model = xLSTM(len(tokenizer), config.embedding_size, config.hidden_size,
              config.num_layers, config.num_blocks, config.dropout,
              config.bidirectional, config.lstm_type)
model.to(device)

# Train the model
optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
criterion = torch.nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)
train(model, train_dataset, optimizer, criterion, config, device)

For more examples and detailed usage, users can refer to the comprehensive documentation provided in the documentation folder of the project.

Running and Testing

To run PyxLSTM effectively, users should follow these straightforward steps:

1. Clone the repository and navigate to the project's directory.
2. Install all necessary dependencies.
3. Execute the unit tests to ensure the proper functioning of various modules using the unittest framework.

git clone https://github.com/muditbhargava66/PyxLSTM.git
cd PyxLSTM
pip install -r requirements.txt
python -m unittest discover tests

This sequence will verify the functionality of the library through several test modules, ensuring that each component of PyxLSTM is in optimal working condition.

Conclusion

PyxLSTM is a promising library for those interested in advanced sequence modeling methodologies. Its flexible architecture, coupled with robust design and documentation, make it an excellent resource for developers and researchers aiming to enhance their projects using the latest advancements in xLSTM technology. For any additional queries or contributions, stakeholders are encouraged to reach out to the maintainer or explore the GitHub repository.