PyxLSTM

欢迎PR GitHub发布最后提交

PyxLSTM是一个Python库，提供了基于Beck等人（2024）的研究论文"xLSTM: 扩展长短期记忆"的高效和可扩展的扩展长短期记忆（xLSTM）架构实现。xLSTM通过引入指数门控、记忆混合和矩阵记忆结构来增强传统LSTM，从而提高序列建模任务的性能和可扩展性。

特性

实现了xLSTM的sLSTM（标量LSTM）和mLSTM（矩阵LSTM）变体
支持前置和后置上投影块结构，实现灵活的模型架构
提供高级模型定义和训练工具，易于使用
包含用于训练、评估和文本生成的脚本
提供数据处理工具和可自定义的数据集类
轻量级和模块化设计，可无缝集成到现有项目中
经过广泛测试和文档记录，确保可靠性和可用性
适用于广泛的序列建模任务，包括语言建模、文本生成等

安装

要安装PyxLSTM，可以使用pip：

pip install PyxLSTM

开发安装

对于包含测试依赖项的开发安装：

pip install PyxLSTM[dev]

或者，您可以克隆仓库并手动安装：

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

使用方法

以下是使用PyxLSTM进行语言建模的基本示例：

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  # 假设train函数在training模块中定义

# 加载配置
config = load_config("path/to/config.yaml")
set_seed(config.seed)
device = get_device()

# 初始化分词器和数据集
tokenizer = Tokenizer(config.vocab_file)
train_dataset = LanguageModelingDataset(config.train_data, tokenizer, config.max_length)

# 创建xLSTM模型
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)

# 训练模型
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)

有关更详细的使用说明和示例，请参阅文档。

代码目录结构

xLSTM/
│
├── xLSTM/
│   ├── __init__.py
│   ├── slstm.py
│   ├── mlstm.py
│   ├── block.py
│   └── model.py
│
├── utils/
│   ├── config.py
│   ├── logging.py
│   └── utils.py
│
├── tests/
│   ├── test_slstm.py  
│   ├── test_mlstm.py
│   ├── test_block.py
│   └── test_model.py
│
├── docs/
│   ├── slstm.md
│   ├── mlstm.md
│   └── training.md
│
├── examples/
│   ├── language_modeling.py
│   └── xLSTM_shape_verification.py
│
├── .gitignore
├── pyproject.toml
├── MANIFEST.in
├── requirements.txt
├── README.md
└── LICENSE

xLSTM/: 主要的Python包，包含实现代码。
- slstm.py: sLSTM模块的实现。
- mlstm.py: mLSTM模块的实现。
- block.py: xLSTM块的实现（上投影前后）。
- model.py: xLSTM模型的高级定义。
utils/: 实用工具模块。
- config.py: 配置管理。
- logging.py: 日志设置。
- utils.py: 杂项实用函数。
tests/: 不同模块的单元测试。
- test_slstm.py: sLSTM模块的测试。
- test_mlstm.py: mLSTM模块的测试。
- test_block.py: xLSTM块的测试。
- test_model.py: 整体xLSTM模型的测试。
docs/: 文档文件。
- README.md: 主要文档文件。
- slstm.md: sLSTM的文档。
- mlstm.md: mLSTM的文档。
- training.md: 训练指南。
.gitignore: Git忽略文件，用于排除不必要的文件/目录。
setup.py: 包安装脚本。
requirements.txt: Python依赖项列表。
README.md: 项目README文件。
LICENSE: 项目许可证文件。

运行和测试代码库

要运行和测试PyxLSTM代码库，请按以下步骤操作：

克隆PyxLSTM仓库：

git clone https://github.com/muditbhargava66/PyxLSTM.git

进入克隆的目录：
```
cd PyxLSTM
```
安装所需依赖：
```
pip install -r requirements.txt
```
运行单元测试：
```
python -m unittest discover tests
```
此命令将运行tests目录中的所有单元测试。它将执行测试文件test_slstm.py、test_mlstm.py、test_block.py和test_model.py。

如果遇到任何问题或有进一步的问题，请参阅PyxLSTM文档或联系维护者寻求帮助。

文档

PyxLSTM的文档可在docs目录中找到。它提供了关于库组件、使用指南和示例的详细信息。

引用

如果您在研究或项目中使用PyxLSTM，请引用原始xLSTM论文：

@article{Beck2024xLSTM,
  title={xLSTM: Extended Long Short-Term Memory},
  author={Beck, Maximilian and Pöppel, Korbinian and Spanring, Markus and Auer, Andreas and Prudnikova, Oleksandra and Kopp, Michael and Klambauer, Günter and Brandstetter, Johannes and Hochreiter, Sepp},
  journal={arXiv preprint arXiv:2405.04517},
  year={2024}
}

论文链接：https://arxiv.org/abs/2405.04517