train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d2eea17d-499a-418a-b19c-541ed2af2166",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Python 3.13 is already installed\n",
      "Pinned `\u001b[36m/.python-version\u001b[39m` to `\u001b[32m3.13\u001b[39m`\n"
     ]
    }
   ],
   "source": [
    "!uv python install 313\n",
    "!uv python pin 313"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fc531160-8579-4899-87ca-040eade771f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using CPython \u001b[36m3.13.11\u001b[39m\u001b[36m\u001b[39m\n",
      "Creating virtual environment at: \u001b[36m.venv\u001b[39m\n",
      "Activate with: \u001b[32msource .venv/bin/activate\u001b[39m\n"
     ]
    }
   ],
   "source": [
    "!uv venv --clear\n",
    "!uv pip install ipykernel\n",
    "!uv run python -m ipykernel install --user --name=uv-kernel --display-name \"Python (uv venv)\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "0ef9c787-54f1-42f5-b6f9-1a18aa694b68",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2K\u001b[2mResolved \u001b[1m80 packages\u001b[0m \u001b[2min 32ms\u001b[0m\u001b[0m                                         \u001b[0m\n",
      "\u001b[2K\u001b[2mInstalled \u001b[1m1 package\u001b[0m \u001b[2min 24ms\u001b[0m\u001b[0m                                 \u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1maccelerate\u001b[0m\u001b[2m==1.14.0\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!uv pip install sentence-transformers huggingface_hub datasets tokenizers accelerate>=1.1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b50e7888-4d6b-48f5-a38f-82cd6df53da7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/.venv/lib/python3.13/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset, DatasetDict\n",
    "from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, SentenceTransformerTrainingArguments, SentenceTransformerModelCardData\n",
    "from sentence_transformers.sentence_transformer.modules import StaticEmbedding\n",
    "from sentence_transformers.sentence_transformer.losses import MultipleNegativesRankingLoss, MatryoshkaLoss\n",
    "from sentence_transformers.sentence_transformer.training_args import BatchSamplers, MultiDatasetBatchSamplers\n",
    "from tokenizers import Tokenizer\n",
    "from huggingface_hub import notebook_login"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8f7eaae2-8f37-4f4c-90d5-a35a2d258e55",
   "metadata": {},
   "outputs": [],
   "source": [
    "notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7cfd6110-7cac-4c3d-be2a-592c24e6fb73",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process__FinModernBERT_pairs(batch):\n",
    "\tanchors = []\n",
    "\tpositives = []\n",
    "\n",
    "\tfor i, score in enumerate(batch[\"score\"]):\n",
    "\t\tif float(score) >= 0.9:\n",
    "\t\t\tanchors.append(batch[\"sentence_a\"][i])\n",
    "\t\t\tpositives.append(batch[\"sentence_b\"][i])\n",
    "\n",
    "\treturn dict(anchor=anchors, positive=positives)\n",
    "FinModernBERT_pairs = load_dataset(\"BatuhanECB/FinModernBERT-pairs-sec-synthetic-v1\", data_dir=\"sts\", split=\"train\", revision=\"refs/convert/parquet\")\n",
    "FinModernBERT_pairs = FinModernBERT_pairs.map(process__FinModernBERT_pairs, batched=True, remove_columns=FinModernBERT_pairs.column_names)\n",
    "\n",
    "def process__Quora_Question_Pairs(batch):\n",
    "\tanchors = []\n",
    "\tpositives = []\n",
    "\n",
    "\tfor i, score in enumerate(batch[\"is_duplicate\"]):\n",
    "\t\tif score:\n",
    "\t\t\tanchors.append(batch[\"question1\"][i])\n",
    "\t\t\tpositives.append(batch[\"question2\"][i])\n",
    "\n",
    "\treturn dict(anchor=anchors, positive=positives)\n",
    "Quora_Question_Pairs = load_dataset(\"Heliosoph/Quora-Question-Pairs\", split=\"train\", revision=\"refs/convert/parquet\")\n",
    "Quora_Question_Pairs = Quora_Question_Pairs.map(process__Quora_Question_Pairs, batched=True, remove_columns=Quora_Question_Pairs.column_names)\n",
    "\n",
    "def process__wildchat_paraphrases(batch):\n",
    "\tanchors = []\n",
    "\tpositives = []\n",
    "\n",
    "\tfor i, prompt in enumerate(batch[\"prompt\"]):\n",
    "\t\tfor para in batch[\"paraphrases\"][i]:\n",
    "\t\t\tanchors.append(prompt)\n",
    "\t\t\tpositives.append(para)\n",
    "\n",
    "\treturn dict(anchor=anchors, positive=positives)\n",
    "wildchat_paraphrases = load_dataset(\"owenkaplinsky/wildchat-paraphrases\", split=\"train\", revision=\"refs/convert/parquet\")\n",
    "wildchat_paraphrases = wildchat_paraphrases.map(process__wildchat_paraphrases, batched=True, remove_columns=wildchat_paraphrases.column_names)\n",
    "\n",
    "def process__ogce1(batch):\n",
    "\tanchors = []\n",
    "\tpositives = []\n",
    "\n",
    "\tfor i, signal in enumerate(batch[\"signal_type\"]):\n",
    "\t\tif signal.startswith(\"gradient_\") or signal in \"word_synonym|definition_example|query_alternate|query_same_concept\":\n",
    "\t\t\tif float(batch[\"label\"][i]) >= 0.85:\n",
    "\t\t\t\tanchors.append(batch[\"text_a\"][i])\n",
    "\t\t\t\tpositives.append(batch[\"text_b\"][i])\n",
    "\n",
    "\treturn dict(anchor=anchors, positives=positives)\n",
    "ogce1 = load_dataset(\"mjbommar/ogbert-v1-contrastive\", split=\"train\", revision=\"refs/convert/parquet\")\n",
    "ogce1 = ogce1.map(process__ogce1, batched=True, remove_columns=ogce1.column_names)\n",
    "\n",
    "def process__ogce13(batch):\n",
    "\tanchors = []\n",
    "\tpositives = []\n",
    "\tnegatives_1 = []\n",
    "\tnegatives_2 = []\n",
    "\n",
    "\tfor i, anchor in enumerate(batch[\"source_sentence\"]):\n",
    "\t\tanchors.append(anchor)\n",
    "\t\tpositives.append(batch[\"near_synonym_sentence\"][i])\n",
    "\t\tnegatives_1.append(batch[\"near_antonym_sentence\"][i])\n",
    "\t\tnegatives_2.append(batch[\"antonym_sentence\"][i])\n",
    "\n",
    "\treturn dict(anchor=anchors, positives=positives, negatives_1=negatives_1, negatives_2=negatives_2)\n",
    "ogce13 = load_dataset(\"mjbommar/opengloss-v1.3-contrastive-examples\", split=\"train\", revision=\"refs/convert/parquet\")\n",
    "ogce13 = ogce13.map(process__ogce13, batched=True, remove_columns=ogce13.column_names)\n",
    "\n",
    "ds_train = DatasetDict({\n",
    "\t\"BatuhanECB/FinModernBERT-pairs-sec-synthetic-v1\": FinModernBERT_pairs,\n",
    "\t\"Heliosoph/Quora-Question-Pairs\": Quora_Question_Pairs,\n",
    "\t\"owenkaplinsky/wildchat-paraphrases\": wildchat_paraphrases,\n",
    "    \"mjbommar/ogbert-v1-contrastive\": ogce1,\n",
    "\t\"mjbommar/opengloss-v1.3-contrastive-examples\": ogce13\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "15a9927e-faa5-426a-b7bb-67b9f6ad1902",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = Tokenizer.from_pretrained(\"poolside/Laguna-S-2.1\")\n",
    "static_embedding = StaticEmbedding(tokenizer, embedding_dim=1024)\n",
    "\n",
    "PurpleStatic = SentenceTransformer(\n",
    "\tmodules=[static_embedding],\n",
    "\tmodel_card_data=SentenceTransformerModelCardData(\n",
    "\t\tlanguage=\"en\",\n",
    "\t\tlicense=\"wtfpl\",\n",
    "\t\tmodel_name=\"PurpleStatic: Static Embeddings\"\n",
    "\t)\n",
    ")\n",
    "\n",
    "base_loss = MultipleNegativesRankingLoss(PurpleStatic)\n",
    "loss = MatryoshkaLoss(PurpleStatic, base_loss, matryoshka_dims=[1024, 768, 512, 256, 128, 64, 32])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec6ecbba-0d42-4189-9067-b5f92b4637c7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='21371' max='21400' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [21371/21400 2:03:29 < 00:12, 2.34 it/s, Epoch 19.97/20]\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>5000</td>\n",
       "      <td>40.791266</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6000</td>\n",
       "      <td>40.154305</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7000</td>\n",
       "      <td>38.176074</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>8000</td>\n",
       "      <td>38.800340</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>9000</td>\n",
       "      <td>38.206828</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10000</td>\n",
       "      <td>37.408781</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>11000</td>\n",
       "      <td>36.580504</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>12000</td>\n",
       "      <td>37.152621</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>13000</td>\n",
       "      <td>36.002770</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>14000</td>\n",
       "      <td>35.222508</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>15000</td>\n",
       "      <td>36.238773</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>16000</td>\n",
       "      <td>35.199824</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>17000</td>\n",
       "      <td>35.261895</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>18000</td>\n",
       "      <td>35.194262</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>19000</td>\n",
       "      <td>36.039648</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>20000</td>\n",
       "      <td>34.257383</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>21000</td>\n",
       "      <td>35.113922</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# we are running on RTX 5090\n",
    "args = SentenceTransformerTrainingArguments(\n",
    "\tnum_train_epochs=20,\n",
    "\tper_device_train_batch_size=4096,\n",
    "\tlearning_rate=1e-4,\n",
    "\twarmup_steps=0.1,\n",
    "\tfp16=False,\n",
    "\tbf16=True,\n",
    "\tbatch_sampler=BatchSamplers.NO_DUPLICATES,\n",
    "\tmulti_dataset_batch_sampler=MultiDatasetBatchSamplers.PROPORTIONAL,\n",
    "\n",
    "\tlogging_steps=1000,\n",
    "    logging_first_step=True,\n",
    "    save_strategy=\"no\"\n",
    ")\n",
    "\n",
    "trainer = SentenceTransformerTrainer(\n",
    "\tmodel=PurpleStatic,\n",
    "\targs=args,\n",
    "\ttrain_dataset=ds_train,\n",
    "\tloss=loss\n",
    ")\n",
    "trainer.train(resume_from_checkpoint=\"trainer_output/checkpoint-4000/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "389fbdee-c86d-48a6-837b-6df7a2ec75cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "PurpleStatic.push_to_hub(\"LocalWisdom/PurpleStatic\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (uv venv)",
   "language": "python",
   "name": "uv-kernel"
  },
  "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.13.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}