SecureBERT_Plus

SecureBERT_Plus项目介绍

SecureBERT_Plus是一种专为网络安全领域设计的基础语言模型，这个模型的前身为SecureBERT。它基于RoBERTa模型，借助更大规模的数据集与强大的计算能力进行了训练。在原SecureBERT模型的基础上，SecureBERT_Plus在屏蔽语言模型（Masked Language Model, MLM）任务上实现了约9%的性能提升。这一改进标志着在理解和表示网络安全相关文本数据方面取得了显著进步。

数据集

SecureBERT_Plus利用了比之前版本多八倍的数据集进行训练，这样的增加使得模型能够更全面地理解和解析网络安全领域的各种文本信息。

如何加载模型

SecureBERT_Plus已上传到Huggingface平台，用户可以直接通过Huggingface的Transformers库来加载和使用该模型。以下是加载模型的简单代码示例：

from transformers import RobertaTokenizer, RobertaModel
import torch

tokenizer = RobertaTokenizer.from_pretrained("ehsanaghaei/SecureBERT_Plus")
model = RobertaModel.from_pretrained("ehsanaghaei/SecureBERT_Plus")

inputs = tokenizer("This is SecureBERT Plus!", return_tensors="pt")
outputs = model(**inputs)

last_hidden_states = outputs.last_hidden_state

屏蔽词填充（MLM）

为了预测文本中被屏蔽的词汇，可以使用以下代码。在文本中插入 <mask> 标记以指示需要预测的词语。

import torch
import transformers
from transformers import RobertaTokenizerFast

tokenizer = RobertaTokenizerFast.from_pretrained("ehsanaghaei/SecureBERT_Plus")
model = transformers.RobertaForMaskedLM.from_pretrained("ehsanaghaei/SecureBERT_Plus")

def predict_mask(sent, tokenizer, model, topk=10, print_results=True):
    token_ids = tokenizer.encode(sent, return_tensors='pt')
    masked_position = (token_ids.squeeze() == tokenizer.mask_token_id).nonzero()
    masked_pos = [mask.item() for mask in masked_position]
    words = []
    with torch.no_grad():
        output = model(token_ids)
    
    last_hidden_state = output[0].squeeze()
    list_of_list = []
    for index, mask_index in enumerate(masked_pos):
        mask_hidden_state = last_hidden_state[mask_index]
        idx = torch.topk(mask_hidden_state, k=topk, dim=0)[1]
        words = [tokenizer.decode(i.item()).strip() for i in idx]
        words = [w.replace(' ', '') for w in words]
        list_of_list.append(words)
        if print_results:
            print("Mask ", "Predictions: ", words)

    return words

while True:
    sent = input("Text here: \t")
    print("SecureBERT: ")
    predict_mask(sent, tokenizer, model)
    print("===========================\n")