python
import json
import logging
import os
import sys
import torch
from peft import LoraConfig, get_peft_model, TaskType, PeftModel
from sklearn.model_selection import train_test_split
from datasets import Dataset # 需要导入 Dataset
from transformers import (
AutoTokenizer,
AutoModelForSequenceClassification,
TrainingArguments,
Trainer,
BitsAndBytesConfig,
DataCollatorWithPadding
)
# 配置日志
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)
class QwenLinuxClassifier:
"""
基于Qwen3-8B模型的文本分类器,使用LoRA技术进行高效微调
Args:
model_dir (str): 模型保存路径,默认为"./models/qwen3_8b_classifier"
Attributes:
base_model_path (str): 基础模型路径
model_save_path (str): 微调模型保存路径
device (torch.device): 计算设备(CUDA/CPU)
label_map (dict): 标签到ID的映射
id_to_label (dict): ID到标签的映射
tokenizer: 分词器对象
model: 模型对象
"""
def __init__(self, model_dir="./models/qwen3_8b_classifier"):
"""
初始化分类器
Args:
model_dir (str): 模型保存目录路径,用于保存微调后的模型权重和配置文件
初始化流程:
1. 设置模型路径配置
2. 配置计算设备(CUDA优先)
3. 定义标签映射关系
4. 加载分词器并配置padding设置
5. 调用load_model()加载或初始化模型
"""
# 1. 路径配置 (适配 Linux)
# 基础模型路径:指向预训练的Qwen3-8B模型目录
self.base_model_path = "./models/Qwen/Qwen3-8B"
# 模型保存路径:使用绝对路径确保Linux环境下的正确性
self.model_save_path = os.path.abspath(model_dir)
# 2. 设备配置
# 自动检测CUDA可用性,优先使用GPU,否则回退到CPU
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 3. 标签映射
# label_map: 将中文标签映射为数字ID,用于模型训练
self.label_map = {"通用知识": 0, "专业咨询": 1}
# id_to_label: 反向映射,用于预测时将数字ID转换回标签名称
self.id_to_label = {v: k for k, v in self.label_map.items()}
# 4. 加载分词器
# 从基础模型路径加载分词器,trust_remote_code=True允许加载自定义代码
self.tokenizer = AutoTokenizer.from_pretrained(
self.base_model_path,
trust_remote_code=True
)
# 设置填充token为结束token(EOS),确保序列长度一致
self.tokenizer.pad_token = self.tokenizer.eos_token
# 设置右侧填充,与大多数因果语言模型的训练方式一致
self.tokenizer.padding_side = "right"
# 模型对象初始化为None,在load_model()中实际加载
self.model = None
# 加载模型权重或初始化新模型
self.load_model()
def load_model(self):
"""
加载模型,优先加载已存在的微调权重,否则加载原始模型并配置LoRA
加载策略:
- 如果model_save_path中存在adapter_config.json,则加载已有的LoRA权重
- 否则加载基础模型并配置新的LoRA适配器
LoRA配置参数说明:
- r (int): LoRA低秩矩阵的秩,控制可训练参数量,越大表达能力越强
- lora_alpha (int): 缩放系数,影响LoRA层的输出幅度
- lora_dropout (float): Dropout率,防止过拟合
- target_modules (list): 应用LoRA的目标层,包括注意力层和MLP层
"""
# QLoRA 配置:4-bit量化配置,大幅降低显存占用
bnb_config = BitsAndBytesConfig(
load_in_4bit=True, # 启用4-bit量化
bnb_4bit_use_double_quant=True, # 启用嵌套量化进一步节省显存
bnb_4bit_quant_type="nf4", # 使用NF4量化类型,精度优于FP4
bnb_4bit_compute_dtype=torch.bfloat16 # 计算时使用bf16,Linux+4090D推荐
)
# 检查是否存在已有的 Adapter 配置文件
adapter_path = os.path.join(self.model_save_path, "adapter_config.json")
if os.path.exists(adapter_path):
# 场景1: 存在微调权重,加载基础模型+Adapter
logger.info(f">>> 发现微调权重,正在加载: {self.model_save_path}")
# 首先加载基础模型(带量化配置)
base_model = AutoModelForSequenceClassification.from_pretrained(
self.base_model_path, # 基础模型路径
attn_implementation="eager", # 使用eager注意力实现,兼容性更好
num_labels=2, # 二分类任务
quantization_config=bnb_config, # 应用4-bit量化
trust_remote_code=True, # 允许加载自定义模型代码
device_map="auto" # 自动分配层到可用设备
)
# 加载预训练的LoRA权重
self.model = PeftModel.from_pretrained(base_model, self.model_save_path)
else:
# 场景2: 不存在微调权重,加载基础模型并配置新LoRA
logger.info(">>> 加载原始模型并配置 LoRA...")
self.model = AutoModelForSequenceClassification.from_pretrained(
self.base_model_path, # 基础模型路径
attn_implementation="eager", # 注意力实现方式
num_labels=2, # 分类类别数
quantization_config=bnb_config, # 量化配置
trust_remote_code=True, # 信任远程代码
device_map="auto" # 自动设备映射
)
# 配置LoRA参数
peft_config = LoraConfig(
task_type=TaskType.SEQ_CLS, # 任务类型:序列分类
inference_mode=False, # 训练模式(非推理模式)
r=16, # LoRA秩:4090D显存充足,可设为16
lora_alpha=32, # 缩放系数:通常设为2*r
lora_dropout=0.05, # Dropout率:防止过拟合
# 目标模块:对哪些层应用LoRA,包含注意力层和MLP层
target_modules=[
"q_proj", "k_proj", "v_proj", "o_proj", # 注意力投影层
"gate_proj", "up_proj", "down_proj" # MLP投影层
]
)
# 将LoRA配置应用到基础模型
self.model = get_peft_model(self.model, peft_config)
# 设置模型的padding token ID,与分词器保持一致
self.model.config.pad_token_id = self.tokenizer.pad_token_id
def preprocess_data(self, texts, labels):
"""
数据预处理:将文本编码为模型输入格式
Args:
texts (list[str]): 文本列表,每个元素是一条待分类的文本字符串
labels (list[str] | list[int]): 对应的标签列表,
可以是字符串标签("通用知识"/"专业咨询")或数字ID(0/1)
Returns:
dict: 包含以下键的字典:
- input_ids (list[list[int]]: 文本的token ID序列
- attention_mask (list[list[int]]): 注意力掩码,1表示真实token,0表示padding
- labels (list[int]): 转换后的数字标签列表
处理流程:
1. 将字符串标签转换为数字ID
2. 使用tokenizer对文本进行编码
3. 将标签添加到编码结果中
"""
# 将标签转换为id: 如果是字符串则从label_map查找,否则直接使用
label_ids = [
self.label_map.get(label, 0) if isinstance(label, str) else label
for label in labels
]
# 分词编码
encodings = self.tokenizer(
texts, # 输入文本列表
truncation=True, # 超长序列截断
padding=False, # 不在此处padding,由DataCollator动态处理
max_length=512, # 最大序列长度
return_tensors=None # 返回Python列表格式,便于转为Dataset
)
# 将标签添加到编码字典中
encodings['labels'] = label_ids
return encodings
def create_dataset(self, encodings):
"""
将编码后的数据转换为 Hugging Face Dataset 格式
Args:
encodings (dict): 编码后的数据字典,包含:
- input_ids: token ID列表的列表
- attention_mask: 注意力掩码列表的列表
- labels: 标签ID列表
Returns:
Dataset: Hugging Face Dataset对象,兼容Trainer API
说明:
Dataset格式是Hugging Face标准数据格式,支持高效的内存映射和批量处理
"""
return Dataset.from_dict(encodings)
def train_model(self, data_file):
"""
训练模型的主要方法
Args:
data_file (str): 训练数据文件路径,要求JSONL格式。
每行是一个JSON对象,必须包含:
- query (str): 输入文本
- label (str|int): 标签,可以是"通用知识"/"专业咨询"或0/1
训练流程:
1. 读取并解析JSONL数据文件
2. 划分训练集和验证集(9:1)
3. 数据预处理和Dataset创建
4. 配置训练参数
5. 初始化Trainer并执行训练
6. 保存微调后的模型
训练参数说明:
- num_train_epochs (int): 训练轮数
- per_device_train_batch_size (int): 每设备训练批次大小
- gradient_accumulation_steps (int): 梯度累积步数,等效batch_size = 4*4=16
- learning_rate (float): 学习率,LoRA通常使用较大学习率(1e-4)
- lr_scheduler_type (str): 学习率调度策略,cosine为余弦退火
- bf16 (bool): 启用bfloat16混合精度,加速训练并节省显存
- gradient_checkpointing (bool): 梯度检查点,时间换显存
"""
# 检查数据文件是否存在
if not os.path.exists(data_file):
logger.error(f"找不到数据集: {data_file}")
return
# 读取JSONL格式数据文件
with open(data_file, "r", encoding="utf-8") as f:
# 逐行解析JSON
data = [json.loads(line) for line in f]
# 提取文本和标签
texts = [item["query"] for item in data] # 输入文本列表
labels = [item["label"] for item in data] # 标签列表
# 划分训练集和验证集,test_size=0.1表示10%作为验证集
train_texts, val_texts, train_labels, val_labels = train_test_split(
texts, # 输入文本
labels, # 对应标签
test_size=0.1, # 验证集比例
random_state=42 # 随机种子,保证可复现
)
# 数据预处理:将文本编码为模型输入格式
train_encodings = self.preprocess_data(train_texts, train_labels)
val_encodings = self.preprocess_data(val_texts, val_labels)
# 创建Hugging Face Dataset对象
train_dataset = self.create_dataset(train_encodings)
val_dataset = self.create_dataset(val_encodings)
# 配置训练参数 (针对 Linux/4090D 优化)
training_args = TrainingArguments(
output_dir="./qwen_checkpoints", # 检查点保存目录
num_train_epochs=5, # 训练总轮数
per_device_train_batch_size=4, # 每设备批次大小
gradient_accumulation_steps=4, # 梯度累积步数
learning_rate=1e-4, # 学习率
lr_scheduler_type="cosine", # 余弦退火学习率调度
logging_steps=10, # 每10步记录日志
eval_strategy="epoch", # 每个epoch结束后评估
save_strategy="epoch", # 每个epoch结束后保存
bf16=True, # 启用bf16混合精度
gradient_checkpointing=True, # 启用梯度检查点节省显存
load_best_model_at_end=True, # 训练结束加载最佳模型
report_to="none" # 禁用第三方报告(WandB等)
)
# 初始化Trainer
trainer = Trainer(
model=self.model, # 待训练模型
args=training_args, # 训练参数
train_dataset=train_dataset, # 训练数据集
eval_dataset=val_dataset, # 验证数据集
data_collator=DataCollatorWithPadding(
tokenizer=self.tokenizer
) # 动态padding的数据收集器
)
# 执行训练
trainer.train()
# 保存微调后的模型权重
self.model.save_pretrained(self.model_save_path)
# 保存分词器(包含配置)
self.tokenizer.save_pretrained(self.model_save_path)
logger.info(f"模型已保存至 {self.model_save_path}")
def predict(self, text):
"""
对单条文本进行预测
Args:
text (str): 待分类的文本字符串,例如"帮我查一下这门课的申请条件"
Returns:
str: 预测的标签名称,"通用知识"或"专业咨询"
预测流程:
1. 将模型设为评估模式(关闭Dropout等训练特定层)
2. 对输入文本进行tokenize并移至计算设备
3. 前向传播获取logits
4. 取概率最大的类别作为预测结果
5. 将数字ID映射回标签名称
"""
# 设置评估模式:禁用Dropout,BatchNorm使用运行统计
self.model.eval()
# 文本编码并移至设备
inputs = self.tokenizer(
text, # 输入文本
return_tensors="pt", # 返回PyTorch张量
padding=True, # 启用padding
truncation=True, # 超长截断
max_length=128 # 预测时使用较短长度即可
).to(self.device) # 移至GPU/CPU
# 禁用梯度计算,节省内存并加速推理
with torch.no_grad():
# 前向传播获取模型输出
outputs = self.model(**inputs)
# 获取分类logits(未归一化的分数)
logits = outputs.logits
# 取概率最大的维度作为预测类别
pred = torch.argmax(logits, dim=-1).item()
# 将数字ID转换回标签名称并返回
return self.id_to_label[pred]
# 启动脚本
if __name__ == "__main__":
# 数据文件路径:Linux建议使用绝对路径或相对于脚本的路径
DATA_PATH = "./dataset/model_generic_1000.jsonl"
# 初始化分类器(会自动加载已有模型或初始化新模型)
classifier = QwenLinuxClassifier()
# 如果数据文件存在,则执行训练
if os.path.exists(DATA_PATH):
classifier.train_model(DATA_PATH)
# 执行预测示例
print(classifier.predict("帮我查一下这门课的申请条件"))
训练模型主要是在Linux云平台上训练
在云端(如 AutoDL、阿里云 PAI、腾讯云等)训练 Qwen-8B 这种级别的模型,不仅是运行一段脚本,更是一套环境配置、数据搬运与后台管理的流程。
第一步:创建实例与环境初始化
在云平台选择镜像时,务必选择 PyTorch 2.x + CUDA 12.1 以上版本的官方镜像。
-
登录终端:使用 SSH 或平台提供的 JupyterLab 终端。
-
创建工作目录:
Bash
bashmkdir -p /root/autodl-tmp/qwen_project # 建议在数据盘(如 autodl-tmp)操作,防止系统盘撑爆 cd /root/autodl-tmp/qwen_project -
安装必要工具:
Bash
bashpip install -U modelscope transformers datasets peft accelerate bitsandbytes scikit-learn # 4090 必装 Flash Attention 2,能节省约 20% 显存并加速,可选择,注意上述代码运行时未设置,请不要运行下面代码 pip install flash-attn --no-build-isolation
第二步:极速下载模型
不要用 git clone,在国内云服务器上,用 ModelScope 是最快的:
Python
python
# 创建一个 download_model.py 并运行
from modelscope import snapshot_download
model_dir = snapshot_download('Qwen/Qwen3-8B', cache_dir='./models')
print(f"模型下载完成,路径: {model_dir}")
第三步:上传你的数据集
-
准备数据 :将你的数据处理成
jsonl格式,例如train.jsonl:JSONL格式:
{"query": "帮我看看这个课程多少钱?", "label": "专业咨询"} {"query": "什么是量子力学?", "label": "通用知识"} -
上传
第四步:启动后台训练 (核心步骤)
云端训练最忌讳直接在 SSH 窗口运行,因为一旦网络波动断开连接,训练就会崩溃。必须使用 tmux 或 nohup。
-
创建并进入虚拟终端:
Bash
tmux new -s train_session -
运行训练脚本:
运行:
python train.py