{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "import os\n",
    "os.environ[\"HF_ENDPOINT\"] = \"https://hf-mirror.com\"\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import pandas as pd\n",
    "import transformers\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments\n",
    "from tqdm import tqdm"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "train_df = pd.read_csv('./train.csv')\n",
    "test_df = pd.read_csv('./test.csv')\n",
    "from sklearn.model_selection import train_test_split\n",
    "train_df, val_df = train_test_split(train_df, test_size=0.1, random_state=42, stratify=train_df['label'])\n",
    "print(f\"Train: {len(train_df)}, Val: {len(val_df)}, Test: {len(test_df)}\")"
   ],
   "id": "f32065752f1597fb",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "train_df.head()",
   "id": "460c6fffc81814df",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-07-02T09:54:33.329021843Z",
     "start_time": "2026-07-02T09:54:22.292806435Z"
    }
   },
   "cell_type": "code",
   "source": [
    "model_name = \"xlm-roberta-base\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=3)"
   ],
   "id": "3d828c6bd763c919",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Loading weights:   0%|          | 0/197 [00:00<?, ?it/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "0b2f498a9bf4470c9b105f34749cff74"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[transformers] \u001B[1mXLMRobertaForSequenceClassification LOAD REPORT\u001B[0m from: xlm-roberta-base\n",
      "Key                         | Status     | \n",
      "----------------------------+------------+-\n",
      "lm_head.bias                | UNEXPECTED | \n",
      "roberta.pooler.dense.weight | UNEXPECTED | \n",
      "lm_head.dense.weight        | UNEXPECTED | \n",
      "lm_head.layer_norm.weight   | UNEXPECTED | \n",
      "roberta.pooler.dense.bias   | UNEXPECTED | \n",
      "lm_head.dense.bias          | UNEXPECTED | \n",
      "lm_head.layer_norm.bias     | UNEXPECTED | \n",
      "classifier.out_proj.bias    | MISSING    | \n",
      "classifier.dense.weight     | MISSING    | \n",
      "classifier.dense.bias       | MISSING    | \n",
      "classifier.out_proj.weight  | MISSING    | \n",
      "\n",
      "Notes:\n",
      "- UNEXPECTED:\tcan be ignored when loading from different task/architecture; not ok if you expect identical arch.\n",
      "- MISSING:\tthose params were newly initialized because missing from the checkpoint. Consider training on your downstream task.\n"
     ]
    }
   ],
   "execution_count": 241
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "embedded_text=tokenizer(\"你好啊\")",
   "id": "a69150070518a5a3",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "tokenizer.decode(embedded_text['input_ids'],skip_special_tokens=True)",
   "id": "56c67212e3124e3e",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "embedded_text=tokenizer([\"你好啊\",\"我是灰太狼\"],[\"我不好\",\"我是红太狼\"])\n",
    "embedded_text=tokenizer(\"你好啊\",return_tensors='pt')"
   ],
   "id": "c89d28ab49d59f24",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "class NliDataset(Dataset):\n",
    "    def __init__(self, df, tokenizer, max_length=128):\n",
    "        self.df = df.reset_index(drop=True)\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_length = max_length\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        row = self.df.iloc[idx]\n",
    "        premise = str(row['premise'])\n",
    "        hypothesis = str(row['hypothesis'])\n",
    "        # 编码文本对，返回 input_ids, attention_mask\n",
    "        encoding = self.tokenizer(\n",
    "            premise,\n",
    "            hypothesis,\n",
    "            truncation=True,\n",
    "            padding='max_length',\n",
    "            max_length=self.max_length,\n",
    "            return_tensors='pt'  # 返回 PyTorch Tensor\n",
    "        )\n",
    "        # 去掉 batch 维度（因为只处理单条）\n",
    "        item = {\n",
    "            'input_ids': encoding['input_ids'].squeeze(0),\n",
    "            'attention_mask': encoding['attention_mask'].squeeze(0)\n",
    "        }\n",
    "        if 'label' in row:\n",
    "            item['labels'] = torch.tensor(row['label'], dtype=torch.long)\n",
    "        return item"
   ],
   "id": "540f99b92c4a059b",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "65bcf49c94fcdfd9",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "batch_size = 16  # 根据显存调整，推荐使用 16 或 32\n",
    "max_length = 128\n",
    "\n",
    "train_dataset = NliDataset(train_df, tokenizer, max_length)\n",
    "val_dataset = NliDataset(val_df, tokenizer, max_length)\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)"
   ],
   "id": "35aa7869a8899205",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "5f70009f93be4960",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from transformers import get_linear_schedule_with_warmup\n",
    "from sklearn.metrics import accuracy_score\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)\n",
    "\n",
    "# 计算总训练步数（用于 warmup）\n",
    "total_steps = len(train_loader) * 3  # 假设训练 3 个 epoch\n",
    "warmup_steps = int(0.1 * total_steps)  # warmup 比例为 10%\n",
    "\n",
    "scheduler = get_linear_schedule_with_warmup(\n",
    "    optimizer,\n",
    "    num_warmup_steps=warmup_steps,\n",
    "    num_training_steps=total_steps\n",
    ")\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model.to(device)"
   ],
   "id": "ffc4d718d1cf64d",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "def evaluate(model, data_loader, device):\n",
    "    model.eval()\n",
    "    all_preds = []\n",
    "    all_labels = []\n",
    "    with torch.no_grad():\n",
    "        for batch in tqdm(data_loader, desc=\"Evaluating\"):\n",
    "            input_ids = batch['input_ids'].to(device)\n",
    "            attention_mask = batch['attention_mask'].to(device)\n",
    "            labels = batch['labels'].to(device)\n",
    "\n",
    "            outputs = model(input_ids, attention_mask=attention_mask)\n",
    "            logits = outputs.logits\n",
    "            preds = torch.argmax(logits, dim=1)\n",
    "\n",
    "            all_preds.extend(preds.cpu().numpy())\n",
    "            all_labels.extend(labels.cpu().numpy())\n",
    "\n",
    "    acc = accuracy_score(all_labels, all_preds)\n",
    "    return acc"
   ],
   "id": "cf187465e1a6ab4b",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "num_epochs = 3\n",
    "best_val_acc = 0.0\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "\n",
    "    progress_bar = tqdm(train_loader, desc=f'Epoch {epoch+1}/{num_epochs}')\n",
    "    for batch in progress_bar:\n",
    "        # 将数据移至设备\n",
    "        input_ids = batch['input_ids'].to(device)\n",
    "        attention_mask = batch['attention_mask'].to(device)\n",
    "        labels = batch['labels'].to(device)\n",
    "\n",
    "        # 前向传播\n",
    "        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)\n",
    "        loss = outputs.loss\n",
    "\n",
    "        # 反向传播\n",
    "        loss.backward()\n",
    "\n",
    "        # 梯度裁剪（防止梯度爆炸）\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n",
    "\n",
    "        # 更新参数\n",
    "        optimizer.step()\n",
    "        scheduler.step()   # 更新学习率\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # 记录损失\n",
    "        total_loss += loss.item()\n",
    "        progress_bar.set_postfix({'loss': loss.item()})\n",
    "\n",
    "    avg_train_loss = total_loss / len(train_loader)\n",
    "    print(f\"Epoch {epoch+1} - Average Train Loss: {avg_train_loss:.4f}\")\n",
    "\n",
    "    # 在每个 epoch 结束后评估验证集\n",
    "    val_acc = evaluate(model, val_loader, device)\n",
    "    print(f\"Epoch {epoch+1} - Validation Accuracy: {val_acc:.4f}\")\n",
    "\n",
    "    # 保存最佳模型\n",
    "    if val_acc > best_val_acc:\n",
    "        best_val_acc = val_acc\n",
    "        torch.save(model.state_dict(), 'best_model.pt')\n",
    "        print(\"Best model saved!\")"
   ],
   "id": "9d75896a5b9cf921",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# 加载最佳模型权重\n",
    "model.load_state_dict(torch.load('best_model.pt'))\n",
    "model.eval()\n",
    "\n",
    "# 构建测试集 Dataset 和 DataLoader（注意测试集没有 label）\n",
    "class TestDataset(Dataset):\n",
    "    def __init__(self, df, tokenizer, max_length=128):\n",
    "        self.df = df.reset_index(drop=True)\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_length = max_length\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        row = self.df.iloc[idx]\n",
    "        premise = str(row['premise'])\n",
    "        hypothesis = str(row['hypothesis'])\n",
    "        encoding = self.tokenizer(\n",
    "            premise,\n",
    "            hypothesis,\n",
    "            truncation=True,\n",
    "            padding='max_length',\n",
    "            max_length=self.max_length,\n",
    "            return_tensors='pt'\n",
    "        )\n",
    "        return {\n",
    "            'input_ids': encoding['input_ids'].squeeze(0),\n",
    "            'attention_mask': encoding['attention_mask'].squeeze(0)\n",
    "        }\n",
    "\n",
    "test_dataset = TestDataset(test_df, tokenizer, max_length)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "# 预测\n",
    "all_preds = []\n",
    "with torch.no_grad():\n",
    "    for batch in tqdm(test_loader, desc=\"Predicting\"):\n",
    "        input_ids = batch['input_ids'].to(device)\n",
    "        attention_mask = batch['attention_mask'].to(device)\n",
    "        outputs = model(input_ids, attention_mask=attention_mask)\n",
    "        logits = outputs.logits\n",
    "        preds = torch.argmax(logits, dim=1)\n",
    "        all_preds.extend(preds.cpu().numpy())\n",
    "\n",
    "# 生成提交文件\n",
    "submission = pd.DataFrame({\n",
    "    'id': test_df['id'],\n",
    "    'label': all_preds\n",
    "})\n",
    "submission.to_csv('submission.csv', index=False)\n",
    "print(\"Submission saved!\")"
   ],
   "id": "df17fe10b2f36fc5",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}