78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
import argparse
|
|
from config import TrainingConfig, WandBConfig
|
|
from data_loader import DataLoader
|
|
from model_handler import ModelHandler
|
|
from trainer import CustomTrainer
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="Train a language model with LoRA")
|
|
|
|
# wandb args
|
|
wandb_group = parser.add_argument_group("Weights & Biases")
|
|
wandb_group.add_argument(
|
|
"--wandb_project", type=str, default="lora-training", help="WandB project name"
|
|
)
|
|
wandb_group.add_argument("--wandb_name", type=str, help="WandB run name")
|
|
wandb_group.add_argument("--wandb_entity", type=str, help="WandB entity/username")
|
|
wandb_group.add_argument(
|
|
"--wandb_tags", type=str, nargs="+", help="WandB tags for the run"
|
|
)
|
|
wandb_group.add_argument("--wandb_notes", type=str, help="Notes for the WandB run")
|
|
# rest
|
|
parser.add_argument(
|
|
"--base_model",
|
|
type=str,
|
|
default="unsloth/Qwen2.5-7B",
|
|
help="Base model to fine-tune",
|
|
)
|
|
parser.add_argument(
|
|
"--dataset", type=str, required=True, help="Path to the training dataset"
|
|
)
|
|
parser.add_argument(
|
|
"--output_dir", type=str, default="/output/", help="Directory to save the model"
|
|
)
|
|
parser.add_argument("--hub_token", type=str, help="Hugging Face Hub token")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
|
|
try:
|
|
wandb_config = WandBConfig(
|
|
enabled=args.wandb_project is not None,
|
|
project=args.wandb_project or "lora-training",
|
|
name=args.wandb_name,
|
|
entity=args.wandb_entity,
|
|
tags=args.wandb_tags,
|
|
notes=args.wandb_notes,
|
|
)
|
|
|
|
config = TrainingConfig(
|
|
base_model=args.base_model, output_dir=args.output_dir, wandb=wandb_config
|
|
)
|
|
|
|
model_handler = ModelHandler(config)
|
|
model, tokenizer = model_handler.setup_model()
|
|
|
|
data_loader = DataLoader(tokenizer, config.data.template)
|
|
dataset = data_loader.load_dataset(args.dataset)
|
|
|
|
trainer_handler = CustomTrainer(config)
|
|
trainer = trainer_handler.create_trainer(model, tokenizer, dataset)
|
|
trainer_stats = trainer_handler.train_and_log(trainer)
|
|
|
|
model.save_pretrained(args.output_dir)
|
|
tokenizer.save_pretrained(args.output_dir)
|
|
|
|
print("Training completed successfully!")
|
|
|
|
except Exception as e:
|
|
print(f"Error during training: {e}")
|
|
raise
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|