trloom 0.1.0__py3-none-any.whl

trloom/__init__.py

@@ -0,0 +1,19 @@
+"""TRLoom: YAML-driven end-to-end fine-tuning on top of Hugging Face TRL."""
+
+from __future__ import annotations
+
+from trloom.config import FineTuneConfig, load_config
+from trloom.job import FineTuneJob, available_methods, run_from_yaml
+from trloom.trainers import list_trainers
+
+__all__ = [
+    "FineTuneConfig",
+    "FineTuneJob",
+    "__version__",
+    "available_methods",
+    "list_trainers",
+    "load_config",
+    "run_from_yaml",
+]
+
+__version__ = "0.1.0"

trloom/__main__.py

@@ -0,0 +1,8 @@
+"""python -m trloom"""
+
+from __future__ import annotations
+
+from trloom.cli import main
+
+if __name__ == "__main__":  # pragma: no cover
+    raise SystemExit(main())

trloom/cli.py

@@ -0,0 +1,148 @@
+"""Command-line interface for TRLoom."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import sys
+from pathlib import Path
+
+
+def _configure_logging(verbose: bool) -> None:
+    level = logging.DEBUG if verbose else logging.INFO
+    logging.basicConfig(
+        level=level,
+        format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+    )
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="trloom",
+        description="YAML-driven fine-tuning on top of Hugging Face TRL.",
+    )
+    parser.add_argument(
+        "--version",
+        action="store_true",
+        help="Print version and exit.",
+    )
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        action="store_true",
+        help="Enable debug logging.",
+    )
+
+    subparsers = parser.add_subparsers(dest="command")
+
+    run_parser = subparsers.add_parser("run", help="Run a fine-tuning job from a YAML config.")
+    run_parser.add_argument("config", type=str, help="Path to the YAML configuration file.")
+    run_parser.add_argument(
+        "--modal",
+        action="store_true",
+        help="Force execution on Modal (overrides modal.enabled in YAML).",
+    )
+    run_parser.add_argument(
+        "--local",
+        action="store_true",
+        help="Force local execution (overrides modal.enabled in YAML).",
+    )
+
+    subparsers.add_parser("methods", help="List TRL training methods available in this environment.")
+
+    validate_parser = subparsers.add_parser("validate", help="Validate a YAML config without training.")
+    validate_parser.add_argument("config", type=str, help="Path to the YAML configuration file.")
+
+    modal_parser = subparsers.add_parser(
+        "modal-script",
+        help="Write a standalone Modal entrypoint script for a YAML config.",
+    )
+    modal_parser.add_argument("config", type=str, help="Path to the YAML configuration file.")
+    modal_parser.add_argument(
+        "-o",
+        "--output",
+        type=str,
+        default=None,
+        help="Destination path for the generated script.",
+    )
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    _configure_logging(getattr(args, "verbose", False))
+
+    if args.version:
+        from trloom import __version__
+
+        print(__version__)
+        return 0
+
+    if args.command is None:
+        parser.print_help()
+        return 0
+
+    if args.command == "methods":
+        from trloom import available_methods
+
+        methods = available_methods(include_experimental=True)
+        for method in methods:
+            print(method)
+        return 0
+
+    if args.command == "validate":
+        from trloom import load_config
+        from trloom.trainers import get_trainer_spec
+
+        config = load_config(args.config)
+        spec = get_trainer_spec(config.method)
+        payload = {
+            "ok": True,
+            "method": config.method,
+            "trainer": spec.trainer_name,
+            "config_class": spec.config_name,
+            "experimental": spec.experimental,
+            "output_dir": str(config.resolved_output_dir()),
+            "wandb_enabled": config.wandb.enabled,
+            "modal_enabled": config.modal.enabled,
+        }
+        print(json.dumps(payload, indent=2))
+        return 0
+
+    if args.command == "modal-script":
+        from trloom.modal_support.runner import write_modal_entrypoint
+
+        dest = write_modal_entrypoint(args.config, destination=args.output)
+        print(dest)
+        return 0
+
+    if args.command == "run":
+        from trloom import run_from_yaml
+
+        if args.modal and args.local:
+            parser.error("Use only one of --modal / --local.")
+        use_modal: bool | None
+        if args.modal:
+            use_modal = True
+        elif args.local:
+            use_modal = False
+        else:
+            use_modal = None
+
+        config_path = Path(args.config)
+        if not config_path.is_file():
+            print(f"Config not found: {config_path}", file=sys.stderr)
+            return 1
+
+        run_from_yaml(config_path, use_modal=use_modal)
+        return 0
+
+    parser.print_help()
+    return 1
+
+
+if __name__ == "__main__":  # pragma: no cover
+    raise SystemExit(main())

trloom/config/__init__.py

@@ -0,0 +1,22 @@
+"""Configuration schema and YAML loading utilities."""
+
+from __future__ import annotations
+
+from trloom.config.loader import load_config, load_config_dict
+from trloom.config.schema import (
+    DatasetConfig,
+    FineTuneConfig,
+    ModalConfig,
+    ModelConfig,
+    WandbConfig,
+)
+
+__all__ = [
+    "DatasetConfig",
+    "FineTuneConfig",
+    "ModalConfig",
+    "ModelConfig",
+    "WandbConfig",
+    "load_config",
+    "load_config_dict",
+]

trloom/config/loader.py

@@ -0,0 +1,36 @@
+"""Load and validate TRLoom YAML configuration files."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from trloom.config.schema import FineTuneConfig
+
+
+def _read_yaml(path: Path) -> dict[str, Any]:
+    with path.open("r", encoding="utf-8") as handle:
+        data = yaml.safe_load(handle)
+    if data is None:
+        raise ValueError(f"Config file is empty: {path}")
+    if not isinstance(data, dict):
+        raise TypeError(f"Config root must be a mapping, got {type(data).__name__}: {path}")
+    return data
+
+
+def load_config_dict(source: str | Path | dict[str, Any]) -> dict[str, Any]:
+    """Return a raw configuration dictionary from a path or mapping."""
+    if isinstance(source, dict):
+        return dict(source)
+    path = Path(source).expanduser().resolve()
+    if not path.is_file():
+        raise FileNotFoundError(f"Config file not found: {path}")
+    return _read_yaml(path)
+
+
+def load_config(source: str | Path | dict[str, Any]) -> FineTuneConfig:
+    """Load and validate a :class:`FineTuneConfig` from YAML or a dict."""
+    data = load_config_dict(source)
+    return FineTuneConfig.model_validate(data)

trloom/config/schema.py

@@ -0,0 +1,185 @@
+"""Pydantic models for TRLoom YAML configuration."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+
+
+class ModelConfig(BaseModel):
+    """Model / PEFT / quantization settings (mirrors TRL ``ModelConfig`` fields)."""
+
+    model_config = ConfigDict(extra="allow")
+
+    model_name_or_path: str
+    model_revision: str = "main"
+    dtype: Literal["auto", "bfloat16", "float16", "float32"] | None = "float32"
+    attn_implementation: str | None = None
+    trust_remote_code: bool = False
+
+    use_peft: bool = False
+    lora_r: int = 16
+    lora_alpha: int = 32
+    lora_dropout: float = 0.05
+    lora_target_modules: list[str] | str | None = None
+    lora_target_parameters: list[str] | None = None
+    lora_modules_to_save: list[str] | None = None
+    lora_task_type: str = "CAUSAL_LM"
+    use_rslora: bool = False
+    use_dora: bool = False
+
+    load_in_8bit: bool = False
+    load_in_4bit: bool = False
+    bnb_4bit_quant_type: Literal["fp4", "nf4"] = "nf4"
+    use_bnb_nested_quant: bool = False
+    bnb_4bit_quant_storage: str | None = None
+
+    @model_validator(mode="after")
+    def _validate_quantization(self) -> ModelConfig:
+        if self.load_in_8bit and self.load_in_4bit:
+            raise ValueError("Cannot enable both load_in_8bit and load_in_4bit.")
+        return self
+
+
+class DatasetSourceConfig(BaseModel):
+    """A single dataset entry (Hub repo id or local path)."""
+
+    model_config = ConfigDict(extra="allow")
+
+    path: str
+    name: str | None = None
+    split: str | None = None
+    data_files: str | list[str] | dict[str, Any] | None = None
+    data_dir: str | None = None
+    streaming: bool = False
+    columns: dict[str, str] | None = None
+    weight: float | None = None
+
+
+class DatasetConfig(BaseModel):
+    """Dataset loading configuration.
+
+    Prefer ``path`` for a single Hub/local dataset, or ``datasets`` for a mixture.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    path: str | None = None
+    name: str | None = None
+    # Optional datasets split selector, e.g. "train" or "train[:64]"
+    split: str | None = None
+    data_files: str | list[str] | dict[str, Any] | None = None
+    data_dir: str | None = None
+    streaming: bool = False
+
+    datasets: list[DatasetSourceConfig] | None = None
+
+    train_split: str = "train"
+    eval_split: str | None = "test"
+    text_column: str | None = None
+    columns: dict[str, str] | None = None
+
+    # Optional Hugging Face datasets kwargs passthrough
+    kwargs: dict[str, Any] = Field(default_factory=dict)
+
+    @model_validator(mode="after")
+    def _require_source(self) -> DatasetConfig:
+        if not self.path and not self.datasets:
+            raise ValueError("Dataset config requires either 'path' or 'datasets'.")
+        return self
+
+
+class WandbConfig(BaseModel):
+    """Weights & Biases logging settings."""
+
+    model_config = ConfigDict(extra="allow")
+
+    enabled: bool = False
+    project: str | None = None
+    entity: str | None = None
+    run_name: str | None = None
+    group: str | None = None
+    tags: list[str] = Field(default_factory=list)
+    notes: str | None = None
+    mode: Literal["online", "offline", "disabled"] | None = None
+    dir: str | None = None
+    job_type: str | None = "train"
+    # Extra keys forwarded to wandb.init
+    init_kwargs: dict[str, Any] = Field(default_factory=dict)
+
+
+class ModalConfig(BaseModel):
+    """Modal Labs remote execution settings."""
+
+    model_config = ConfigDict(extra="allow")
+
+    enabled: bool = False
+    app_name: str = "trloom"
+    gpu: str = "T4"
+    timeout: int = 60 * 60 * 4
+    cpu: float | None = None
+    memory: int | None = None
+    volume_name: str = "trloom-outputs"
+    volume_mount: str = "/outputs"
+    # Names of Modal Secrets to attach (create with `modal secret create ...`).
+    # Leave empty for public Hub models/datasets that need no tokens.
+    secrets: list[str] = Field(default_factory=list)
+    pip_packages: list[str] = Field(default_factory=list)
+    python_version: str = "3.11"
+    region: str | None = None
+    # How to get trloom into the Modal image:
+    # - local: mount the local installed package (best for development)
+    # - git: pip install from git_url (best for reproducibility)
+    # - pypi: pip install trloom (once published)
+    install_source: Literal["local", "git", "pypi"] = "local"
+    git_url: str = "git+https://github.com/saqlain2204/trloom.git"
+    # If set, copy training output from the Modal volume back to this local path
+    download_dir: str | None = None
+
+
+class FineTuneConfig(BaseModel):
+    """Root configuration for a TRLoom fine-tuning job."""
+
+    model_config = ConfigDict(extra="allow")
+
+    method: str = Field(
+        ...,
+        description="TRL training method, e.g. sft, dpo, grpo, kto, reward, rloo.",
+    )
+    model: ModelConfig
+    dataset: DatasetConfig
+    training: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Keyword arguments forwarded to the TRL *Config class "
+        "(SFTConfig, DPOConfig, …).",
+    )
+    wandb: WandbConfig = Field(default_factory=WandbConfig)
+    modal: ModalConfig = Field(default_factory=ModalConfig)
+
+    # Optional extras
+    reward_funcs: list[str] | str | None = None
+    trainer_kwargs: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Extra kwargs passed directly to the TRL Trainer constructor.",
+    )
+    push_to_hub: bool = False
+    hub_model_id: str | None = None
+    seed: int | None = 42
+
+    @field_validator("method")
+    @classmethod
+    def _normalize_method(cls, value: str) -> str:
+        return value.strip().lower().replace("-", "_").replace(" ", "_")
+
+    @field_validator("reward_funcs", mode="before")
+    @classmethod
+    def _coerce_reward_funcs(cls, value: Any) -> Any:
+        if isinstance(value, str):
+            return [value]
+        return value
+
+    def resolved_output_dir(self) -> Path:
+        output_dir = self.training.get("output_dir", "./outputs")
+        return Path(output_dir).expanduser().resolve()

trloom/data/__init__.py

@@ -0,0 +1,7 @@
+"""Dataset loading helpers for Hub and local sources."""
+
+from __future__ import annotations
+
+from trloom.data.dataset import load_train_eval_datasets
+
+__all__ = ["load_train_eval_datasets"]

trloom/data/dataset.py

@@ -0,0 +1,210 @@
+"""Load training datasets from the Hugging Face Hub or local files."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any
+
+from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict, concatenate_datasets, load_dataset
+
+from trloom.config.schema import DatasetConfig, DatasetSourceConfig
+
+logger = logging.getLogger(__name__)
+
+DatasetLike = Dataset | DatasetDict | IterableDataset | IterableDatasetDict
+
+
+def _is_local_path(path: str) -> bool:
+    candidate = Path(path).expanduser()
+    return candidate.exists() or path.startswith((".", "/", "~", "\\")) or ":\\" in path
+
+
+def _infer_builder(path: str) -> str | None:
+    """Infer a datasets builder name from a local file extension."""
+    suffix = Path(path).suffix.lower()
+    mapping = {
+        ".json": "json",
+        ".jsonl": "json",
+        ".csv": "csv",
+        ".tsv": "csv",
+        ".parquet": "parquet",
+        ".txt": "text",
+        ".arrow": "arrow",
+    }
+    return mapping.get(suffix)
+
+
+def _rename_columns(dataset: Any, columns: dict[str, str] | None) -> Any:
+    if not columns:
+        return dataset
+    rename_map = {src: dst for src, dst in columns.items() if src in getattr(dataset, "column_names", [])}
+    if not rename_map:
+        return dataset
+    return dataset.rename_columns(rename_map)
+
+
+def _load_single_source(
+    source: DatasetSourceConfig | DatasetConfig,
+    *,
+    default_split: str | None = None,
+) -> DatasetLike:
+    path = source.path
+    if path is None:
+        raise ValueError("Dataset source is missing 'path'.")
+
+    kwargs: dict[str, Any] = {}
+    if isinstance(source, DatasetConfig):
+        kwargs.update(source.kwargs)
+    # Allow extras from DatasetSourceConfig / DatasetConfig
+    extras = source.model_extra or {}
+    for key, value in extras.items():
+        if key not in {"columns", "weight"}:
+            kwargs.setdefault(key, value)
+
+    name = source.name
+    data_files = source.data_files
+    data_dir = source.data_dir
+    streaming = bool(source.streaming)
+    split = source.split if getattr(source, "split", None) else default_split
+
+    if data_files is not None:
+        kwargs["data_files"] = data_files
+    if data_dir is not None:
+        kwargs["data_dir"] = data_dir
+
+    if _is_local_path(path):
+        local = Path(path).expanduser().resolve()
+        if local.is_dir():
+            # load_from_disk for saved DatasetDict / Dataset
+            try:
+                from datasets import load_from_disk
+
+                logger.info("Loading local dataset from disk: %s", local)
+                dataset = load_from_disk(str(local))
+            except Exception:
+                logger.info("Falling back to load_dataset for directory: %s", local)
+                dataset = load_dataset(str(local), name=name, split=split, streaming=streaming, **kwargs)
+        else:
+            builder = _infer_builder(str(local))
+            if builder is None and data_files is None:
+                raise ValueError(
+                    f"Unsupported local dataset file type for '{local}'. "
+                    "Use json/jsonl/csv/tsv/parquet/txt/arrow or provide data_files."
+                )
+            logger.info("Loading local dataset file via builder=%s path=%s", builder, local)
+            load_path = builder or str(local)
+            file_kwargs = dict(kwargs)
+            if builder and "data_files" not in file_kwargs:
+                file_kwargs["data_files"] = str(local)
+            dataset = load_dataset(load_path, name=name, split=split, streaming=streaming, **file_kwargs)
+    else:
+        logger.info("Loading Hub dataset: %s (name=%s, split=%s)", path, name, split)
+        dataset = load_dataset(path, name=name, split=split, streaming=streaming, **kwargs)
+
+    columns = getattr(source, "columns", None)
+    if columns and hasattr(dataset, "column_names"):
+        dataset = _rename_columns(dataset, columns)
+    elif columns and isinstance(dataset, (DatasetDict, IterableDatasetDict)):
+        dataset = type(dataset)({k: _rename_columns(v, columns) for k, v in dataset.items()})
+
+    return dataset
+
+
+def _as_dataset_dict(dataset: DatasetLike, train_split: str) -> DatasetDict | IterableDatasetDict:
+    if isinstance(dataset, (DatasetDict, IterableDatasetDict)):
+        return dataset
+    # Single split dataset — wrap under train_split key
+    if isinstance(dataset, IterableDataset):
+        return IterableDatasetDict({train_split: dataset})
+    return DatasetDict({train_split: dataset})
+
+
+def _mix_datasets(
+    sources: list[DatasetSourceConfig],
+    *,
+    train_split: str,
+) -> DatasetDict | IterableDatasetDict:
+    loaded: list[Dataset | IterableDataset] = []
+    weights: list[float] = []
+    streaming = False
+
+    for source in sources:
+        ds = _load_single_source(source, default_split=source.split or train_split)
+        if isinstance(ds, (DatasetDict, IterableDatasetDict)):
+            if train_split not in ds:
+                raise KeyError(
+                    f"Mixture source '{source.path}' has no split '{train_split}'. "
+                    f"Available: {list(ds.keys())}"
+                )
+            part = ds[train_split]
+        else:
+            part = ds
+        if isinstance(part, IterableDataset):
+            streaming = True
+        loaded.append(part)
+        weights.append(float(source.weight) if source.weight is not None else 1.0)
+
+    if streaming:
+        # Interleave for streaming mixtures; fall back to concatenate when weights are equal.
+        from datasets import interleave_datasets
+
+        mixed = interleave_datasets(loaded, probabilities=_normalize(weights), seed=42)
+        return IterableDatasetDict({train_split: mixed})
+
+    if len(set(weights)) == 1:
+        mixed = concatenate_datasets(loaded)  # type: ignore[arg-type]
+    else:
+        from datasets import interleave_datasets
+
+        mixed = interleave_datasets(loaded, probabilities=_normalize(weights), seed=42)
+    return DatasetDict({train_split: mixed})
+
+
+def _normalize(weights: list[float]) -> list[float]:
+    total = sum(weights)
+    if total <= 0:
+        raise ValueError("Dataset mixture weights must sum to a positive value.")
+    return [w / total for w in weights]
+
+
+def load_train_eval_datasets(
+    config: DatasetConfig,
+) -> tuple[Dataset | IterableDataset, Dataset | IterableDataset | None]:
+    """Load train (and optional eval) datasets from a :class:`DatasetConfig`."""
+    if config.datasets:
+        dataset_dict = _mix_datasets(config.datasets, train_split=config.train_split)
+    else:
+        assert config.path is not None
+        raw = _load_single_source(config)
+        dataset_dict = _as_dataset_dict(raw, config.train_split)
+        if config.columns:
+            if isinstance(dataset_dict, (DatasetDict, IterableDatasetDict)):
+                dataset_dict = type(dataset_dict)(
+                    {k: _rename_columns(v, config.columns) for k, v in dataset_dict.items()}
+                )
+
+    if config.train_split not in dataset_dict:
+        available = list(dataset_dict.keys())
+        raise KeyError(
+            f"Train split '{config.train_split}' not found in dataset. Available splits: {available}"
+        )
+
+    train_dataset = dataset_dict[config.train_split]
+    eval_dataset: Dataset | IterableDataset | None = None
+    if config.eval_split and config.eval_split in dataset_dict:
+        eval_dataset = dataset_dict[config.eval_split]
+
+    if config.text_column and hasattr(train_dataset, "column_names"):
+        if config.text_column not in train_dataset.column_names:
+            raise KeyError(
+                f"text_column '{config.text_column}' not found. "
+                f"Columns: {train_dataset.column_names}"
+            )
+
+    logger.info(
+        "Loaded train dataset (%s)%s",
+        type(train_dataset).__name__,
+        f" and eval split '{config.eval_split}'" if eval_dataset is not None else "",
+    )
+    return train_dataset, eval_dataset