eihab2342/code-efficiency

1# src/preprocessing/tokenizer.py2"""3Wrapper على tokenizer CodeT5.4Input format:  "optimize: <slow_code>"5Output format: <fast_code>6"""7from transformers import AutoTokenizer8from src.utils.config import PRETRAINED_MODEL, MAX_INPUT_LEN, MAX_TARGET_LEN9from src.utils.logger import get_logger10 11log = get_logger("tokenizer")12 13 14def get_tokenizer():15    log.info(f"تحميل tokenizer: {PRETRAINED_MODEL}")16    return AutoTokenizer.from_pretrained(PRETRAINED_MODEL)17 18 19def encode_sample(tokenizer, slow_code: str, fast_code: str | None = None) -> dict:20    """21    يحوّل sample واحد لـ tensors جاهزة للموديل.22    لو fast_code=None → inference mode (بدون labels).23    """24    inputs = tokenizer(25        f"optimize: {slow_code}",26        max_length=MAX_INPUT_LEN,27        padding="max_length",28        truncation=True,29        return_tensors="pt",30    )31 32    if fast_code is None:33        return inputs34 35    with tokenizer.as_target_tokenizer():36        targets = tokenizer(37            fast_code,38            max_length=MAX_TARGET_LEN,39            padding="max_length",40            truncation=True,41            return_tensors="pt",42        )43 44    labels = targets["input_ids"]45    labels[labels == tokenizer.pad_token_id] = -100   # ignore padding في الـ loss46    inputs["labels"] = labels47    return inputs48