vespaembed 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

vespaembed/tasks/__init__.py

@@ -0,0 +1,23 @@
+# Tasks module - imports are performed when this module is loaded
+# This registers all tasks with the Registry
+#
+# Available tasks (by data format):
+# 1. pairs - Text pairs for semantic search (anchor, positive)
+# 2. triplets - Text triplets with hard negatives (anchor, positive, negative)
+# 3. similarity - Text pairs with similarity scores (sentence1, sentence2, score)
+# 4. tsdae - Unlabeled text for unsupervised learning (text)
+#
+# Matryoshka is a training option (--matryoshka flag) that can be enabled for any task except TSDAE.
+from vespaembed.tasks.base import BaseTask
+from vespaembed.tasks.pairs import PairsTask
+from vespaembed.tasks.similarity import SimilarityTask
+from vespaembed.tasks.triplets import TripletsTask
+from vespaembed.tasks.tsdae import TSDAETask
+
+__all__ = [
+    "BaseTask",
+    "PairsTask",
+    "SimilarityTask",
+    "TripletsTask",
+    "TSDAETask",
+]

vespaembed/tasks/base.py

@@ -0,0 +1,144 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+from datasets import Dataset
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.training_args import BatchSamplers
+
+
+class BaseTask(ABC):
+    """Base class for all training tasks."""
+
+    # Task metadata
+    name: str = ""
+    description: str = ""
+
+    # Column configuration
+    expected_columns: list[str] = []  # Required columns
+    optional_columns: list[str] = []  # Optional columns (included if present)
+    column_aliases: dict[str, list[str]] = {}
+
+    # Loss configuration
+    loss_options: list[str] = []  # Available loss variants
+    default_loss: str = ""  # Default loss variant
+
+    # Batch sampler (NO_DUPLICATES for in-batch negative losses)
+    batch_sampler: BatchSamplers = BatchSamplers.BATCH_SAMPLER
+
+    def __init__(self, loss_variant: str | None = None):
+        """Initialize task with optional loss variant.
+
+        Args:
+            loss_variant: Which loss variant to use (if task supports multiple)
+        """
+        # Label encoding mappings (set by prepare_dataset if task has labels)
+        self._label_to_idx: dict[str, int] | None = None
+        self._idx_to_label: dict[int, str] | None = None
+
+        # Loss variant selection
+        if loss_variant:
+            if self.loss_options and loss_variant not in self.loss_options:
+                raise ValueError(f"Unknown loss variant '{loss_variant}'. Options: {self.loss_options}")
+            self._loss_variant = loss_variant
+        else:
+            self._loss_variant = self.default_loss
+
+        # Track which optional columns were found
+        self._found_optional_columns: list[str] = []
+
+    @property
+    def loss_variant(self) -> str:
+        """Return the selected loss variant."""
+        return self._loss_variant
+
+    @property
+    def found_optional_columns(self) -> list[str]:
+        """Return list of optional columns that were found in the dataset."""
+        return self._found_optional_columns
+
+    @property
+    def label_to_idx(self) -> dict[str, int] | None:
+        """Return label to index mapping, or None if task has no labels."""
+        return self._label_to_idx
+
+    @property
+    def idx_to_label(self) -> dict[int, str] | None:
+        """Return index to label mapping, or None if task has no labels."""
+        return self._idx_to_label
+
+    @property
+    def num_labels(self) -> int | None:
+        """Return number of labels, or None if task has no labels."""
+        if self._label_to_idx is None:
+            return None
+        return len(self._label_to_idx)
+
+    def get_label_config(self) -> dict | None:
+        """Return label configuration in HuggingFace format.
+
+        Returns:
+            Dict with id2label, label2id, and num_labels, or None if no labels
+        """
+        if self._label_to_idx is None:
+            return None
+        return {
+            "id2label": {str(k): v for k, v in self._idx_to_label.items()},
+            "label2id": self._label_to_idx,
+            "num_labels": len(self._label_to_idx),
+        }
+
+    @abstractmethod
+    def get_loss(self, model: SentenceTransformer, **kwargs) -> Any:
+        """Return configured loss function from sentence-transformers.
+
+        Args:
+            model: The SentenceTransformer model
+            **kwargs: Additional loss configuration
+
+        Returns:
+            Loss function instance
+        """
+        raise NotImplementedError
+
+    def prepare_dataset(self, dataset: Dataset) -> Dataset:
+        """Normalize column names and reorder for the loss function.
+
+        Args:
+            dataset: Input dataset
+
+        Returns:
+            Prepared dataset with normalized columns
+        """
+        # 1. Rename aliased columns to canonical names
+        for canonical, aliases in self.column_aliases.items():
+            for alias in aliases:
+                if alias in dataset.column_names and canonical not in dataset.column_names:
+                    dataset = dataset.rename_column(alias, canonical)
+                    break
+
+        # 2. Validate required columns exist
+        missing = set(self.expected_columns) - set(dataset.column_names)
+        if missing:
+            available = ", ".join(sorted(dataset.column_names))
+            raise ValueError(
+                f"Missing required columns for task '{self.name}': {missing}. " f"Available columns: {available}"
+            )
+
+        # 3. Find which optional columns are present
+        self._found_optional_columns = [col for col in self.optional_columns if col in dataset.column_names]
+
+        # 4. Select and reorder columns (required + found optional)
+        columns_to_select = self.expected_columns + self._found_optional_columns
+        return dataset.select_columns(columns_to_select)
+
+    @abstractmethod
+    def get_evaluator(self, eval_dataset: Dataset) -> Any:
+        """Return appropriate sentence-transformers evaluator for this task.
+
+        Args:
+            eval_dataset: Evaluation dataset (already prepared)
+
+        Returns:
+            Evaluator instance
+        """
+        raise NotImplementedError

vespaembed/tasks/pairs.py

@@ -0,0 +1,91 @@
+from datasets import Dataset
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.evaluation import InformationRetrievalEvaluator
+from sentence_transformers.losses import (
+    CachedGISTEmbedLoss,
+    CachedMultipleNegativesRankingLoss,
+    GISTEmbedLoss,
+    MultipleNegativesRankingLoss,
+    MultipleNegativesSymmetricRankingLoss,
+)
+from sentence_transformers.training_args import BatchSamplers
+
+from vespaembed.core.registry import Registry
+from vespaembed.tasks.base import BaseTask
+
+
+@Registry.register_task("pairs")
+class PairsTask(BaseTask):
+    """Text pairs task for semantic search and retrieval.
+
+    Use this when you have query-document pairs without explicit negatives.
+    The loss uses in-batch negatives for contrastive learning.
+
+    Data format:
+    - anchor: Query/question text
+    - positive: Relevant document/answer
+
+    Loss variants:
+    - mnr: MultipleNegativesRankingLoss (default) - standard in-batch negatives
+    - mnr_symmetric: Bidirectional ranking - use if you need "given answer, find query"
+    - gist: GISTEmbedLoss - uses guide model to filter false negatives
+    - cached_mnr: Cached version - allows larger effective batch sizes
+    - cached_gist: Cached GIST - combines both benefits
+
+    Tip: Use mine_hard_negatives() to create triplets for better results.
+    """
+
+    name = "pairs"
+    description = "Text pairs for semantic search (anchor, positive)"
+
+    expected_columns = ["anchor", "positive"]
+    column_aliases = {
+        "anchor": ["query", "question", "sent1", "sentence1", "text1"],
+        "positive": ["document", "answer", "pos", "sent2", "sentence2", "text2"],
+    }
+
+    # Loss variants
+    loss_options = ["mnr", "mnr_symmetric", "gist", "cached_mnr", "cached_gist"]
+    default_loss = "mnr"
+
+    batch_sampler = BatchSamplers.NO_DUPLICATES
+
+    def get_loss(self, model: SentenceTransformer, **kwargs) -> MultipleNegativesRankingLoss:
+        """Return the selected loss variant."""
+        guide_model = kwargs.pop("guide_model", None)
+        mini_batch_size = kwargs.pop("mini_batch_size", 32)
+
+        if self._loss_variant == "mnr":
+            return MultipleNegativesRankingLoss(model, **kwargs)
+
+        elif self._loss_variant == "mnr_symmetric":
+            return MultipleNegativesSymmetricRankingLoss(model, **kwargs)
+
+        elif self._loss_variant == "gist":
+            if guide_model is None:
+                guide_model = model
+            return GISTEmbedLoss(model, guide=guide_model, **kwargs)
+
+        elif self._loss_variant == "cached_mnr":
+            return CachedMultipleNegativesRankingLoss(model, mini_batch_size=mini_batch_size, **kwargs)
+
+        elif self._loss_variant == "cached_gist":
+            if guide_model is None:
+                guide_model = model
+            return CachedGISTEmbedLoss(model, guide=guide_model, mini_batch_size=mini_batch_size, **kwargs)
+
+        else:
+            return MultipleNegativesRankingLoss(model, **kwargs)
+
+    def get_evaluator(self, eval_dataset: Dataset):
+        """Return InformationRetrievalEvaluator for pair data."""
+        queries = {str(i): text for i, text in enumerate(eval_dataset["anchor"])}
+        corpus = {str(i): text for i, text in enumerate(eval_dataset["positive"])}
+        relevant_docs = {str(i): {str(i)} for i in range(len(eval_dataset))}
+
+        return InformationRetrievalEvaluator(
+            queries=queries,
+            corpus=corpus,
+            relevant_docs=relevant_docs,
+            name="pairs-eval",
+        )

vespaembed/tasks/similarity.py

@@ -0,0 +1,84 @@
+from datasets import Dataset
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator
+from sentence_transformers.losses import AnglELoss, CoSENTLoss, CosineSimilarityLoss
+from sentence_transformers.training_args import BatchSamplers
+
+from vespaembed.core.registry import Registry
+from vespaembed.tasks.base import BaseTask
+
+
+@Registry.register_task("similarity")
+class SimilarityTask(BaseTask):
+    """Text pairs with similarity scores task.
+
+    Use this when you have pairs of sentences with continuous similarity
+    scores (e.g., 0.0 to 1.0 or 0 to 5). Common for STS (Semantic Textual
+    Similarity) benchmarks.
+
+    Data format:
+    - sentence1: First sentence
+    - sentence2: Second sentence
+    - score: Similarity score (will be normalized to 0-1 range)
+
+    Loss variants (similar performance, pick one):
+    - cosine: CosineSimilarityLoss (default) - simple and effective
+    - cosent: CoSENTLoss - ranking-based, from CoSENT paper
+    - angle: AnglELoss - angle-optimized, from AnglE paper
+
+    According to papers: AnglE >= CoSENT >= Cosine, but results are often similar.
+    """
+
+    name = "similarity"
+    description = "Text pairs with similarity scores (STS-style)"
+
+    expected_columns = ["sentence1", "sentence2", "score"]
+    column_aliases = {
+        "sentence1": ["sent1", "text1", "anchor", "query"],
+        "sentence2": ["sent2", "text2", "positive", "document"],
+        "score": ["similarity", "label", "sim_score"],
+    }
+
+    # Loss variants
+    loss_options = ["cosine", "cosent", "angle"]
+    default_loss = "cosine"
+
+    batch_sampler = BatchSamplers.BATCH_SAMPLER
+
+    def prepare_dataset(self, dataset: Dataset) -> Dataset:
+        """Prepare dataset and normalize scores to 0-1 range if needed."""
+        dataset = super().prepare_dataset(dataset)
+
+        # Check if scores need normalization (e.g., 0-5 scale to 0-1)
+        scores = dataset["score"]
+        max_score = max(scores)
+
+        if max_score > 1.0:
+            # Normalize to 0-1 range
+            dataset = dataset.map(lambda x: {"score": x["score"] / max_score})
+
+        return dataset
+
+    def get_loss(self, model: SentenceTransformer, **kwargs):
+        """Return the selected loss variant."""
+        if self._loss_variant == "cosine":
+            return CosineSimilarityLoss(model, **kwargs)
+
+        elif self._loss_variant == "cosent":
+            return CoSENTLoss(model, **kwargs)
+
+        elif self._loss_variant == "angle":
+            return AnglELoss(model, **kwargs)
+
+        else:
+            # Fallback to default
+            return CosineSimilarityLoss(model, **kwargs)
+
+    def get_evaluator(self, eval_dataset: Dataset) -> EmbeddingSimilarityEvaluator:
+        """Return EmbeddingSimilarityEvaluator."""
+        return EmbeddingSimilarityEvaluator(
+            sentences1=eval_dataset["sentence1"],
+            sentences2=eval_dataset["sentence2"],
+            scores=eval_dataset["score"],
+            name="similarity-eval",
+        )

vespaembed/tasks/triplets.py

@@ -0,0 +1,90 @@
+from datasets import Dataset
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.evaluation import TripletEvaluator
+from sentence_transformers.losses import (
+    CachedGISTEmbedLoss,
+    CachedMultipleNegativesRankingLoss,
+    GISTEmbedLoss,
+    MultipleNegativesRankingLoss,
+    MultipleNegativesSymmetricRankingLoss,
+)
+from sentence_transformers.training_args import BatchSamplers
+
+from vespaembed.core.registry import Registry
+from vespaembed.tasks.base import BaseTask
+
+
+@Registry.register_task("triplets")
+class TripletsTask(BaseTask):
+    """Text triplets task for semantic search and retrieval.
+
+    Use this when you have query-document pairs WITH explicit hard negatives.
+    The loss uses both explicit negatives and in-batch negatives.
+
+    Data format:
+    - anchor: Query/question text
+    - positive: Relevant document/answer
+    - negative: Hard negative (irrelevant but similar document)
+
+    Loss variants:
+    - mnr: MultipleNegativesRankingLoss (default) - explicit + in-batch negatives
+    - mnr_symmetric: Bidirectional ranking - use if you need "given answer, find query"
+    - gist: GISTEmbedLoss - uses guide model to filter false negatives
+    - cached_mnr: Cached version - allows larger effective batch sizes
+    - cached_gist: Cached GIST - combines both benefits
+
+    Tip: Hard negatives significantly improve model quality. Use mine_hard_negatives()
+    to generate them from pairs data.
+    """
+
+    name = "triplets"
+    description = "Text triplets for semantic search (anchor, positive, negative)"
+
+    expected_columns = ["anchor", "positive", "negative"]
+    column_aliases = {
+        "anchor": ["query", "question", "sent1", "sentence1", "text1", "premise"],
+        "positive": ["document", "answer", "pos", "sent2", "sentence2", "text2", "entailment"],
+        "negative": ["neg", "hard_negative", "sent3", "sentence3", "text3", "contradiction"],
+    }
+
+    # Loss variants (same as pairs)
+    loss_options = ["mnr", "mnr_symmetric", "gist", "cached_mnr", "cached_gist"]
+    default_loss = "mnr"
+
+    batch_sampler = BatchSamplers.NO_DUPLICATES
+
+    def get_loss(self, model: SentenceTransformer, **kwargs) -> MultipleNegativesRankingLoss:
+        """Return the selected loss variant."""
+        guide_model = kwargs.pop("guide_model", None)
+        mini_batch_size = kwargs.pop("mini_batch_size", 32)
+
+        if self._loss_variant == "mnr":
+            return MultipleNegativesRankingLoss(model, **kwargs)
+
+        elif self._loss_variant == "mnr_symmetric":
+            return MultipleNegativesSymmetricRankingLoss(model, **kwargs)
+
+        elif self._loss_variant == "gist":
+            if guide_model is None:
+                guide_model = model
+            return GISTEmbedLoss(model, guide=guide_model, **kwargs)
+
+        elif self._loss_variant == "cached_mnr":
+            return CachedMultipleNegativesRankingLoss(model, mini_batch_size=mini_batch_size, **kwargs)
+
+        elif self._loss_variant == "cached_gist":
+            if guide_model is None:
+                guide_model = model
+            return CachedGISTEmbedLoss(model, guide=guide_model, mini_batch_size=mini_batch_size, **kwargs)
+
+        else:
+            return MultipleNegativesRankingLoss(model, **kwargs)
+
+    def get_evaluator(self, eval_dataset: Dataset):
+        """Return TripletEvaluator for triplet data."""
+        return TripletEvaluator(
+            anchors=eval_dataset["anchor"],
+            positives=eval_dataset["positive"],
+            negatives=eval_dataset["negative"],
+            name="triplets-eval",
+        )

vespaembed/tasks/tsdae.py

@@ -0,0 +1,102 @@
+import random
+
+from datasets import Dataset
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.losses import DenoisingAutoEncoderLoss
+from sentence_transformers.training_args import BatchSamplers
+
+from vespaembed.core.registry import Registry
+from vespaembed.tasks.base import BaseTask
+
+
+def _add_noise(text: str, del_ratio: float = 0.6) -> str:
+    """Add noise to text by randomly deleting words.
+
+    Args:
+        text: Input text string
+        del_ratio: Probability of keeping each word (default 0.6 = 60% kept)
+
+    Returns:
+        Noisy version of the text with some words randomly deleted
+    """
+    words = text.split()
+    if not words:
+        return text
+    # Keep words with probability del_ratio
+    kept_words = [word for word in words if random.random() < del_ratio]
+    if len(kept_words) == 0:
+        # Keep at least one random word
+        return random.choice(words)
+    return " ".join(kept_words)
+
+
+@Registry.register_task("tsdae")
+class TSDAETask(BaseTask):
+    """TSDAE (Transformer-based Sequential Denoising Auto-Encoder) task.
+
+    Unsupervised training that learns embeddings by reconstructing
+    corrupted input sentences. Useful for domain adaptation when
+    you only have unlabeled text.
+
+    Data format:
+    - text/sentence: Raw text sentences (no labels needed)
+
+    The task automatically adds noise by randomly deleting ~40% of words
+    from the input text. The model learns to reconstruct the original
+    text from the corrupted version.
+
+    Reference: https://arxiv.org/abs/2104.06979
+    """
+
+    name = "tsdae"
+    description = "TSDAE - unsupervised domain adaptation with denoising auto-encoder"
+
+    expected_columns = ["text"]
+    column_aliases = {
+        "text": ["sentence", "sentences", "content", "input"],
+    }
+
+    batch_sampler = BatchSamplers.BATCH_SAMPLER
+
+    def prepare_dataset(self, dataset: Dataset) -> Dataset:
+        """Prepare dataset for TSDAE training with noise.
+
+        Adds noise to text by randomly deleting words (40% deletion rate).
+        The model learns to reconstruct the original text from the noisy input.
+
+        Args:
+            dataset: Input dataset with 'text' column
+
+        Returns:
+            Dataset with 'anchor' (noisy) and 'positive' (original) columns
+        """
+        # First, apply base class normalization (handles column aliases)
+        dataset = super().prepare_dataset(dataset)
+
+        # Create noisy versions by randomly deleting words
+        # anchor = noisy text (input), positive = original text (target)
+        dataset = dataset.map(
+            lambda x: {"anchor": _add_noise(x["text"]), "positive": x["text"]},
+            desc="Adding noise to text",
+        )
+
+        # Remove the original text column as we now have anchor/positive
+        dataset = dataset.remove_columns(["text"])
+
+        return dataset
+
+    def get_loss(self, model: SentenceTransformer, **kwargs) -> DenoisingAutoEncoderLoss:
+        """Return DenoisingAutoEncoderLoss for unsupervised training."""
+        return DenoisingAutoEncoderLoss(
+            model=model,
+            tie_encoder_decoder=True,
+            **kwargs,
+        )
+
+    def get_evaluator(self, eval_dataset: Dataset):
+        """Return None - TSDAE has no intrinsic evaluator.
+
+        Evaluation should be done on downstream tasks (e.g., STS, retrieval)
+        that match your target use case.
+        """
+        return None