Upload folder using huggingface_hub

src/imrnns/api.py

@@ -7,7 +7,7 @@ from .beir_data import load_beir_source
 from .caching import build_cache
 from .checkpoints import default_checkpoint_name, load_model, save_checkpoint
 from .data import ContrastiveCachedDataset, load_cached_split
-from .encoders import resolve_encoder_spec
+from .encoders import encoder_storage_key, resolve_encoder_spec
 from .evaluation import evaluate_model
 from .model import IMRNN, ModelConfig
 from .training import TrainingConfig, train_model
@@ -124,7 +124,7 @@ def train(
         k_values=[k],
     )
 
-    checkpoint_stem = encoder or encoder_spec.key
+    checkpoint_stem = encoder_storage_key(encoder or encoder_spec.key)
     checkpoint_path = output_dir / default_checkpoint_name(checkpoint_stem, dataset)
     metadata = {
         "encoder": checkpoint_stem,

src/imrnns/assets.py

@@ -5,7 +5,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
 
-from .encoders import normalize_encoder_name
+from .encoders import encoder_storage_key, normalize_encoder_name
 
 
 @dataclass(frozen=True)
@@ -43,10 +43,10 @@ def discover_cached_embeddings(assets_root: Path) -> list[AssetMatch]:
             assets.append(AssetMatch(encoder=encoder, dataset=dataset, path=entry))
             continue
 
-        match = re.fullmatch(r"cache_(mini|minilm|e5|mpnet)_(.+)", entry.name)
+        match = re.fullmatch(r"cache_(.+)_(.+)", entry.name)
         if match:
             encoder, dataset = match.groups()
-            encoder = normalize_encoder_name(encoder)
+            encoder = encoder_storage_key(encoder)
             assets.append(AssetMatch(encoder=encoder, dataset=dataset, path=entry))
     return assets
 
@@ -68,34 +68,39 @@ def discover_repo_checkpoints(repo_root: Path) -> list[AssetMatch]:
     if not base_dir.exists():
         return assets
     for entry in sorted(base_dir.rglob("*.pt")):
-        match = re.fullmatch(r"imrnns-(minilm|e5)-(.+)\.pt", entry.name)
-        if not match:
+        encoder = encoder_storage_key(entry.parent.name)
+        prefix = f"imrnns-{entry.parent.name}-"
+        if not entry.name.startswith(prefix) or not entry.name.endswith(".pt"):
             continue
-        encoder, dataset = match.groups()
-        if encoder == "minilm":
-            encoder = "mini"
+        dataset = entry.name.removeprefix(prefix).removesuffix(".pt")
         assets.append(AssetMatch(encoder=encoder, dataset=dataset, path=entry))
     return assets
 
 
 def resolve_cache_dir(assets_root: Path, encoder: str, dataset: str) -> Path:
-    encoder = normalize_encoder_name(encoder)
+    encoder = encoder_storage_key(encoder)
     dataset = dataset.lower()
     for asset in discover_cached_embeddings(assets_root):
         if asset.encoder == encoder and asset.dataset.lower() == dataset:
             return asset.path
+    direct = assets_root / f"cache_{encoder}_{dataset}"
+    if direct.exists():
+        return direct
     raise FileNotFoundError(
         f"No cached embeddings found for encoder='{encoder}' dataset='{dataset}' under {assets_root}"
     )
 
 
 def resolve_checkpoint_path(assets_root: Path, encoder: str, dataset: str) -> Optional[Path]:
-    encoder = normalize_encoder_name(encoder)
+    encoder = encoder_storage_key(encoder)
     dataset = dataset.lower()
     for asset in discover_repo_checkpoints(package_root()):
         if asset.encoder == encoder and asset.dataset.lower() == dataset:
             return asset.path
     for asset in discover_checkpoints(assets_root):
-        if asset.encoder == encoder and asset.dataset.lower() == dataset:
+        if encoder_storage_key(asset.encoder) == encoder and asset.dataset.lower() == dataset:
             return asset.path
+    direct = assets_root / f"imrnns-{encoder}-{dataset}.pt"
+    if direct.exists():
+        return direct
     return None

src/imrnns/checkpoints.py

@@ -6,14 +6,12 @@ from typing import Any
 
 import torch
 
-from .encoders import normalize_encoder_name
+from .encoders import encoder_storage_key, normalize_encoder_name
 from .model import IMRNN, ModelConfig
 
 
 def default_checkpoint_name(encoder: str, dataset: str) -> str:
-    normalized = normalize_encoder_name(encoder)
-    display = "minilm" if normalized == "mini" else normalized
-    return f"imrnns-{display}-{dataset}.pt"
+    return f"imrnns-{encoder_storage_key(encoder)}-{dataset}.pt"
 
 
 def sanitize_legacy_state_dict(state_dict: dict[str, Any]) -> dict[str, Any]:

src/imrnns/cli.py

@@ -17,7 +17,7 @@ from .beir_data import load_beir_source
 from .caching import build_cache
 from .checkpoints import default_checkpoint_name, load_model, save_checkpoint
 from .data import ContrastiveCachedDataset, load_cached_split
-from .encoders import normalize_encoder_name, resolve_encoder_spec
+from .encoders import encoder_storage_key, normalize_encoder_name, resolve_encoder_spec
 from .evaluation import evaluate_model
 from .model import IMRNN, ModelConfig
 from .training import TrainingConfig, train_model
@@ -25,6 +25,7 @@ from .training import TrainingConfig, train_model
 
 def _add_common_args(parser: argparse.ArgumentParser) -> None:
     parser.add_argument("--assets-root", type=Path, default=default_assets_root())
+    parser.add_argument("--datasets-dir", type=Path)
     parser.add_argument("--encoder")
     parser.add_argument("--encoder-model-name")
     parser.add_argument("--embedding-dim", type=int)
@@ -46,9 +47,8 @@ def _resolve_encoder_spec(args: argparse.Namespace):
 
 def _encoder_label(args: argparse.Namespace, encoder_spec) -> str:
     if args.encoder:
-        normalized = normalize_encoder_name(args.encoder)
-        return "minilm" if normalized == "mini" else normalized
-    return encoder_spec.key.replace("/", "-")
+        return encoder_storage_key(args.encoder)
+    return encoder_storage_key(encoder_spec.key)
 
 
 def _command_list_assets(args: argparse.Namespace) -> int:
@@ -76,7 +76,7 @@ def _load_training_inputs(args: argparse.Namespace):
     encoder_spec = _resolve_encoder_spec(args)
     encoder_label = _encoder_label(args, encoder_spec)
     cache_dir = args.cache_dir or resolve_cache_dir(args.assets_root, encoder_label, args.dataset)
-    datasets_dir = args.assets_root / "datasets"
+    datasets_dir = args.datasets_dir or (args.assets_root / "datasets")
     beir_source = load_beir_source(args.dataset, datasets_dir=datasets_dir, max_queries=args.max_queries)
     train_split = load_cached_split(cache_dir, "train", beir_source, encoder_spec, args.device)
     val_split = load_cached_split(cache_dir, "val", beir_source, encoder_spec, args.device)
@@ -96,7 +96,7 @@ def _command_cache(args: argparse.Namespace) -> int:
         dataset_name=args.dataset,
         encoder_spec=encoder_spec,
         cache_dir=cache_dir,
-        datasets_dir=args.assets_root / "datasets",
+        datasets_dir=args.datasets_dir or (args.assets_root / "datasets"),
         device=args.device,
         batch_size=args.batch_size,
         num_negatives=args.num_negatives,
@@ -194,7 +194,7 @@ def _command_evaluate(args: argparse.Namespace) -> int:
             f"No checkpoint found for encoder='{encoder_label}' dataset='{args.dataset}'. Provide --checkpoint."
         )
 
-    datasets_dir = args.assets_root / "datasets"
+    datasets_dir = args.datasets_dir or (args.assets_root / "datasets")
     beir_source = load_beir_source(args.dataset, datasets_dir=datasets_dir, max_queries=args.max_queries)
     test_split = load_cached_split(cache_dir, "test", beir_source, encoder_spec, args.device)
     model, metadata, missing, unexpected = load_model(

src/imrnns/encoders.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 from typing import Optional
+import re
 
 
 @dataclass(frozen=True)
@@ -47,6 +48,13 @@ def normalize_encoder_name(name: str) -> str:
     return aliases.get(key, key)
 
 
+def encoder_storage_key(name: str) -> str:
+    normalized = normalize_encoder_name(name)
+    if normalized == "mini":
+        return "minilm"
+    return re.sub(r"[^a-z0-9._-]+", "-", normalized.lower()).strip("-")
+
+
 def get_encoder_spec(name: str) -> EncoderSpec:
     key = normalize_encoder_name(name)
     if key not in ENCODER_SPECS: