# %% 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!")