{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2026-07-05T15:05:59.049967188Z",
     "start_time": "2026-07-05T15:05:56.833387637Z"
    }
   },
   "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": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-07-05T15:05:59.150289671Z",
     "start_time": "2026-07-05T15:05:59.051997256Z"
    }
   },
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train: 10908, Val: 1212, Test: 5195\n"
     ]
    }
   ],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-07-05T15:05:59.184794043Z",
     "start_time": "2026-07-05T15:05:59.156968452Z"
    }
   },
   "cell_type": "code",
   "source": "train_df.head()",
   "id": "460c6fffc81814df",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "              id                                            premise  \\\n",
       "9989  6417481c8d  यद्यपि हम आज के समय में अल कायदा के साथ केएसएए...   \n",
       "3880  5c5ca34cf6             'Upload him into his body? What body?'   \n",
       "8559  1d3c28ecff  yeah and then about every five years you have ...   \n",
       "6316  22e2a4903d  θέλω να πω ότι υπήρχε είχα, είχα το ρολόι μου ...   \n",
       "762   307016c21f  Yet, in the mouths of the white townsfolk of S...   \n",
       "\n",
       "                                             hypothesis lang_abv language  \\\n",
       "9989  हर व्यक्ति केएसएम को हमेशा अल कायदा के बराबर म...       hi    Hindi   \n",
       "3880                I don't think he has a body at all.       en  English   \n",
       "8559  You have to dig them up every five years, thro...       en  English   \n",
       "6316      Ήταν ενοχλητικό όταν κάλυψε τα παπούτσια μου.       el    Greek   \n",
       "762   White townsfolk in Salisbury, N.C. are easily ...       en  English   \n",
       "\n",
       "      label  \n",
       "9989      2  \n",
       "3880      1  \n",
       "8559      1  \n",
       "6316      1  \n",
       "762       1  "
      ],
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>premise</th>\n",
       "      <th>hypothesis</th>\n",
       "      <th>lang_abv</th>\n",
       "      <th>language</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>9989</th>\n",
       "      <td>6417481c8d</td>\n",
       "      <td>यद्यपि हम आज के समय में अल कायदा के साथ केएसएए...</td>\n",
       "      <td>हर व्यक्ति केएसएम को हमेशा अल कायदा के बराबर म...</td>\n",
       "      <td>hi</td>\n",
       "      <td>Hindi</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3880</th>\n",
       "      <td>5c5ca34cf6</td>\n",
       "      <td>'Upload him into his body? What body?'</td>\n",
       "      <td>I don't think he has a body at all.</td>\n",
       "      <td>en</td>\n",
       "      <td>English</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8559</th>\n",
       "      <td>1d3c28ecff</td>\n",
       "      <td>yeah and then about every five years you have ...</td>\n",
       "      <td>You have to dig them up every five years, thro...</td>\n",
       "      <td>en</td>\n",
       "      <td>English</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6316</th>\n",
       "      <td>22e2a4903d</td>\n",
       "      <td>θέλω να πω ότι υπήρχε είχα, είχα το ρολόι μου ...</td>\n",
       "      <td>Ήταν ενοχλητικό όταν κάλυψε τα παπούτσια μου.</td>\n",
       "      <td>el</td>\n",
       "      <td>Greek</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>762</th>\n",
       "      <td>307016c21f</td>\n",
       "      <td>Yet, in the mouths of the white townsfolk of S...</td>\n",
       "      <td>White townsfolk in Salisbury, N.C. are easily ...</td>\n",
       "      <td>en</td>\n",
       "      <td>English</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 4
  },
  {
   "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
}
