| """Fine-tuning entry point for CORTEX AI. |
| |
| This script is written to run on GPU hardware. It is NOT executed in this |
| repository: the sandbox has 4 CPU cores, ~15 GiB of RAM and no accelerator, and |
| the checkpoint is 893 GB. Treat it as the documented path, not a verified run. |
| |
| Usage:: |
| |
| python -m cortex_ai.training.train --config train_config.json |
| python -m cortex_ai.training.train --dataset data/cortex_sft.jsonl --epochs 1 |
| |
| Requires: torch, transformers, peft, datasets |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| from .config import TrainConfig |
| from .dataset import dataset_stats, iter_encoded, load_jsonl |
|
|
|
|
| def build_arg_parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser(description="Affinage de CORTEX AI") |
| parser.add_argument("--config", type=str, default="", help="fichier JSON de configuration") |
| parser.add_argument("--dataset", type=str, default="", help="dataset JSONL") |
| parser.add_argument("--output-dir", type=str, default="", help="dossier de sortie") |
| parser.add_argument("--epochs", type=int, default=0) |
| parser.add_argument("--dry-run", action="store_true", help="valider les donnees sans entrainer") |
| return parser |
|
|
|
|
| def main(argv: list[str] | None = None) -> int: |
| args = build_arg_parser().parse_args(argv) |
|
|
| config = TrainConfig.load(args.config) if args.config else TrainConfig() |
| if args.dataset: |
| config.dataset_path = args.dataset |
| if args.output_dir: |
| config.output_dir = args.output_dir |
| if args.epochs: |
| config.epochs = args.epochs |
|
|
| examples = load_jsonl(config.dataset_path) |
| stats = dataset_stats(examples) |
| print("=== Donnees ===") |
| for key, value in stats.items(): |
| print(f" {key:20s} {value}") |
|
|
| |
| encoded = list(iter_encoded(examples, config.thinking_mode)) |
| print(f" {'exemples encodes':20s} {len(encoded)}") |
| print(f" {'prompt le plus long':20s} {max(len(e) for e in encoded)} caracteres") |
|
|
| if args.dry_run: |
| print("\n--dry-run : validation terminee, aucun entrainement lance.") |
| return 0 |
|
|
| try: |
| import torch |
| from datasets import Dataset |
| from peft import LoraConfig, get_peft_model |
| from transformers import ( |
| AutoModelForCausalLM, |
| AutoTokenizer, |
| DataCollatorForLanguageModeling, |
| Trainer, |
| TrainingArguments, |
| ) |
| except ImportError as exc: |
| print( |
| f"\nDependance manquante : {exc.name}.\n" |
| "Installez torch, transformers, peft et datasets, puis relancez.\n" |
| "Sur cette machine (sans GPU, 893 Go de checkpoint), l'entrainement " |
| "n'est pas executable : utilisez --dry-run pour valider les donnees.", |
| file=sys.stderr, |
| ) |
| return 2 |
|
|
| if not torch.cuda.is_available(): |
| print( |
| "\nAucun GPU detecte. L'affinage de ce checkpoint exige plusieurs " |
| "accelerateurs. Utilisez --dry-run pour valider les donnees.", |
| file=sys.stderr, |
| ) |
| return 2 |
|
|
| print("\n=== Chargement du modele ===") |
| tokenizer = AutoTokenizer.from_pretrained(config.model_id) |
| model = AutoModelForCausalLM.from_pretrained( |
| config.model_id, torch_dtype=torch.bfloat16, device_map="auto" |
| ) |
|
|
| if config.method == "lora": |
| model = get_peft_model( |
| model, |
| LoraConfig( |
| r=config.lora.r, |
| lora_alpha=config.lora.alpha, |
| lora_dropout=config.lora.dropout, |
| target_modules=config.lora.target_modules, |
| task_type="CAUSAL_LM", |
| ), |
| ) |
| model.print_trainable_parameters() |
|
|
| dataset = Dataset.from_dict({"text": encoded}) |
|
|
| def tokenize(batch): |
| return tokenizer(batch["text"], truncation=True, max_length=config.max_seq_length) |
|
|
| dataset = dataset.map(tokenize, batched=True, remove_columns=["text"]) |
|
|
| training_args = TrainingArguments( |
| output_dir=config.output_dir, |
| num_train_epochs=config.epochs, |
| per_device_train_batch_size=config.batch_size, |
| gradient_accumulation_steps=config.gradient_accumulation_steps, |
| learning_rate=config.learning_rate, |
| warmup_ratio=config.warmup_ratio, |
| weight_decay=config.weight_decay, |
| logging_steps=config.logging_steps, |
| save_steps=config.save_steps, |
| bf16=True, |
| seed=config.seed, |
| report_to=[], |
| ) |
|
|
| trainer = Trainer( |
| model=model, |
| args=training_args, |
| train_dataset=dataset, |
| data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False), |
| ) |
| trainer.train() |
| trainer.save_model(config.output_dir) |
| tokenizer.save_pretrained(config.output_dir) |
| print(f"\nModele enregistre dans {config.output_dir}") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
