eihab2342/code-efficiency

1# src/preprocessing/dataset.py2"""3PyTorch Dataset — يلف الـ samples ويجهزها للـ DataLoader4"""5import torch6from torch.utils.data import Dataset7from transformers import PreTrainedTokenizer8from src.preprocessing.data_loader import Sample9from src.utils.config import MAX_INPUT_LEN, MAX_TARGET_LEN10from typing import List11 12 13class CodeOptDataset(Dataset):14 15    def __init__(self, samples: List[Sample], tokenizer: PreTrainedTokenizer):16        self.samples   = samples17        self.tokenizer = tokenizer18 19    def __len__(self) -> int:20        return len(self.samples)21 22    def __getitem__(self, idx: int) -> dict:23        s = self.samples[idx]24 25        model_inputs = self.tokenizer(26            f"optimize: {s.slow_code}",27            max_length=MAX_INPUT_LEN,28            padding="max_length",29            truncation=True,30        )31 32        with self.tokenizer.as_target_tokenizer():33            labels = self.tokenizer(34                s.fast_code,35                max_length=MAX_TARGET_LEN,36                padding="max_length",37                truncation=True,38            )39 40        label_ids = labels["input_ids"]41        # ignore padding في الـ cross-entropy loss42        label_ids = [43            l if l != self.tokenizer.pad_token_id else -10044            for l in label_ids45        ]46 47        return {48            "input_ids":      torch.tensor(model_inputs["input_ids"]),49            "attention_mask": torch.tensor(model_inputs["attention_mask"]),50            "labels":         torch.tensor(label_ids),51        }52