forked from Eshe/competition-vd
57 lines
2.3 KiB
Python
Executable File
57 lines
2.3 KiB
Python
Executable File
import torch
|
|
import torch.nn as nn
|
|
from transformers import RobertaModel, RobertaTokenizer
|
|
|
|
class CodeBERTClassifier(nn.Module):
|
|
def __init__(self, num_labels):
|
|
super(CodeBERTClassifier, self).__init__()
|
|
self.codebert = RobertaModel.from_pretrained("microsoft/codebert-base")
|
|
self.dropout = nn.Dropout(0.1)
|
|
self.classifier = nn.Linear(self.codebert.config.hidden_size, num_labels)
|
|
|
|
def forward(self, input_ids, attention_mask):
|
|
outputs = self.codebert(input_ids=input_ids, attention_mask=attention_mask)
|
|
pooled_output = outputs[1]
|
|
pooled_output = self.dropout(pooled_output)
|
|
logits = self.classifier(pooled_output)
|
|
return logits
|
|
|
|
class GNNLayer(nn.Module):
|
|
def __init__(self, in_features, out_features):
|
|
super(GNNLayer, self).__init__()
|
|
self.linear = nn.Linear(in_features, out_features)
|
|
|
|
def forward(self, x, adj):
|
|
x = self.linear(x)
|
|
x = torch.matmul(adj, x)
|
|
return torch.relu(x)
|
|
|
|
class GNNClassifier(nn.Module):
|
|
def __init__(self, input_dim, hidden_dim, num_classes):
|
|
super(GNNClassifier, self).__init__()
|
|
self.gnn1 = GNNLayer(input_dim, hidden_dim)
|
|
self.gnn2 = GNNLayer(hidden_dim, num_classes)
|
|
|
|
def forward(self, x, adj):
|
|
x = self.gnn1(x, adj)
|
|
x = self.gnn2(x, adj)
|
|
return x
|
|
|
|
class MultitaskModel(nn.Module):
|
|
def __init__(self, num_vulnerability_types, num_cwe_types):
|
|
super(MultitaskModel, self).__init__()
|
|
self.codebert = RobertaModel.from_pretrained("microsoft/codebert-base")
|
|
self.dropout = nn.Dropout(0.1)
|
|
self.vulnerability_classifier = nn.Linear(self.codebert.config.hidden_size, num_vulnerability_types)
|
|
self.cwe_classifier = nn.Linear(self.codebert.config.hidden_size, num_cwe_types)
|
|
|
|
def forward(self, input_ids, attention_mask):
|
|
outputs = self.codebert(input_ids=input_ids, attention_mask=attention_mask)
|
|
pooled_output = outputs[1]
|
|
pooled_output = self.dropout(pooled_output)
|
|
vulnerability_logits = self.vulnerability_classifier(pooled_output)
|
|
cwe_logits = self.cwe_classifier(pooled_output)
|
|
return vulnerability_logits, cwe_logits
|
|
|
|
def get_tokenizer():
|
|
return RobertaTokenizer.from_pretrained("microsoft/codebert-base") |