基于NLP的恶意网页识别

引言

欢迎阅读《 基于NLP的恶意网页识别》，在前三篇中，我们已经使用PaddleNLP进行了恶意网页的分类，包括使用文本分类模型和预训练模型Fine-tune。本篇文章将着重优化模型，处理HTML标签提取结果不理想的情况，并最终将训练好的模型部署成可用的Python应用程序。

项目目录回顾

在前三篇文章中，我们已经完成了以下内容：

1. 使用PaddleNLP的文本分类模型进行简单的正常网页与被黑网页的二分类。
2. 使用PaddleNLP的预训练模型Fine-tune，提高HTML网页内容处理结果的判断准确率。
3. 进行正常网页与恶意网页的二分类，提取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:continuef_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:continuef_w.write(line)

这段代码将原始的训练集和验证集中包含<script>标签内信息不完整的样本去除，得到新的训练集train_list2.txt和验证集eval_list2.txt。

使用预训练模型Fine-tune

接下来，我们使用PaddleNLP进行预训练模型Fine-tune，以优化HTML标签提取结果的预测准确率。首先，我们需要定义自定义数据集和加载预训练模型。

from paddlenlp.datasets import load_datasetdef 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_dataloaderbatch_size = 128
max_seq_length = 64trans_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.1num_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 = batchlogits = 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 += 1if 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 tagstext = ','.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链接、图片信息等，加入到系统中，进一步提升恶意网页识别的准确性。