teapottrainer 0.1.0__tar.gz

teapottrainer-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Zakery Clarke
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

teapottrainer-0.1.0/MANIFEST.in

@@ -0,0 +1 @@
+prune tests

teapottrainer-0.1.0/PKG-INFO

@@ -0,0 +1,95 @@
+Metadata-Version: 2.4
+Name: teapottrainer
+Version: 0.1.0
+Summary: A modular, config-driven library for training LLMs.
+Author-email: Zakery Clarke <zakerytclarke@gmail.com>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/zakerytclarke/teapot_trainer
+Project-URL: Repository, https://github.com/zakerytclarke/teapot_trainer
+Project-URL: Issues, https://github.com/zakerytclarke/teapot_trainer/issues
+Keywords: llm,training,transformers,pytorch
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.1
+Requires-Dist: transformers>=4.40
+Requires-Dist: datasets>=2.14
+Requires-Dist: pydantic>=2
+Requires-Dist: tensorboard>=2.14
+Requires-Dist: tqdm>=4.66
+Requires-Dist: huggingface_hub>=0.20
+Provides-Extra: dev
+Requires-Dist: pytest>=8; extra == "dev"
+Requires-Dist: build; extra == "dev"
+Requires-Dist: twine; extra == "dev"
+Dynamic: license-file
+
+# teapottrainer
+
+`teapottrainer` is a config-driven library for pretraining and supervised fine-tuning of causal language models. You describe the model, the dataset mix, and the token budgets. `TeapotTrainer` runs pretraining, then supervised training, and can resume from the last checkpoint.
+
+```bash
+pip install teapottrainer
+```
+
+## Train
+
+```python
+from teapottrainer import (
+    DatasetSpec,
+    FINEWEB_EDU,
+    GSM8K,
+    ModelConfig,
+    TeapotTrainer,
+    TrainingConfig,
+)
+
+config = TrainingConfig(
+    model=ModelConfig(
+        model_type="llama",
+        hidden_size=768,
+        layers=12,
+        heads=12,
+        kv_heads=4,
+        intermediate=2048,
+        seq_len=2048,
+    ),
+    pretrain_datasets=[
+        DatasetSpec(source=FINEWEB_EDU, target_tokens=1_000_000_000),
+    ],
+    sft_datasets=[
+        DatasetSpec(source=GSM8K, target_tokens=4_000_000),
+    ],
+    tokenizer="HuggingFaceTB/SmolLM2-135M",
+    batch_size=4,
+    grad_accum=8,
+    lr=5e-4,
+)
+
+trainer = TeapotTrainer(config)
+trainer.train(fresh=True)
+```
+
+`train(fresh=False)` continues from `checkpoints/resume` in the run directory: weights, optimizer, and the sample cursor. Pass `init_from` to start from an existing Hugging Face model or a local checkpoint instead of randomly initialized weights.
+
+Set `sft_at_checkpoints=True` to fine-tune a copy of the base model at each pretrain eval and save that copy beside the base weights. Pretraining keeps going from the base model.
+
+Runs are written under `./runs/`. Loss and task scores go to TensorBoard and to `metrics.jsonl` in the run directory.
+
+## Datasets
+
+Preset sources live on the package: `FINEWEB_EDU`, `DCLM_EDU`, `STACK_EDU`, `FINEMATH`, `GSM8K`, `SQUAD_V2`, `HOTPOT`, `SYNTHQA`, `TEAPOT_CHAT`, the SmolTalk subsets, and the reading, tool-calling, and instruction-following evals. A `DatasetSpec` assigns a source a token budget and a mix weight.
+
+To add a source, build a `DatasetSourceConfig` with a `chat_template_fn` that returns a prompt and an answer. The trainer supervises the answer tokens.
+
+## Develop
+
+The `train.py` script in this repository is one example run. It is not part of the installed package.
+
+```bash
+pip install -e ".[dev]"
+pytest
+```

teapottrainer-0.1.0/README.md

@@ -0,0 +1,66 @@
+# teapottrainer
+
+`teapottrainer` is a config-driven library for pretraining and supervised fine-tuning of causal language models. You describe the model, the dataset mix, and the token budgets. `TeapotTrainer` runs pretraining, then supervised training, and can resume from the last checkpoint.
+
+```bash
+pip install teapottrainer
+```
+
+## Train
+
+```python
+from teapottrainer import (
+    DatasetSpec,
+    FINEWEB_EDU,
+    GSM8K,
+    ModelConfig,
+    TeapotTrainer,
+    TrainingConfig,
+)
+
+config = TrainingConfig(
+    model=ModelConfig(
+        model_type="llama",
+        hidden_size=768,
+        layers=12,
+        heads=12,
+        kv_heads=4,
+        intermediate=2048,
+        seq_len=2048,
+    ),
+    pretrain_datasets=[
+        DatasetSpec(source=FINEWEB_EDU, target_tokens=1_000_000_000),
+    ],
+    sft_datasets=[
+        DatasetSpec(source=GSM8K, target_tokens=4_000_000),
+    ],
+    tokenizer="HuggingFaceTB/SmolLM2-135M",
+    batch_size=4,
+    grad_accum=8,
+    lr=5e-4,
+)
+
+trainer = TeapotTrainer(config)
+trainer.train(fresh=True)
+```
+
+`train(fresh=False)` continues from `checkpoints/resume` in the run directory: weights, optimizer, and the sample cursor. Pass `init_from` to start from an existing Hugging Face model or a local checkpoint instead of randomly initialized weights.
+
+Set `sft_at_checkpoints=True` to fine-tune a copy of the base model at each pretrain eval and save that copy beside the base weights. Pretraining keeps going from the base model.
+
+Runs are written under `./runs/`. Loss and task scores go to TensorBoard and to `metrics.jsonl` in the run directory.
+
+## Datasets
+
+Preset sources live on the package: `FINEWEB_EDU`, `DCLM_EDU`, `STACK_EDU`, `FINEMATH`, `GSM8K`, `SQUAD_V2`, `HOTPOT`, `SYNTHQA`, `TEAPOT_CHAT`, the SmolTalk subsets, and the reading, tool-calling, and instruction-following evals. A `DatasetSpec` assigns a source a token budget and a mix weight.
+
+To add a source, build a `DatasetSourceConfig` with a `chat_template_fn` that returns a prompt and an answer. The trainer supervises the answer tokens.
+
+## Develop
+
+The `train.py` script in this repository is one example run. It is not part of the installed package.
+
+```bash
+pip install -e ".[dev]"
+pytest
+```

teapottrainer-0.1.0/pyproject.toml

@@ -0,0 +1,51 @@
+[build-system]
+requires = ["setuptools>=77"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "teapottrainer"
+dynamic = ["version"]
+authors = [
+    { name = "Zakery Clarke", email = "zakerytclarke@gmail.com" },
+]
+description = "A modular, config-driven library for training LLMs."
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE"]
+requires-python = ">=3.10"
+keywords = ["llm", "training", "transformers", "pytorch"]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Operating System :: OS Independent",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "torch>=2.1",
+    "transformers>=4.40",
+    "datasets>=2.14",
+    "pydantic>=2",
+    "tensorboard>=2.14",
+    "tqdm>=4.66",
+    "huggingface_hub>=0.20",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=8", "build", "twine"]
+
+[project.urls]
+Homepage = "https://github.com/zakerytclarke/teapot_trainer"
+Repository = "https://github.com/zakerytclarke/teapot_trainer"
+Issues = "https://github.com/zakerytclarke/teapot_trainer/issues"
+
+[tool.setuptools]
+packages = ["teapottrainer"]
+
+[tool.setuptools.dynamic]
+version = { attr = "teapottrainer.__version__" }
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+# train.py lives at the repo root and is not installed. The pytest console
+# script does not put that directory on sys.path, so the tests that import it
+# need it added explicitly.
+pythonpath = ["."]

teapottrainer-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

teapottrainer-0.1.0/teapottrainer/__init__.py

@@ -0,0 +1,148 @@
+from .config import TrainingConfig, ModelConfig, DatasetSpec
+from .trainer import TeapotTrainer
+from .datasets import (
+    DatasetSourceConfig,
+    DatasetRegistry,
+    FINEWEB_EDU,
+    DCLM_EDU,
+    STACK_EDU,
+    FINEMATH,
+    INFIWEBMATH,
+    COSMOPEDIA_V2,
+    SMOLTALK,
+    SMOLTALK_MAGPIE,
+    SMOLTALK_EVERYDAY,
+    SMOLTALK_OPENHERMES,
+    SMOLTALK_CONSTRAINTS,
+    SMOLTALK_SYSTEMCHATS,
+    SMOLTALK_SUMMARIZE,
+    SMOLTALK_REWRITE,
+    SMOLTALK_APIGEN,
+    SMOLTALK_APIGEN_DECISION,
+    SMOLTALK_APIGEN_ALIASED,
+    SMOLTALK_APIGEN_MISSING,
+    SMOLTALK_APIGEN_UNRELATED,
+    SMOLTALK_METAMATH,
+    SMOLTALK_MAGPIE_TEST,
+    SMOLTALK_CONSTRAINTS_TEST,
+    SMOLTALK_APIGEN_TEST,
+    SMOLTALK_APIGEN_DECISION_TEST,
+    SMOLTALK_APIGEN_ALIASED_TEST,
+    GSM8K,
+    GSM8K_TEST,
+    WHEN2CALL,
+    WHEN2CALL_TEST,
+    WHEN2CALL_ASK,
+    WHEN2CALL_REFUSE,
+    SQUAD_V2,
+    SQUAD_V2_VAL,
+    HOTPOT,
+    HOTPOT_VAL,
+    SYNTHQA,
+    SYNTHQA_TEST,
+    SYNTHQA_ANSWERABLE,
+    SYNTHQA_UNANSWERABLE,
+    TEAPOT_CHAT,
+    SQUAD_V2_ANSWERABLE,
+    SQUAD_V2_UNANSWERABLE,
+    SQUAD_V2_CONTRASTIVE,
+    HOTPOT_CONTRASTIVE,
+    SYNTHQA_CONTRASTIVE,
+    IFEVAL,
+    MT_BENCH,
+    ALPACA_EVAL,
+    ARENA_HARD,
+    ARC_EASY,
+    ARC_EASY_VAL,
+    HELLASWAG,
+    HELLASWAG_VAL,
+    PIQA,
+    PIQA_VAL,
+    WINOGRANDE,
+    WINOGRANDE_VAL,
+    COQA,
+    COQA_VAL,
+    DROP,
+    DROP_VAL,
+    HERMES_FUNCTION_CALLING,
+    BFCL_SIMPLE,
+    BFCL_MULTIPLE,
+    BFCL_PARALLEL,
+    BFCL_IRRELEVANCE,
+)
+
+__version__ = "0.1.0"
+__all__ = [
+    "TeapotTrainer",
+    "TrainingConfig",
+    "ModelConfig",
+    "DatasetSpec",
+    "DatasetSourceConfig",
+    "DatasetRegistry",
+    "FINEWEB_EDU",
+    "DCLM_EDU",
+    "STACK_EDU",
+    "FINEMATH",
+    "INFIWEBMATH",
+    "COSMOPEDIA_V2",
+    "SMOLTALK",
+    "SMOLTALK_MAGPIE",
+    "SMOLTALK_EVERYDAY",
+    "SMOLTALK_OPENHERMES",
+    "SMOLTALK_CONSTRAINTS",
+    "SMOLTALK_SYSTEMCHATS",
+    "SMOLTALK_SUMMARIZE",
+    "SMOLTALK_REWRITE",
+    "SMOLTALK_APIGEN",
+    "SMOLTALK_APIGEN_DECISION",
+    "SMOLTALK_APIGEN_ALIASED",
+    "SMOLTALK_APIGEN_MISSING",
+    "SMOLTALK_APIGEN_UNRELATED",
+    "SMOLTALK_METAMATH",
+    "SMOLTALK_MAGPIE_TEST",
+    "SMOLTALK_CONSTRAINTS_TEST",
+    "SMOLTALK_APIGEN_TEST",
+    "SMOLTALK_APIGEN_DECISION_TEST",
+    "SMOLTALK_APIGEN_ALIASED_TEST",
+    "GSM8K",
+    "GSM8K_TEST",
+    "WHEN2CALL",
+    "WHEN2CALL_TEST",
+    "WHEN2CALL_ASK",
+    "WHEN2CALL_REFUSE",
+    "SQUAD_V2",
+    "SQUAD_V2_VAL",
+    "HOTPOT",
+    "HOTPOT_VAL",
+    "SYNTHQA",
+    "SYNTHQA_TEST",
+    "SYNTHQA_ANSWERABLE",
+    "SYNTHQA_UNANSWERABLE",
+    "TEAPOT_CHAT",
+    "SQUAD_V2_ANSWERABLE",
+    "SQUAD_V2_UNANSWERABLE",
+    "SQUAD_V2_CONTRASTIVE",
+    "HOTPOT_CONTRASTIVE",
+    "SYNTHQA_CONTRASTIVE",
+    "IFEVAL",
+    "MT_BENCH",
+    "ALPACA_EVAL",
+    "ARENA_HARD",
+    "ARC_EASY",
+    "ARC_EASY_VAL",
+    "HELLASWAG",
+    "HELLASWAG_VAL",
+    "PIQA",
+    "PIQA_VAL",
+    "WINOGRANDE",
+    "WINOGRANDE_VAL",
+    "COQA",
+    "COQA_VAL",
+    "DROP",
+    "DROP_VAL",
+    "HERMES_FUNCTION_CALLING",
+    "BFCL_SIMPLE",
+    "BFCL_MULTIPLE",
+    "BFCL_PARALLEL",
+    "BFCL_IRRELEVANCE",
+]

teapottrainer-0.1.0/teapottrainer/config.py

@@ -0,0 +1,54 @@
+from pydantic import BaseModel, Field, ConfigDict
+from typing import Optional, List, Any, Union
+from .datasets import DatasetSourceConfig, FINEWEB_EDU
+
+class ModelConfig(BaseModel):
+    model_type: str = "llama"
+    hidden_size: int = 768
+    layers: int = Field(12, alias="num_hidden_layers")
+    heads: int = Field(12, alias="num_attention_heads")
+    kv_heads: int = Field(4, alias="num_key_value_heads")
+    intermediate: int = Field(3072, alias="intermediate_size")
+    seq_len: int = Field(2048, alias="max_position_embeddings")
+    
+    custom_config_cls: Optional[Any] = None
+    custom_model_cls: Optional[Any] = None
+
+    model_config = ConfigDict(populate_by_name=True, arbitrary_types_allowed=True)
+
+
+class DatasetSpec(BaseModel):
+    source: Union[DatasetSourceConfig, str] = FINEWEB_EDU
+    target_tokens: int = 1_000_000_000
+    weight: float = 1.0
+
+    model_config = ConfigDict(populate_by_name=True, arbitrary_types_allowed=True)
+
+
+class TrainingConfig(BaseModel):
+    model: ModelConfig = Field(default_factory=ModelConfig, alias="model_config")
+    
+    pretrain_datasets: List[DatasetSpec] = Field(
+        default_factory=lambda: [DatasetSpec(source=FINEWEB_EDU, target_tokens=2_000_000_000)]
+    )
+    sft_datasets: List[DatasetSpec] = Field(default_factory=list)
+    eval_datasets: List[DatasetSpec] = Field(default_factory=list)
+    
+    tokenizer: str = "Qwen/Qwen2.5-0.5B"
+    # Hugging Face id or local directory. None builds a new model from the config.
+    init_from: Optional[str] = None
+    batch_size: int = 4
+    grad_accum: int = 32
+    lr: float = 5e-4
+    weight_decay: float = 0.1
+    compile: bool = True
+    logging_dir: Optional[str] = None
+    
+    eval_every_percent: float = 5.0
+    # At each pretrain eval checkpoint, train a copy on the full SFT mix and
+    # save it beside the base weights. The pretrain model keeps training.
+    sft_at_checkpoints: bool = False
+    holdout_ratio: float = 0.05
+    seed: int = 42
+
+    model_config = ConfigDict(populate_by_name=True)