examples/Angle-ATEC.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a2a75be0",
   "metadata": {},
   "source": [
    "# AnglE-optimized Text Embeddings\n",
    "\n",
    "> It is Angle 📐, not Angel 👼.\n",
    "\n",
    "🔥 基于 AnglE 开箱即用的文本向量库，支持中英双语，可用于文本相似度计算、检索召回、匹配等场景。代码基于 🤗transformers 构建，提供易用的微调接口，可在 3090Ti、 4090 等消费级 GPU 上微调 LLaMA-7B 模型，支持多卡分布式训练。\n",
    "\n",
    "\n",
    "<a href=\"https://arxiv.org/abs/2309.12871\">\n",
    "    <img src=\"https://img.shields.io/badge/Arxiv-2306.06843-yellow.svg?style=flat-square\" alt=\"https://arxiv.org/abs/2309.12871\" />\n",
    "</a>\n",
    "<a href=\"https://pypi.org/project/angle_emb/\">\n",
    "    <img src=\"https://img.shields.io/pypi/v/angle_emb?style=flat-square\" alt=\"PyPI version\" />\n",
    "</a>\n",
    "\n",
    "<a href=\"https://pypi.org/project/angle_emb/\">\n",
    "    <img src=\"https://img.shields.io/pypi/dm/angle_emb?style=flat-square\" alt=\"PyPI Downloads\" />\n",
    "</a>\n",
    "<a href=\"http://makeapullrequest.com\">\n",
    "    <img src=\"https://img.shields.io/badge/PRs-welcome-brightgreen.svg?style=flat-square\" alt=\"http://makeapullrequest.com\" />\n",
    "</a>\n",
    "\n",
    "[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/angle-optimized-text-embeddings/semantic-textual-similarity-on-sick-r-1)](https://paperswithcode.com/sota/semantic-textual-similarity-on-sick-r-1?p=angle-optimized-text-embeddings)\n",
    "[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/angle-optimized-text-embeddings/semantic-textual-similarity-on-sts16)](https://paperswithcode.com/sota/semantic-textual-similarity-on-sts16?p=angle-optimized-text-embeddings)\n",
    "[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/angle-optimized-text-embeddings/semantic-textual-similarity-on-sts15)](https://paperswithcode.com/sota/semantic-textual-similarity-on-sts15?p=angle-optimized-text-embeddings)\n",
    "[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/angle-optimized-text-embeddings/semantic-textual-similarity-on-sts14)](https://paperswithcode.com/sota/semantic-textual-similarity-on-sts14?p=angle-optimized-text-embeddings)\n",
    "[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/angle-optimized-text-embeddings/semantic-textual-similarity-on-sts13)](https://paperswithcode.com/sota/semantic-textual-similarity-on-sts13?p=angle-optimized-text-embeddings)\n",
    "[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/angle-optimized-text-embeddings/semantic-textual-similarity-on-sts12)](https://paperswithcode.com/sota/semantic-textual-similarity-on-sts12?p=angle-optimized-text-embeddings)\n",
    "[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/angle-optimized-text-embeddings/semantic-textual-similarity-on-sts-benchmark)](https://paperswithcode.com/sota/semantic-textual-similarity-on-sts-benchmark?p=angle-optimized-text-embeddings)\n",
    "\n",
    "\n",
    "\n",
    "如果你有使用我们的代码及预训练模型，欢迎给我们三连，三连方式为：\n",
    "1) 给本项目 GitHub 加个 star\n",
    "2) 粘贴以下引用信息到你 paper 的 bibtex\n",
    "3) 在你的 paper 正文中引用\n",
    "\n",
    "```bibtex\n",
    "@article{li2023angle,\n",
    "  title={AnglE-Optimized Text Embeddings},\n",
    "  author={Li, Xianming and Li, Jing},\n",
    "  journal={arXiv preprint arXiv:2309.12871},\n",
    "  year={2023}\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17cee9a7",
   "metadata": {},
   "source": [
    "# 1. 安装依赖"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "73327e2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Collecting angle-emb\n",
      "  Downloading angle_emb-0.1.1-py3-none-any.whl (12 kB)\n",
      "Requirement already satisfied: peft in /home/jupyter-sean/.local/lib/python3.9/site-packages (from angle-emb) (0.5.0)\n",
      "Collecting transformers>=4.32.1\n",
      "  Downloading transformers-4.34.1-py3-none-any.whl (7.7 MB)\n",
      "     |████████████████████████████████| 7.7 MB 23.4 MB/s            \n",
      "\u001b[?25hRequirement already satisfied: boltons in /home/jupyter-sean/.local/lib/python3.9/site-packages (from angle-emb) (23.0.0)\n",
      "Requirement already satisfied: datasets in /home/jupyter-sean/.local/lib/python3.9/site-packages (from angle-emb) (2.13.1)\n",
      "Requirement already satisfied: prettytable in /home/jupyter-sean/.local/lib/python3.9/site-packages (from angle-emb) (3.9.0)\n",
      "Collecting bitsandbytes\n",
      "  Using cached bitsandbytes-0.41.1-py3-none-any.whl (92.6 MB)\n",
      "Collecting tokenizers<0.15,>=0.14\n",
      "  Downloading tokenizers-0.14.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.8 MB)\n",
      "     |████████████████████████████████| 3.8 MB 70.0 MB/s            \n",
      "\u001b[?25hRequirement already satisfied: regex!=2019.12.17 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from transformers>=4.32.1->angle-emb) (2023.6.3)\n",
      "Requirement already satisfied: requests in /opt/tljh/user/lib/python3.9/site-packages (from transformers>=4.32.1->angle-emb) (2.26.0)\n",
      "Requirement already satisfied: filelock in /home/jupyter-sean/.local/lib/python3.9/site-packages (from transformers>=4.32.1->angle-emb) (3.12.0)\n",
      "Collecting huggingface-hub<1.0,>=0.16.4\n",
      "  Downloading huggingface_hub-0.18.0-py3-none-any.whl (301 kB)\n",
      "     |████████████████████████████████| 301 kB 76.6 MB/s            \n",
      "\u001b[?25hRequirement already satisfied: tqdm>=4.27 in /opt/tljh/user/lib/python3.9/site-packages (from transformers>=4.32.1->angle-emb) (4.62.3)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/tljh/user/lib/python3.9/site-packages (from transformers>=4.32.1->angle-emb) (22.0)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from transformers>=4.32.1->angle-emb) (6.0)\n",
      "Requirement already satisfied: numpy>=1.17 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from transformers>=4.32.1->angle-emb) (1.24.3)\n",
      "Requirement already satisfied: safetensors>=0.3.1 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from transformers>=4.32.1->angle-emb) (0.3.1)\n",
      "Requirement already satisfied: pandas in /home/jupyter-sean/.local/lib/python3.9/site-packages (from datasets->angle-emb) (2.0.3)\n",
      "Requirement already satisfied: multiprocess in /home/jupyter-sean/.local/lib/python3.9/site-packages (from datasets->angle-emb) (0.70.14)\n",
      "Requirement already satisfied: xxhash in /home/jupyter-sean/.local/lib/python3.9/site-packages (from datasets->angle-emb) (3.2.0)\n",
      "Requirement already satisfied: dill<0.3.7,>=0.3.0 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from datasets->angle-emb) (0.3.6)\n",
      "Requirement already satisfied: aiohttp in /home/jupyter-sean/.local/lib/python3.9/site-packages (from datasets->angle-emb) (3.8.4)\n",
      "Requirement already satisfied: fsspec[http]>=2021.11.1 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from datasets->angle-emb) (2023.5.0)\n",
      "Requirement already satisfied: pyarrow>=8.0.0 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from datasets->angle-emb) (12.0.1)\n",
      "Requirement already satisfied: psutil in /opt/tljh/user/lib/python3.9/site-packages (from peft->angle-emb) (5.9.4)\n",
      "Requirement already satisfied: accelerate in /home/jupyter-sean/.local/lib/python3.9/site-packages (from peft->angle-emb) (0.23.0)\n",
      "Requirement already satisfied: torch>=1.13.0 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from peft->angle-emb) (2.0.1)\n",
      "Requirement already satisfied: wcwidth in /opt/tljh/user/lib/python3.9/site-packages (from prettytable->angle-emb) (0.2.6)\n",
      "Requirement already satisfied: charset-normalizer<4.0,>=2.0 in /opt/tljh/user/lib/python3.9/site-packages (from aiohttp->datasets->angle-emb) (2.0.0)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from aiohttp->datasets->angle-emb) (4.0.2)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from aiohttp->datasets->angle-emb) (1.9.2)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from aiohttp->datasets->angle-emb) (1.3.3)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /opt/tljh/user/lib/python3.9/site-packages (from aiohttp->datasets->angle-emb) (22.1.0)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from aiohttp->datasets->angle-emb) (1.3.1)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from aiohttp->datasets->angle-emb) (6.0.4)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from huggingface-hub<1.0,>=0.16.4->transformers>=4.32.1->angle-emb) (4.7.1)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/tljh/user/lib/python3.9/site-packages (from requests->transformers>=4.32.1->angle-emb) (2022.12.7)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/tljh/user/lib/python3.9/site-packages (from requests->transformers>=4.32.1->angle-emb) (3.1)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /opt/tljh/user/lib/python3.9/site-packages (from requests->transformers>=4.32.1->angle-emb) (1.26.7)\n",
      "  Using cached huggingface_hub-0.17.3-py3-none-any.whl (295 kB)\n",
      "Requirement already satisfied: nvidia-nccl-cu11==2.14.3 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from torch>=1.13.0->peft->angle-emb) (2.14.3)\n",
      "Requirement already satisfied: nvidia-cusolver-cu11==11.4.0.1 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from torch>=1.13.0->peft->angle-emb) (11.4.0.1)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu11==11.7.99 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from torch>=1.13.0->peft->angle-emb) (11.7.99)\n",
      "Requirement already satisfied: nvidia-cusparse-cu11==11.7.4.91 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from torch>=1.13.0->peft->angle-emb) (11.7.4.91)\n",
      "Requirement already satisfied: jinja2 in /opt/tljh/user/lib/python3.9/site-packages (from torch>=1.13.0->peft->angle-emb) (3.1.2)\n",
      "Requirement already satisfied: nvidia-cufft-cu11==10.9.0.58 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from torch>=1.13.0->peft->angle-emb) (10.9.0.58)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu11==11.7.99 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from torch>=1.13.0->peft->angle-emb) (11.7.99)\n",
      "Requirement already satisfied: nvidia-cublas-cu11==11.10.3.66 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from torch>=1.13.0->peft->angle-emb) (11.10.3.66)\n",
      "Requirement already satisfied: nvidia-nvtx-cu11==11.7.91 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from torch>=1.13.0->peft->angle-emb) (11.7.91)\n",
      "Requirement already satisfied: sympy in /home/jupyter-sean/.local/lib/python3.9/site-packages (from torch>=1.13.0->peft->angle-emb) (1.12)\n",
      "Requirement already satisfied: triton==2.0.0 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from torch>=1.13.0->peft->angle-emb) (2.0.0)\n",
      "Requirement already satisfied: nvidia-curand-cu11==10.2.10.91 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from torch>=1.13.0->peft->angle-emb) (10.2.10.91)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu11==11.7.101 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from torch>=1.13.0->peft->angle-emb) (11.7.101)\n",
      "Requirement already satisfied: nvidia-cudnn-cu11==8.5.0.96 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from torch>=1.13.0->peft->angle-emb) (8.5.0.96)\n",
      "Requirement already satisfied: networkx in /home/jupyter-sean/.local/lib/python3.9/site-packages (from torch>=1.13.0->peft->angle-emb) (3.1)\n",
      "Requirement already satisfied: setuptools in /opt/tljh/user/lib/python3.9/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch>=1.13.0->peft->angle-emb) (58.2.0)\n",
      "Requirement already satisfied: wheel in /opt/tljh/user/lib/python3.9/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch>=1.13.0->peft->angle-emb) (0.37.0)\n",
      "Requirement already satisfied: lit in /home/jupyter-sean/.local/lib/python3.9/site-packages (from triton==2.0.0->torch>=1.13.0->peft->angle-emb) (16.0.6)\n",
      "Requirement already satisfied: cmake in /home/jupyter-sean/.local/lib/python3.9/site-packages (from triton==2.0.0->torch>=1.13.0->peft->angle-emb) (3.26.4)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: tzdata>=2022.1 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from pandas->datasets->angle-emb) (2023.3)\n",
      "Requirement already satisfied: pytz>=2020.1 in /opt/tljh/user/lib/python3.9/site-packages (from pandas->datasets->angle-emb) (2022.6)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /opt/tljh/user/lib/python3.9/site-packages (from pandas->datasets->angle-emb) (2.8.2)\n",
      "Requirement already satisfied: six>=1.5 in /opt/tljh/user/lib/python3.9/site-packages (from python-dateutil>=2.8.2->pandas->datasets->angle-emb) (1.16.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/tljh/user/lib/python3.9/site-packages (from jinja2->torch>=1.13.0->peft->angle-emb) (2.1.1)\n",
      "Requirement already satisfied: mpmath>=0.19 in /home/jupyter-sean/.local/lib/python3.9/site-packages (from sympy->torch>=1.13.0->peft->angle-emb) (1.3.0)\n",
      "Installing collected packages: huggingface-hub, tokenizers, transformers, bitsandbytes, angle-emb\n",
      "  Attempting uninstall: huggingface-hub\n",
      "    Found existing installation: huggingface-hub 0.15.1\n",
      "    Uninstalling huggingface-hub-0.15.1:\n",
      "      Successfully uninstalled huggingface-hub-0.15.1\n",
      "\u001b[33m  WARNING: The script huggingface-cli is installed in '/home/jupyter-sean/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\n",
      "  Attempting uninstall: tokenizers\n",
      "    Found existing installation: tokenizers 0.13.3\n",
      "    Uninstalling tokenizers-0.13.3:\n",
      "      Successfully uninstalled tokenizers-0.13.3\n",
      "  Attempting uninstall: transformers\n",
      "    Found existing installation: transformers 4.29.2\n",
      "    Uninstalling transformers-4.29.2:\n",
      "      Successfully uninstalled transformers-4.29.2\n",
      "\u001b[33m  WARNING: The script transformers-cli is installed in '/home/jupyter-sean/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\n",
      "Successfully installed angle-emb-0.1.1 bitsandbytes-0.41.1 huggingface-hub-0.17.3 tokenizers-0.14.1 transformers-4.34.1\n"
     ]
    }
   ],
   "source": [
    "!pip install -U angle-emb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "20db9d14",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7fb2e68b4f50>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "import random\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "\n",
    "seed = 42\n",
    "os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e08c647c",
   "metadata": {},
   "source": [
    "# 2. 加载数据\n",
    "\n",
    "需要封装成 datasets.Dataset 格式，必须包含 `text1`, `text2`, `label` 三列，`label` 列是数值类型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "edd9914f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset nli_zh (/home/jupyter-sean/.cache/huggingface/datasets/shibing624___nli_zh/ATEC/1.0.0/65b555276ee420c801e1c9eb830db959e37f42fa60c68c8b07a4448b8c436706)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f253ec2f6ae04def84d218730370f8ac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "ds = load_dataset('shibing624/nli_zh', 'ATEC')\n",
    "ds = ds.rename_column('sentence1', 'text1')\n",
    "ds = ds.rename_column('sentence2', 'text2')\n",
    "ds = ds.select_columns([\"text1\", \"text2\", \"label\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6358dab",
   "metadata": {},
   "source": [
    "# 3. 加载模型训练\n",
    "\n",
    "参数主要调整 loss_kwargs，请大家搜索参数，各参数含义参照 Paper: https://arxiv.org/abs/2309.12871\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1a2d4a34",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached shuffled indices for dataset at /home/jupyter-sean/.cache/huggingface/datasets/shibing624___nli_zh/ATEC/1.0.0/65b555276ee420c801e1c9eb830db959e37f42fa60c68c8b07a4448b8c436706/cache-bc417f8e3c8845d1.arrow\n",
      "Loading cached processed dataset at /home/jupyter-sean/.cache/huggingface/datasets/shibing624___nli_zh/ATEC/1.0.0/65b555276ee420c801e1c9eb830db959e37f42fa60c68c8b07a4448b8c436706/cache-f15bff99645a64a2_*_of_00008.arrow\n",
      "Loading cached processed dataset at /home/jupyter-sean/.cache/huggingface/datasets/shibing624___nli_zh/ATEC/1.0.0/65b555276ee420c801e1c9eb830db959e37f42fa60c68c8b07a4448b8c436706/cache-cea1adf8fcf578ef_*_of_00008.arrow\n",
      "Loading cached processed dataset at /home/jupyter-sean/.cache/huggingface/datasets/shibing624___nli_zh/ATEC/1.0.0/65b555276ee420c801e1c9eb830db959e37f42fa60c68c8b07a4448b8c436706/cache-e2aa4a5813639a58_*_of_00008.arrow\n"
     ]
    }
   ],
   "source": [
    "from angle_emb import AnglE, AngleDataTokenizer\n",
    "\n",
    "angle = AnglE.from_pretrained('hfl/chinese-roberta-wwm-ext', max_length=128, pooling_strategy='cls').cuda()\n",
    "\n",
    "train_ds = ds['train'].shuffle().map(AngleDataTokenizer(angle.tokenizer, angle.max_length), num_proc=8)\n",
    "valid_ds = ds['validation'].map(AngleDataTokenizer(angle.tokenizer, angle.max_length), num_proc=8)\n",
    "test_ds = ds['test'].map(AngleDataTokenizer(angle.tokenizer, angle.max_length), num_proc=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30e2d1d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You're using a BertTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='4867' max='4885' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [4867/4885 10:57 < 00:02, 7.40 it/s, Epoch 4.98/5]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>500</td>\n",
       "      <td>13.415200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1000</td>\n",
       "      <td>12.894200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1500</td>\n",
       "      <td>12.172400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2000</td>\n",
       "      <td>12.194400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2500</td>\n",
       "      <td>11.538200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3000</td>\n",
       "      <td>11.391500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3500</td>\n",
       "      <td>10.889600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4000</td>\n",
       "      <td>10.797600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4500</td>\n",
       "      <td>10.341900</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluate: 313it [00:14, 21.63it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "new best corrcoef!\n",
      "save to ckpts/atec/best-checkpoint\n",
      "corrcoef: 0.4922603377898468, accuracy: 0.8564, best corrcoef: 0.4922603377898468\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluate: 313it [00:14, 21.59it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "new best corrcoef!\n",
      "save to ckpts/atec/best-checkpoint\n",
      "corrcoef: 0.5103606507039081, accuracy: 0.8622, best corrcoef: 0.5103606507039081\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluate: 313it [00:14, 21.81it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "new best corrcoef!\n",
      "save to ckpts/atec/best-checkpoint\n",
      "corrcoef: 0.5121040463618336, accuracy: 0.8621, best corrcoef: 0.5121040463618336\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluate: 313it [00:14, 21.80it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "new best corrcoef!\n",
      "save to ckpts/atec/best-checkpoint\n",
      "corrcoef: 0.5130152790494402, accuracy: 0.8646, best corrcoef: 0.5130152790494402\n"
     ]
    }
   ],
   "source": [
    "angle.fit(\n",
    "    train_ds=train_ds,\n",
    "    valid_ds=test_ds,\n",
    "    output_dir='ckpts/atec',\n",
    "    batch_size=64,\n",
    "    epochs=5,\n",
    "    learning_rate=3e-5,\n",
    "    save_steps=1000,\n",
    "    eval_steps=1000,\n",
    "    warmup_steps=0,\n",
    "    gradient_accumulation_steps=1,\n",
    "    loss_kwargs={\n",
    "        'w1': 1.0,\n",
    "        'w2': 5.0,\n",
    "        'w3': 1.0,\n",
    "        'cosine_tau': 20,\n",
    "        'ibn_tau': 20,\n",
    "        'angle_tau': 1.0\n",
    "    },\n",
    "    fp16=True,\n",
    "    logging_steps=500\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98f224a1",
   "metadata": {},
   "source": [
    "# 4. 评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "946f09b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load best checkpoint and evaluate\n",
    "\n",
    "angle = AnglE.from_pretrained('hfl/chinese-roberta-wwm-ext', pretrained_model_path='ckpts/atec/best-checkpoint').cuda()\n",
    "angle.evaluate(test_ds, device=angle.device)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}