add custom models

.gitignore

@@ -205,3 +205,6 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+
+# Embedding Bench custom models
+custom_models.json

app.py

@@ -14,9 +14,18 @@ from corpus import build_corpus
 from dataset_config import DATASET_PRESETS, DatasetConfig
 from evals.quality import evaluate_quality
 from evals.speed import evaluate_speed
-from models import REGISTRY, ModelConfig
+from models import (
+    REGISTRY,
+    VALID_BACKENDS,
+    ModelConfig,
+    load_custom_models_from_file,
+    register_model,
+    save_custom_model_to_file,
+)
 from wrapper import load_model
 
+load_custom_models_from_file()
+
 # ---------------------------------------------------------------------------
 # Page config & custom CSS
 # ---------------------------------------------------------------------------
@@ -114,6 +123,40 @@ selected_models = st.sidebar.multiselect(
     label_visibility="collapsed",
 )
 
+with st.sidebar.expander("➕ Add Custom Model"):
+    with st.form("add_model_form", clear_on_submit=True):
+        new_key = st.text_input("Registry key", placeholder="my-model")
+        new_name = st.text_input("Display name", placeholder="My Custom Model")
+        new_model_id = st.text_input("HuggingFace model ID", placeholder="org/model-name")
+        new_backend = st.selectbox("Backend", sorted(VALID_BACKENDS))
+        new_gguf_file = st.text_input(
+            "GGUF filename (gguf backend only)", value="", placeholder="model.gguf"
+        )
+        new_is_baseline = st.checkbox("Mark as baseline", value=False)
+        new_persist = st.checkbox("Save to disk", value=False,
+                                  help="Persist to custom_models.json so it loads next session")
+        submitted = st.form_submit_button("Add Model", use_container_width=True)
+    if submitted:
+        if not new_key or not new_name or not new_model_id:
+            st.sidebar.error("Key, name, and model ID are required.")
+        elif new_backend == "gguf" and not new_gguf_file:
+            st.sidebar.error("GGUF filename is required for gguf backend.")
+        else:
+            cfg = ModelConfig(
+                name=new_name,
+                model_id=new_model_id,
+                is_baseline=new_is_baseline,
+                backend=new_backend,
+                gguf_file=new_gguf_file or None,
+            )
+            try:
+                register_model(new_key, cfg)
+                if new_persist:
+                    save_custom_model_to_file(new_key, cfg)
+                st.rerun()
+            except ValueError as e:
+                st.sidebar.error(str(e))
+
 st.sidebar.markdown("**Datasets**")
 available_datasets = list(DATASET_PRESETS.keys())
 selected_datasets = st.sidebar.multiselect(

bench.py

@@ -5,7 +5,7 @@ import argparse
 from corpus import build_corpus
 from dataset_config import DATASET_PRESETS, DatasetConfig
 from evals import evaluate_memory, evaluate_quality, evaluate_speed
-from models import REGISTRY
+from models import REGISTRY, ModelConfig, load_custom_models_from_file, register_model
 from report import print_report
 from wrapper import load_model
 
@@ -18,9 +18,15 @@ def main(argv: list[str] | None = None) -> None:
     parser.add_argument(
         "--models",
         nargs="+",
-        default=list(REGISTRY.keys()),
-        choices=list(REGISTRY.keys()),
-        help="Models to benchmark (default: all)",
+        default=None,
+        help="Models to benchmark (default: all registered)",
+    )
+    parser.add_argument(
+        "--add-model",
+        action="append",
+        default=[],
+        metavar="KEY:NAME:MODEL_ID:BACKEND[:GGUF_FILE]",
+        help="Register a custom model. Can be repeated.",
     )
     parser.add_argument("--corpus-size", type=int, default=1000)
     parser.add_argument("--batch-size", type=int, default=64)
@@ -61,6 +67,28 @@ def main(argv: list[str] | None = None) -> None:
 
     args = parser.parse_args(argv)
 
+    # Load persisted custom models and register any --add-model entries
+    load_custom_models_from_file()
+    for spec in args.add_model:
+        parts = spec.split(":")
+        if len(parts) < 4:
+            parser.error(f"--add-model requires KEY:NAME:MODEL_ID:BACKEND, got: {spec}")
+        key, name, model_id, backend = parts[0], parts[1], parts[2], parts[3]
+        gguf_file = parts[4] if len(parts) > 4 else None
+        try:
+            register_model(key, ModelConfig(
+                name=name, model_id=model_id, backend=backend, gguf_file=gguf_file,
+            ))
+        except ValueError as e:
+            parser.error(str(e))
+
+    if args.models is None:
+        args.models = list(REGISTRY.keys())
+    else:
+        for k in args.models:
+            if k not in REGISTRY:
+                parser.error(f"Unknown model key: '{k}'. Available: {list(REGISTRY.keys())}")
+
     # Build list of dataset configs
     if args.dataset:
         # Custom dataset overrides presets

models.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
 
 
 @dataclass
@@ -43,3 +45,34 @@ REGISTRY: dict[str, ModelConfig] = {
     #     backend="libembedding",
     # ),
 }
+
+VALID_BACKENDS = {"sbert", "fastembed", "libembedding", "gguf"}
+CUSTOM_MODELS_PATH = Path(__file__).parent / "custom_models.json"
+
+
+def register_model(key: str, config: ModelConfig) -> None:
+    if key in REGISTRY:
+        raise ValueError(f"Model key '{key}' already exists in registry")
+    if config.backend not in VALID_BACKENDS:
+        raise ValueError(f"Invalid backend '{config.backend}'. Must be one of: {VALID_BACKENDS}")
+    REGISTRY[key] = config
+
+
+def load_custom_models_from_file(path: Path = CUSTOM_MODELS_PATH) -> None:
+    if not path.exists():
+        return
+    with open(path) as f:
+        entries = json.load(f)
+    for key, fields in entries.items():
+        if key not in REGISTRY:
+            REGISTRY[key] = ModelConfig(**fields)
+
+
+def save_custom_model_to_file(key: str, config: ModelConfig, path: Path = CUSTOM_MODELS_PATH) -> None:
+    existing: dict = {}
+    if path.exists():
+        with open(path) as f:
+            existing = json.load(f)
+    existing[key] = asdict(config)
+    with open(path, "w") as f:
+        json.dump(existing, f, indent=2)