[tool.poetry]
name = "pico-train"
version = "1.0.0"
description = "A minimalistic framework for transparently training language models and storing comprehensive checkpoints for in-depth learning dynamics research"
authors = ["Richard Diehl Martinez <richard@picolm.io>"]
license = "Apache 2.0"
readme = "README.md"
packages = [{include = "src"}]

[tool.poetry.scripts]
train = "scripts.train:main"

[tool.poetry.dependencies]
python = "^3.10,<3.13"
lightning = "^2.4.0"
click = "^8.1.7"
wandb = "^0.18.1"
huggingface-hub = {extras = ["cli"], version = "^0.25.1"}
datasets = "^3.0.1,<3.2.0"
transformers = "^4.45.2"
pre-commit = "^4.0.1"
torch = "^2.5.1"
evaluate = "^0.4.3"
deepspeed = "^0.16.2"
rich = "^13.9.4"

[tool.poetry.group.dev.dependencies]
ipykernel = "^6.29.5"
jupyter = "^1.1.1"

[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"