AURASHIELDV3 / data /preprocessing.py
ANJU DEVI
AuraShield V3 Docker deployment
6960b79
Raw
History Blame Contribute Delete
1.95 kB
import re
def clean_code(code):
"""
Basic code cleaning: removes empty lines and trailing whitespace.
"""
lines = code.split('\n')
cleaned_lines = [line.rstrip() for line in lines if line.strip()]
return '\n'.join(cleaned_lines)
def create_pairs(data):
"""
Extracts vulnerable and secure code pairs.
Returns a list of tuples: (vulnerable_code, secure_code)
"""
pairs = []
for item in data:
vuln = item.get('vulnerable_code', '')
secure = item.get('secure_code', '')
if vuln and secure:
pairs.append((vuln, secure))
return pairs
def prepare_training_data(data):
"""
Formats data for LLM Fine-tuning (Prompt-Response format).
"""
training_samples = []
for item in data:
vuln_code = item.get('vulnerable_code', '')
secure_code = item.get('secure_code', '')
prompt = f"### Instruction:\nFix the following vulnerable code securely:\n\n### Input:\n{vuln_code}\n\n### Response:\n"
completion = f"{secure_code}<|endoftext|>"
training_samples.append({
"prompt": prompt,
"completion": completion
})
return training_samples
import hashlib
def check_data_leakage(train_set, test_set):
"""
Checks if any code snippets from the train set appear in the test set.
"""
train_hashes = {hashlib.md5(item['code'].strip().encode()).hexdigest() for item in train_set}
leaked_count = 0
for item in test_set:
test_hash = hashlib.md5(item['code'].strip().encode()).hexdigest()
if test_hash in train_hashes:
leaked_count += 1
return leaked_count == 0, leaked_count
def tokenize_data(tokenizer, data_samples, max_length=512):
"""
Tokenizes data for model training.
(Requires a tokenizer from transformers)
"""
# Placeholder for actual tokenization logic
pass