eunJ/codebert_vulnerability_detector_multi

CodeBERT Vulnerability Detector (Multi-class)

C/C++ 코드의 취약점을 탐지하는 다중 클래스 분류 모델입니다.

모델 정보

• 기반 모델: microsoft/codebert-base
• 분류 클래스: 4개 (CWE-79, CWE-89, CWE-119, 기타)
• 입력: C/C++ 소스 코드 텍스트

사용 방법

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# 모델 로드
model_name = "eunJ/codebert_vulnerability_detector_multi"
tokenizer = AutoTokenizer.from_pretrained(model_name)  
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# 코드 분석
code = '''
char buffer[100];
gets(buffer);  
'''

inputs = tokenizer(code, return_tensors="pt", max_length=512, truncation=True)
with torch.no_grad():
    outputs = model(**inputs)
    predictions = torch.softmax(outputs.logits, dim=-1)
    predicted_class = torch.argmax(predictions)

print(f"예측 클래스: {predicted_class.item()}")

클래스 레이블

• 0: CWE-79 (Cross-site Scripting)
• 1: CWE-89 (SQL Injection)
• 2: CWE-119 (Buffer Overflow)
• 3: CWE-Other (기타)