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

vespaembed/core/trainer.py

@@ -0,0 +1,573 @@
+import json
+import os
+from pathlib import Path
+from typing import Callable, Optional
+
+from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer
+from sentence_transformers.training_args import SentenceTransformerTrainingArguments
+from transformers import TrainerCallback
+
+from vespaembed.core.config import TrainingConfig
+from vespaembed.core.registry import Registry
+from vespaembed.datasets.loader import load_dataset
+from vespaembed.utils.logging import logger
+
+# Get HuggingFace token from environment (fallback for loading private models)
+HF_TOKEN_ENV = os.environ.get("HF_TOKEN")
+
+
+class ProgressCallback(TrainerCallback):
+    """Callback to report training progress like tqdm."""
+
+    def __init__(self, callback: Callable[[dict], None]):
+        self.callback = callback
+        self.start_time = None
+        self.total_steps = 0
+        self.total_epochs = 0
+
+    def on_train_begin(self, args, state, control, **kwargs):
+        """Called at the beginning of training."""
+        import time
+
+        self.start_time = time.time()
+        self.total_steps = state.max_steps
+        self.total_epochs = args.num_train_epochs
+
+        if self.callback:
+            self.callback(
+                {
+                    "type": "train_start",
+                    "total_steps": self.total_steps,
+                    "total_epochs": self.total_epochs,
+                }
+            )
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        """Called when the trainer logs metrics."""
+        import time
+
+        if logs and self.callback:
+            current_step = state.global_step
+            elapsed = time.time() - self.start_time if self.start_time else 0
+
+            # Calculate progress
+            progress_pct = (current_step / self.total_steps * 100) if self.total_steps > 0 else 0
+            steps_per_sec = current_step / elapsed if elapsed > 0 else 0
+            remaining_steps = self.total_steps - current_step
+            eta_seconds = remaining_steps / steps_per_sec if steps_per_sec > 0 else 0
+
+            progress = {
+                "type": "progress",
+                "epoch": state.epoch,
+                "total_epochs": self.total_epochs,
+                "step": current_step,
+                "total_steps": self.total_steps,
+                "progress_pct": progress_pct,
+                "loss": logs.get("loss"),
+                "learning_rate": logs.get("learning_rate"),
+                "steps_per_sec": steps_per_sec,
+                "elapsed_seconds": elapsed,
+                "eta_seconds": eta_seconds,
+            }
+            self.callback(progress)
+
+    def on_train_end(self, args, state, control, **kwargs):
+        """Called at the end of training."""
+        import time
+
+        if self.callback:
+            elapsed = time.time() - self.start_time if self.start_time else 0
+            self.callback(
+                {
+                    "type": "train_end",
+                    "total_steps": state.global_step,
+                    "elapsed_seconds": elapsed,
+                }
+            )
+
+
+class VespaEmbedTrainer:
+    """High-level trainer that wraps SentenceTransformerTrainer."""
+
+    def __init__(
+        self,
+        config: TrainingConfig,
+        progress_callback: Optional[Callable[[dict], None]] = None,
+    ):
+        """Initialize the trainer.
+
+        Args:
+            config: Training configuration
+            progress_callback: Optional callback for progress updates
+        """
+        self.config = config
+        self.model = None
+        self.task = None
+        self.progress_callback = progress_callback
+
+    def _load_model(self) -> SentenceTransformer:
+        """Load the model with optional LoRA and Unsloth support.
+
+        Supports four modes:
+        1. Standard: SentenceTransformer only
+        2. Standard + LoRA: SentenceTransformer with PEFT adapter
+        3. Unsloth: FastSentenceTransformer (faster training)
+        4. Unsloth + LoRA: FastSentenceTransformer with LoRA via get_peft_model
+        """
+        use_unsloth = self.config.unsloth.enabled
+
+        if use_unsloth:
+            return self._load_unsloth_model()
+        else:
+            return self._load_standard_model()
+
+    def _load_standard_model(self) -> SentenceTransformer:
+        """Load model with standard SentenceTransformer, optionally with LoRA."""
+        logger.info(f"Loading model: {self.config.base_model}")
+        model = SentenceTransformer(self.config.base_model, token=HF_TOKEN_ENV)
+
+        # Set max_seq_length if specified
+        if self.config.max_seq_length:
+            model.max_seq_length = self.config.max_seq_length
+            logger.info(f"Set max_seq_length: {self.config.max_seq_length}")
+        else:
+            logger.info(f"Using model's default max_seq_length: {model.max_seq_length}")
+
+        # Note: gradient_checkpointing is handled by SentenceTransformerTrainingArguments
+        # The Trainer will automatically enable it on the model during training
+
+        # Add LoRA adapter if enabled
+        if self.config.lora.enabled:
+            try:
+                from peft import LoraConfig, TaskType
+
+                logger.info(
+                    f"Adding LoRA adapter: r={self.config.lora.r}, "
+                    f"alpha={self.config.lora.alpha}, dropout={self.config.lora.dropout}, "
+                    f"target_modules={self.config.lora.target_modules}"
+                )
+
+                peft_config = LoraConfig(
+                    task_type=TaskType.FEATURE_EXTRACTION,
+                    r=self.config.lora.r,
+                    lora_alpha=self.config.lora.alpha,
+                    lora_dropout=self.config.lora.dropout,
+                    target_modules=self.config.lora.target_modules,
+                )
+                model.add_adapter(peft_config)
+
+            except ImportError:
+                raise ImportError("PEFT not installed. Install with: pip install peft")
+
+        return model
+
+    def _load_unsloth_model(self) -> SentenceTransformer:
+        """Load model with Unsloth for faster training."""
+        import torch
+
+        try:
+            from unsloth import FastSentenceTransformer
+        except ImportError:
+            raise ImportError("Unsloth not installed. Install with: pip install unsloth")
+
+        # Auto-detect max_seq_length from model if not specified
+        max_seq_length = self.config.max_seq_length
+        if max_seq_length is None:
+            # Load model temporarily to get max_seq_length
+            temp_model = SentenceTransformer(self.config.base_model, token=HF_TOKEN_ENV)
+            max_seq_length = temp_model.max_seq_length
+            del temp_model
+            logger.info(f"Auto-detected max_seq_length: {max_seq_length}")
+        else:
+            logger.info(f"Using specified max_seq_length: {max_seq_length}")
+
+        # Determine dtype based on precision settings
+        # BF16 is preferred for Unsloth when available (better for training)
+        if self.config.training.bf16:
+            dtype = torch.bfloat16
+            logger.info("Using BF16 precision for Unsloth")
+        elif self.config.training.fp16:
+            dtype = torch.float16
+            logger.info("Using FP16 precision for Unsloth")
+        else:
+            dtype = None  # Let Unsloth auto-detect
+            logger.info("Using auto-detected precision for Unsloth")
+
+        # Full finetuning when Unsloth is enabled but LoRA is not
+        full_finetuning = not self.config.lora.enabled
+
+        # Use Unsloth's optimized gradient checkpointing ("unsloth") when enabled
+        # This is faster and uses less memory than standard PyTorch GC
+        # Must be passed to from_pretrained for full finetuning, or get_peft_model for LoRA
+        use_gc = "unsloth" if self.config.gradient_checkpointing else False
+
+        logger.info(f"Loading model with Unsloth: {self.config.base_model}")
+        if full_finetuning:
+            logger.info(f"Full finetuning mode (gradient_checkpointing={use_gc})")
+        model = FastSentenceTransformer.from_pretrained(
+            model_name=self.config.base_model,
+            max_seq_length=max_seq_length,
+            full_finetuning=full_finetuning,
+            dtype=dtype,
+            use_gradient_checkpointing=use_gc if full_finetuning else False,
+            token=HF_TOKEN_ENV,
+        )
+
+        # Add LoRA if enabled
+        if self.config.lora.enabled:
+            logger.info(
+                f"Applying Unsloth LoRA: r={self.config.lora.r}, "
+                f"alpha={self.config.lora.alpha}, target_modules={self.config.lora.target_modules}"
+            )
+
+            model = FastSentenceTransformer.get_peft_model(
+                model,
+                r=self.config.lora.r,
+                lora_alpha=self.config.lora.alpha,
+                lora_dropout=self.config.lora.dropout,
+                target_modules=self.config.lora.target_modules,
+                bias="none",
+                use_gradient_checkpointing=use_gc,
+                random_state=3407,
+                use_rslora=False,
+                loftq_config=None,
+                task_type="FEATURE_EXTRACTION",
+            )
+
+        return model
+
+    def _log_training_config(self):
+        """Log training configuration parameters."""
+        logger.info("=" * 60)
+        logger.info("Training Configuration")
+        logger.info("=" * 60)
+
+        # Model & Task
+        logger.info(f"  Base Model:      {self.config.base_model}")
+        logger.info(f"  Task:            {self.config.task}")
+        if self.config.loss_variant:
+            logger.info(f"  Loss Variant:    {self.config.loss_variant}")
+
+        # Data
+        logger.info(f"  Training Data:   {self.config.data.train}")
+        if self.config.data.eval:
+            logger.info(f"  Eval Data:       {self.config.data.eval}")
+
+        # Training hyperparameters
+        t = self.config.training
+        logger.info(f"  Epochs:          {t.epochs}")
+        logger.info(f"  Batch Size:      {t.batch_size}")
+        logger.info(f"  Learning Rate:   {t.learning_rate}")
+        logger.info(f"  Optimizer:       {t.optimizer}")
+        logger.info(f"  Scheduler:       {t.scheduler}")
+        logger.info(f"  Warmup Ratio:    {t.warmup_ratio}")
+        logger.info(f"  Weight Decay:    {t.weight_decay}")
+
+        # Precision
+        if t.bf16:
+            logger.info("  Precision:       BF16")
+        elif t.fp16:
+            logger.info("  Precision:       FP16")
+        else:
+            logger.info("  Precision:       FP32")
+
+        # Optional features
+        if self.config.max_seq_length:
+            logger.info(f"  Max Seq Length:  {self.config.max_seq_length}")
+        if self.config.gradient_checkpointing:
+            logger.info("  Grad Checkpoint: Enabled")
+        if t.gradient_accumulation_steps > 1:
+            logger.info(f"  Grad Accum:      {t.gradient_accumulation_steps}")
+
+        # LoRA
+        if self.config.lora.enabled:
+            logger.info(f"  LoRA:            r={self.config.lora.r}, alpha={self.config.lora.alpha}")
+
+        # Unsloth
+        if self.config.unsloth.enabled:
+            logger.info(f"  Unsloth:         Enabled (save: {self.config.unsloth.save_method})")
+
+        # Matryoshka
+        if self.config.matryoshka_dims:
+            logger.info(f"  Matryoshka:      {self.config.matryoshka_dims}")
+
+        # Output
+        logger.info(f"  Output Dir:      {self.config.output.dir}")
+        if self.config.output.push_to_hub:
+            logger.info(f"  Push to Hub:     {self.config.output.hf_username}")
+
+        logger.info("=" * 60)
+
+    def train(self) -> SentenceTransformer:
+        """Run the training process.
+
+        Returns:
+            Trained SentenceTransformer model
+        """
+        # Log configuration
+        self._log_training_config()
+
+        # 1. Load model
+        self.model = self._load_model()
+
+        # 2. Get task (pass task-specific params if applicable)
+        task_cls = Registry.get_task(self.config.task)
+        # Handle loss_variant as either enum or string
+        loss_variant = self.config.loss_variant
+        if loss_variant is not None:
+            loss_variant = loss_variant.value if hasattr(loss_variant, "value") else loss_variant
+
+        if loss_variant:
+            self.task = task_cls(loss_variant=loss_variant)
+        else:
+            self.task = task_cls()
+
+        loss_info = f" (loss: {self.task.loss_variant})" if self.task.loss_variant else ""
+        logger.info(f"Using task: {self.task.name} - {self.task.description}{loss_info}")
+
+        # 3. Load and prepare training data
+        logger.info(f"Loading training data: {self.config.data.train}")
+        train_data = load_dataset(
+            self.config.data.train,
+            subset=self.config.data.subset,
+            split=self.config.data.split or "train",
+        )
+        train_data = self.task.prepare_dataset(train_data)
+        logger.info(f"Training samples: {len(train_data)}")
+
+        # 4. Load and prepare evaluation data (optional)
+        eval_data = None
+        evaluator = None
+        if self.config.data.eval:
+            logger.info(f"Loading evaluation data: {self.config.data.eval}")
+            # Use eval_split if specified (for HF datasets), otherwise use default split detection
+            eval_split = self.config.data.eval_split
+            eval_data = load_dataset(
+                self.config.data.eval,
+                subset=self.config.data.subset,  # Use same subset as training
+                split=eval_split,
+            )
+            eval_data = self.task.prepare_dataset(eval_data)
+            evaluator = self.task.get_evaluator(eval_data)
+            logger.info(f"Evaluation samples: {len(eval_data)}")
+
+        # 5. Create loss function
+        loss = self.task.get_loss(self.model)
+
+        # Wrap with MatryoshkaLoss if dimensions specified (not supported for TSDAE)
+        if self.config.matryoshka_dims:
+            if self.config.task == "tsdae":
+                raise ValueError("Matryoshka is not supported with TSDAE (uses decoder architecture)")
+            from sentence_transformers.losses import MatryoshkaLoss
+
+            logger.info(f"Wrapping with MatryoshkaLoss: {self.config.matryoshka_dims}")
+            loss = MatryoshkaLoss(self.model, loss, matryoshka_dims=self.config.matryoshka_dims)
+
+        # 6. Create output directory
+        output_dir = Path(self.config.output.dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        # 7. Training arguments
+        logger.info(f"Optimizer: {self.config.training.optimizer}, Scheduler: {self.config.training.scheduler}")
+        args = SentenceTransformerTrainingArguments(
+            output_dir=str(output_dir),
+            num_train_epochs=self.config.training.epochs,
+            per_device_train_batch_size=self.config.training.batch_size,
+            per_device_eval_batch_size=self.config.training.batch_size,
+            learning_rate=self.config.training.learning_rate,
+            warmup_ratio=self.config.training.warmup_ratio,
+            weight_decay=self.config.training.weight_decay,
+            fp16=self.config.training.fp16,
+            bf16=self.config.training.bf16,
+            optim=self.config.training.optimizer,
+            lr_scheduler_type=self.config.training.scheduler,
+            gradient_checkpointing=self.config.gradient_checkpointing,
+            batch_sampler=self.task.batch_sampler,
+            eval_strategy="steps" if evaluator else "no",
+            eval_steps=self.config.training.eval_steps if evaluator else None,
+            save_strategy="steps",
+            save_steps=self.config.training.save_steps,
+            save_total_limit=self.config.output.save_total_limit,
+            logging_steps=self.config.training.logging_steps,
+            gradient_accumulation_steps=self.config.training.gradient_accumulation_steps,
+            load_best_model_at_end=True if evaluator else False,
+            report_to="tensorboard",
+            logging_dir=str(output_dir / "logs"),
+        )
+
+        # 8. Create trainer
+        callbacks = []
+        if self.progress_callback:
+            callbacks.append(ProgressCallback(self.progress_callback))
+
+        trainer = SentenceTransformerTrainer(
+            model=self.model,
+            args=args,
+            train_dataset=train_data,
+            eval_dataset=eval_data,
+            loss=loss,
+            evaluator=evaluator,
+            callbacks=callbacks if callbacks else None,
+        )
+
+        # 9. Train
+        logger.info("Starting training...")
+        trainer.train()
+
+        # 10. Save final model
+        final_path = output_dir / "final"
+        logger.info(f"Saving model to: {final_path}")
+        self._save_model(final_path)
+
+        # 11. Add label mappings to config.json if task has labels (HuggingFace convention)
+        label_config = self.task.get_label_config()
+        if label_config:
+            config_path = final_path / "config.json"
+            if config_path.exists():
+                with open(config_path) as f:
+                    config = json.load(f)
+                config.update(label_config)
+                with open(config_path, "w") as f:
+                    json.dump(config, f, indent=2)
+                logger.info(f"Added label mappings to config.json ({label_config['num_labels']} labels)")
+
+        # 12. Push to hub if configured (always private)
+        if self.config.output.push_to_hub and self.config.output.hf_username:
+            if not HF_TOKEN_ENV:
+                logger.warning("HF_TOKEN environment variable not set, skipping push to hub")
+            else:
+                # Construct repo name from username and project directory name
+                project_name = Path(self.config.output.dir).name
+                repo_id = f"{self.config.output.hf_username}/{project_name}"
+                logger.info(f"Pushing to HuggingFace Hub (private): {repo_id}")
+                self._push_to_hub(repo_id)
+
+        logger.success("Training completed!")
+        return self.model
+
+    def _save_model(self, path: Path) -> None:
+        """Save the model based on configuration.
+
+        Handles different save methods for standard, LoRA, and Unsloth models.
+        """
+        path_str = str(path)
+
+        # Add vespaembed tag to model card metadata
+        if hasattr(self.model, "model_card_data") and self.model.model_card_data is not None:
+            self.model.model_card_data.add_tags("vespaembed")
+
+        if self.config.unsloth.enabled:
+            save_method = self.config.unsloth.save_method
+
+            if save_method == "lora":
+                # Save only LoRA adapters
+                logger.info("Saving LoRA adapters only")
+                self.model.save_pretrained(path_str)
+            elif save_method == "merged_16bit":
+                # Merge and save as FP16
+                logger.info("Saving merged model (FP16)")
+                self.model.save_pretrained_merged(
+                    path_str,
+                    tokenizer=self.model.tokenizer,
+                    save_method="merged_16bit",
+                )
+            elif save_method == "merged_4bit":
+                # Merge and save as 4-bit
+                logger.info("Saving merged model (4-bit)")
+                self.model.save_pretrained_merged(
+                    path_str,
+                    tokenizer=self.model.tokenizer,
+                    save_method="merged_4bit",
+                )
+        else:
+            # Standard or LoRA (PEFT) save
+            self.model.save_pretrained(path_str)
+
+        # Add vespaembed mention to README.md
+        self._add_vespaembed_to_readme(path)
+
+    def _add_vespaembed_to_readme(self, path: Path) -> None:
+        """Add vespaembed mention to the README.md file.
+
+        This method is idempotent - it will not add duplicate mentions if called multiple times.
+        """
+        readme_path = path / "README.md"
+        if not readme_path.exists():
+            return
+
+        content = readme_path.read_text(encoding="utf-8")
+
+        # Check if vespaembed mention already exists (idempotency)
+        if "github.com/vespaai-playground/vespaembed" in content:
+            return
+
+        # Insert vespaembed mention after the first heading
+        vespaembed_mention = (
+            "\n> This model was trained using " "[vespaembed](https://github.com/vespaai-playground/vespaembed).\n"
+        )
+
+        # Find the first heading and insert after it
+        lines = content.split("\n")
+        new_lines = []
+        inserted = False
+
+        for line in lines:
+            new_lines.append(line)
+            # Insert after the first markdown heading (starting with #)
+            if not inserted and line.startswith("# ") and not line.startswith("# For reference"):
+                new_lines.append(vespaembed_mention)
+                inserted = True
+
+        if inserted:
+            readme_path.write_text("\n".join(new_lines), encoding="utf-8")
+
+    def _push_to_hub(self, repo_id: str) -> None:
+        """Push the model to HuggingFace Hub.
+
+        Handles different push methods for standard, LoRA, and Unsloth models.
+        """
+        if self.config.unsloth.enabled:
+            save_method = self.config.unsloth.save_method
+
+            if save_method == "lora":
+                # Push only LoRA adapters
+                self.model.push_to_hub(repo_id, token=HF_TOKEN_ENV, private=True)
+            else:
+                # Push merged model
+                self.model.push_to_hub_merged(
+                    repo_id,
+                    tokenizer=self.model.tokenizer,
+                    save_method=save_method,
+                    token=HF_TOKEN_ENV,
+                    private=True,
+                )
+        else:
+            # Standard or LoRA (PEFT) save
+            self.model.push_to_hub(repo_id, token=HF_TOKEN_ENV, private=True)
+
+        # Upload the modified README.md with vespaembed mention
+        self._upload_readme_to_hub(repo_id)
+
+    def _upload_readme_to_hub(self, repo_id: str) -> None:
+        """Upload the modified README.md to HuggingFace Hub.
+
+        This uploads the local README.md (which contains the vespaembed mention)
+        to overwrite the auto-generated one on the hub.
+        """
+        from huggingface_hub import HfApi
+
+        # Get the local README.md path from the final output directory
+        output_dir = Path(self.config.output.dir)
+        readme_path = output_dir / "final" / "README.md"
+
+        if not readme_path.exists():
+            logger.warning(f"README.md not found at {readme_path}, skipping upload")
+            return
+
+        api = HfApi()
+        api.upload_file(
+            path_or_fileobj=str(readme_path),
+            path_in_repo="README.md",
+            repo_id=repo_id,
+            token=HF_TOKEN_ENV,
+        )

vespaembed/datasets/__init__.py

@@ -0,0 +1,3 @@
+from vespaembed.datasets.loader import load_dataset
+
+__all__ = ["load_dataset"]

vespaembed/datasets/formats/__init__.py

@@ -0,0 +1,5 @@
+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
+
+__all__ = ["load_csv", "load_hf_dataset", "load_jsonl"]

vespaembed/datasets/formats/csv.py

@@ -0,0 +1,15 @@
+import pandas as pd
+from datasets import Dataset
+
+
+def load_csv(path: str) -> Dataset:
+    """Load a CSV file as a HuggingFace Dataset.
+
+    Args:
+        path: Path to CSV file
+
+    Returns:
+        HuggingFace Dataset
+    """
+    df = pd.read_csv(path)
+    return Dataset.from_pandas(df)

vespaembed/datasets/formats/huggingface.py

@@ -0,0 +1,34 @@
+import os
+from typing import Optional
+
+from datasets import Dataset
+from datasets import load_dataset as hf_load_dataset
+
+# Get HuggingFace token from environment
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+
+def load_hf_dataset(
+    name: str,
+    subset: Optional[str] = None,
+    split: str = "train",
+) -> Dataset:
+    """Load a dataset from the HuggingFace Hub.
+
+    Args:
+        name: Dataset name (e.g., "sentence-transformers/all-nli")
+        subset: Dataset subset/configuration (optional)
+        split: Dataset split (default: "train")
+
+    Returns:
+        HuggingFace Dataset
+
+    Note:
+        Uses HF_TOKEN environment variable for authentication if set.
+    """
+    if subset:
+        dataset = hf_load_dataset(name, subset, split=split, token=HF_TOKEN)
+    else:
+        dataset = hf_load_dataset(name, split=split, token=HF_TOKEN)
+
+    return dataset

vespaembed/datasets/formats/jsonl.py

@@ -0,0 +1,26 @@
+import json
+
+from datasets import Dataset
+
+
+def load_jsonl(path: str) -> Dataset:
+    """Load a JSONL file as a HuggingFace Dataset.
+
+    Args:
+        path: Path to JSONL file
+
+    Returns:
+        HuggingFace Dataset
+    """
+    records = []
+
+    with open(path) as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                records.append(json.loads(line))
+
+    if not records:
+        raise ValueError(f"No records found in {path}")
+
+    return Dataset.from_list(records)