| "script": "# /// script\n# dependencies = [\n# \"sentence-transformers>=3.0\",\n# \"datasets>=2.20\",\n# \"torch>=2.4\",\n# \"accelerate>=0.33\",\n# \"transformers>=4.44\",\n# \"peft>=0.12\",\n# \"huggingface_hub>=0.24\",\n# ]\n# ///\n\nimport json\nimport os\nfrom pathlib import Path\n\nfrom datasets import Dataset\nfrom huggingface_hub import hf_hub_download\nfrom sentence_transformers import (\n SentenceTransformer,\n SentenceTransformerTrainer,\n SentenceTransformerTrainingArguments,\n losses,\n)\n\nDATASET_REPO = os.environ[\"DATASET_REPO\"]\nHUB_MODEL_ID = os.environ[\"HUB_MODEL_ID\"]\nBASE_MODEL = os.environ.get(\"BASE_MODEL\", \"BAAI/bge-m3\")\n\n\ndef read_jsonl(filename):\n path = Path(hf_hub_download(repo_id=DATASET_REPO, filename=filename, repo_type=\"dataset\"))\n rows = []\n with path.open(\"r\", encoding=\"utf-8\") as fh:\n for line in fh:\n line = line.strip()\n if line:\n rows.append(json.loads(line))\n return rows\n\n\ndef to_training_rows(rows):\n return [\n {\n \"anchor\": \"Represent this sentence for searching relevant passages: \" + row[\"query\"],\n \"positive\": row[\"positive_text\"],\n \"negative\": row[\"negative_text\"],\n }\n for row in rows\n ]\n\n\ntrain_rows = read_jsonl(\"train_triplets.jsonl\")\ndev_rows = read_jsonl(\"dev_triplets.jsonl\")\ntrain_dataset = Dataset.from_list(to_training_rows(train_rows))\neval_dataset = Dataset.from_list(to_training_rows(dev_rows)) if dev_rows else None\n\nmodel = SentenceTransformer(BASE_MODEL)\n\nfrom peft import LoraConfig, TaskType\nmodel.add_adapter(\n LoraConfig(\n task_type=TaskType.FEATURE_EXTRACTION,\n r=16,\n lora_alpha=32,\n lora_dropout=0.05,\n bias=\"none\",\n )\n)\n\ncached_loss = getattr(losses, \"CachedMultipleNegativesRankingLoss\", None)\nbase_loss = cached_loss(model=model) if cached_loss is not None else losses.MultipleNegativesRankingLoss(model=model)\ntrain_loss = losses.MatryoshkaLoss(model=model, loss=base_loss, matryoshka_dims=[1024, 256])\n\nargs = SentenceTransformerTrainingArguments(\n output_dir=\"vinci-v0-output\",\n num_train_epochs=1,\n per_device_train_batch_size=16,\n per_device_eval_batch_size=16,\n gradient_accumulation_steps=4,\n learning_rate=2e-05,\n seed=42,\n data_seed=42,\n eval_strategy=\"epoch\" if eval_dataset is not None else \"no\",\n save_strategy=\"epoch\",\n save_total_limit=2,\n logging_steps=10,\n report_to=\"none\",\n run_name=\"vinci-v0-gpu_16gb_lora\",\n push_to_hub=True,\n hub_model_id=HUB_MODEL_ID,\n hub_strategy=\"end\",\n max_steps=20,\n)\n\ntrainer = SentenceTransformerTrainer(\n model=model,\n args=args,\n train_dataset=train_dataset,\n eval_dataset=eval_dataset,\n loss=train_loss,\n)\ntrainer.train()\nmodel.save_pretrained(\"vinci-v0-output\")\nmodel.push_to_hub(HUB_MODEL_ID, exist_ok=True)\n", |
