pranavKHF/code_vulnerability_detection

1from transformers import RobertaTokenizer2from torch.utils.data import Dataset, DataLoader3import torch4import json5from pathlib import Path6 7 8class VulnerabilityDataset(Dataset):9    """PyTorch dataset for vulnerability detection"""10 11    def __init__(self, data_path, tokenizer, max_length=512):12        self.tokenizer = tokenizer13        self.max_length = max_length14 15        self.data = []16        data_path = Path(data_path)17 18        if not data_path.exists():19            raise FileNotFoundError(f"Dataset file not found: {data_path}")20 21        with open(data_path, "r", encoding="utf-8") as f:22            for line in f:23                line = line.strip()24                if line:25                    self.data.append(json.loads(line))26 27        print(f"{data_path.name}: {len(self.data)} samples")28 29    def __len__(self):30        return len(self.data)31 32    def __getitem__(self, idx):33        sample = self.data[idx]34 35        code = sample["func"]          # confirmed correct36        label = sample["target"]       # confirmed correct (0/1)37 38        encoding = self.tokenizer(39            code,40            truncation=True,41            padding="max_length",42            max_length=self.max_length,43            return_tensors="pt"44        )45 46        return {47            "input_ids": encoding["input_ids"].squeeze(0),48            "attention_mask": encoding["attention_mask"].squeeze(0),49            "labels": torch.tensor(label, dtype=torch.long)50        }51 52 53def load_tokenizer(model_name="Salesforce/codet5-base"):54    print(f"Tokenizer: {model_name}")55    return RobertaTokenizer.from_pretrained(model_name)56 57 58def create_dataloader(59    train_path,60    valid_path,61    test_path,62    tokenizer,63    batch_size=8,64    max_length=512,65    num_workers=2,66):67    train_dataset = VulnerabilityDataset(train_path, tokenizer, max_length)68    valid_dataset = VulnerabilityDataset(valid_path, tokenizer, max_length)69    test_dataset = VulnerabilityDataset(test_path, tokenizer, max_length)70 71    if len(train_dataset) == 0:72        raise RuntimeError(f"No samples found in {train_path}")73 74    train_loader = DataLoader(75        train_dataset,76        batch_size=batch_size,77        shuffle=True,78        num_workers=num_workers,79        pin_memory=True,80        persistent_workers=True81    )82 83    valid_loader = DataLoader(84        valid_dataset,85        batch_size=batch_size,86        shuffle=False,87        num_workers=num_workers,88        pin_memory=True,89        persistent_workers=True90    )91 92    test_loader = DataLoader(93        test_dataset,94        batch_size=batch_size,95        shuffle=False,96        num_workers=num_workers,97        pin_memory=True,98        persistent_workers=True99    )100 101    return train_loader, valid_loader, test_loader102