embed-train 1.0.0__py3-none-any.whl

embed_train/__init__.py

@@ -0,0 +1,37 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+import yaml
+
+from embed_train.exceptions import InvalidRunnerClassError
+from embed_train.utils import load_class
+from retrievalbase.mixins import FromConfigMixin
+
+if TYPE_CHECKING:
+    from embed_train.settings import RunnerSettings
+
+
+class Runner[TCRunner: "RunnerSettings"](FromConfigMixin[TCRunner], ABC):
+    def __init__(self, config: TCRunner):
+        self.config = config
+
+    @abstractmethod
+    def run(self) -> None:
+        raise NotImplementedError()
+
+    @classmethod
+    def _validate_klass(self, symbol: Any) -> type["Runner[TCRunner]"]:
+        if not isinstance(symbol, type):
+            raise InvalidRunnerClassError(f"Loaded object is not a class: {symbol!r}")
+        # 2️⃣ Ensure it is a Runner subclass
+        if not issubclass(symbol, Runner):
+            raise InvalidRunnerClassError(f"Class {symbol.__module__}.{symbol.__name__} is not a subclass of Runner")
+        return symbol
+
+    @classmethod
+    def get_runner(cls, config_path: str) -> "Runner[TCRunner]":
+        with open(config_path) as f:
+            yaml_config = yaml.safe_load(f)
+            symbol = load_class(yaml_config["module_path"])
+            klass: type[Runner[TCRunner]] = cls._validate_klass(symbol)
+            return klass.from_settings()

embed_train/constants.py

@@ -0,0 +1,3 @@
+RANDOM_SEED = 42
+CONFIG_PATH = "/config/config.yaml"
+TRUST_REMOTE_CODE = True

embed_train/exceptions.py

@@ -0,0 +1,31 @@
+class EmbedTrainError(Exception):
+    """Base exception for all embed_train errors."""
+
+
+# -------------------------
+# Runtime errors
+# -------------------------
+class EmbedTrainRuntimeError(RuntimeError, EmbedTrainError):
+    """Runtime errors during execution (training, loading, etc.)."""
+
+
+class MissingContextError(EmbedTrainRuntimeError):
+    """Raised when required context is missing."""
+
+
+# -------------------------
+# Type errors
+# -------------------------
+class EmbedTrainTypeError(TypeError, EmbedTrainError):
+    """Invalid type provided."""
+
+
+class InvalidRunnerClassError(EmbedTrainTypeError):
+    """Runner class is not valid or does not inherit expected base."""
+
+
+# -------------------------
+# Value errors
+# -------------------------
+class EmbedTrainValueError(ValueError, EmbedTrainError):
+    """Invalid value provided in configuration or runtime."""

embed_train/models/__init__.py

@@ -0,0 +1,61 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Self, cast
+
+from torch import nn
+from transformers import PreTrainedModel
+
+from embed_train.utils import load_checkpoint_state_dict
+from retrievalbase.mixins import FromConfigMixin
+
+if TYPE_CHECKING:
+    from embed_train.settings import ModelSettings
+
+
+class Model[TCModel: "ModelSettings"](FromConfigMixin[TCModel], ABC):
+    def __init__(self, config: TCModel):
+        self.config = config
+
+    @abstractmethod
+    def to_hf_model(self) -> PreTrainedModel:  # inplace
+        raise NotImplementedError()
+
+    def save(
+        self,
+        repo_dir: str | Path,
+        push: bool = False,
+        **push_kwargs: Any,
+    ) -> None:
+        repo_dir_path = Path(repo_dir) if isinstance(repo_dir, str) else repo_dir
+        self.to_hf_model().save_pretrained(repo_dir_path)
+        if push:
+            self.to_hf_model().push_to_hub(
+                repo_id=repo_dir_path.name,
+                **push_kwargs,
+            )
+
+    def to(self, device: str) -> Self:
+        cast(nn.Module, self.to_hf_model()).to(device)
+        return self
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        config: TCModel,
+        checkpoint: str,
+        device: str,
+        strict: bool = True,
+    ) -> "Model[TCModel]":
+        model = cls(config)
+        hf_model = model.to_hf_model()
+        checkpoint_path = Path(checkpoint)
+        if checkpoint_path.is_dir():
+            state_dict_path = checkpoint_path / "model.safetensors"
+            if not state_dict_path.exists():
+                raise FileNotFoundError(f"No model.safetensors found in {checkpoint_path}")
+            state_dict = load_checkpoint_state_dict(state_dict_path)
+        else:
+            state_dict = load_checkpoint_state_dict(checkpoint_path)
+        hf_model.load_state_dict(state_dict, strict=strict)
+        hf_model.to(device)  # ty: ignore[invalid-argument-type]
+        return model

embed_train/push_to_hf/__init__.py

@@ -0,0 +1,131 @@
+# finetune_embedder/src/runners/push_to_hf.py
+import inspect
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from huggingface_hub import HfApi
+
+from embed_train import Runner
+from embed_train.models import Model
+from embed_train.utils import load_class
+
+_logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from embed_train.settings import PushToHFRunnerSettings
+
+
+class PushToHFRunner[TCPushToHFRunner: "PushToHFRunnerSettings[Any]"](Runner[TCPushToHFRunner]):
+    def __init__(self, config: TCPushToHFRunner) -> None:
+        super().__init__(config)
+
+    def _infer_model_name(self) -> str:
+        name: str = self.wrapper_cls.__name__
+        if not name.endswith("ModelWrapper"):
+            raise ValueError(f"Model wrapper class must end with 'ModelWrapper', got {name}")
+        return name.removesuffix("ModelWrapper").lower()
+
+    def _find_model_source_dir(self) -> Path:
+        module_file = Path(inspect.getfile(self.wrapper_cls)).resolve()
+        models_dir = module_file.parent
+        return models_dir
+
+    def _collect_model_files(self) -> list[Path]:
+        model_name = self._infer_model_name()
+        models_dir = self._find_model_source_dir()
+
+        model_dir = models_dir / model_name
+        if not model_dir.exists():
+            raise FileNotFoundError(f"Model directory not found: {model_dir}")
+
+        files: list[Path] = []
+
+        for pattern in (
+            "modeling_*.py",
+            "configuration_*.py",
+            "vllm_modeling_*.py",
+            "vllm_configuration_*.py",
+        ):
+            files.extend(model_dir.glob(pattern))
+
+        if not files:
+            raise RuntimeError(f"No model source files found in {model_dir}")
+
+        return sorted(set(files))
+
+    def _save_model_source_files(self) -> None:
+        files = self._collect_model_files()
+        for src in files:
+            dst = Path(self.config.hf.repo) / src.name
+            dst.write_text(src.read_text(encoding="utf-8"), encoding="utf-8")
+
+    def _load_model(self) -> Model[Any]:
+        _logger.info(f"Loading model class from {self.config.model.module_path}")
+        model_cls: Model[Any] = load_class(self.config.model.module_path)
+        return model_cls.from_checkpoint(
+            config=self.config.model,
+            checkpoint=self.config.checkpoint_path,
+            device=self.config.device,
+        )
+
+    def _load_wrapper(self) -> None:
+        _logger.info(f"Loading model wrapper → {self.config.model.module_path}")
+        self.wrapper_cls = load_class(self.config.model.module_path)
+
+    def _load_model_from_checkpoint(self) -> None:
+        _logger.info(f"Loading model from checkpoint → {self.config.checkpoint_path}")
+        self.model: Model[Any] = self.wrapper_cls.from_checkpoint(
+            config=self.config.model,
+            checkpoint=self.config.checkpoint_path,
+            device=self.config.device,
+        )
+
+    def _save_model_locally(self) -> None:
+        _logger.info("Saving model locally → %s", self.config.hf.repo)
+        self.model.save(
+            repo_dir=self.config.hf.repo,
+            push=False,  # we control pushing
+            create_repo=self.config.create_repo,
+            revision=self.config.hf.revision,
+            private=self.config.hf.private,
+            commit_message=self.config.hf.commit_message,
+        )
+
+    def _ensure_hf_repo_exists(self) -> None:
+        if not self.config.create_repo:
+            return
+        _logger.info(
+            f"Ensuring Hugging Face repo exists | repo={self.config.hf.repo} | private={self.config.hf.private}"
+        )
+        api = HfApi()
+        api.create_repo(
+            repo_id=self.config.hf.repo,
+            repo_type="model",
+            private=self.config.hf.private,
+            exist_ok=True,  # safe for retries / CI
+        )
+
+    def _push_repo_to_hf(self) -> None:
+        _logger.info(
+            f"Pushing full repository folder to Hugging Face | repo={self.config.hf.repo} | revision={self.config.hf.revision}"
+        )
+        api = HfApi()
+        api.upload_folder(
+            folder_path=self.config.hf.repo,
+            repo_id=self.config.hf.repo,
+            repo_type="model",
+            revision=self.config.hf.revision,
+            commit_message=self.config.hf.commit_message,
+        )
+        _logger.info("Repository successfully pushed to Hugging Face")
+
+    def run(self) -> None:
+        self._load_wrapper()
+        self._load_model_from_checkpoint()
+        self._save_model_locally()
+        self._save_model_source_files()
+        if self.config.push:
+            self._ensure_hf_repo_exists()
+            self._push_repo_to_hf()
+        _logger.info("PushToHFRunner completed successfully")

embed_train/settings.py

@@ -0,0 +1,202 @@
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Literal, Self
+
+from pydantic import Field, model_validator
+from pydantic_settings import BaseSettings
+
+from embed_train.constants import TRUST_REMOTE_CODE
+from embed_train.exceptions import EmbedTrainValueError
+from retrievalbase.dataset.settings import HuggingFaceDatasetAdaptaterSettings
+from retrievalbase.settings import FromConfigMixinSettings
+
+if TYPE_CHECKING:
+    from retrievalbase.connector.settings import DatasetConnectorSettings
+    from retrievalbase.evaluation.settings import ProcessorSettings
+
+
+class TorchDatasetSettings[TCDatasetConnector: "DatasetConnectorSettings"](FromConfigMixinSettings):
+    dataset_connector: TCDatasetConnector
+
+
+class TokenizerSettings(BaseSettings):
+    name: str
+    padding: str | bool
+    truncation: bool
+    max_length: int
+    trust_remote_code: bool = TRUST_REMOTE_CODE
+
+
+class SamplerSettings(BaseSettings):
+    module_path: str
+    seed: int | None
+
+
+class DistanceSettings(FromConfigMixinSettings):
+    pass
+
+
+class StepRangeDistanceSettings(DistanceSettings):
+    pass
+
+
+class PositiveSamplerSettings[TCDistance: "DistanceSettings"](SamplerSettings):
+    max_step_distance: int
+    distance: TCDistance
+    k: int
+
+
+class CollateFnSettings[TCProcessor: "ProcessorSettings"](BaseSettings):
+    module_path: str
+    tokenizer: TokenizerSettings
+    processor: TCProcessor
+
+
+class InBatchPositiveCollateFnSettings[TCProcessor: "ProcessorSettings"](CollateFnSettings[TCProcessor]):
+    pass
+
+
+class MultiPositiveInBatchCollateFnSettings[TCProcessor: "ProcessorSettings"](CollateFnSettings[TCProcessor]):
+    n_pos: int
+
+
+class ModelSettings(FromConfigMixinSettings):
+    pass
+
+
+class LossSettings(FromConfigMixinSettings):
+    pass
+
+
+class TrainerSettings(FromConfigMixinSettings):
+    data_dir: Path
+
+
+class PyTorchTrainerSettings[
+    TCModel: "ModelSettings",
+    TCLoss: "LossSettings",
+    TCTorchDataset: "TorchDatasetSettings[Any]",
+    TCCollateFn: "CollateFnSettings[Any]",
+](TrainerSettings):
+    model: TCModel
+    torch_dataset: TCTorchDataset
+    collate_fn: TCCollateFn
+    train_frac: float
+    num_epochs: int
+    batch_size: int
+    shuffle: bool
+    lr: float
+    device: str
+    save_every: int
+    drop_last: bool
+    loss: TCLoss
+    resume_from: str | None = None
+
+
+class ContrastiveLossSettings(LossSettings):
+    temperature: float
+
+
+class QueryMultiPositiveDatasetSettings[TCDatasetConnector: "DatasetConnectorSettings"](
+    TorchDatasetSettings[TCDatasetConnector]
+):
+    pass
+
+
+class InBatchNegativeContrastiveLossSettings(ContrastiveLossSettings):
+    pass
+
+
+class MultiPositiveContrastiveLossSettings(ContrastiveLossSettings):
+    n_pos: int
+
+
+class RunnerSettings(FromConfigMixinSettings):
+    pass
+
+
+class HFSettings(BaseSettings):
+    repo: str
+    revision: str | None
+    private: bool
+    commit_message: str | None
+
+
+class PushToHFRunnerSettings[TCModel: "ModelSettings"](RunnerSettings):
+    checkpoint_path: str = Field(init=False)
+    device: str = Field(init=False)
+    push: bool = Field(init=False)
+    create_repo: bool = Field(init=False)
+    hf: HFSettings = Field(init=False)
+    model: TCModel = Field(init=False)
+
+
+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
+    precision_recall_at_k: list[int]
+    mrr_at_k: list[int]
+    ndcg_at_k: list[int]
+    batch_size: int
+
+
+class SentenceTransformerLoss(LossSettings):
+    kwargs: dict[str, Any]
+
+
+class SentenceTransformersTrainerSettings[
+    TCDatasetConnector: "DatasetConnectorSettings",
+    TCModel: "ModelSettings",
+](TrainerSettings):
+    hf_dataset: HuggingFaceDatasetAdaptaterSettings[TCDatasetConnector]
+    model: TCModel
+    tokenizer: TokenizerSettings
+    loss: SentenceTransformerLoss
+    pooling: Literal["cls", "mean_tokens", "max_tokens"]
+    batch_size: int
+    num_epochs: int
+    lr: float
+    warmup_ratio: float
+    eval_steps: int
+    save_steps: int
+    logging_steps: int
+    fp16: bool
+    train_frac: float
+    resume_from: str | None
+    evaluation: EvalutationSettings
+    trust_remote_code: bool = TRUST_REMOTE_CODE
+
+    # -------------------------
+    # CROSS-FIELD VALIDATION
+    # -------------------------
+    @model_validator(mode="after")
+    def _validate_steps_consistency(self) -> Self:
+        # 1️⃣ ordering
+        if not (self.logging_steps <= self.eval_steps <= self.save_steps):
+            raise EmbedTrainValueError(
+                "Expected logging_steps ≤ eval_steps ≤ save_steps, got "
+                f"{self.logging_steps} ≤ {self.eval_steps} ≤ {self.save_steps}"
+            )
+
+        # 2️⃣ alignment (CRITICAL)
+        if self.save_steps % self.eval_steps != 0:
+            raise EmbedTrainValueError(
+                "save_steps must be a multiple of eval_steps "
+                f"(got save_steps={self.save_steps}, eval_steps={self.eval_steps})"
+            )
+
+        return self

embed_train/train/__init__.py

@@ -0,0 +1,17 @@
+from typing import TYPE_CHECKING, Any
+
+from embed_train import Runner
+from embed_train.train.trainers import Trainer
+from embed_train.utils import load_class
+
+if TYPE_CHECKING:
+    from embed_train.settings import TrainRunnerSettings
+
+
+class TrainRunner[TCTrainRunner: "TrainRunnerSettings[Any]"](Runner[TCTrainRunner]):
+    def __init__(self, config: TCTrainRunner):
+        super().__init__(config)
+        self.trainer: Trainer[Any] = load_class(self.config.trainer.module_path).from_config(self.config.trainer)
+
+    def run(self) -> None:
+        self.trainer.train()

embed_train/train/dataset/__init__.py

@@ -0,0 +1,109 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, cast
+
+import torch
+from datasets import Dataset as HFDataset  # type: ignore[import-untyped]
+from torch.utils.data import Dataset as TorchDatasetBase
+from transformers import AutoTokenizer, PreTrainedTokenizer
+
+from embed_train.utils import load_class
+from retrievalbase.dataset import TextDataset
+from retrievalbase.enums import EmbeddingPurposeEnum
+from retrievalbase.evaluation import Processor
+from retrievalbase.mixins import FromConfigMixin
+
+_logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from embed_train.settings import CollateFnSettings, TorchDatasetSettings
+
+
+class CollateFn[TCCollateFn: "CollateFnSettings[Any]", T: dict[str, Any]](ABC):
+    def __init__(self, config: TCCollateFn, context: dict[str, Any] | None) -> None:
+        self.config = config
+        self.context = context
+        self.tokenizer: PreTrainedTokenizer = self._load_tokenizer()
+        self.processor: Processor[Any] = self._load_processor()
+
+    @abstractmethod
+    def _process_batch(self, batch: list[T]) -> tuple[list[str], list[str]]:
+        raise NotImplementedError()
+
+    def _load_tokenizer(self) -> PreTrainedTokenizer:
+        tokenizer = AutoTokenizer.from_pretrained(  # nosec CWE-494
+            self.config.tokenizer.name,
+            trust_remote_code=self.config.tokenizer.trust_remote_code,
+        )
+        return cast(PreTrainedTokenizer, tokenizer)
+
+    def _load_processor(self) -> Processor[Any]:
+        processor: Processor[Any] = load_class(self.config.processor.module_path).from_config(self.config.processor)
+        return processor
+
+    def __call__(self, batch: list[T]) -> tuple[torch.Tensor, torch.Tensor]:
+        queries, candidates = self._process_batch(batch)
+        return self._post_process(queries, candidates)
+
+    def _post_process(self, queries: list[str], candidates: list[str]) -> tuple[torch.Tensor, torch.Tensor]:
+        queries = [
+            self.processor(
+                text=query,
+                purpose=cast(EmbeddingPurposeEnum, EmbeddingPurposeEnum.QUERY),
+            )
+            for query in queries
+        ]
+        candidates = [
+            self.processor(
+                text=candidate,
+                purpose=cast(EmbeddingPurposeEnum, EmbeddingPurposeEnum.DOCUMENT),
+            )
+            for candidate in candidates
+        ]
+        return self._tokenizer_qc(queries, candidates)
+
+    def _tokenizer_qc(self, queries: list[str], candidates: list[str]) -> tuple[torch.Tensor, torch.Tensor]:
+        q_tok = self.tokenizer(
+            queries,
+            padding=self.config.tokenizer.padding,
+            truncation=self.config.tokenizer.truncation,
+            max_length=self.config.tokenizer.max_length,
+            return_tensors="pt",
+        )
+        c_tok = self.tokenizer(
+            candidates,
+            padding=self.config.tokenizer.padding,
+            truncation=self.config.tokenizer.truncation,
+            max_length=self.config.tokenizer.max_length,
+            return_tensors="pt",
+        )
+        return q_tok, c_tok
+
+
+class TorchDataset[TCTorchDataset: "TorchDatasetSettings[Any]", T: dict[str, Any]](
+    FromConfigMixin[TCTorchDataset],
+    TorchDatasetBase[T],
+    ABC,
+):
+    def __init__(self, config: TCTorchDataset) -> None:
+        self.config = config
+        self.dataset = self._load_dataset()
+
+    @abstractmethod
+    def __len__(self) -> int:
+        raise NotImplementedError()
+
+    def _load_dataset(self) -> TextDataset[Any]:
+        dataset: TextDataset[Any] = (
+            load_class(self.config.dataset_connector.module_path).from_config(self.config.dataset_connector).load_text()
+        )
+        _logger.info(f"Instantiating dataset | class={dataset.__class__.__name__} |")
+        _logger.info(f"Dataset loaded | type={type(dataset).__name__} | size={len(dataset)}")
+        return dataset
+
+    def to_hf_dataset(self) -> HFDataset:
+        """
+        Convert *this* TorchDataset (via __getitem__) to a HuggingFace Dataset.
+        """
+        rows: list[dict[str, Any]] = [self[i] for i in range(len(self))]
+        return HFDataset.from_list(rows)

embed_train/train/dataset/collate.py

@@ -0,0 +1,52 @@
+import random
+from typing import Any
+
+from embed_train.settings import (
+    InBatchPositiveCollateFnSettings,
+    MultiPositiveInBatchCollateFnSettings,
+)
+from embed_train.train.dataset import CollateFn
+
+
+class InBatchPositiveCollateFn(CollateFn[InBatchPositiveCollateFnSettings, dict[str, str]]):
+    def __init__(
+        self,
+        config: InBatchPositiveCollateFnSettings,
+        context: dict[str, Any] | None,
+    ) -> None:
+        super().__init__(config, context)
+
+    def _process_batch(
+        self,
+        batch: list[dict[str, str]],
+    ) -> tuple[list[str], list[str]]:
+        queries: list[str] = []
+        positives: list[str] = []
+
+        for item in batch:
+            queries.append(item["query"])
+            positives.append(random.choice(item["positives"]))  # nosec B311
+        return queries, positives
+
+
+class MultiPositiveInBatchCollateFn(CollateFn[MultiPositiveInBatchCollateFnSettings, dict[str, str]]):
+    def __init__(
+        self,
+        config: MultiPositiveInBatchCollateFnSettings,
+        context: dict[str, Any] | None,
+    ) -> None:
+        super().__init__(config, context)
+
+    def _process_batch(
+        self,
+        batch: list[dict[str, str]],
+    ) -> tuple[list[str], list[str]]:
+        queries: list[str] = []
+        passages: list[str] = []
+        for item in batch:
+            query = item["query"]
+            positives = item["positives"]
+            sampled_positives = [random.choice(positives) for _ in range(self.config.n_pos)]  # nosec B311
+            queries.append(query)
+            passages.extend(sampled_positives)
+        return queries, passages

embed_train/train/dataset/sampling/__init__.py

@@ -0,0 +1,46 @@
+import random
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+from embed_train.settings import StepRangeDistanceSettings
+from retrievalbase.mixins import FromConfigMixin
+
+if TYPE_CHECKING:
+    from embed_train.settings import DistanceSettings, SamplerSettings
+
+
+class Sampler[TCSampler: "SamplerSettings", T: dict[str, Any]](ABC):
+    RANDOM_SEED = 42
+
+    def __init__(self, config: TCSampler, context: dict[str, Any] | None) -> None:
+        self.config = config
+        self.context = context
+        self.rng = random.Random(self.config.seed)  # nosec B311
+
+    @abstractmethod
+    def sample(self, item: T, **kwargs: Any) -> list[str]:
+        pass
+
+    def _init_components(self) -> None:
+        return None
+
+
+class Distance[TCDistance: "DistanceSettings"](FromConfigMixin[TCDistance], ABC):
+    def __init__(self, config: TCDistance) -> None:
+        self.config = config
+
+    @abstractmethod
+    def __call__(self, a: tuple[int, int], b: tuple[int, int]) -> int:
+        raise NotImplementedError
+
+
+class StepRangeDistance(Distance[StepRangeDistanceSettings]):
+    def __init__(self, config: StepRangeDistanceSettings) -> None:
+        super().__init__(config)
+
+    def __call__(self, a: tuple[int, int], b: tuple[int, int]) -> int:
+        if a[1] < b[0]:
+            return b[0] - a[1]
+        if b[1] < a[0]:
+            return a[0] - b[1]
+        return 0