207 lines
6.6 KiB
Python
207 lines
6.6 KiB
Python
# %%
|
||
import os
|
||
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
|
||
import torch
|
||
from torch.utils.data import Dataset, DataLoader
|
||
import pandas as pd
|
||
import transformers
|
||
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
|
||
from tqdm import tqdm
|
||
# %%
|
||
train_df = pd.read_csv('./train.csv')
|
||
test_df = pd.read_csv('./test.csv')
|
||
from sklearn.model_selection import train_test_split
|
||
train_df, val_df = train_test_split(train_df, test_size=0.1, random_state=42, stratify=train_df['label'])
|
||
print(f"Train: {len(train_df)}, Val: {len(val_df)}, Test: {len(test_df)}")
|
||
# %%
|
||
train_df.head()
|
||
# %%
|
||
model_name = "xlm-roberta-base"
|
||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=3)
|
||
# %%
|
||
embedded_text=tokenizer("你好啊")
|
||
# %%
|
||
tokenizer.decode(embedded_text['input_ids'],skip_special_tokens=True)
|
||
# %%
|
||
embedded_text=tokenizer(["你好啊","我是灰太狼"],["我不好","我是红太狼"])
|
||
embedded_text=tokenizer("你好啊",return_tensors='pt')
|
||
# %%
|
||
class NliDataset(Dataset):
|
||
def __init__(self, df, tokenizer, max_length=128):
|
||
self.df = df.reset_index(drop=True)
|
||
self.tokenizer = tokenizer
|
||
self.max_length = max_length
|
||
|
||
def __len__(self):
|
||
return len(self.df)
|
||
|
||
def __getitem__(self, idx):
|
||
row = self.df.iloc[idx]
|
||
premise = str(row['premise'])
|
||
hypothesis = str(row['hypothesis'])
|
||
# 编码文本对，返回 input_ids, attention_mask
|
||
encoding = self.tokenizer(
|
||
premise,
|
||
hypothesis,
|
||
truncation=True,
|
||
padding='max_length',
|
||
max_length=self.max_length,
|
||
return_tensors='pt' # 返回 PyTorch Tensor
|
||
)
|
||
# 去掉 batch 维度（因为只处理单条）
|
||
item = {
|
||
'input_ids': encoding['input_ids'].squeeze(0),
|
||
'attention_mask': encoding['attention_mask'].squeeze(0)
|
||
}
|
||
if 'label' in row:
|
||
item['labels'] = torch.tensor(row['label'], dtype=torch.long)
|
||
return item
|
||
# %%
|
||
|
||
# %%
|
||
batch_size = 16 # 根据显存调整，推荐使用 16 或 32
|
||
max_length = 128
|
||
|
||
train_dataset = NliDataset(train_df, tokenizer, max_length)
|
||
val_dataset = NliDataset(val_df, tokenizer, max_length)
|
||
|
||
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
|
||
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
|
||
# %%
|
||
|
||
# %%
|
||
from transformers import get_linear_schedule_with_warmup
|
||
from sklearn.metrics import accuracy_score
|
||
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
|
||
|
||
# 计算总训练步数（用于 warmup）
|
||
total_steps = len(train_loader) * 3 # 假设训练 3 个 epoch
|
||
warmup_steps = int(0.1 * total_steps) # warmup 比例为 10%
|
||
|
||
scheduler = get_linear_schedule_with_warmup(
|
||
optimizer,
|
||
num_warmup_steps=warmup_steps,
|
||
num_training_steps=total_steps
|
||
)
|
||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
model.to(device)
|
||
# %%
|
||
def evaluate(model, data_loader, device):
|
||
model.eval()
|
||
all_preds = []
|
||
all_labels = []
|
||
with torch.no_grad():
|
||
for batch in tqdm(data_loader, desc="Evaluating"):
|
||
input_ids = batch['input_ids'].to(device)
|
||
attention_mask = batch['attention_mask'].to(device)
|
||
labels = batch['labels'].to(device)
|
||
|
||
outputs = model(input_ids, attention_mask=attention_mask)
|
||
logits = outputs.logits
|
||
preds = torch.argmax(logits, dim=1)
|
||
|
||
all_preds.extend(preds.cpu().numpy())
|
||
all_labels.extend(labels.cpu().numpy())
|
||
|
||
acc = accuracy_score(all_labels, all_preds)
|
||
return acc
|
||
# %%
|
||
num_epochs = 3
|
||
best_val_acc = 0.0
|
||
|
||
for epoch in range(num_epochs):
|
||
model.train()
|
||
total_loss = 0
|
||
|
||
progress_bar = tqdm(train_loader, desc=f'Epoch {epoch+1}/{num_epochs}')
|
||
for batch in progress_bar:
|
||
# 将数据移至设备
|
||
input_ids = batch['input_ids'].to(device)
|
||
attention_mask = batch['attention_mask'].to(device)
|
||
labels = batch['labels'].to(device)
|
||
|
||
# 前向传播
|
||
outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
|
||
loss = outputs.loss
|
||
|
||
# 反向传播
|
||
loss.backward()
|
||
|
||
# 梯度裁剪（防止梯度爆炸）
|
||
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
|
||
|
||
# 更新参数
|
||
optimizer.step()
|
||
scheduler.step() # 更新学习率
|
||
optimizer.zero_grad()
|
||
|
||
# 记录损失
|
||
total_loss += loss.item()
|
||
progress_bar.set_postfix({'loss': loss.item()})
|
||
|
||
avg_train_loss = total_loss / len(train_loader)
|
||
print(f"Epoch {epoch+1} - Average Train Loss: {avg_train_loss:.4f}")
|
||
|
||
# 在每个 epoch 结束后评估验证集
|
||
val_acc = evaluate(model, val_loader, device)
|
||
print(f"Epoch {epoch+1} - Validation Accuracy: {val_acc:.4f}")
|
||
|
||
# 保存最佳模型
|
||
if val_acc > best_val_acc:
|
||
best_val_acc = val_acc
|
||
torch.save(model.state_dict(), 'best_model.pt')
|
||
print("Best model saved!")
|
||
# %%
|
||
# 加载最佳模型权重
|
||
model.load_state_dict(torch.load('best_model.pt'))
|
||
model.eval()
|
||
|
||
# 构建测试集 Dataset 和 DataLoader（注意测试集没有 label）
|
||
class TestDataset(Dataset):
|
||
def __init__(self, df, tokenizer, max_length=128):
|
||
self.df = df.reset_index(drop=True)
|
||
self.tokenizer = tokenizer
|
||
self.max_length = max_length
|
||
|
||
def __len__(self):
|
||
return len(self.df)
|
||
|
||
def __getitem__(self, idx):
|
||
row = self.df.iloc[idx]
|
||
premise = str(row['premise'])
|
||
hypothesis = str(row['hypothesis'])
|
||
encoding = self.tokenizer(
|
||
premise,
|
||
hypothesis,
|
||
truncation=True,
|
||
padding='max_length',
|
||
max_length=self.max_length,
|
||
return_tensors='pt'
|
||
)
|
||
return {
|
||
'input_ids': encoding['input_ids'].squeeze(0),
|
||
'attention_mask': encoding['attention_mask'].squeeze(0)
|
||
}
|
||
|
||
test_dataset = TestDataset(test_df, tokenizer, max_length)
|
||
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
|
||
|
||
# 预测
|
||
all_preds = []
|
||
with torch.no_grad():
|
||
for batch in tqdm(test_loader, desc="Predicting"):
|
||
input_ids = batch['input_ids'].to(device)
|
||
attention_mask = batch['attention_mask'].to(device)
|
||
outputs = model(input_ids, attention_mask=attention_mask)
|
||
logits = outputs.logits
|
||
preds = torch.argmax(logits, dim=1)
|
||
all_preds.extend(preds.cpu().numpy())
|
||
|
||
# 生成提交文件
|
||
submission = pd.DataFrame({
|
||
'id': test_df['id'],
|
||
'label': all_preds
|
||
})
|
||
submission.to_csv('submission.csv', index=False)
|
||
print("Submission saved!") |