Spaces:
Sleeping
Sleeping
| 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 | |