kan-gpt

kan-gpt

The PyTorch implementation of Generative Pre-trained Transformers (GPTs) using Kolmogorov-Arnold Networks (KANs) for language modeling

Stars: 663

Visit
 screenshot

The KAN-GPT repository is a PyTorch implementation of Generative Pre-trained Transformers (GPTs) using Kolmogorov-Arnold Networks (KANs) for language modeling. It provides a model for generating text based on prompts, with a focus on improving performance compared to traditional MLP-GPT models. The repository includes scripts for training the model, downloading datasets, and evaluating model performance. Development tasks include integrating with other libraries, testing, and documentation.

README:

KAN-GPT

PyPI - Downloads PyPI - Version codecov CI GitHub License

The PyTorch implementation of Generative Pre-trained Transformers (GPTs) using Kolmogorov-Arnold Networks (KANs) for language modeling

Install it from PyPI

pip install kan_gpt

Citation

If you find our work useful cite us!

@misc{GANESH2024KANGPT,
  author       = {Aditya Nalgunda Ganesh},
  title        = {KAN-GPT: The PyTorch implementation of Generative Pre-trained Transformers (GPTs) using Kolmogorov-Arnold Networks (KANs) for language modeling},
  year         = {2024},
  month        = {May},
  note         = {Release 1.0.0, 9th May 2024},
  url          = {https://github.com/AdityaNG/kan-gpt/}
}

Usage

Refer to the KAN_GPT.ipynb and kan_gpt/prompt.py for usage examples. The following is an outline of how to use the model:

from kan_gpt.model import GPT
from transformers import GPT2Tokenizer

model_config = GPT.get_default_config()
model_config.model_type = "gpt2"
model_config.vocab_size = 50257
model_config.block_size = 1024
model = GPT(model_config)

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

prompt = "Bangalore is often described as the "

prompt_encoded = tokenizer.encode(
  text=prompt, add_special_tokens=False
)

x = torch.tensor(prompt_encoded).unsqueeze(0)

model.eval()
y = model.generate(x, 50)  # sample 50 tokens

result = tokenizer.decode(y[0])

print(result)

# Bangalore is often described as the Silicon Valley of India.
# The city has witnessed rapid growth in the past two decades.....

Setup for Development

# Download Repo
git clone https://github.com/AdityaNG/kan-gpt
cd kan-gpt
git pull

# Download Dataset
./scripts/download_webtext.sh
./scripts/download_tinyshakespeare.sh

# Install dependencies for development
pip install -r requirements.txt
pip install -e .

Train

Use the following dummy script to make sure everything is working as expected

WANDB_MODE=offline CUDA_VISIBLE_DEVICE="" python3 -m kan_gpt.train --architecture MLP --batch_size 1 --dummy_dataset --device cpu --max_iters 200
WANDB_MODE=offline CUDA_VISIBLE_DEVICE="" python3 -m kan_gpt.train --architecture KAN --batch_size 1 --dummy_dataset --device cpu --max_iters 200

Then make use of the training script

python -m kan_gpt.train

Prompt

You can prompt the model to produce text as follows

python -m kan_gpt.prompt --prompt "Bangalore is often described as the " --model_path (checkpoint)

Results

We train and compare KAN-GPT with an equivalent MLP-GPT model on the Tiny Shakespeare dataset. We observe that the KAN-GPT performs slightly better than the MLP-GPT. We are looking into further experiments to dive deeper. The results are shown below:

Metrics
results_loss results_cross_entropy results_perplexity

TODOs

  • [x] Integrate minGPT and pykan
  • [x] Dataset downloading script for WebText
  • [x] PyTorch Dataset parser for WebText
  • [x] PyTorch Dataset parser for tinyshakespeare
  • [x] Mini training POC for KAN-GPT
    • [x] Integrate KAN training logic from KAN.train_kan
    • [x] Train a dummy batch w/o any memory issues
  • [x] Mini training POC for MLP-GPT
  • [x] Train MLP-GPT on the webtext dataset as a baseline
  • [x] Train KAN-GPT on the webtext dataset as a baseline
  • [x] Metrics comparing KAN-GPT and MLP-GPT
  • [x] Auto Save checkpoints
  • [x] Auto Save checkpoints to W&B
  • [ ] Auto Download model weights from git / huggingface
  • [x] W&B hyperparam sweep script
  • [x] Script to load checkpoint in interactive mode
  • [ ] Reduce requrements.txt constraints
  • [ ] Define pydantic model for training and sweep args
  • [ ] Pruning the package, get rid of unused code
  • [ ] Training script to PyTorch Lighting
  • [x] Documentation: mkdocs gh-deploy
  • [x] Integrate with efficient-kan
  • [x] Test Cases
    • [x] KAN: Forward-Backward test
    • [x] GPT: Forward-Backward test
    • [x] KAN_GPT: Forward-Backward test
    • [x] EFFICIENT_KAN: Forward-Backward test

Development

Read the CONTRIBUTING.md file.

References

For Tasks:

Click tags to check more tools for each tasks

For Jobs:

Alternative AI tools for kan-gpt

Similar Open Source Tools

For similar tasks

For similar jobs