Maaac/CodeLLaMA-Linux-BugFix

1from pydriller import Repository2import os3import json4from tqdm import tqdm5import re6 7REPO_PATH = '../linux'8OUTPUT_FILE = './output/linux_bugfix_dataset.jsonl'9 10TEST_MODE = False  # Set to False to process the full repository11 12BUGFIX_KEYWORDS = [13    'fix', 'bug', 'leak', 'null', 'overflow', 'error', 'failure',14    'crash', 'panic', 'memory', 'race', 'deadlock', 'corruption',15    'security', 'vulnerability', 'exploit', 'buffer', 'stack'16]17 18def is_bugfix_commit(msg):19    msg_lower = msg.lower()20    return any(keyword in msg_lower for keyword in BUGFIX_KEYWORDS)21 22def extract_instruction_from_commit_msg(msg):23    lines = msg.strip().splitlines()24    for line in lines:25        line = line.strip()26        if len(line) < 5 or not any(c.isalpha() for c in line):27            continue28        if line.lower().startswith((29            '[patch]', 'signed-off-by', 'reviewed-by', 'tested-by', 'ack',30            'reported-by', 'cc:', 'co-authored-by', 'patchwork-id',31            'suggested-by', 'fixes:', 'link:', 'cherry picked from commit'32        )):33            continue34        return line35    return msg.strip().splitlines()[0] if msg.strip() else "fix"36 37def extract_code_context(code, line_number, context_lines=10):38    if not code:39        return ""40    lines = code.split('\n')41    start = max(0, line_number - context_lines)42    end = min(len(lines), line_number + context_lines)43    return '\n'.join(lines[start:end])44 45def extract_diff_context(diff_text, context_lines=5):46    if not diff_text:47        return ""48    lines = diff_text.split('\n')49    change_lines = [i for i, line in enumerate(lines) if line.startswith('+') or line.startswith('-')]50    if not change_lines:51        return diff_text52    start = max(0, change_lines[0] - context_lines)53    end = min(len(lines), change_lines[-1] + context_lines + 1)54    return '\n'.join(lines[start:end])55 56def create_dataset_entry(original_code, commit_msg, diff_code):57    return {58        "input": {59            "original code": original_code.strip(),60            "instruction": extract_instruction_from_commit_msg(commit_msg)61        },62        "output": {63            "diff codes": diff_code.strip()64        }65    }66 67def process_commit(commit):68    entries = []69    if not is_bugfix_commit(commit.msg):70        return entries71 72    for mod in commit.modified_files:73        if not mod.new_path or not mod.new_path.endswith(('.c', '.h')):74            continue75        if mod.change_type.name != "MODIFY":76            continue77        if not mod.diff or not mod.source_code_before:78            continue79 80        focused_diff = extract_diff_context(mod.diff)81 82        diff_lines = mod.diff.split('\n')83        line_numbers = []84        for line in diff_lines:85            if line.startswith('@@'):86                match = re.search(r'@@ -(\d+),?\d* \+\d+,?\d* @@', line)87                if match:88                    line_numbers.append(int(match.group(1)))89 90        if line_numbers:91            focused_code = extract_code_context(mod.source_code_before, line_numbers[0])92        else:93            focused_code = '\n'.join(mod.source_code_before.split('\n')[:50])94 95        entry = create_dataset_entry(96            original_code=focused_code,97            commit_msg=commit.msg,98            diff_code=focused_diff99        )100        entries.append(entry)101 102    return entries103 104def main():105    if not os.path.exists(REPO_PATH):106        print(f"\u274c Repository not found at: {REPO_PATH}")107        return108 109    os.makedirs('./output', exist_ok=True)110 111    print("\ud83d\udd0d Building Linux kernel bug-fix dataset...")112    print(f"\ud83d\udcc1 Repository: {REPO_PATH}")113    print(f"\ud83d\udcce Output: {OUTPUT_FILE}")114 115    output_file = OUTPUT_FILE.replace('.jsonl', '_test.jsonl') if TEST_MODE else OUTPUT_FILE116 117    repo = Repository(REPO_PATH)118    dataset_entries = []119    processed_commits = 0120    total_commits = 0121    bugfix_commits = 0122 123    for commit in tqdm(repo.traverse_commits(), desc="Processing commits"):124        total_commits += 1125        if TEST_MODE and MAX_COMMITS_TEST and total_commits > MAX_COMMITS_TEST:126            break127        if is_bugfix_commit(commit.msg):128            bugfix_commits += 1129            entries = process_commit(commit)130            if entries:131                dataset_entries.extend(entries)132                processed_commits += 1133                if TEST_MODE:134                    print(f"\n\ud83d\udd0d Bug-fix commit {processed_commits}: {commit.hash[:8]}")135                    print(f"\ud83d\udcdd Message: {extract_instruction_from_commit_msg(commit.msg)}")136                    print(f"\ud83d\udcca Files: {len(entries)} entries extracted")137                    print(f"\ud83d\udcc1 Files: {[mod.new_path for mod in commit.modified_files if mod.new_path and mod.new_path.endswith(('.c', '.h'))]}")138 139    with open(output_file, 'w', encoding='utf-8') as f:140        for entry in dataset_entries:141            f.write(json.dumps(entry, ensure_ascii=False) + '\n')142 143    print(f"\n\u2705 Dataset creation completed!")144    print(f"\ud83d\udcca Total commits processed: {total_commits}")145    print(f"\ud83d\udc1b Bug-fix commits found: {bugfix_commits}")146    print(f"\ud83d\udcdd Commits with valid entries: {processed_commits}")147    print(f"\ud83d\udcdd Total dataset entries: {len(dataset_entries)}")148    print(f"\ud83d\udcce Saved to: {output_file}")149 150    if dataset_entries:151        print(f"\n\ud83d\udccb Sample dataset entry:")152        sample = dataset_entries[0]153        print(json.dumps(sample, indent=2, ensure_ascii=False)[:800] + "...")154        print(f"\n\ud83d\udcc1 Dataset structure:")155        print(f"   - Input: original code + instruction")156        print(f"   - Output: diff codes")157        print(f"   - Format: JSONL (one JSON object per line)")158 159if __name__ == "__main__":160    main()