環境準備
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)