pyproject.toml

[project]
name = "litgpt"
version = "0.3.0.dev0"
description = "Hackable implementation of state-of-the-art open-source LLMs"
authors = [
    { name = "Lightning AI", email = "contact@lightning.ai" },
]
readme = "README.md"
license = { file = "LICENSE" }

dependencies = [
    "torch==2.3.1",
    "lightning>=2.3.3",
    "jsonargparse[signatures]>=4.27.6",
]

[project.urls]
homepage = "https://github.com/lightning-AI/litgpt"
documentation = "https://github.com/lightning-AI/litgpt/tutorials"

[project.scripts]
litgpt = "litgpt.__main__:main"

[project.optional-dependencies]
test = [
    "pytest",
    "pytest-rerunfailures",
    "pytest-timeout",
    "transformers>=4.38.0",
    "einops",
    "protobuf",
    "lightning-thunder; python_version >= '3.10'",
]
all = [
    "bitsandbytes==0.42.0",      # quantization
    "sentencepiece",             # llama-based models
    "tokenizers",                # pythia, falcon, redpajama
    "datasets",                  # eval
    "requests",                  # litgpt.data
    "litdata",                   # litgpt.data
    "zstandard",                 # litgpt.data.prepare_slimpajama.py
    "pandas",                    # litgpt.data.prepare_starcoder.py
    "pyarrow",                   # litgpt.data.prepare_starcoder.py
    "torchmetrics",              # litgpt.pretrain
    "safetensors",               # download
    "huggingface_hub[hf_transfer]>=0.21.0",  # download
    "transformers", #tokenizer
    "wandb",
    "mosaicml-streaming",
    "xformers",
    "einops",
]
flashattn = [
    "flash-attn",
]

[build-system]
requires = [
    "setuptools",
    "wheel",
]
build-backend = "setuptools.build_meta"

[tool.setuptools.packages.find]
include = [
    "litgpt",
    "litgpt.*",
]
exclude = []

[tool.setuptools.package-data]
litgpt = [
    "LICENSE",
    "README.md",
]