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

{embed_train-3.0.0 → embed_train-3.2.0}/AGENTS.md

@@ -36,13 +36,15 @@ Key modules:
 - `src/embed_train/train/trainers/hf/__init__.py`
   SentenceTransformers-based training path with `InformationRetrievalEvaluator`.
 - `src/embed_train/train/dataset/__init__.py`
-  Base `TorchDataset` and `CollateFn` abstractions.
+  Base `TorchDataset`, `HardNegativeMiner`, and `CollateFn` abstractions.
 - `src/embed_train/train/dataset/collate.py`
-  Built-in collate functions for in-batch positive training.
+  Built-in collate functions for in-batch positive and hard-negative training.
+- `src/embed_train/train/dataset/hard_negatives.py`
+  SentenceTransformers-backed hard-negative mining implementation.
 - `src/embed_train/train/dataset/torch_datasets.py`
-  Built-in grouped and flattened query/positive dataset views.
+  Built-in grouped, flattened query/positive, and hard-negative dataset views.
 - `src/embed_train/train/trainers/torch/loss.py`
-  Built-in contrastive losses.
+  Built-in contrastive losses, including hard-negative candidate ranking.
 - `src/embed_train/push_to_hf/__init__.py`
   `PushToHFRunner` for checkpoint restore, local repo export, and HF upload.
 - `src/embed_train/utils.py`
@@ -102,6 +104,9 @@ For most changes, inspect the matching implementation and tests together:
 - Preserve the row contracts expected by built-in collate functions:
   - grouped format: `{"query": str, "positives": list[str]}`
   - flattened format: `{"query": str, "positive": str}`
+  - hard-negative format: `{"query": str, "positive": str, "negative": str}` or `negative_<n>` columns
+- Keep hard-negative mining isolated behind `HardNegativeMiner`/`HardNegativeDataset` unless a task explicitly changes the abstraction boundary.
+- Preserve the hard-negative candidate order expected by `HardNegativeContrastiveLoss`: one positive first, followed by one or more negatives for the same query.
 - When changing row formats, update the collate functions and tests together.
 - Keep tokenizer behavior explicit through settings rather than hidden defaults.
 

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

@@ -1,3 +1,17 @@
+# [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)
+
+
+### Features
+
+* add hard negative training ([4d79e25](https://gitlab.com/efysent/agentic-core/embed-train/commit/4d79e25877f7483dd7c6f13a97b1dd846c998c9a))
+
 # [3.0.0](https://gitlab.com/efysent/agentic-core/embed-train/compare/v2.0.0...v3.0.0) (2026-05-09)
 
 

embed_train-3.0.0/README.md → embed_train-3.2.0/PKG-INFO

@@ -1,3 +1,17 @@
+Metadata-Version: 2.4
+Name: embed-train
+Version: 3.2.0
+Author-email: jalal <jalalkhaldi3@gmail.com>
+Requires-Python: <3.13,>=3.11
+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==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` is a config-driven library for training, evaluating, checkpointing, and publishing embedding models.
@@ -61,9 +75,10 @@ src/embed_train/
 ├── train/
 │   ├── __init__.py                # TrainRunner
 │   ├── dataset/
-│   │   ├── __init__.py            # Base TorchDataset and CollateFn abstractions
-│   │   ├── collate.py             # Built-in in-batch positive collate functions
-│   │   ├── torch_datasets.py      # Built-in query/positive dataset views
+│   │   ├── __init__.py            # Base TorchDataset, HardNegativeMiner, and CollateFn abstractions
+│   │   ├── collate.py             # Built-in in-batch and hard-negative collate functions
+│   │   ├── hard_negatives.py      # SentenceTransformers hard-negative miner
+│   │   ├── torch_datasets.py      # Built-in query/positive and hard-negative dataset views
 │   │   └── sampling/
 │   │       └── samplers.py        # Built-in positive sampler(s)
 │   └── trainers/
@@ -107,7 +122,8 @@ The base class already provides:
 
 The PyTorch flow separates concerns clearly:
 
-- `TorchDataset`: how rows are loaded and exposed
+- `TorchDataset`: how rows are loaded, optionally mined, and exposed
+- `HardNegativeMiner`: how query/positive rows can be expanded with mined negatives
 - `CollateFn`: how rows become query/document text pairs and tokenized tensors
 - `Processor`: text normalization or preprocessing, provided by `retrievalbase`
 
@@ -124,8 +140,12 @@ The current repository includes these reusable implementations:
 
 - `embed_train.train.dataset.torch_datasets.QueryMultiPositiveDataset`
 - `embed_train.train.dataset.torch_datasets.QueryPositiveDataset`
+- `embed_train.train.dataset.torch_datasets.HardNegativeDataset`
+- `embed_train.train.dataset.hard_negatives.SentenceTransformerHardNegativeMiner`
+- `embed_train.train.dataset.collate.HardNegativeCollateFn`
 - `embed_train.train.dataset.collate.InBatchNegativeCollateFn`
 - `embed_train.train.dataset.collate.MultiPositiveInBatchCollateFn`
+- `embed_train.train.trainers.torch.loss.HardNegativeContrastiveLoss`
 - `embed_train.train.trainers.torch.loss.InBatchNegativeContrastiveLoss`
 - `embed_train.train.trainers.torch.loss.MultiPositiveContrastiveLoss`
 - `embed_train.train.trainers.hf.SentenceTransformersTrainer`
@@ -147,6 +167,20 @@ What the built-in trainer does:
 - logs to TensorBoard
 - saves checkpoints to `data_dir/checkpoints/...`
 
+For hard-negative training, configure the trainer with:
+
+- `torch_dataset`: `embed_train.train.dataset.torch_datasets.HardNegativeDataset`
+- `torch_dataset.hard_negative_miner`: `embed_train.train.dataset.hard_negatives.SentenceTransformerHardNegativeMiner`
+- `collate_fn`: `embed_train.train.dataset.collate.HardNegativeCollateFn`
+- `loss`: `embed_train.train.trainers.torch.loss.HardNegativeContrastiveLoss`
+
+`HardNegativeDataset` converts rows with `metadata.query` and `page_content` into a Hugging Face
+`Dataset`, mines negatives with `sentence_transformers.util.mine_hard_negatives`, and then exposes
+the mined rows through the normal `TorchDataset` interface. `HardNegativeCollateFn` expects each row
+to contain `query`, `positive`, and either `negative` or numbered `negative_<n>` fields. It emits one
+positive followed by that row's negatives, which is the candidate layout required by
+`HardNegativeContrastiveLoss`.
+
 ### SentenceTransformers Training
 
 Use this when your data is naturally represented as a Hugging Face `Dataset` and you want a standard SentenceTransformers training path with IR evaluation.
@@ -204,10 +238,11 @@ This is the right place to define pooling, projection heads, shared or separate
 
 Subclass `TorchDataset` when you need a different row shape or data-loading strategy.
 
-The built-in datasets show two common patterns:
+The built-in datasets show these common patterns:
 
 - grouped query -> many positives
 - flattened query -> single positive
+- flattened query -> single positive plus mined hard negatives
 
 ### Add a Custom Collate Function
 

embed_train-3.0.0/PKG-INFO → embed_train-3.2.0/README.md

@@ -1,16 +1,3 @@
-Metadata-Version: 2.4
-Name: embed-train
-Version: 3.0.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: retrievalbase<2.0.0,>=1.0.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
-Description-Content-Type: text/markdown
-
 # embed-train
 
 `embed-train` is a config-driven library for training, evaluating, checkpointing, and publishing embedding models.
@@ -74,9 +61,10 @@ src/embed_train/
 ├── train/
 │   ├── __init__.py                # TrainRunner
 │   ├── dataset/
-│   │   ├── __init__.py            # Base TorchDataset and CollateFn abstractions
-│   │   ├── collate.py             # Built-in in-batch positive collate functions
-│   │   ├── torch_datasets.py      # Built-in query/positive dataset views
+│   │   ├── __init__.py            # Base TorchDataset, HardNegativeMiner, and CollateFn abstractions
+│   │   ├── collate.py             # Built-in in-batch and hard-negative collate functions
+│   │   ├── hard_negatives.py      # SentenceTransformers hard-negative miner
+│   │   ├── torch_datasets.py      # Built-in query/positive and hard-negative dataset views
 │   │   └── sampling/
 │   │       └── samplers.py        # Built-in positive sampler(s)
 │   └── trainers/
@@ -120,7 +108,8 @@ The base class already provides:
 
 The PyTorch flow separates concerns clearly:
 
-- `TorchDataset`: how rows are loaded and exposed
+- `TorchDataset`: how rows are loaded, optionally mined, and exposed
+- `HardNegativeMiner`: how query/positive rows can be expanded with mined negatives
 - `CollateFn`: how rows become query/document text pairs and tokenized tensors
 - `Processor`: text normalization or preprocessing, provided by `retrievalbase`
 
@@ -137,8 +126,12 @@ The current repository includes these reusable implementations:
 
 - `embed_train.train.dataset.torch_datasets.QueryMultiPositiveDataset`
 - `embed_train.train.dataset.torch_datasets.QueryPositiveDataset`
+- `embed_train.train.dataset.torch_datasets.HardNegativeDataset`
+- `embed_train.train.dataset.hard_negatives.SentenceTransformerHardNegativeMiner`
+- `embed_train.train.dataset.collate.HardNegativeCollateFn`
 - `embed_train.train.dataset.collate.InBatchNegativeCollateFn`
 - `embed_train.train.dataset.collate.MultiPositiveInBatchCollateFn`
+- `embed_train.train.trainers.torch.loss.HardNegativeContrastiveLoss`
 - `embed_train.train.trainers.torch.loss.InBatchNegativeContrastiveLoss`
 - `embed_train.train.trainers.torch.loss.MultiPositiveContrastiveLoss`
 - `embed_train.train.trainers.hf.SentenceTransformersTrainer`
@@ -160,6 +153,20 @@ What the built-in trainer does:
 - logs to TensorBoard
 - saves checkpoints to `data_dir/checkpoints/...`
 
+For hard-negative training, configure the trainer with:
+
+- `torch_dataset`: `embed_train.train.dataset.torch_datasets.HardNegativeDataset`
+- `torch_dataset.hard_negative_miner`: `embed_train.train.dataset.hard_negatives.SentenceTransformerHardNegativeMiner`
+- `collate_fn`: `embed_train.train.dataset.collate.HardNegativeCollateFn`
+- `loss`: `embed_train.train.trainers.torch.loss.HardNegativeContrastiveLoss`
+
+`HardNegativeDataset` converts rows with `metadata.query` and `page_content` into a Hugging Face
+`Dataset`, mines negatives with `sentence_transformers.util.mine_hard_negatives`, and then exposes
+the mined rows through the normal `TorchDataset` interface. `HardNegativeCollateFn` expects each row
+to contain `query`, `positive`, and either `negative` or numbered `negative_<n>` fields. It emits one
+positive followed by that row's negatives, which is the candidate layout required by
+`HardNegativeContrastiveLoss`.
+
 ### SentenceTransformers Training
 
 Use this when your data is naturally represented as a Hugging Face `Dataset` and you want a standard SentenceTransformers training path with IR evaluation.
@@ -217,10 +224,11 @@ This is the right place to define pooling, projection heads, shared or separate
 
 Subclass `TorchDataset` when you need a different row shape or data-loading strategy.
 
-The built-in datasets show two common patterns:
+The built-in datasets show these common patterns:
 
 - grouped query -> many positives
 - flattened query -> single positive
+- flattened query -> single positive plus mined hard negatives
 
 ### Add a Custom Collate Function
 

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

@@ -1,6 +1,6 @@
 [project]
 name = "embed-train"
-version = "3.0.0"
+version = "3.2.0"
 description = ""
 authors = [
     { name = "jalal", email = "jalalkhaldi3@gmail.com" }
@@ -9,12 +9,13 @@ 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",
-    "retrievalbase>=1.0.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",
 ]
 
 [build-system]

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

@@ -57,6 +57,10 @@ class MultiPositiveInBatchCollateFnSettings[TCProcessor: "ProcessorSettings"](Co
     n_pos: int
 
 
+class HardNegativeCollateFnSettings[TCProcessor: "ProcessorSettings"](CollateFnSettings[TCProcessor]):
+    pass
+
+
 class ModelSettings(FromConfigMixinSettings):
     pass
 
@@ -69,6 +73,41 @@ class TrainerSettings(FromConfigMixinSettings):
     data_dir: Path
 
 
+class HardNegativeMinerSettings(FromConfigMixinSettings):
+    pass
+
+
+class SentenceTransformerHardNegativeMinerSettings(HardNegativeMinerSettings):
+    model_name_or_path: str
+    cross_encoder_model_name_or_path: str | None = None
+    tokenizer: TokenizerSettings
+    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
+    range_max: int | None = None
+    max_score: float | None = None
+    min_score: float | None = None
+    absolute_margin: float | None = None
+    relative_margin: float | None = None
+    num_negatives: int = 3
+    sampling_strategy: Literal["random", "top"] = "top"
+    query_prompt_name: str | None = None
+    query_prompt: str | None = None
+    corpus_prompt_name: str | None = None
+    corpus_prompt: str | None = None
+    include_positives: bool = False
+    output_format: Literal["triplet", "n-tuple", "labeled-pair", "labeled-list"] = "n-tuple"
+    output_scores: bool = False
+    batch_size: int = 32
+    faiss_batch_size: int = 16384
+    use_faiss: bool = False
+    use_multi_process: bool | list[str] = False
+    verbose: bool = True
+    cache_folder: str | None = None
+    trust_remote_code: bool = TRUST_REMOTE_CODE
+
+
 class PyTorchTrainerSettings[
     TCModel: "ModelSettings",
     TCLoss: "LossSettings",
@@ -107,6 +146,13 @@ class QueryPositiveDatasetSettings[TCDatasetConnector: "DatasetConnectorSettings
     pass
 
 
+class HardNegativeDatasetSettings[
+    TCDatasetConnector: "DatasetConnectorSettings",
+    TCHardNegativeMiner: "HardNegativeMinerSettings",
+](TorchDatasetSettings[TCDatasetConnector]):
+    hard_negative_miner: TCHardNegativeMiner
+
+
 class InBatchNegativeContrastiveLossSettings(ContrastiveLossSettings):
     pass
 
@@ -115,6 +161,10 @@ class MultiPositiveContrastiveLossSettings(ContrastiveLossSettings):
     n_pos: int
 
 
+class HardNegativeContrastiveLossSettings(ContrastiveLossSettings):
+    pass
+
+
 class RunnerSettings(FromConfigMixinSettings):
     pass
 
@@ -139,17 +189,6 @@ class TrainRunnerSettings[TCTrainer: "TrainerSettings"](RunnerSettings):
     trainer: TCTrainer
 
 
-class HardNegativesSettings(BaseSettings):
-    range_min: int
-    range_max: int
-    max_score: float
-    relative_margin: float
-    num_negatives: int
-    sampling_strategy: Literal["random", "top"]
-    batch_size: int
-    use_faiss: bool
-
-
 class EvalutationSettings(BaseSettings):
     query_column: str
     document_column: str

{embed_train-3.0.0 → embed_train-3.2.0}/src/embed_train/train/dataset/__init__.py

@@ -16,7 +16,7 @@ from retrievalbase.mixins import FromConfigMixin
 _logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
-    from embed_train.settings import CollateFnSettings, TorchDatasetSettings
+    from embed_train.settings import CollateFnSettings, HardNegativeMinerSettings, TorchDatasetSettings
 
 
 class CollateFn[TCCollateFn: "CollateFnSettings[Any]", T: dict[str, Any]](ABC):
@@ -107,3 +107,12 @@ class TorchDataset[TCTorchDataset: "TorchDatasetSettings[Any]", T: dict[str, Any
         """
         rows: list[dict[str, Any]] = [self[i] for i in range(len(self))]
         return HFDataset.from_list(rows)
+
+
+class HardNegativeMiner[TCHardNegativeMiner: "HardNegativeMinerSettings"](FromConfigMixin[TCHardNegativeMiner], ABC):
+    def __init__(self, config: TCHardNegativeMiner) -> None:
+        self.config = config
+
+    @abstractmethod
+    def mine(self, dataset: HFDataset) -> HFDataset:
+        raise NotImplementedError

{embed_train-3.0.0 → embed_train-3.2.0}/src/embed_train/train/dataset/collate.py

@@ -2,6 +2,7 @@ import random
 from typing import Any
 
 from embed_train.settings import (
+    HardNegativeCollateFnSettings,
     InBatchNegativeCollateFnSettings,
     MultiPositiveInBatchCollateFnSettings,
 )
@@ -50,3 +51,36 @@ class MultiPositiveInBatchCollateFn(CollateFn[MultiPositiveInBatchCollateFnSetti
             queries.append(query)
             passages.extend(sampled_positives)
         return queries, passages
+
+
+class HardNegativeCollateFn(CollateFn[HardNegativeCollateFnSettings, dict[str, Any]]):
+    def __init__(
+        self,
+        config: HardNegativeCollateFnSettings,
+        context: dict[str, Any] | None,
+    ) -> None:
+        super().__init__(config, context)
+
+    def _process_batch(
+        self,
+        batch: list[dict[str, Any]],
+    ) -> tuple[list[str], list[str]]:
+        queries: list[str] = []
+        passages: list[str] = []
+
+        for item in batch:
+            queries.append(item["query"])
+            passages.append(item["positive"])
+            passages.extend(self._negative_passages(item))
+
+        return queries, passages
+
+    def _negative_passages(self, item: dict[str, Any]) -> list[str]:
+        if "negative" in item:
+            return [item["negative"]]
+
+        negative_keys = sorted(
+            (key for key in item if key.startswith("negative_")),
+            key=lambda key: int(key.rsplit("_", 1)[1]),
+        )
+        return [item[key] for key in negative_keys]

embed_train-3.2.0/src/embed_train/train/dataset/hard_negatives.py

@@ -0,0 +1,136 @@
+import gc
+import logging
+
+import torch
+from datasets import Dataset as HFDataset  # type: ignore[import-untyped]
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.cross_encoder import CrossEncoder
+from sentence_transformers.sentence_transformer.modules import Pooling, Transformer
+from sentence_transformers.util import mine_hard_negatives
+
+from embed_train.settings import SentenceTransformerHardNegativeMinerSettings
+from embed_train.train.dataset import HardNegativeMiner
+
+_logger = logging.getLogger(__name__)
+
+
+class SentenceTransformerHardNegativeMiner(HardNegativeMiner[SentenceTransformerHardNegativeMinerSettings]):
+    def __init__(self, config: "SentenceTransformerHardNegativeMinerSettings") -> None:
+        super().__init__(config)
+
+    def _load_sentence_transformer(self) -> SentenceTransformer:
+        trust = self.config.trust_remote_code or self.config.tokenizer.trust_remote_code
+
+        _logger.info(
+            f"Loading SentenceTransformer | "
+            f"model={self.config.model_name_or_path} | "
+            f"pooling={self.config.pooling} | "
+            f"max_length={self.config.tokenizer.max_length}"
+        )
+
+        transformer = Transformer(
+            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_kwargs={
+                "trust_remote_code": trust,
+            },
+            config_kwargs={
+                "trust_remote_code": trust,
+            },
+            processor_kwargs={
+                "trust_remote_code": trust,
+                "padding": self.config.tokenizer.padding,
+                "truncation": self.config.tokenizer.truncation,
+                "model_max_length": self.config.tokenizer.max_length,
+            },
+        )
+
+        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_embedding_dimension()}")
+
+        return model
+
+    def _load_cross_encoder(self) -> "CrossEncoder | None":
+        if not self.config.cross_encoder_model_name_or_path:
+            _logger.info("No CrossEncoder configured")
+            return None
+
+        _logger.info(f"Loading CrossEncoder | model={self.config.cross_encoder_model_name_or_path}")
+
+        model = CrossEncoder(
+            self.config.cross_encoder_model_name_or_path,
+            trust_remote_code=self.config.trust_remote_code,
+        )
+
+        _logger.info("CrossEncoder loaded successfully")
+
+        return model
+
+    def mine(self, dataset: HFDataset) -> HFDataset:
+        _logger.info(
+            f"Starting hard negative mining | "
+            f"dataset_size={len(dataset)} | "
+            f"num_negatives={self.config.num_negatives} | "
+            f"sampling_strategy={self.config.sampling_strategy} | "
+            f"use_faiss={self.config.use_faiss}"
+        )
+
+        model = self._load_sentence_transformer()
+        cross_encoder = self._load_cross_encoder()
+
+        try:
+            mined = mine_hard_negatives(
+                dataset=dataset,
+                model=model,
+                anchor_column_name=self.config.anchor_column_name,
+                positive_column_name=self.config.positive_column_name,
+                cross_encoder=cross_encoder,
+                range_min=self.config.range_min,
+                range_max=self.config.range_max,
+                max_score=self.config.max_score,
+                min_score=self.config.min_score,
+                absolute_margin=self.config.absolute_margin,
+                relative_margin=self.config.relative_margin,
+                num_negatives=self.config.num_negatives,
+                sampling_strategy=self.config.sampling_strategy,
+                query_prompt_name=self.config.query_prompt_name,
+                query_prompt=self.config.query_prompt,
+                corpus_prompt_name=self.config.corpus_prompt_name,
+                corpus_prompt=self.config.corpus_prompt,
+                include_positives=self.config.include_positives,
+                output_format=self.config.output_format,
+                output_scores=self.config.output_scores,
+                batch_size=self.config.batch_size,
+                faiss_batch_size=self.config.faiss_batch_size,
+                use_faiss=self.config.use_faiss,
+                use_multi_process=self.config.use_multi_process,
+                verbose=self.config.verbose,
+                cache_folder=self.config.cache_folder,
+            )
+
+            _logger.info(f"Hard negative mining completed successfully | mined_dataset_size={len(mined)}")
+        finally:
+            _logger.info("Cleaning GPU memory after hard negative mining")
+            del model
+            if cross_encoder is not None:
+                del cross_encoder
+            gc.collect()
+            if torch.cuda.is_available():
+                allocated_before = torch.cuda.memory_allocated() / 1024**3
+                reserved_before = torch.cuda.memory_reserved() / 1024**3
+                torch.cuda.empty_cache()
+                torch.cuda.ipc_collect()
+                allocated_after = torch.cuda.memory_allocated() / 1024**3
+                reserved_after = torch.cuda.memory_reserved() / 1024**3
+                _logger.info(
+                    f"CUDA memory cleanup completed | "
+                    f"allocated_before={allocated_before:.2f}GB | "
+                    f"reserved_before={reserved_before:.2f}GB | "
+                    f"allocated_after={allocated_after:.2f}GB | "
+                    f"reserved_after={reserved_after:.2f}GB"
+                )
+        return mined

{embed_train-3.0.0 → embed_train-3.2.0}/src/embed_train/train/dataset/torch_datasets.py

@@ -1,9 +1,16 @@
 from typing import Any
 
 import polars as pl
+from datasets import Dataset as HFDataset  # type: ignore[import-untyped]
 
-from embed_train.settings import QueryMultiPositiveDatasetSettings, QueryPositiveDatasetSettings
+from embed_train.settings import (
+    HardNegativeDatasetSettings,
+    QueryMultiPositiveDatasetSettings,
+    QueryPositiveDatasetSettings,
+)
 from embed_train.train.dataset import TorchDataset
+from embed_train.train.dataset.hard_negatives import HardNegativeMiner
+from embed_train.utils import load_class
 
 
 class QueryMultiPositiveDataset(TorchDataset[QueryMultiPositiveDatasetSettings[Any], dict[str, Any]]):
@@ -71,3 +78,36 @@ class QueryPositiveDataset(
             "query": self._queries[index],
             "positive": positive,
         }
+
+
+class HardNegativeDataset(
+    TorchDataset[
+        HardNegativeDatasetSettings[Any, Any],
+        dict[str, Any],
+    ]
+):
+    def __init__(
+        self,
+        config: HardNegativeDatasetSettings[Any, Any],
+    ) -> None:
+        super().__init__(config)
+        input_dataset = self._to_query_positive_hf_dataset()
+        miner = self._load_hard_negative_miner()
+        self._dataset = miner.mine(input_dataset)
+
+    def _to_query_positive_hf_dataset(self) -> HFDataset:
+        rows = self.dataset.polars.select(
+            pl.col("metadata").struct.field("query").alias("query"),
+            pl.col("page_content").alias("positive"),
+        )
+        return HFDataset.from_list(rows.to_dicts())
+
+    def _load_hard_negative_miner(self) -> HardNegativeMiner[Any]:
+        miner_cls = load_class(self.config.hard_negative_miner.module_path)
+        return miner_cls.from_config(self.config.hard_negative_miner)
+
+    def __len__(self) -> int:
+        return len(self._dataset)
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        return self._dataset[index]

{embed_train-3.0.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.0.0 → embed_train-3.2.0}/src/embed_train/train/trainers/torch/__init__.py

@@ -189,7 +189,7 @@ class PyTorchTrainer[TCPyTorchTrainer: "PyTorchTrainerSettings[Any, Any, Any, An
             self.model.load_state_dict(ckpt["model_state"])
             self.optimizer.load_state_dict(ckpt["optimizer_state"])
             _logger.info(f"Resuming training from epoch {ckpt['epoch']}")
-            return ckpt["epoch"]  # type:ignore[no-any-return]
+            return ckpt["epoch"]
         return 0
 
     def run_epoch(