AI red teamer （人工智能红队）系列32 – 人工智能信息安全应用 – 从零构建影评分类模型

数据介绍

Maas 等人（2011 年）推出的 IMDB 数据集收集了从互联网电影数据库中提取的电影评论，并为情感分析添加了注释。它包括 50,000 篇评论，平均分为训练集和测试集，其精心策划的正面和负面例子的混合使研究人员能够对各种自然语言处理技术进行基准测试和改进。IMDB 数据集对后来开发基于向量的单词表示法的工作产生了影响，并且仍然是评估情感分类任务中分类性能和模型架构的常用基准资源(Maas et al., 2011).

目标

训练一个可以预测电影评论是正面（1）还是负面（）的模型，从此处下载 数据集。

思路

利用TF-IDF向量化技术将原始文本评论转换为模型可理解的数值特征。构建包含TF-IDF向量器和逻辑回归分类器的机器学习管道，并在训练数据上对其进行拟合，最后得到模型。

下载解压数据集

import requests
import zipfile
import io
import os

# 下载数据集
url = "https://academy.hackthebox.com/storage/modules/292/skills_assessment_data.zip"
response = requests.get(url)

if response.status_code == 200:
    print("下载成功")

    # 解压数据集
    with zipfile.ZipFile(io.BytesIO(response.content)) as z:
        z.extractall("skills_assessment_data")
        print("解压成功")

    # 列出解压后的文件
    extracted_files = os.listdir("skills_assessment_data")
    print("解压的文件:", extracted_files)
else:
    print("下载失败")

导入必要模块

import pandas as pd
import numpy as np
import json
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report, accuracy_score
import joblib

• pandas 和 numpy 是数据处理的常用工具。
• json 用于加载和解析JSON格式的数据文件。
• sklearn.model_selection.train_test_split 用于数据分割。
• sklearn.feature_extraction.text.TfidfVectorizer 将文本数据转换成数值特征向量。
• sklearn.linear_model.LogisticRegression 是我们选择的分类算法。
• sklearn.metrics.classification_report 和 accuracy_score 用于评估模型性能。
• joblib 库用于高效地保存和加载训练好的模型。

加载和准备数据

import json

# 加载JSON数据，并指定UTF-8编码
with open("skills_assessment_data/train.json", "r", encoding="utf-8") as f:
    train_data = json.load(f)

with open("skills_assessment_data/test.json", "r", encoding="utf-8") as f:
    test_data = json.load(f)

# 提取特征 (X) 和标签 (y)
X_train = [sample["text"] for sample in train_data]
y_train = [sample["label"] for sample in train_data]

X_test = [sample["text"] for sample in test_data]
y_test = [sample["label"] for sample in test_data]

构建和训练机器学习管道

from sklearn.pipeline import Pipeline
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression

# 创建机器学习管道
pipeline = Pipeline([
    ("vectorizer", TfidfVectorizer()),  # 第一步：文本向量化
    ("classifier", LogisticRegression(max_iter=1000)) # 第二步：逻辑回归分类器，增加迭代次数以确保收敛
])

# 在训练数据上拟合管道
pipeline.fit(X_train, y_train)

模型评估

from sklearn.metrics import classification_report, accuracy_score

# 使用训练好的管道对测试集进行预测
y_pred = pipeline.predict(X_test)

# 打印模型准确率
print("准确率 (Accuracy):", accuracy_score(y_test, y_pred))

# 打印详细的分类报告（包含精确率、召回率、F1分数等）
print("分类报告 (Classification Report):\n", classification_report(y_test, y_pred))

保存模型

import joblib

# 保存完整的、已训练好的模型
joblib.dump(pipeline, "skills_assessment.joblib")