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

vespaembed/__init__.py

@@ -1 +1 @@
-__version__ = "0.0.1"
+__version__ = "0.0.2"

vespaembed/cli/__init__.py

@@ -0,0 +1,17 @@
+from abc import ABC, abstractmethod
+from argparse import ArgumentParser
+
+
+class BaseCommand(ABC):
+    """Base class for all CLI commands."""
+
+    @staticmethod
+    @abstractmethod
+    def register_subcommand(parser: ArgumentParser):
+        """Register the subcommand with argparse."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def execute(self):
+        """Execute the command."""
+        raise NotImplementedError

vespaembed/cli/commands/__init__.py

@@ -0,0 +1,7 @@
+from vespaembed.cli.commands.evaluate import EvaluateCommand
+from vespaembed.cli.commands.export import ExportCommand
+from vespaembed.cli.commands.info import InfoCommand
+from vespaembed.cli.commands.serve import ServeCommand
+from vespaembed.cli.commands.train import TrainCommand
+
+__all__ = ["EvaluateCommand", "ExportCommand", "InfoCommand", "ServeCommand", "TrainCommand"]

vespaembed/cli/commands/evaluate.py

@@ -0,0 +1,85 @@
+from argparse import ArgumentParser, Namespace
+
+from vespaembed.cli import BaseCommand
+from vespaembed.core.registry import Registry
+from vespaembed.datasets.loader import load_dataset
+from vespaembed.models.loader import load_model
+from vespaembed.utils.logging import logger
+
+
+def evaluate_command_factory(args: Namespace) -> "EvaluateCommand":
+    """Factory function for EvaluateCommand."""
+    return EvaluateCommand(
+        model_path=args.model,
+        data_path=args.data,
+        task=args.task,
+    )
+
+
+class EvaluateCommand(BaseCommand):
+    """Evaluate a trained embedding model."""
+
+    @staticmethod
+    def register_subcommand(parser: ArgumentParser):
+        """Register the evaluate subcommand."""
+        eval_parser = parser.add_parser(
+            "evaluate",
+            help="Evaluate a trained embedding model",
+        )
+
+        eval_parser.add_argument(
+            "--model",
+            type=str,
+            required=True,
+            help="Path to trained model",
+        )
+        eval_parser.add_argument(
+            "--data",
+            type=str,
+            required=True,
+            help="Path to evaluation data",
+        )
+        eval_parser.add_argument(
+            "--task",
+            type=str,
+            required=True,
+            choices=["mnr", "triplet", "contrastive", "sts", "nli", "tsdae", "matryoshka"],
+            help="Task type (determines evaluator)",
+        )
+
+        eval_parser.set_defaults(func=evaluate_command_factory)
+
+    def __init__(self, model_path: str, data_path: str, task: str):
+        self.model_path = model_path
+        self.data_path = data_path
+        self.task = task
+
+    def execute(self):
+        """Execute the evaluate command."""
+        # Load model
+        logger.info(f"Loading model: {self.model_path}")
+        model = load_model(self.model_path)
+
+        # Load and prepare data
+        logger.info(f"Loading evaluation data: {self.data_path}")
+        task_cls = Registry.get_task(self.task)
+        task = task_cls()
+
+        eval_data = load_dataset(self.data_path)
+        eval_data = task.prepare_dataset(eval_data)
+
+        # Create evaluator
+        evaluator = task.get_evaluator(eval_data)
+
+        if evaluator is None:
+            logger.warning(f"No evaluator available for task: {self.task}")
+            return
+
+        # Run evaluation
+        logger.info("Running evaluation...")
+        results = evaluator(model)
+
+        # Print results
+        logger.success("Evaluation Results:")
+        for key, value in results.items():
+            logger.print(f"  {key}: {value:.4f}")

vespaembed/cli/commands/export.py

@@ -0,0 +1,86 @@
+from argparse import ArgumentParser, Namespace
+
+from vespaembed.cli import BaseCommand
+from vespaembed.models.export import export_model, push_to_hub
+from vespaembed.models.loader import load_model
+from vespaembed.utils.logging import logger
+
+
+def export_command_factory(args: Namespace) -> "ExportCommand":
+    """Factory function for ExportCommand."""
+    return ExportCommand(
+        model_path=args.model,
+        output_path=args.output,
+        format=args.format,
+        hub_id=args.hub_id,
+    )
+
+
+class ExportCommand(BaseCommand):
+    """Export a trained model to different formats."""
+
+    @staticmethod
+    def register_subcommand(parser: ArgumentParser):
+        """Register the export subcommand."""
+        export_parser = parser.add_parser(
+            "export",
+            help="Export a trained model to different formats",
+        )
+
+        export_parser.add_argument(
+            "--model",
+            type=str,
+            required=True,
+            help="Path to trained model",
+        )
+        export_parser.add_argument(
+            "--output",
+            type=str,
+            default=None,
+            help="Output path for exported model",
+        )
+        export_parser.add_argument(
+            "--format",
+            type=str,
+            default="onnx",
+            choices=["onnx"],
+            help="Export format (default: onnx)",
+        )
+        export_parser.add_argument(
+            "--hub-id",
+            type=str,
+            default=None,
+            help="Push to HuggingFace Hub with this repo ID",
+        )
+
+        export_parser.set_defaults(func=export_command_factory)
+
+    def __init__(
+        self,
+        model_path: str,
+        output_path: str = None,
+        format: str = "onnx",
+        hub_id: str = None,
+    ):
+        self.model_path = model_path
+        self.output_path = output_path or f"{model_path}_exported"
+        self.format = format
+        self.hub_id = hub_id
+
+    def execute(self):
+        """Execute the export command."""
+        # Load model
+        logger.info(f"Loading model: {self.model_path}")
+        model = load_model(self.model_path)
+
+        # Export
+        if self.format:
+            logger.info(f"Exporting to {self.format}: {self.output_path}")
+            export_path = export_model(model, self.output_path, self.format)
+            logger.success(f"Model exported to: {export_path}")
+
+        # Push to Hub
+        if self.hub_id:
+            logger.info(f"Pushing to HuggingFace Hub: {self.hub_id}")
+            url = push_to_hub(model, self.hub_id)
+            logger.success(f"Model pushed to: {url}")

vespaembed/cli/commands/info.py

@@ -0,0 +1,52 @@
+from argparse import ArgumentParser, Namespace
+
+# Import tasks to register them
+import vespaembed.tasks  # noqa: F401
+from vespaembed.cli import BaseCommand
+from vespaembed.core.registry import Registry
+from vespaembed.utils.logging import logger
+
+
+def info_command_factory(args: Namespace) -> "InfoCommand":
+    """Factory function for InfoCommand."""
+    return InfoCommand(show_tasks=args.tasks)
+
+
+class InfoCommand(BaseCommand):
+    """Show information about available tasks."""
+
+    @staticmethod
+    def register_subcommand(parser: ArgumentParser):
+        """Register the info subcommand."""
+        info_parser = parser.add_parser(
+            "info",
+            help="Show information about available tasks",
+        )
+
+        info_parser.add_argument(
+            "--tasks",
+            action="store_true",
+            help="List available tasks",
+        )
+
+        info_parser.set_defaults(func=info_command_factory)
+
+    def __init__(self, show_tasks: bool = True):
+        self.show_tasks = show_tasks
+
+    def execute(self):
+        """Execute the info command."""
+        if self.show_tasks or True:  # Default to showing tasks
+            self._show_tasks()
+
+    def _show_tasks(self):
+        """Display available tasks."""
+        logger.print("\n[bold]Available Tasks:[/bold]\n")
+
+        tasks_info = Registry.get_task_info()
+
+        for task in tasks_info:
+            logger.print(f"  [cyan]{task['name']}[/cyan]")
+            logger.print(f"    {task['description']}")
+            logger.print(f"    Expected columns: {', '.join(task['expected_columns'])}")
+            logger.print("")

vespaembed/cli/commands/serve.py

@@ -0,0 +1,49 @@
+from argparse import ArgumentParser, Namespace
+
+from vespaembed.cli import BaseCommand
+from vespaembed.utils.logging import logger
+
+
+def serve_command_factory(args: Namespace) -> "ServeCommand":
+    """Factory function for ServeCommand."""
+    return ServeCommand(host=args.host, port=args.port)
+
+
+class ServeCommand(BaseCommand):
+    """Start the web UI server."""
+
+    @staticmethod
+    def register_subcommand(parser: ArgumentParser):
+        """Register the serve subcommand."""
+        serve_parser = parser.add_parser(
+            "serve",
+            help="Start the web UI server",
+        )
+
+        serve_parser.add_argument(
+            "--host",
+            type=str,
+            default="127.0.0.1",
+            help="Host to bind to (default: 127.0.0.1)",
+        )
+        serve_parser.add_argument(
+            "--port",
+            type=int,
+            default=8000,
+            help="Port to bind to (default: 8000)",
+        )
+
+        serve_parser.set_defaults(func=serve_command_factory)
+
+    def __init__(self, host: str = "127.0.0.1", port: int = 8000):
+        self.host = host
+        self.port = port
+
+    def execute(self):
+        """Execute the serve command."""
+        import uvicorn
+
+        from vespaembed.web.app import app
+
+        logger.info(f"Starting VespaEmbed web UI at http://{self.host}:{self.port}")
+        uvicorn.run(app, host=self.host, port=self.port)

vespaembed/cli/commands/train.py

@@ -0,0 +1,267 @@
+import time
+from argparse import ArgumentParser, Namespace
+from pathlib import Path
+from typing import Optional
+
+# Import tasks to register them
+import vespaembed.tasks  # noqa: F401
+from vespaembed.cli import BaseCommand
+from vespaembed.core.config import TrainingConfig, load_config_from_yaml
+from vespaembed.core.trainer import VespaEmbedTrainer
+from vespaembed.utils.logging import logger
+
+# Projects directory
+PROJECTS_DIR = Path.home() / ".vespaembed" / "projects"
+
+
+def train_command_factory(args: Namespace) -> "TrainCommand":
+    """Factory function for TrainCommand."""
+    return TrainCommand(
+        config_path=args.config,
+        data=args.data,
+        task=args.task,
+        base_model=args.base_model,
+        project=args.project,
+        eval_data=args.eval_data,
+        epochs=args.epochs,
+        batch_size=args.batch_size,
+        learning_rate=args.learning_rate,
+        optimizer=args.optimizer,
+        scheduler=args.scheduler,
+        unsloth=args.unsloth,
+        matryoshka=args.matryoshka,
+        matryoshka_dims=args.matryoshka_dims,
+        subset=args.subset,
+        split=args.split,
+    )
+
+
+class TrainCommand(BaseCommand):
+    """Train or fine-tune an embedding model."""
+
+    @staticmethod
+    def register_subcommand(parser: ArgumentParser):
+        """Register the train subcommand."""
+        train_parser = parser.add_parser(
+            "train",
+            help="Train or fine-tune an embedding model",
+        )
+
+        # Config file (alternative to CLI args)
+        train_parser.add_argument(
+            "--config",
+            type=str,
+            default=None,
+            help="Path to YAML configuration file",
+        )
+
+        # Required arguments (unless using config file)
+        train_parser.add_argument(
+            "--data",
+            type=str,
+            default=None,
+            help="Path to training data (CSV, JSONL, or HF dataset)",
+        )
+        train_parser.add_argument(
+            "--task",
+            type=str,
+            default=None,
+            choices=["pairs", "triplets", "similarity", "tsdae"],
+            help="Training task type",
+        )
+        train_parser.add_argument(
+            "--base-model",
+            type=str,
+            default=None,
+            help="Base model name or path",
+        )
+
+        # Optional arguments
+        train_parser.add_argument(
+            "--project",
+            type=str,
+            default=None,
+            help="Project name. Output saved to ~/.vespaembed/projects/<name>/",
+        )
+        train_parser.add_argument(
+            "--eval-data",
+            type=str,
+            default=None,
+            help="Path to evaluation data",
+        )
+        train_parser.add_argument(
+            "--epochs",
+            type=int,
+            default=3,
+            help="Number of training epochs (default: 3)",
+        )
+        train_parser.add_argument(
+            "--batch-size",
+            type=int,
+            default=32,
+            help="Batch size (default: 32)",
+        )
+        train_parser.add_argument(
+            "--learning-rate",
+            type=float,
+            default=2e-5,
+            help="Learning rate (default: 2e-5)",
+        )
+        train_parser.add_argument(
+            "--optimizer",
+            type=str,
+            default="adamw_torch",
+            choices=["adamw_torch", "adamw_torch_fused", "adamw_8bit", "adafactor", "sgd"],
+            help="Optimizer type (default: adamw_torch)",
+        )
+        train_parser.add_argument(
+            "--scheduler",
+            type=str,
+            default="linear",
+            choices=["linear", "cosine", "cosine_with_restarts", "constant", "constant_with_warmup", "polynomial"],
+            help="Learning rate scheduler (default: linear)",
+        )
+        train_parser.add_argument(
+            "--unsloth",
+            action="store_true",
+            help="Use Unsloth for faster training",
+        )
+        train_parser.add_argument(
+            "--matryoshka",
+            action="store_true",
+            help="Enable Matryoshka embeddings (multi-dimensional)",
+        )
+        train_parser.add_argument(
+            "--matryoshka-dims",
+            type=str,
+            default="768,512,256,128,64",
+            help="Matryoshka dimensions, comma-separated (default: 768,512,256,128,64)",
+        )
+
+        # HuggingFace dataset options
+        train_parser.add_argument(
+            "--subset",
+            type=str,
+            default=None,
+            help="HuggingFace dataset subset",
+        )
+        train_parser.add_argument(
+            "--split",
+            type=str,
+            default=None,
+            help="HuggingFace dataset split",
+        )
+
+        train_parser.set_defaults(func=train_command_factory)
+
+    def __init__(
+        self,
+        config_path: Optional[str] = None,
+        data: Optional[str] = None,
+        task: Optional[str] = None,
+        base_model: Optional[str] = None,
+        project: Optional[str] = None,
+        eval_data: Optional[str] = None,
+        epochs: int = 3,
+        batch_size: int = 32,
+        learning_rate: float = 2e-5,
+        optimizer: str = "adamw_torch",
+        scheduler: str = "linear",
+        unsloth: bool = False,
+        matryoshka: bool = False,
+        matryoshka_dims: str = "768,512,256,128,64",
+        subset: Optional[str] = None,
+        split: Optional[str] = None,
+    ):
+        self.config_path = config_path
+        self.data = data
+        self.task = task
+        self.base_model = base_model
+        self.project = project
+        self.eval_data = eval_data
+        self.epochs = epochs
+        self.batch_size = batch_size
+        self.learning_rate = learning_rate
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.unsloth = unsloth
+        self.matryoshka = matryoshka
+        self.matryoshka_dims = matryoshka_dims
+        self.subset = subset
+        self.split = split
+
+    def _generate_project_name(self) -> str:
+        """Generate a random project name."""
+        import random
+        import string
+
+        return "".join(random.choices(string.ascii_lowercase + string.digits, k=8))
+
+    def _resolve_output_dir(self, project_name: str) -> Path:
+        """Resolve output directory from project name."""
+        PROJECTS_DIR.mkdir(parents=True, exist_ok=True)
+        output_dir = PROJECTS_DIR / project_name
+        if output_dir.exists():
+            # Append timestamp to make unique
+            timestamp = int(time.time())
+            output_dir = PROJECTS_DIR / f"{project_name}-{timestamp}"
+        output_dir.mkdir(parents=True, exist_ok=True)
+        return output_dir
+
+    def execute(self):
+        """Execute the train command."""
+        # Load config from file or build from CLI args
+        if self.config_path:
+            logger.info(f"Loading config from: {self.config_path}")
+            config = load_config_from_yaml(self.config_path)
+        else:
+            # Validate required arguments
+            if not self.data:
+                raise ValueError("--data is required (or use --config)")
+            if not self.task:
+                raise ValueError("--task is required (or use --config)")
+            if not self.base_model:
+                raise ValueError("--base-model is required (or use --config)")
+
+            # Generate project name if not provided
+            project_name = self.project or self._generate_project_name()
+            output_dir = self._resolve_output_dir(project_name)
+
+            logger.info(f"Project: {project_name}")
+            logger.info(f"Output: {output_dir}")
+
+            # Parse matryoshka dimensions if enabled
+            matryoshka_dims = None
+            if self.matryoshka:
+                if self.task == "tsdae":
+                    raise ValueError("Matryoshka is not supported with TSDAE task")
+                matryoshka_dims = [int(d.strip()) for d in self.matryoshka_dims.split(",") if d.strip()]
+                logger.info(f"Matryoshka enabled with dimensions: {matryoshka_dims}")
+
+            # Build config from CLI args
+            config = TrainingConfig(
+                task=self.task,
+                base_model=self.base_model,
+                data={
+                    "train": self.data,
+                    "eval": self.eval_data,
+                    "subset": self.subset,
+                    "split": self.split,
+                },
+                training={
+                    "epochs": self.epochs,
+                    "batch_size": self.batch_size,
+                    "learning_rate": self.learning_rate,
+                    "optimizer": self.optimizer,
+                    "scheduler": self.scheduler,
+                },
+                output={
+                    "dir": str(output_dir),
+                },
+                unsloth=self.unsloth,
+                matryoshka_dims=matryoshka_dims,
+            )
+
+        # Create and run trainer
+        trainer = VespaEmbedTrainer(config)
+        trainer.train()

vespaembed/cli/vespaembed.py

@@ -0,0 +1,55 @@
+import argparse
+
+from vespaembed.cli.commands.evaluate import EvaluateCommand
+from vespaembed.cli.commands.export import ExportCommand
+from vespaembed.cli.commands.info import InfoCommand
+from vespaembed.cli.commands.serve import ServeCommand
+from vespaembed.cli.commands.train import TrainCommand
+
+
+def main():
+    """Main entry point for vespaembed CLI."""
+    parser = argparse.ArgumentParser(
+        prog="vespaembed",
+        description="VespaEmbed - No-code training for embedding models",
+        usage="vespaembed [<command>] [<args>]",
+    )
+
+    # Global arguments
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="127.0.0.1",
+        help="Host for web UI (default: 127.0.0.1)",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8000,
+        help="Port for web UI (default: 8000)",
+    )
+
+    # Subcommands
+    commands_parser = parser.add_subparsers(dest="command")
+
+    # Register all commands
+    TrainCommand.register_subcommand(commands_parser)
+    EvaluateCommand.register_subcommand(commands_parser)
+    ExportCommand.register_subcommand(commands_parser)
+    ServeCommand.register_subcommand(commands_parser)
+    InfoCommand.register_subcommand(commands_parser)
+
+    args = parser.parse_args()
+
+    # If no command specified, launch web UI
+    if not hasattr(args, "func") or args.func is None:
+        command = ServeCommand(host=args.host, port=args.port)
+        command.execute()
+    else:
+        # Execute the specified command
+        command = args.func(args)
+        command.execute()
+
+
+if __name__ == "__main__":
+    main()

vespaembed/core/__init__.py

@@ -0,0 +1,2 @@
+# Core module - imports are lazy to avoid requiring heavy dependencies
+# Use explicit imports: from vespaembed.core.config import TrainingConfig