File size: 2,089 Bytes
a006332 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | """Command line training entrypoint for Delta Ultra Mini."""
from __future__ import annotations
import argparse
import json
import logging
import os
import sys
from pathlib import Path
from typing import Any
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from delta.trainer import train
logging.basicConfig(level=os.getenv("DELTA_LOG_LEVEL", "INFO").upper())
logger = logging.getLogger(__name__)
def parse_args() -> argparse.Namespace:
"""Parse training arguments."""
parser = argparse.ArgumentParser(description="Train Delta Ultra Mini.")
parser.add_argument("--data_path", required=True, help="Directory or file containing .txt/.jsonl training data.")
parser.add_argument("--output_dir", required=True, help="Output directory for checkpoints.")
parser.add_argument("--epochs", type=float, default=1.0, help="Number of training epochs.")
parser.add_argument("--batch_size", type=int, default=2, help="Per-device train batch size.")
parser.add_argument("--resume_from_checkpoint", default=None, help="Checkpoint path or true to resume latest.")
parser.add_argument("--tokenizer_path", default=None, help="Path to tokenizer.json.")
parser.add_argument("--config_path", default="configs/ultra_mini.json", help="Model config JSON.")
return parser.parse_args()
def main() -> None:
"""Run Trainer-based model training."""
args = parse_args()
with Path(args.config_path).open("r", encoding="utf-8") as handle:
model_config: dict[str, Any] = json.load(handle)
output_dir = Path(args.output_dir)
config = {
"data_path": args.data_path,
"output_dir": str(output_dir),
"epochs": args.epochs,
"batch_size": args.batch_size,
"resume_from_checkpoint": args.resume_from_checkpoint,
"tokenizer_path": args.tokenizer_path or str(output_dir / "tokenizer.json"),
"model": model_config,
}
train(config)
logger.info("Training complete.")
if __name__ == "__main__":
main()
|