eihab2342/code-efficiency

1# src/preprocessing/data_loader.py2"""3يحمّل داتاست PIE4Perf من الملفات المحلية.4 5الملفات المطلوبة في data/raw/:6    train.jsonl   للتدريب7    val.jsonl     للـ validation8    test.jsonl    للـ evaluation9 10تشغيل:11    py -3.11 -m src.preprocessing.data_loader12"""13 14import json15from pathlib import Path16from dataclasses import dataclass17from typing import List, Tuple18from src.utils.config import DATA_RAW_DIR19from src.utils.logger import get_logger20 21log = get_logger("data_loader")22 23 24@dataclass25class Sample:26    slow_code: str27    fast_code: str28    speedup: float29 30    def is_valid(self) -> bool:31        return (32            bool(self.slow_code.strip())33            and bool(self.fast_code.strip())34            and self.speedup > 035        )36 37 38def _parse(record: dict) -> "Sample | None":39    try:40        slow = (record.get("input")41                or record.get("code_v0_no_empty_lines")42                or record.get("slow_code") or "")43 44        fast = (record.get("target")45                or record.get("code_v1_no_empty_lines")46                or record.get("fast_code") or "")47 48        cpu0 = record.get("cpu_time_v0")49        cpu1 = record.get("cpu_time_v1")50        if cpu0 and cpu1 and float(cpu1) > 0:51            speedup = float(cpu0) / float(cpu1)52        elif record.get("improvement_frac"):53            speedup = 1 + float(record["improvement_frac"]) / 10054        else:55            speedup = 1.056 57        s = Sample(slow_code=slow.strip(), fast_code=fast.strip(), speedup=speedup)58        return s if s.is_valid() else None59    except Exception:60        return None61 62 63def load_raw(path: Path) -> List[Sample]:64    """يقرأ ملف JSONL واحد ويرجع list من الـ samples"""65    if not path.exists():66        log.warning(f"الملف مش موجود: {path.name}")67        return []68 69    samples = []70    with open(path, encoding="utf-8") as f:71        for line in f:72            line = line.strip()73            if not line:74                continue75            try:76                record = json.loads(line)77                s = _parse(record)78                if s:79                    samples.append(s)80            except json.JSONDecodeError:81                continue82 83    log.info(f"تحميل {path.name}  ←  {len(samples)} sample")84    return samples85 86 87def load_dataset() -> Tuple[List[Sample], List[Sample], List[Sample]]:88    """89    يقرأ train.jsonl و val.jsonl و test.jsonl من data/raw/90    ويرجع (train, val, test)91    """92    train = load_raw(DATA_RAW_DIR / "train.jsonl")93    val   = load_raw(DATA_RAW_DIR / "val.jsonl")94    test  = load_raw(DATA_RAW_DIR / "test.jsonl")95 96    # لو مفيش داتا خالص استخدم dummy97    if not train and not val and not test:98        log.warning("مفيش داتا في data/raw/  ← بستخدم dummy samples")99        dummy = _dummy()100        n = len(dummy)101        return dummy[:int(n*0.8)], dummy[int(n*0.8):int(n*0.9)], dummy[int(n*0.9):]102 103    log.info(f"Dataset: {len(train)} train / {len(val)} val / {len(test)} test")104    return train, val, test105 106 107def _dummy() -> List[Sample]:108    return [109        Sample("result=[]\nfor x in a:\n    result.append(x*2)",110               "result=[x*2 for x in a]", 1.3),111        Sample("s=''\nfor w in words:\n    s+=w+' '",112               "s=' '.join(words)", 5.0),113        Sample("def f(x=[]):\n    x.append(1)\n    return x",114               "def f(x=None):\n    if x is None: x=[]\n    x.append(1)\n    return x", 1.0),115    ]116 117 118if __name__ == "__main__":119    train, val, test = load_dataset()120    print(f"\nTrain: {len(train)}  Val: {len(val)}  Test: {len(test)}")121    if train:122        print("\nمثال من Train:")123        print("SLOW:", train[0].slow_code[:80])124        print("FAST:", train[0].fast_code[:80])125        print("Speedup:", round(train[0].speedup, 2))