环境准备
https://www.doubao.com/thread/w5e26d6401c003bb2
执行培训脚本
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification, AdamW
import numpy as np# 自定义数据集类
class SentimentDataset(Dataset):def __init__(self, texts, labels, tokenizer, max_length):self.texts = textsself.labels = labelsself.tokenizer = tokenizerself.max_length = max_lengthdef __len__(self):return len(self.texts)def __getitem__(self, idx):text = self.texts[idx]label = self.labels[idx]encoding = self.tokenizer.encode_plus(text,add_special_tokens=True,max_length=self.max_length,padding='max_length',truncation=True,return_tensors='pt')return {'input_ids': encoding['input_ids'].flatten(),'attention_mask': encoding['attention_mask'].flatten(),'labels': torch.tensor(label, dtype=torch.long)}# 新增问答数据
qa_texts = ["李白是那个朝代的诗人？", "地球的卫星是什么？", "中国的首都是哪里？", "水的化学式是什么？", "苹果公司的创始人是谁？"]
qa_labels = [1, 2, 3, 4, 5] # 为每个问题分配一个唯一的标签# 合并数据
all_texts = qa_texts
all_labels = qa_labels# 初始化 tokenizer
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')# 创建数据集和数据加载器
dataset = SentimentDataset(all_texts, all_labels, tokenizer, max_length=128)
dataloader = DataLoader(dataset, batch_size=2, shuffle=True)# 初始化模型，注意 num_labels 需要根据总标签数调整
model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=8)# 定义优化器，加入 weight_decay 进行正则化
optimizer = AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)# 训练模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)# 早停策略相关参数
best_loss = float('inf')
patience = 3
counter = 0for epoch in range(20): # 增加训练轮数，但结合早停策略total_loss = 0model.train()for batch in dataloader:input_ids = batch['input_ids'].to(device)attention_mask = batch['attention_mask'].to(device)labels = batch['labels'].to(device)optimizer.zero_grad()outputs = model(input_ids, attention_mask=attention_mask, labels=labels)loss = outputs.lossloss.backward()optimizer.step()total_loss += loss.item()avg_loss = total_loss / len(dataloader)print(f'Epoch {epoch + 1}, Loss: {avg_loss}')# 早停策略if avg_loss < best_loss:best_loss = avg_losscounter = 0else:counter += 1if counter >= patience:print("Early stopping triggered!")break# 处理提问的函数
def handle_query(model, tokenizer, query_text, device):model.eval()with torch.no_grad():encoding = tokenizer(query_text, return_tensors='pt', padding=True, truncation=True, max_length=128).to(device)logits = model(**encoding).logitspredictions = np.argmax(logits.cpu().numpy(), axis=1)return predictions# 情感标签映射
label_mapping = {1: "唐朝",2: "月球",3: "北京",4: "H2O",5: "史蒂夫·乔布斯、史蒂夫·沃兹尼亚克和罗恩·韦恩"
}# 针对新增的5条数据提问
query_text = ["李白是那个朝代的诗人？", "地球的卫星是什么？", "中国的首都是哪里？", "水的化学式是什么？", "苹果公司的创始人是谁？"]
result = handle_query(model, tokenizer, query_text, device)
readable_result = [label_mapping[pred] for pred in result]
print("Query result:", readable_result)