import wandb
import yaml
from transformers import Trainer, TrainingArguments
from utils.monitor import measure_resources
from data.datasets import load_and_tokenize_data
from models.pert_model import get_pert_model

# Charger la configuration
with open('config/config.yaml', 'r') as f:
    config = yaml.safe_load(f)

# Initialiser wandb
wandb.init(project=config['wandb']['project'], entity=config['wandb']['entity'])

# Charger les données
train_dataset, test_dataset = load_and_tokenize_data(config)

# Charger le modèle
model = get_pert_model(config)

# Définir les arguments de formation
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=config['training']['num_epochs'],
    per_device_train_batch_size=config['training']['batch_size'],
    per_device_eval_batch_size=config['training']['batch_size'],
    evaluation_strategy='epoch',
    save_steps=10_000,
    save_total_limit=2,
    logging_dir='./logs',
    logging_steps=10,
)

# Créer le Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
)

# Mesurer les ressources et entraîner le modèle
measure_resources(trainer, "PERT (Adapters)")