基于NLP的恶意网页识别
基于NLP的恶意网页识别
基于NLP的恶意网页识别
引言
欢迎阅读《 基于NLP的恶意网页识别》，在前三篇中，我们已经使用PaddleNLP进行了恶意网页的分类，包括使用文本分类模型和预训练模型Fine-tune。本篇文章将着重优化模型，处理HTML标签提取结果不理想的情况，并最终将训练好的模型部署成可用的Python应用程序。
项目目录回顾
在前三篇文章中，我们已经完成了以下内容：
- 使用PaddleNLP的文本分类模型进行简单的正常网页与被黑网页的二分类。
- 使用PaddleNLP的预训练模型Fine-tune，提高HTML网页内容处理结果的判断准确率。
- 进行正常网页与恶意网页的二分类，提取HTML标签信息判断网页是否正常。
本篇文章将以第四篇为基础，继续优化HTML标签提取结果，训练并评估模型，并最终将模型部署成可用的Python应用程序。
优化HTML标签提取结果
在之前的训练中，我们发现有些样本的HTML标签提取结果不够理想，主要集中在<script>标签内的信息清理不完整。为了解决这个问题，我们可以设计逻辑判断，将这部分内容留给下一个流程（比如提取文本信息或人工核验）。通过观察样本，我们发现这部分内容大多是恶意网页，所以即使略过也不会对模型的准确性产生太大影响。接下来，我们将优化标签提取结果。
with open("train_list.txt", "r", encoding="utf-8") as f:
lines = f.readlines()
with open("train_list2.txt", "w", encoding="utf-8") as f_w:
for line in lines:
if ";" in line or "+" in line:
continue
f_w.write(line)
with open("eval_list.txt", "r", encoding="utf-8") as f:
lines = f.readlines()
with open("eval_list2.txt", "w", encoding="utf-8") as f_w:
for line in lines:
if ";" in line or "+" in line:
continue
f_w.write(line)
这段代码将原始的训练集和验证集中包含<script>标签内信息不完整的样本去除，得到新的训练集train_list2.txt和验证集eval_list2.txt。
使用预训练模型Fine-tune
接下来，我们使用PaddleNLP进行预训练模型Fine-tune，以优化HTML标签提取结果的预测准确率。首先，我们需要定义自定义数据集和加载预训练模型。
from paddlenlp.datasets import load_dataset
def read(data_path):
with open(data_path, 'r', encoding='utf-8') as f:
for line in f:
line = line.strip('n').split('t')
words = ''.join(line[:-1])
labels = line[-1]
yield {'text': words, 'label': labels}
train_ds = load_dataset(read, data_path='train_list2.txt', lazy=False)
dev_ds = load_dataset(read, data_path='eval_list2.txt', lazy=False)
# 手动添加标签列表
train_ds.label_list = ['0', '1']
dev_ds.label_list = ['0', '1']
这段代码定义了自定义数据集，并手动添加了标签列表。接下来，我们加载预训练模型和tokenizer。
MODEL_NAME = "ernie-2.0-large-en"
ernie_model = paddlenlp.transformers.ErnieModel.from_pretrained(MODEL_NAME)
model = paddlenlp.transformers.ErnieForSequenceClassification.from_pretrained(MODEL_NAME, num_classes=len(train_ds.label_list))
tokenizer = paddlenlp.transformers.ErnieTokenizer.from_pretrained(MODEL_NAME)
数据处理和模型训练
我们使用PaddleNLP提供的数据处理和模型训练的接口，定义了数据处理函数和模型训练的参数。
from functools import partial
from paddlenlp.data import Stack, Tuple, Pad
from utils import convert_example, create_dataloader
batch_size = 128
max_seq_length = 64
trans_func = partial(
convert_example,
tokenizer=tokenizer,
max_seq_length=max_seq_length)
batchify_fn = lambda samples, fn=Tuple(
Pad(axis=0, pad_val=tokenizer.pad_token_id),
Pad(axis=0, pad_val=tokenizer.pad_token_type_id),
Stack(dtype="int64")): [data for data in fn(samples)]
train_data_loader = create_dataloader(
train_ds,
mode='train',
batch_size=batch_size,
batchify_fn=batchify_fn,
trans_fn=trans_func)
dev_data_loader = create_dataloader(
dev_ds,
mode='dev',
batch_size=batch_size,
batchify_fn=batchify_fn,
trans_fn=trans_func)
learning_rate = 5e-6
epochs = 5
warmup_proportion = 0.1
weight_decay = 0.1
num_training_steps = len(train_data_loader) * epochs
lr_scheduler = LinearDecayWithWarmup(learning_rate, num_training_steps, warmup_proportion)
optimizer = paddle.optimizer.AdamW(
learning_rate=lr_scheduler,
parameters=model.parameters(),
weight_decay=weight_decay,
apply_decay_param_fun=lambda x: x in [
p.name for n, p in model.named_parameters()
if not any(nd in n for nd in ["bias", "norm"])
])
criterion = paddle.nn.loss.CrossEntropyLoss()
metric = paddle.metric.Accuracy()
模型训练与评估
我们使用PaddleNLP提供的训练和评估接口，进行模型的训练与评估。同时，使用VisualDL进行可视化记录。
global_step = 0
for epoch in range(
1, epochs + 1):
with LogWriter(logdir="./visualdl") as writer:
for step, batch in enumerate(train_data_loader, start=1):
input_ids, segment_ids, labels = batch
logits = model(input_ids, segment_ids)
loss = criterion(logits, labels)
probs = F.softmax(logits, axis=1)
correct = metric.compute(probs, labels)
metric.update(correct)
acc = metric.accumulate()
global_step += 1
if global_step % 50 == 0:
print("global step %d, epoch: %d, batch: %d, loss: %.5f, acc: %.5f" % (
global_step, epoch, step, loss, acc))
writer.add_scalar(tag="loss", step=global_step, value=loss)
writer.add_scalar(tag="acc", step=global_step, value=acc)
loss.backward()
optimizer.step()
lr_scheduler.step()
optimizer.clear_grad()
evaluate(model, criterion, metric, dev_data_loader)
model.save_pretrained('/home/aistudio/checkpoint')
tokenizer.save_pretrained('/home/aistudio/checkpoint')
模型导出
训练完成后，我们可以将模型导出为静态图参数，以便后续部署使用。
state_dict = paddle.load('/home/aistudio/checkpoint/model_state.pdparams')
model.set_dict(state_dict)
model.eval()
model = paddle.jit.to_static(
model,
input_spec=[
paddle.static.InputSpec(
shape=[None, None], dtype="int64"),
paddle.static.InputSpec(
shape=[None, None], dtype="int64")
])
paddle.jit.save(model, '/home/aistudio/static_graph_params')
部署与预测
最后，我们将训练好的模型导出并进行部署，以便进行预测。这里我们假设已经准备好了一个HTML页面的内容，可以使用BeautifulSoup进行解析，提取HTML标签信息，然后使用训练好的模型进行预测。
html = BeautifulSoup(open('sample.html'), 'html.parser', from_encoding='utf-8')
def read_tags(text):
tags = []
class MyHTMLParser2(HTMLParser):
def handle_endtag(self, tag):
tags.append(tag)
parser = MyHTMLParser2()
parser.feed(text)
return tags
text = ','.join(read_tags(str(html.get_text)))
with open('sample.txt', 'w', encoding='utf-8') as f:
f.write(text)
!python predict.py --model_file=static_graph_params.pdmodel --params_file=static_graph_params.pdiparams
以上代码将HTML页面的标签信息提取并保存到sample.txt文件中，然后使用训练好的模型进行预测。
总结
在本文中，我们通过优化HTML标签提取结果，使用PaddleNLP进行预训练模型Fine-tune，最终将训练好的模型导出并部署成可用的Python应用程序。这一系列步骤构建了一个完整的恶意网页识别系统，可以帮助企业更好地保护用户免受网络攻击。在未来的工作中，我们可以考虑将网页内容的其他组成部分，如URL链接、图片信息等，加入到系统中，进一步提升恶意网页识别的准确性。