embed-train 3.1.0__tar.gz → 3.2.0__tar.gz

{embed_train-3.1.0 → embed_train-3.2.0}/CHANGELOG.md

@@ -1,3 +1,10 @@
+# [3.2.0](https://gitlab.com/efysent/agentic-core/embed-train/compare/v3.1.0...v3.2.0) (2026-05-12)
+
+
+### Features
+
+* remove sentence transformers deprecations ([581032b](https://gitlab.com/efysent/agentic-core/embed-train/commit/581032b97a0bfc50236cc0fa64a857bc65f20314))
+
 # [3.1.0](https://gitlab.com/efysent/agentic-core/embed-train/compare/v3.0.0...v3.1.0) (2026-05-12)
 
 

{embed_train-3.1.0 → embed_train-3.2.0}/PKG-INFO

@@ -1,14 +1,15 @@
 Metadata-Version: 2.4
 Name: embed-train
-Version: 3.1.0
+Version: 3.2.0
 Author-email: jalal <jalalkhaldi3@gmail.com>
 Requires-Python: <3.13,>=3.11
-Requires-Dist: accelerate<2.0.0,>=1.13.0
-Requires-Dist: datasets<5.0.0,>=4.5.0
+Requires-Dist: accelerate==1.13.0
+Requires-Dist: datasets==4.8.4
 Requires-Dist: retrievalbase<3.0.0,>=2.1.0
-Requires-Dist: sentence-transformers<6.0.0,>=5.1.2
-Requires-Dist: tensorboard<3.0.0,>=2.20.0
-Requires-Dist: torch<3.0.0,>=2.9.0
+Requires-Dist: sentence-transformers==5.4.1
+Requires-Dist: tensorboard==2.20.0
+Requires-Dist: torch==2.11.0
+Requires-Dist: transformers==4.57.6
 Description-Content-Type: text/markdown
 
 # embed-train

{embed_train-3.1.0 → embed_train-3.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "embed-train"
-version = "3.1.0"
+version = "3.2.0"
 description = ""
 authors = [
     { name = "jalal", email = "jalalkhaldi3@gmail.com" }
@@ -9,11 +9,12 @@ readme = "README.md"
 requires-python = ">=3.11,<3.13"
 
 dependencies = [
-    "torch>=2.9.0,<3.0.0",
-    "sentence-transformers>=5.1.2,<6.0.0",
-    "datasets>=4.5.0,<5.0.0",
-    "tensorboard>=2.20.0,<3.0.0",
-    "accelerate>=1.13.0,<2.0.0",
+    "torch==2.11.0",
+    "sentence-transformers==5.4.1",
+    "transformers==4.57.6",
+    "datasets==4.8.4",
+    "tensorboard==2.20.0",
+    "accelerate==1.13.0",
     "retrievalbase>=2.1.0,<3.0.0",
 ]
 

{embed_train-3.1.0 → embed_train-3.2.0}/src/embed_train/settings.py

@@ -81,7 +81,7 @@ class SentenceTransformerHardNegativeMinerSettings(HardNegativeMinerSettings):
     model_name_or_path: str
     cross_encoder_model_name_or_path: str | None = None
     tokenizer: TokenizerSettings
-    pooling: Literal["cls", "mean_tokens", "max_tokens"]
+    pooling: Literal["cls", "max", "mean", "mean_sqrt_len_tokens", "weightedmean", "lasttoken"]
     anchor_column_name: str = "query"
     positive_column_name: str = "positive"
     range_min: int = 0

{embed_train-3.1.0 → embed_train-3.2.0}/src/embed_train/train/dataset/hard_negatives.py

@@ -32,13 +32,13 @@ class SentenceTransformerHardNegativeMiner(HardNegativeMiner[SentenceTransformer
             model_name_or_path=self.config.model_name_or_path,
             max_seq_length=self.config.tokenizer.max_length,
             tokenizer_name_or_path=self.config.tokenizer.name,
-            model_args={
+            model_kwargs={
                 "trust_remote_code": trust,
             },
-            config_args={
+            config_kwargs={
                 "trust_remote_code": trust,
             },
-            tokenizer_args={
+            processor_kwargs={
                 "trust_remote_code": trust,
                 "padding": self.config.tokenizer.padding,
                 "truncation": self.config.tokenizer.truncation,
@@ -46,18 +46,11 @@ class SentenceTransformerHardNegativeMiner(HardNegativeMiner[SentenceTransformer
             },
         )
 
-        pooling = Pooling(
-            transformer.get_word_embedding_dimension(),
-            pooling_mode_mean_tokens=self.config.pooling == "mean_tokens",
-            pooling_mode_cls_token=self.config.pooling == "cls",
-            pooling_mode_max_tokens=self.config.pooling == "max_tokens",
-        )
+        pooling = Pooling(transformer.get_embedding_dimension(), pooling_mode=self.config.pooling)
 
         model = SentenceTransformer(modules=[transformer, pooling])
 
-        _logger.info(
-            f"SentenceTransformer loaded successfully | embedding_dim={transformer.get_word_embedding_dimension()}"
-        )
+        _logger.info(f"SentenceTransformer loaded successfully | embedding_dim={transformer.get_embedding_dimension()}")
 
         return model
 

{embed_train-3.1.0 → embed_train-3.2.0}/src/embed_train/train/trainers/hf/__init__.py

@@ -64,7 +64,7 @@ class SentenceTransformersTrainer[TCHFTrainRunner: "SentenceTransformersTrainerS
         _logger.info("Starting SentenceTransformers training...")
         trainer.train()
 
-    def _get_warmup_steps(self, dataset: Dataset) -> float:
+    def _get_warmup_steps(self, dataset: Dataset) -> int:
         train_size = len(dataset)
         steps_per_epoch = train_size // self.config.per_device_train_batch_size
         total_steps = steps_per_epoch * self.config.num_epochs
@@ -96,13 +96,13 @@ class SentenceTransformersTrainer[TCHFTrainRunner: "SentenceTransformersTrainerS
             model_name_or_path=model_path,
             max_seq_length=self.config.tokenizer.max_length,
             tokenizer_name_or_path=self.config.tokenizer.name,
-            model_args={
+            model_kwargs={
                 "trust_remote_code": trust,
             },
-            config_args={
+            config_kwargs={
                 "trust_remote_code": trust,
             },
-            tokenizer_args={
+            processor_kwargs={
                 "trust_remote_code": trust,
                 "padding": self.config.tokenizer.padding,
                 "truncation": self.config.tokenizer.truncation,
@@ -110,7 +110,7 @@ class SentenceTransformersTrainer[TCHFTrainRunner: "SentenceTransformersTrainerS
             },
         )
         pooling = Pooling(
-            transformer.get_word_embedding_dimension(),
+            transformer.get_embedding_dimension(),
             pooling_mode_mean_tokens=self.config.pooling == "mean_tokens",
             pooling_mode_cls_token=self.config.pooling == "cls",
             pooling_mode_max_tokens=self.config.pooling == "max_tokens",

{embed_train-3.1.0 → embed_train-3.2.0}/tests/fixtures/components.py

@@ -264,7 +264,7 @@ def build_hard_negative_miner_settings(**overrides: Any) -> SentenceTransformerH
         "module_path": "embed_train.train.dataset.hard_negatives.SentenceTransformerHardNegativeMiner",
         "model_name_or_path": "dummy-miner",
         "tokenizer": build_tokenizer_settings(),
-        "pooling": "mean_tokens",
+        "pooling": "mean",
         "range_min": 0,
         "range_max": 2,
         "relative_margin": 0.1,

{embed_train-3.1.0 → embed_train-3.2.0}/tests/unit/test_push_to_hf.py

@@ -158,6 +158,29 @@ def test_save_model_locally_passes_expected_arguments(tmp_path: Path) -> None:
     ]
 
 
+def test_ensure_hf_repo_exists_creates_repo_when_enabled(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    config = build_push_to_hf_runner_settings(tmp_path, create_repo=True)
+    runner = PushToHFRunner(config)
+    payloads: list[dict[str, object]] = []
+
+    class DummyApi:
+        def create_repo(self, **kwargs: object) -> None:
+            payloads.append(kwargs)
+
+    monkeypatch.setattr("embed_train.push_to_hf.HfApi", DummyApi)
+
+    runner._ensure_hf_repo_exists()
+
+    assert payloads == [
+        {
+            "repo_id": config.hf.repo,
+            "repo_type": "model",
+            "private": config.hf.private,
+            "exist_ok": True,
+        }
+    ]
+
+
 def test_push_repo_to_hf_calls_upload_folder(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
     config = build_push_to_hf_runner_settings(tmp_path)
     runner = PushToHFRunner(config)

{embed_train-3.1.0 → embed_train-3.2.0}/tests/unit/test_train/test_collate.py

@@ -69,3 +69,15 @@ def test_hard_negative_collate_processes_negative_columns(monkeypatch) -> None:
 
     assert cast(Any, q_tok)["input_ids"].shape == torch.Size([2, 3])
     assert cast(Any, c_tok)["input_ids"].shape == torch.Size([6, 3])
+
+
+def test_hard_negative_collate_processes_single_negative(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "embed_train.train.dataset.AutoTokenizer.from_pretrained", lambda *args, **kwargs: DummyTokenizer()
+    )
+    collate = HardNegativeCollateFn(build_hard_negative_collate_settings(), context=None)
+
+    q_tok, c_tok = collate(cast(Any, [{"query": "q1", "positive": "p1", "negative": "n1"}]))
+
+    assert cast(Any, q_tok)["input_ids"].shape == torch.Size([1, 3])
+    assert cast(Any, c_tok)["input_ids"].shape == torch.Size([2, 3])

{embed_train-3.1.0 → embed_train-3.2.0}/tests/unit/test_train/test_dataset.py

@@ -2,13 +2,16 @@ from __future__ import annotations
 
 from typing import Any, cast
 
+import pytest
 import torch
+from datasets import Dataset
 
-from embed_train.train.dataset import CollateFn
+from embed_train.train.dataset import CollateFn, HardNegativeMiner
 from tests.fixtures.components import (
     DummyDatasetConnector,
     DummyTokenizer,
     build_collate_settings,
+    build_hard_negative_miner_settings,
     build_multi_positive_torch_dataset_settings,
 )
 
@@ -62,3 +65,14 @@ def test_torch_dataset_loads_runtime_dataset_and_converts_to_hf() -> None:
     assert len(dataset) == 1
     assert hf_dataset[0]["query"] == "query-a"
     assert hf_dataset[0]["positives"] == ["doc-1", "doc-2"]
+
+
+def test_hard_negative_miner_base_mine_raises_not_implemented() -> None:
+    class ConcreteHardNegativeMiner(HardNegativeMiner):
+        def mine(self, dataset: Dataset) -> Dataset:
+            return super().mine(dataset)
+
+    miner = ConcreteHardNegativeMiner(build_hard_negative_miner_settings())
+
+    with pytest.raises(NotImplementedError):
+        miner.mine(Dataset.from_list([]))

embed_train-3.2.0/tests/unit/test_train/test_hard_negatives.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+from typing import Any
+
+from datasets import Dataset
+
+from embed_train.train.dataset.hard_negatives import SentenceTransformerHardNegativeMiner
+from tests.fixtures.components import build_hard_negative_miner_settings
+
+
+class DummyTransformer:
+    calls: list[dict[str, Any]] = []
+
+    def __init__(self, **kwargs: Any) -> None:
+        self.calls.append(kwargs)
+
+    def get_embedding_dimension(self) -> int:
+        return 8
+
+
+class DummyPooling:
+    calls: list[dict[str, Any]] = []
+
+    def __init__(self, embedding_dimension: int, *, pooling_mode: str) -> None:
+        self.calls.append({"embedding_dimension": embedding_dimension, "pooling_mode": pooling_mode})
+
+
+class DummySentenceTransformer:
+    calls: list[list[object]] = []
+
+    def __init__(self, *, modules: list[object]) -> None:
+        self.calls.append(modules)
+
+
+class DummyCrossEncoder:
+    calls: list[dict[str, Any]] = []
+
+    def __init__(self, model_name_or_path: str, *, trust_remote_code: bool) -> None:
+        self.calls.append(
+            {
+                "model_name_or_path": model_name_or_path,
+                "trust_remote_code": trust_remote_code,
+            }
+        )
+
+
+def test_sentence_transformer_hard_negative_miner_mines_with_config(
+    monkeypatch,
+) -> None:
+    captured: dict[str, Any] = {}
+    cleanup_calls: list[str] = []
+    mined = Dataset.from_list([{"query": "q1", "positive": "p1", "negative": "n1"}])
+
+    def fake_mine_hard_negatives(**kwargs: Any) -> Dataset:
+        captured.update(kwargs)
+        return mined
+
+    class DummyCuda:
+        @staticmethod
+        def is_available() -> bool:
+            return True
+
+        @staticmethod
+        def memory_allocated() -> int:
+            cleanup_calls.append("allocated")
+            return 1024**3
+
+        @staticmethod
+        def memory_reserved() -> int:
+            cleanup_calls.append("reserved")
+            return 2 * 1024**3
+
+        @staticmethod
+        def empty_cache() -> None:
+            cleanup_calls.append("empty_cache")
+
+        @staticmethod
+        def ipc_collect() -> None:
+            cleanup_calls.append("ipc_collect")
+
+    monkeypatch.setattr("embed_train.train.dataset.hard_negatives.Transformer", DummyTransformer)
+    monkeypatch.setattr("embed_train.train.dataset.hard_negatives.Pooling", DummyPooling)
+    monkeypatch.setattr("embed_train.train.dataset.hard_negatives.SentenceTransformer", DummySentenceTransformer)
+    monkeypatch.setattr("embed_train.train.dataset.hard_negatives.CrossEncoder", DummyCrossEncoder)
+    monkeypatch.setattr("embed_train.train.dataset.hard_negatives.mine_hard_negatives", fake_mine_hard_negatives)
+    monkeypatch.setattr("embed_train.train.dataset.hard_negatives.torch.cuda", DummyCuda)
+
+    config = build_hard_negative_miner_settings(
+        cross_encoder_model_name_or_path="dummy-reranker",
+        trust_remote_code=True,
+    )
+    miner = SentenceTransformerHardNegativeMiner(config)
+
+    result = miner.mine(Dataset.from_list([{"query": "q1", "positive": "p1"}]))
+
+    assert result is mined
+    assert DummyTransformer.calls[-1]["model_name_or_path"] == "dummy-miner"
+    assert DummyTransformer.calls[-1]["processor_kwargs"]["model_max_length"] == config.tokenizer.max_length
+    assert DummyPooling.calls[-1] == {"embedding_dimension": 8, "pooling_mode": "mean"}
+    assert len(DummySentenceTransformer.calls[-1]) == 2
+    assert DummyCrossEncoder.calls[-1] == {
+        "model_name_or_path": "dummy-reranker",
+        "trust_remote_code": True,
+    }
+    assert captured["anchor_column_name"] == "query"
+    assert captured["positive_column_name"] == "positive"
+    assert captured["cross_encoder"] is not None
+    assert captured["num_negatives"] == 2
+    assert captured["sampling_strategy"] == "top"
+    assert "empty_cache" in cleanup_calls
+    assert "ipc_collect" in cleanup_calls
+
+
+def test_sentence_transformer_hard_negative_miner_skips_cross_encoder() -> None:
+    miner = SentenceTransformerHardNegativeMiner(build_hard_negative_miner_settings())
+
+    assert miner._load_cross_encoder() is None

{embed_train-3.1.0 → embed_train-3.2.0}/tests/unit/test_train/test_hf_trainer.py

@@ -19,7 +19,7 @@ class DummyTransformer:
     def __init__(self, **kwargs):
         self.kwargs = kwargs
 
-    def get_word_embedding_dimension(self) -> int:
+    def get_embedding_dimension(self) -> int:
         return 16
 
 

{embed_train-3.1.0 → embed_train-3.2.0}/tests/unit/test_train/test_loss.py

@@ -58,3 +58,26 @@ def test_hard_negative_contrastive_loss_rejects_invalid_candidate_layout() -> No
 
     with pytest.raises(EmbedTrainValueError, match="one positive and at least one negative"):
         loss(q_emb, c_emb)
+
+
+def test_hard_negative_contrastive_loss_returns_scalar_for_valid_layout() -> None:
+    loss = HardNegativeContrastiveLoss(
+        HardNegativeContrastiveLossSettings(
+            module_path="embed_train.train.trainers.torch.loss.HardNegativeContrastiveLoss",
+            temperature=0.5,
+        )
+    )
+    q_emb = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+    c_emb = torch.tensor(
+        [
+            [1.0, 0.0],
+            [0.0, 1.0],
+            [0.0, 1.0],
+            [1.0, 0.0],
+        ]
+    )
+
+    result = loss(q_emb, c_emb)
+
+    assert result.ndim == 0
+    assert result.item() >= 0

embed_train-3.2.0/tests/unit/test_train/test_torch_datasets.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import pytest
+from datasets import Dataset
+
+from tests.fixtures.components import (
+    DummyDatasetConnector,
+    build_hard_negative_torch_dataset_settings,
+    build_multi_positive_torch_dataset_settings,
+    build_torch_dataset_settings,
+)
+from tests.fixtures.data import build_query_rows
+
+
+def test_query_multi_positive_dataset_groups_rows() -> None:
+    from embed_train.train.dataset.torch_datasets import QueryMultiPositiveDataset
+
+    DummyDatasetConnector.rows = build_query_rows()
+    dataset = QueryMultiPositiveDataset(build_multi_positive_torch_dataset_settings())
+
+    assert len(dataset) == 2
+    grouped = {item["query"]: item["positives"] for item in (dataset[i] for i in range(len(dataset)))}
+    assert grouped["query-a"] == ["doc-1", "doc-2"]
+    assert grouped["query-b"] == ["doc-3"]
+
+
+def test_query_positive_dataset_flattens_rows() -> None:
+    from embed_train.train.dataset.torch_datasets import QueryPositiveDataset
+
+    DummyDatasetConnector.rows = build_query_rows()
+    dataset = QueryPositiveDataset(build_torch_dataset_settings())
+
+    assert len(dataset) == 3
+    rows = {tuple(item.values()) for item in (dataset[i] for i in range(len(dataset)))}
+    assert ("query-a", "doc-1") in rows
+    assert ("query-a", "doc-2") in rows
+    assert ("query-b", "doc-3") in rows
+
+
+def test_query_positive_dataset_rejects_empty_positive() -> None:
+    from embed_train.train.dataset.torch_datasets import QueryPositiveDataset
+
+    DummyDatasetConnector.rows = [
+        {
+            "page_content": "",
+            "metadata": {"query": "query-a", "step_range": (0, 1), "section": "alpha"},
+        }
+    ]
+    dataset = QueryPositiveDataset(build_torch_dataset_settings())
+
+    with pytest.raises(ValueError, match="No positive passage found for index 0"):
+        dataset[0]
+
+
+class DummyHardNegativeMiner:
+    seen_dataset: Dataset | None = None
+
+    @classmethod
+    def from_config(cls, _config: object) -> DummyHardNegativeMiner:
+        return cls()
+
+    def mine(self, dataset: Dataset) -> Dataset:
+        self.__class__.seen_dataset = dataset
+        return Dataset.from_list(
+            [
+                {"query": "query-a", "positive": "doc-1", "negative": "doc-2"},
+                {"query": "query-b", "positive": "doc-3", "negative": "doc-1"},
+            ]
+        )
+
+
+def test_hard_negative_dataset_mines_and_exposes_rows(monkeypatch: pytest.MonkeyPatch) -> None:
+    from embed_train.train.dataset.torch_datasets import HardNegativeDataset
+
+    DummyDatasetConnector.rows = build_query_rows()
+    monkeypatch.setattr("embed_train.train.dataset.torch_datasets.load_class", lambda _: DummyHardNegativeMiner)
+
+    dataset = HardNegativeDataset(build_hard_negative_torch_dataset_settings())
+
+    assert len(dataset) == 2
+    assert dataset[0] == {"query": "query-a", "positive": "doc-1", "negative": "doc-2"}
+    assert DummyHardNegativeMiner.seen_dataset is not None
+    assert DummyHardNegativeMiner.seen_dataset[0] == {"query": "query-a", "positive": "doc-1"}

{embed_train-3.1.0 → embed_train-3.2.0}/uv.lock

@@ -110,15 +110,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" },
 ]
 
-[[package]]
-name = "annotated-doc"
-version = "0.0.4"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/57/ba/046ceea27344560984e26a590f90bc7f4a75b06701f653222458922b558c/annotated_doc-0.0.4.tar.gz", hash = "sha256:fbcda96e87e9c92ad167c2e53839e57503ecfda18804ea28102353485033faa4", size = 7288, upload-time = "2025-11-10T22:07:42.062Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/1e/d3/26bf1008eb3d2daa8ef4cacc7f3bfdc11818d111f7e2d0201bc6e3b49d45/annotated_doc-0.0.4-py3-none-any.whl", hash = "sha256:571ac1dc6991c450b25a9c2d84a3705e2ae7a53467b5d111c24fa8baabbed320", size = 5303, upload-time = "2025-11-10T22:07:40.673Z" },
-]
-
 [[package]]
 name = "annotated-types"
 version = "0.7.0"
@@ -284,18 +275,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/db/8f/61959034484a4a7c527811f4721e75d02d653a35afb0b6054474d8185d4c/charset_normalizer-3.4.7-py3-none-any.whl", hash = "sha256:3dce51d0f5e7951f8bb4900c257dad282f49190fdbebecd4ba99bcc41fef404d", size = 61958, upload-time = "2026-04-02T09:28:37.794Z" },
 ]
 
-[[package]]
-name = "click"
-version = "8.3.2"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "colorama", marker = "sys_platform == 'win32'" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/57/75/31212c6bf2503fdf920d87fee5d7a86a2e3bcf444984126f13d8e4016804/click-8.3.2.tar.gz", hash = "sha256:14162b8b3b3550a7d479eafa77dfd3c38d9dc8951f6f69c78913a8f9a7540fd5", size = 302856, upload-time = "2026-04-03T19:14:45.118Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/e4/20/71885d8b97d4f3dde17b1fdb92dbd4908b00541c5a3379787137285f602e/click-8.3.2-py3-none-any.whl", hash = "sha256:1924d2c27c5653561cd2cae4548d1406039cb79b858b747cfea24924bbc1616d", size = 108379, upload-time = "2026-04-03T19:14:43.505Z" },
-]
-
 [[package]]
 name = "colorama"
 version = "0.4.6"
@@ -459,7 +438,7 @@ wheels = [
 
 [[package]]
 name = "embed-train"
-version = "3.1.0"
+version = "3.2.0"
 source = { editable = "." }
 dependencies = [
     { name = "accelerate" },
@@ -468,6 +447,7 @@ dependencies = [
     { name = "sentence-transformers" },
     { name = "tensorboard" },
     { name = "torch" },
+    { name = "transformers" },
 ]
 
 [package.dev-dependencies]
@@ -484,12 +464,13 @@ dev = [
 
 [package.metadata]
 requires-dist = [
-    { name = "accelerate", specifier = ">=1.13.0,<2.0.0" },
-    { name = "datasets", specifier = ">=4.5.0,<5.0.0" },
+    { name = "accelerate", specifier = "==1.13.0" },
+    { name = "datasets", specifier = "==4.8.4" },
     { name = "retrievalbase", specifier = ">=2.1.0,<3.0.0" },
-    { name = "sentence-transformers", specifier = ">=5.1.2,<6.0.0" },
-    { name = "tensorboard", specifier = ">=2.20.0,<3.0.0" },
-    { name = "torch", specifier = ">=2.9.0,<3.0.0" },
+    { name = "sentence-transformers", specifier = "==5.4.1" },
+    { name = "tensorboard", specifier = "==2.20.0" },
+    { name = "torch", specifier = "==2.11.0" },
+    { name = "transformers", specifier = "==4.57.6" },
 ]
 
 [package.metadata.requires-dev]
@@ -702,22 +683,21 @@ http2 = [
 
 [[package]]
 name = "huggingface-hub"
-version = "1.11.0"
+version = "0.36.2"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "filelock" },
     { name = "fsspec" },
-    { name = "hf-xet", marker = "platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" },
-    { name = "httpx" },
+    { name = "hf-xet", marker = "platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" },
     { name = "packaging" },
     { name = "pyyaml" },
+    { name = "requests" },
     { name = "tqdm" },
-    { name = "typer" },
     { name = "typing-extensions" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/dc/89/e7aa12d8a6b9259bed10671abb25ae6fa437c0f88a86ecbf59617bae7759/huggingface_hub-1.11.0.tar.gz", hash = "sha256:15fb3713c7f9cdff7b808a94fd91664f661ab142796bb48c9cd9493e8d166278", size = 761749, upload-time = "2026-04-16T13:07:39.73Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/7c/b7/8cb61d2eece5fb05a83271da168186721c450eb74e3c31f7ef3169fa475b/huggingface_hub-0.36.2.tar.gz", hash = "sha256:1934304d2fb224f8afa3b87007d58501acfda9215b334eed53072dd5e815ff7a", size = 649782, upload-time = "2026-02-06T09:24:13.098Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/37/02/4f3f8997d1ea7fe0146b343e5e14bd065fa87af790d07e5576d31b31cc18/huggingface_hub-1.11.0-py3-none-any.whl", hash = "sha256:42a6de0afbfeb5e022222d36398f029679db4eb4778801aafda32257ae9131ab", size = 645499, upload-time = "2026-04-16T13:07:37.716Z" },
+    { url = "https://files.pythonhosted.org/packages/a8/af/48ac8483240de756d2438c380746e7130d1c6f75802ef22f3c6d49982787/huggingface_hub-0.36.2-py3-none-any.whl", hash = "sha256:48f0c8eac16145dfce371e9d2d7772854a4f591bcb56c9cf548accf531d54270", size = 566395, upload-time = "2026-02-06T09:24:11.133Z" },
 ]
 
 [[package]]
@@ -2085,15 +2065,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/e1/e3/c164c88b2e5ce7b24d667b9bd83589cf4f3520d97cad01534cd3c4f55fdb/setuptools-81.0.0-py3-none-any.whl", hash = "sha256:fdd925d5c5d9f62e4b74b30d6dd7828ce236fd6ed998a08d81de62ce5a6310d6", size = 1062021, upload-time = "2026-02-06T21:10:37.175Z" },
 ]
 
-[[package]]
-name = "shellingham"
-version = "1.5.4"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/58/15/8b3609fd3830ef7b27b655beb4b4e9c62313a4e8da8c676e142cc210d58e/shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de", size = 10310, upload-time = "2023-10-24T04:13:40.426Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" },
-]
-
 [[package]]
 name = "six"
 version = "1.17.0"
@@ -2279,22 +2250,23 @@ wheels = [
 
 [[package]]
 name = "transformers"
-version = "5.5.4"
+version = "4.57.6"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
+    { name = "filelock" },
     { name = "huggingface-hub" },
     { name = "numpy" },
     { name = "packaging" },
     { name = "pyyaml" },
     { name = "regex" },
+    { name = "requests" },
     { name = "safetensors" },
     { name = "tokenizers" },
     { name = "tqdm" },
-    { name = "typer" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/a5/1e/1e244ab2ab50a863e6b52cc55761910567fa532b69a6740f6e99c5fdbd98/transformers-5.5.4.tar.gz", hash = "sha256:2e67cadba81fc7608cc07c4dd54f524820bc3d95b1cabd0ef3db7733c4f8b82e", size = 8227649, upload-time = "2026-04-13T16:55:55.181Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/c4/35/67252acc1b929dc88b6602e8c4a982e64f31e733b804c14bc24b47da35e6/transformers-4.57.6.tar.gz", hash = "sha256:55e44126ece9dc0a291521b7e5492b572e6ef2766338a610b9ab5afbb70689d3", size = 10134912, upload-time = "2026-01-16T10:38:39.284Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/29/fb/162a66789c65e5afa3b051309240c26bf37fbc8fea285b4546ae747995a2/transformers-5.5.4-py3-none-any.whl", hash = "sha256:0bd6281b82966fe5a7a16f553ea517a9db1dee6284d7cb224dfd88fc0dd1c167", size = 10236696, upload-time = "2026-04-13T16:55:51.497Z" },
+    { url = "https://files.pythonhosted.org/packages/03/b8/e484ef633af3887baeeb4b6ad12743363af7cce68ae51e938e00aaa0529d/transformers-4.57.6-py3-none-any.whl", hash = "sha256:4c9e9de11333ddfe5114bc872c9f370509198acf0b87a832a0ab9458e2bd0550", size = 11993498, upload-time = "2026-01-16T10:38:31.289Z" },
 ]
 
 [[package]]
@@ -2332,21 +2304,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/88/39/bca669095ccf0a400af941fdf741578d4c2d6719f1b7f10e6dbec10aa862/ty-0.0.31-py3-none-win_arm64.whl", hash = "sha256:e9cb15fad26545c6a608f40f227af3a5513cb376998ca6feddd47ca7d93ffafa", size = 10590392, upload-time = "2026-04-15T15:47:57.968Z" },
 ]
 
-[[package]]
-name = "typer"
-version = "0.24.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "annotated-doc" },
-    { name = "click" },
-    { name = "rich" },
-    { name = "shellingham" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/f5/24/cb09efec5cc954f7f9b930bf8279447d24618bb6758d4f6adf2574c41780/typer-0.24.1.tar.gz", hash = "sha256:e39b4732d65fbdcde189ae76cf7cd48aeae72919dea1fdfc16593be016256b45", size = 118613, upload-time = "2026-02-21T16:54:40.609Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/4a/91/48db081e7a63bb37284f9fbcefda7c44c277b18b0e13fbc36ea2335b71e6/typer-0.24.1-py3-none-any.whl", hash = "sha256:112c1f0ce578bfb4cab9ffdabc68f031416ebcc216536611ba21f04e9aa84c9e", size = 56085, upload-time = "2026-02-21T16:54:41.616Z" },
-]
-
 [[package]]
 name = "types-pyyaml"
 version = "6.0.12.20260408"

embed_train-3.1.0/tests/unit/test_train/test_torch_datasets.py

@@ -1,33 +0,0 @@
-from __future__ import annotations
-
-from tests.fixtures.components import (
-    DummyDatasetConnector,
-    build_multi_positive_torch_dataset_settings,
-    build_torch_dataset_settings,
-)
-from tests.fixtures.data import build_query_rows
-
-
-def test_query_multi_positive_dataset_groups_rows() -> None:
-    from embed_train.train.dataset.torch_datasets import QueryMultiPositiveDataset
-
-    DummyDatasetConnector.rows = build_query_rows()
-    dataset = QueryMultiPositiveDataset(build_multi_positive_torch_dataset_settings())
-
-    assert len(dataset) == 2
-    grouped = {item["query"]: item["positives"] for item in (dataset[i] for i in range(len(dataset)))}
-    assert grouped["query-a"] == ["doc-1", "doc-2"]
-    assert grouped["query-b"] == ["doc-3"]
-
-
-def test_query_positive_dataset_flattens_rows() -> None:
-    from embed_train.train.dataset.torch_datasets import QueryPositiveDataset
-
-    DummyDatasetConnector.rows = build_query_rows()
-    dataset = QueryPositiveDataset(build_torch_dataset_settings())
-
-    assert len(dataset) == 3
-    rows = {tuple(item.values()) for item in (dataset[i] for i in range(len(dataset)))}
-    assert ("query-a", "doc-1") in rows
-    assert ("query-a", "doc-2") in rows
-    assert ("query-b", "doc-3") in rows