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

vespaembed/datasets/loader.py

@@ -0,0 +1,80 @@
+from pathlib import Path
+from typing import Optional
+
+from datasets import Dataset
+
+from vespaembed.datasets.formats.csv import load_csv
+from vespaembed.datasets.formats.huggingface import load_hf_dataset
+from vespaembed.datasets.formats.jsonl import load_jsonl
+
+
+def load_dataset(
+    path: str,
+    subset: Optional[str] = None,
+    split: Optional[str] = None,
+) -> Dataset:
+    """Load a dataset from various sources.
+
+    Supports:
+    - CSV files (.csv)
+    - JSONL files (.jsonl)
+    - HuggingFace datasets (org/dataset-name)
+
+    Args:
+        path: Path to file or HuggingFace dataset name
+        subset: HuggingFace dataset subset (optional)
+        split: HuggingFace dataset split (optional, defaults to "train")
+
+    Returns:
+        HuggingFace Dataset object
+
+    Raises:
+        ValueError: If file format is not supported
+        FileNotFoundError: If file does not exist
+    """
+    path_obj = Path(path)
+
+    # Check if it's a local file
+    if path_obj.exists():
+        suffix = path_obj.suffix.lower()
+
+        if suffix == ".csv":
+            return load_csv(path)
+        elif suffix in (".jsonl", ".json"):
+            return load_jsonl(path)
+        else:
+            raise ValueError(f"Unsupported file format: {suffix}. " "Supported formats: .csv, .jsonl")
+
+    # Check if it looks like a HuggingFace dataset
+    if "/" in path or not path_obj.suffix:
+        return load_hf_dataset(path, subset=subset, split=split or "train")
+
+    # File doesn't exist and doesn't look like HF dataset
+    raise FileNotFoundError(
+        f"File not found: {path}. " "Provide a valid file path or HuggingFace dataset name (e.g., 'org/dataset-name')."
+    )
+
+
+def preview_dataset(dataset: Dataset, num_samples: int = 5) -> list[dict]:
+    """Preview a dataset by returning the first N samples.
+
+    Args:
+        dataset: Dataset to preview
+        num_samples: Number of samples to return
+
+    Returns:
+        List of sample dictionaries
+    """
+    return [dataset[i] for i in range(min(num_samples, len(dataset)))]
+
+
+def get_columns(dataset: Dataset) -> list[str]:
+    """Get column names from a dataset.
+
+    Args:
+        dataset: Dataset to inspect
+
+    Returns:
+        List of column names
+    """
+    return dataset.column_names

vespaembed/db.py

@@ -0,0 +1,176 @@
+import json
+import shutil
+import sqlite3
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+from vespaembed.enums import RunStatus
+
+# Default database location
+DEFAULT_DB_DIR = Path.home() / ".vespaembed"
+DEFAULT_DB_PATH = DEFAULT_DB_DIR / "vespaembed.db"
+
+
+def get_db_path() -> Path:
+    """Get the database path, creating directory if needed."""
+    DEFAULT_DB_DIR.mkdir(parents=True, exist_ok=True)
+    return DEFAULT_DB_PATH
+
+
+def get_connection() -> sqlite3.Connection:
+    """Get a database connection."""
+    conn = sqlite3.connect(get_db_path(), check_same_thread=False)
+    conn.row_factory = sqlite3.Row
+    return conn
+
+
+def init_db():
+    """Initialize the database schema."""
+    conn = get_connection()
+    cursor = conn.cursor()
+
+    cursor.execute("""
+        CREATE TABLE IF NOT EXISTS runs (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            status TEXT NOT NULL DEFAULT 'pending',
+            pid INTEGER,
+            config TEXT,
+            project_name TEXT,
+            output_dir TEXT,
+            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+            error_message TEXT
+        )
+    """)
+
+    conn.commit()
+    conn.close()
+
+
+def create_run(config: dict, project_name: str, output_dir: str) -> int:
+    """Create a new training run."""
+    conn = get_connection()
+    cursor = conn.cursor()
+
+    cursor.execute(
+        """
+        INSERT INTO runs (status, config, project_name, output_dir)
+        VALUES (?, ?, ?, ?)
+    """,
+        (RunStatus.PENDING.value, json.dumps(config), project_name, output_dir),
+    )
+
+    run_id = cursor.lastrowid
+    conn.commit()
+    conn.close()
+
+    return run_id
+
+
+def update_run_status(
+    run_id: int,
+    status: RunStatus,
+    pid: Optional[int] = None,
+    error_message: Optional[str] = None,
+):
+    """Update a run's status."""
+    conn = get_connection()
+    cursor = conn.cursor()
+
+    if pid is not None:
+        cursor.execute(
+            """
+            UPDATE runs SET status = ?, pid = ?, updated_at = ?
+            WHERE id = ?
+        """,
+            (status.value, pid, datetime.now(), run_id),
+        )
+    elif error_message is not None:
+        cursor.execute(
+            """
+            UPDATE runs SET status = ?, error_message = ?, updated_at = ?
+            WHERE id = ?
+        """,
+            (status.value, error_message, datetime.now(), run_id),
+        )
+    else:
+        cursor.execute(
+            """
+            UPDATE runs SET status = ?, updated_at = ?
+            WHERE id = ?
+        """,
+            (status.value, datetime.now(), run_id),
+        )
+
+    conn.commit()
+    conn.close()
+
+
+def get_run(run_id: int) -> Optional[dict]:
+    """Get a run by ID."""
+    conn = get_connection()
+    cursor = conn.cursor()
+
+    cursor.execute("SELECT * FROM runs WHERE id = ?", (run_id,))
+    row = cursor.fetchone()
+    conn.close()
+
+    if row:
+        return dict(row)
+    return None
+
+
+def get_all_runs() -> list[dict]:
+    """Get all runs."""
+    conn = get_connection()
+    cursor = conn.cursor()
+
+    cursor.execute("SELECT * FROM runs ORDER BY created_at DESC")
+    rows = cursor.fetchall()
+    conn.close()
+
+    return [dict(row) for row in rows]
+
+
+def get_active_run() -> Optional[dict]:
+    """Get the currently active (running or pending) run."""
+    conn = get_connection()
+    cursor = conn.cursor()
+
+    cursor.execute(
+        """
+        SELECT * FROM runs
+        WHERE status IN (?, ?)
+        ORDER BY created_at DESC
+        LIMIT 1
+    """,
+        (RunStatus.PENDING.value, RunStatus.RUNNING.value),
+    )
+
+    row = cursor.fetchone()
+    conn.close()
+
+    if row:
+        return dict(row)
+    return None
+
+
+def delete_run(run_id: int, delete_files: bool = True):
+    """Delete a run and optionally its output files."""
+    run = get_run(run_id)
+
+    if run and delete_files and run.get("output_dir"):
+        output_path = Path(run["output_dir"])
+        if output_path.exists():
+            shutil.rmtree(output_path, ignore_errors=True)
+
+    conn = get_connection()
+    cursor = conn.cursor()
+    cursor.execute("DELETE FROM runs WHERE id = ?", (run_id,))
+    conn.commit()
+    conn.close()
+
+
+# Initialize database on import
+init_db()

vespaembed/enums.py

@@ -0,0 +1,58 @@
+from enum import Enum
+
+
+class TaskType(str, Enum):
+    """Supported training tasks.
+
+    Tasks are organized by data format:
+    - pairs: Text pairs for semantic search (anchor, positive)
+    - triplets: Text triplets with hard negatives (anchor, positive, negative)
+    - similarity: Text pairs with similarity scores
+    - tsdae: Unlabeled text for unsupervised learning
+
+    Note: Matryoshka is a training option (--matryoshka flag), not a separate task.
+    """
+
+    PAIRS = "pairs"
+    TRIPLETS = "triplets"
+    SIMILARITY = "similarity"
+    TSDAE = "tsdae"
+
+
+class LossVariant(str, Enum):
+    """Available loss function variants.
+
+    For pairs task:
+    - mnr: MultipleNegativesRankingLoss (default, recommended)
+    - mnr_symmetric: Bidirectional ranking
+    - gist: GISTEmbedLoss with guide model
+    - cached_mnr: Cached version for larger batches
+    - cached_gist: Cached GIST
+
+    For similarity task:
+    - cosine: CosineSimilarityLoss (default)
+    - cosent: CoSENTLoss
+    - angle: AnglELoss
+    """
+
+    # Pairs task variants
+    MNR = "mnr"
+    MNR_SYMMETRIC = "mnr_symmetric"
+    GIST = "gist"
+    CACHED_MNR = "cached_mnr"
+    CACHED_GIST = "cached_gist"
+
+    # Similarity task variants
+    COSINE = "cosine"
+    COSENT = "cosent"
+    ANGLE = "angle"
+
+
+class RunStatus(str, Enum):
+    """Training run status."""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    STOPPED = "stopped"
+    ERROR = "error"

vespaembed/evaluation/__init__.py

@@ -0,0 +1,3 @@
+from vespaembed.evaluation.factory import create_evaluator
+
+__all__ = ["create_evaluator"]

vespaembed/evaluation/factory.py

@@ -0,0 +1,86 @@
+from typing import Any, Optional
+
+from datasets import Dataset
+from sentence_transformers.evaluation import (
+    BinaryClassificationEvaluator,
+    EmbeddingSimilarityEvaluator,
+    InformationRetrievalEvaluator,
+    TripletEvaluator,
+)
+
+
+def create_evaluator(
+    task: str,
+    eval_dataset: Dataset,
+    name: str = "eval",
+) -> Optional[Any]:
+    """Create an appropriate evaluator based on task type.
+
+    Args:
+        task: Task name (mnr, triplet, contrastive, sts, nli, tsdae, matryoshka)
+        eval_dataset: Prepared evaluation dataset
+        name: Evaluator name
+
+    Returns:
+        Evaluator instance or None
+    """
+    if task == "mnr":
+        return _create_ir_evaluator(eval_dataset, name)
+    elif task == "triplet":
+        return _create_triplet_evaluator(eval_dataset, name)
+    elif task == "contrastive":
+        return _create_binary_evaluator(eval_dataset, name)
+    elif task == "sts":
+        return _create_similarity_evaluator(eval_dataset, name)
+    elif task == "nli":
+        return _create_similarity_evaluator(eval_dataset, name)
+    elif task == "tsdae":
+        return None  # TSDAE has no intrinsic evaluator
+    elif task == "matryoshka":
+        return _create_ir_evaluator(eval_dataset, name)
+    else:
+        return None
+
+
+def _create_ir_evaluator(dataset: Dataset, name: str) -> InformationRetrievalEvaluator:
+    """Create Information Retrieval evaluator."""
+    queries = {str(i): text for i, text in enumerate(dataset["anchor"])}
+    corpus = {str(i): text for i, text in enumerate(dataset["positive"])}
+    relevant_docs = {str(i): {str(i)} for i in range(len(dataset))}
+
+    return InformationRetrievalEvaluator(
+        queries=queries,
+        corpus=corpus,
+        relevant_docs=relevant_docs,
+        name=name,
+    )
+
+
+def _create_triplet_evaluator(dataset: Dataset, name: str) -> TripletEvaluator:
+    """Create Triplet evaluator."""
+    return TripletEvaluator(
+        anchors=dataset["anchor"],
+        positives=dataset["positive"],
+        negatives=dataset["negative"],
+        name=name,
+    )
+
+
+def _create_binary_evaluator(dataset: Dataset, name: str) -> BinaryClassificationEvaluator:
+    """Create Binary Classification evaluator."""
+    return BinaryClassificationEvaluator(
+        sentences1=dataset["sentence1"],
+        sentences2=dataset["sentence2"],
+        labels=dataset["label"],
+        name=name,
+    )
+
+
+def _create_similarity_evaluator(dataset: Dataset, name: str) -> EmbeddingSimilarityEvaluator:
+    """Create Embedding Similarity evaluator."""
+    return EmbeddingSimilarityEvaluator(
+        sentences1=dataset["sentence1"],
+        sentences2=dataset["sentence2"],
+        scores=dataset["score"],
+        name=name,
+    )

vespaembed/models/__init__.py

@@ -0,0 +1,4 @@
+from vespaembed.models.export import export_model
+from vespaembed.models.loader import load_model
+
+__all__ = ["export_model", "load_model"]

vespaembed/models/export.py

@@ -0,0 +1,89 @@
+from pathlib import Path
+from typing import Optional
+
+from sentence_transformers import SentenceTransformer
+
+
+def export_model(
+    model: SentenceTransformer,
+    output_path: str,
+    format: str = "onnx",
+) -> str:
+    """Export a model to a different format.
+
+    Args:
+        model: SentenceTransformer model
+        output_path: Output directory or file path
+        format: Export format ("onnx")
+
+    Returns:
+        Path to exported model
+
+    Raises:
+        ValueError: If format is not supported
+    """
+    output_path = Path(output_path)
+
+    if format.lower() == "onnx":
+        return _export_onnx(model, output_path)
+    else:
+        raise ValueError(f"Unsupported export format: {format}. Supported: onnx")
+
+
+def _export_onnx(model: SentenceTransformer, output_path: Path) -> str:
+    """Export model to ONNX format.
+
+    Args:
+        model: SentenceTransformer model
+        output_path: Output directory
+
+    Returns:
+        Path to ONNX model
+    """
+    try:
+        import onnx  # noqa: F401
+    except ImportError:
+        raise ImportError("ONNX not installed. Install with: pip install vespaembed[onnx]")
+
+    output_path.mkdir(parents=True, exist_ok=True)
+    onnx_path = output_path / "model.onnx"
+
+    # Use sentence-transformers built-in ONNX export if available
+    # Otherwise fall back to manual export
+    try:
+        model.save(str(output_path), model_name_or_path="model.onnx", create_model_card=False)
+    except Exception:
+        # Manual export via transformers
+        from transformers import AutoTokenizer
+
+        tokenizer = AutoTokenizer.from_pretrained(model[0].auto_model.config._name_or_path)
+
+        # Export the transformer part
+        model[0].auto_model.save_pretrained(output_path)
+        tokenizer.save_pretrained(output_path)
+
+    return str(onnx_path)
+
+
+def push_to_hub(
+    model: SentenceTransformer,
+    repo_id: str,
+    commit_message: Optional[str] = None,
+    private: bool = False,
+) -> str:
+    """Push model to HuggingFace Hub.
+
+    Args:
+        model: SentenceTransformer model
+        repo_id: Repository ID (e.g., "username/model-name")
+        commit_message: Commit message
+        private: Whether to create a private repository
+
+    Returns:
+        URL of the model on HuggingFace Hub
+    """
+    return model.push_to_hub(
+        repo_id=repo_id,
+        commit_message=commit_message or "Upload model via vespaembed",
+        private=private,
+    )

vespaembed/models/loader.py

@@ -0,0 +1,25 @@
+from sentence_transformers import SentenceTransformer
+
+
+def load_model(model_name_or_path: str, use_unsloth: bool = False) -> SentenceTransformer:
+    """Load a sentence transformer model.
+
+    Args:
+        model_name_or_path: Model name from HuggingFace Hub or local path
+        use_unsloth: Whether to use Unsloth for faster inference
+
+    Returns:
+        SentenceTransformer model
+    """
+    if use_unsloth:
+        try:
+            from unsloth import FastSentenceTransformer
+
+            return FastSentenceTransformer.from_pretrained(
+                model_name_or_path,
+                for_inference=True,
+            )
+        except ImportError:
+            raise ImportError("Unsloth not installed. Install with: pip install vespaembed[unsloth]")
+
+    return SentenceTransformer(model_name_or_path)