eihab2342/code-efficiency

1# src/model/trainer.py2"""3Fine-tuning loop لـ CodeT5 على PIE4Perf.4 5تشغيل:6    python -m src.model.trainer7 8المخرجات:9    models/codet5-finetuned/   (weights + tokenizer)10"""11import torch12from torch.utils.data import DataLoader13from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, get_linear_schedule_with_warmup14from torch.optim import AdamW15from tqdm import tqdm16 17from src.preprocessing.data_loader import load_dataset18from src.preprocessing.dataset import CodeOptDataset19from src.utils.config import (20    PRETRAINED_MODEL, FINETUNED_MODEL,21    BATCH_SIZE, EPOCHS, LEARNING_RATE22)23from src.utils.logger import get_logger24 25log = get_logger("trainer")26 27 28def train():29    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")30    log.info(f"Device: {device}")31 32    # ── 1. Data ───────────────────────────────────────33    train_samples, val_samples, _ = load_dataset()34 35    tokenizer = AutoTokenizer.from_pretrained(PRETRAINED_MODEL)36 37    train_ds = CodeOptDataset(train_samples, tokenizer)38    val_ds   = CodeOptDataset(val_samples,   tokenizer)39 40    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True)41    val_loader   = DataLoader(val_ds,   batch_size=BATCH_SIZE)42 43    # ── 2. Model ──────────────────────────────────────44    log.info(f"تحميل الموديل: {PRETRAINED_MODEL}")45    model = AutoModelForSeq2SeqLM.from_pretrained(PRETRAINED_MODEL).to(device)46 47    # ── 3. Optimizer + Scheduler ──────────────────────48    optimizer = AdamW(model.parameters(), lr=LEARNING_RATE)49    total_steps = len(train_loader) * EPOCHS50    scheduler = get_linear_schedule_with_warmup(51        optimizer,52        num_warmup_steps=int(0.1 * total_steps),53        num_training_steps=total_steps,54    )55 56    best_val_loss = float("inf")57 58    # ── 4. Training loop ──────────────────────────────59    for epoch in range(1, EPOCHS + 1):60        # Train61        model.train()62        train_loss = 0.063        for batch in tqdm(train_loader, desc=f"Epoch {epoch}/{EPOCHS} [train]"):64            batch = {k: v.to(device) for k, v in batch.items()}65            outputs = model(**batch)66            loss = outputs.loss67            loss.backward()68            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)69            optimizer.step()70            scheduler.step()71            optimizer.zero_grad()72            train_loss += loss.item()73 74        avg_train = train_loss / len(train_loader)75 76        # Validation77        model.eval()78        val_loss = 0.079        with torch.no_grad():80            for batch in tqdm(val_loader, desc=f"Epoch {epoch}/{EPOCHS} [val]"):81                batch = {k: v.to(device) for k, v in batch.items()}82                outputs = model(**batch)83                val_loss += outputs.loss.item()84 85        avg_val = val_loss / len(val_loader)86        log.info(f"Epoch {epoch}: train_loss={avg_train:.4f}  val_loss={avg_val:.4f}")87 88        # Save best89        if avg_val < best_val_loss:90            best_val_loss = avg_val91            FINETUNED_MODEL.mkdir(parents=True, exist_ok=True)92            model.save_pretrained(FINETUNED_MODEL)93            tokenizer.save_pretrained(FINETUNED_MODEL)94            log.info(f"✅ Saved best model → {FINETUNED_MODEL}")95 96    log.info("Training done!")97 98 99if __name__ == "__main__":100    train()101