Upload train.ipynb

train.ipynb

@@ -0,0 +1,370 @@
+{
+ "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
+}