Views
No views yet
microsoft/unixcoder-baseDetectVul/devign1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2import torch
3
4# Load the fine-tuned model
5tokenizer = AutoTokenizer.from_pretrained("microsoft/unixcoder-base")
6model = AutoModelForSequenceClassification.from_pretrained("mahdin70/unixcoder-code-vulnerability-detector")
7
8# Sample code snippet
9code_snippet = """
10void process(char *input) {
11 char buffer[50];
12 strcpy(buffer, input); // Potential buffer overflow
13}
14"""
15
16# Tokenize the input
17inputs = tokenizer(code_snippet, return_tensors="pt", truncation=True, padding="max_length", max_length=512)
18
19# Run inference
20with torch.no_grad():
21 outputs = model(**inputs)
22 predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
23 predicted_label = torch.argmax(predictions, dim=1).item()
24
25# Output the result
26print("Vulnerable Code" if predicted_label == 1 else "Safe Code")DetectVul/devign0 (Safe), 1 (Vulnerable)| Metric | Score |
|---|---|
| Train Loss | 0.4835 |
| Evaluation Loss | 0.6855 |
| Accuracy | 68.34% |
| F1 Score | 62.14% |
| Precision | 69.18% |
| Recall | 56.40% |
| Factor | Value |
|---|---|
| GPU Used | 2x T4 GPU |
| Training Time | ~1 hour |