[project]
name = "needle"
version = "0.1.0"
requires-python = ">=3.11"
dependencies = [
    "jax",
    "jaxlib",
    "flax",
    "optax",
    "datasets",
    "huggingface_hub",
    "transformers",
    "wandb",
    "sentencepiece",
    "scipy",
    "gcsfs",
    "tqdm",
    "google-genai",
]

[project.optional-dependencies]
tpu = ["jax[tpu]"]
gpu = ["jax[cuda12]"]

[project.scripts]
needle = "needle.cli:main"

[tool.setuptools.packages.find]
include = ["needle*"]

[build-system]
requires = ["setuptools>=68.0"]
build-backend = "setuptools.build_meta"