AstrAI
A lightweight Transformer training & inference framework
## 📖 Table of Contents
English
- [Features](#features)
- [Quick Start](#quick-start)
- [Documentation](#documentation)
- [Contributing](#contributing)
- [Community](#community)
- [License](#license)
---
## English
### Features
- 🚀 **High Performance**: Optimized for both training and inference with efficient parallelization.
- 🔧 **Flexible**: Support for seq/sft/dpo training, customizable model architectures.
- 💡 **Easy to Use**: Simple API with comprehensive examples and demos.
- 📦 **Lightweight**: Minimal dependencies, easy to deploy.
- 🔬 **Research‑Friendly**: Modular design, easy to experiment with new ideas.
- 🤗 **HuggingFace Integration**: Compatible with HuggingFace models and datasets.
### Quick Start
#### Installation
```bash
git clone https://github.com/ViperEkura/AstrAI.git
cd AstrAI
pip install -e .
```
For development dependencies:
```bash
pip install -e ".[dev]"
```
#### Train a Model
```bash
python tools/train.py \
--train_type=seq \
--data_root_path=/path/to/dataset \
--param_path=/path/to/param_path
```
#### Generate Text
```bash
python tools/generate.py --param_path=/path/to/param_path
```
#### Demo
Check out the demos in the `demo/` folder:
```bash
# Download pre‑processed data (required before running demos)
python demo/download.py
# Interactive streaming chat
python demo/stream_chat.py
# Batch generation
python demo/generate_batch.py
# Auto‑regressive generation
python demo/generate_ar.py
```
Watch a video walkthrough on [bilibili](https://www.bilibili.com/video/BV1z5RPYHEkd).
### Documentation
| Document | Description |
|----------|-------------|
| [Parameter Guide](./assets/docs/params.md) | Training & inference parameters |
| [Design Document](./assets/docs/design.md) | Framework architecture & module design |
| [Data Flow](./assets/docs/dataflow.md) | Data processing pipeline details |
| [Model Introduction](./assets/docs/introduction.md) | Model architecture & technical details |
### Contributing
We welcome contributions! Please see our [Contributing Guidelines](CONTRIBUTING.md) for details.
1. Fork the repository.
2. Create a feature branch.
3. Commit your changes.
4. Open a Pull Request.
For major changes, please open an issue first to discuss what you would like to change.
### Community
- **GitHub Issues**: [Issue Tracker](https://github.com/ViperEkura/AstrAI/issues)
- **Discussions**: [GitHub Discussions](https://github.com/ViperEkura/AstrAI/discussions)
- **HuggingFace**: [Model Hub](https://huggingface.co/ViperEk/AstrAI)
### License
This project is licensed under the [GPL-3.0 License](LICENSE).
---
A lightweight Transformer framework designed for both high performance and ease of use.
