teapottrainer 0.1.0__py3-none-any.whl

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/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)