rat-embed 0.1.0__tar.gz

rat_embed-0.1.0/.gitignore

@@ -0,0 +1,14 @@
+# Python
+__pycache__/
+*.py[cod]
+.venv/
+*.egg-info/
+
+# Data (generated at runtime)
+data/*.npy
+data/*.json
+
+# OS
+.DS_Store
+*.Zone.Identifier
+Thumbs.db

rat_embed-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 sojir
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

rat_embed-0.1.0/PKG-INFO

@@ -0,0 +1,99 @@
+Metadata-Version: 2.4
+Name: rat-embed
+Version: 0.1.0
+Summary: Relative Anchor Translation: zero-shot embedding space translation via similarity profiles
+Project-URL: Homepage, https://github.com/jiro-prog/rat-experiment
+Project-URL: Repository, https://github.com/jiro-prog/rat-experiment
+Project-URL: Paper, https://zenodo.org/records/19401277
+Project-URL: Issues, https://github.com/jiro-prog/rat-experiment/issues
+Author: So Jiro
+License-Expression: MIT
+License-File: LICENSE
+Keywords: cross-model retrieval,embedding,model translation,relative representation,zero-shot
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: numpy>=1.24
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Requires-Dist: ruff>=0.1; extra == 'dev'
+Provides-Extra: models
+Requires-Dist: sentence-transformers>=2.2; extra == 'models'
+Description-Content-Type: text/markdown
+
+# RAT — Relative Anchor Translation
+
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.19401277.svg)](https://zenodo.org/records/19401277)
+
+Zero-shot embedding space translation using relative distances to shared anchors. No additional training required.
+
+## Install
+
+```bash
+pip install -e .            # core (numpy only)
+pip install -e ".[models]"  # + sentence-transformers for fit()
+pip install -e ".[dev]"     # + pytest, ruff
+```
+
+## Quick Start
+
+```python
+import numpy as np
+from rat import RATranslator
+
+# 1. Prepare anchor embeddings from both models (K anchors, L2-normalized)
+anchor_a = ...  # (K, D_a) from model A
+anchor_b = ...  # (K, D_b) from model B
+
+# 2. Fit the translator
+translator = RATranslator(kernel="poly").fit_embeddings(anchor_a, anchor_b)
+
+# 3. Transform & retrieve
+query_emb = ...  # (N, D_a) from model A
+db_emb = ...     # (M, D_b) from model B
+results = translator.retrieve(query_emb, db_emb, top_k=10)
+# results["indices"]  → (N, 10) nearest neighbor indices
+# results["scores"]   → (N, 10) cosine similarity scores
+```
+
+Or transform individually for more control:
+
+```python
+q_rel = translator.transform(query_emb, "a")              # query side (no z-score)
+d_rel = translator.transform(db_emb, "b")                 # db side (z-score applied)
+d_rel = translator.transform(db_emb, "b", role="query")   # override: skip z-score
+```
+
+## Advanced: RATHub (multi-model)
+
+```python
+from rat import RATHub
+
+hub = RATHub(kernel="poly")
+hub.set_anchors("minilm", anchor_minilm)      # (K, 384)
+hub.set_anchors("e5", anchor_e5)              # (K, 1024)
+hub.set_anchors("bge", anchor_bge)            # (K, 384)
+
+# Transform from any model
+q = hub.transform("minilm", query_emb, role="query")
+d = hub.transform("e5", db_emb, role="db")
+
+# Or use retrieve directly
+results = hub.retrieve(query_emb, db_emb, "minilm", "e5", top_k=10)
+```
+
+## Paper
+
+See the [Zenodo record](https://zenodo.org/records/19401277) for the full experiment report.
+
+Experiment reproduction code is in `experiments/` (unchanged from the original research).
+
+## License
+
+MIT

rat_embed-0.1.0/README.md

@@ -0,0 +1,70 @@
+# RAT — Relative Anchor Translation
+
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.19401277.svg)](https://zenodo.org/records/19401277)
+
+Zero-shot embedding space translation using relative distances to shared anchors. No additional training required.
+
+## Install
+
+```bash
+pip install -e .            # core (numpy only)
+pip install -e ".[models]"  # + sentence-transformers for fit()
+pip install -e ".[dev]"     # + pytest, ruff
+```
+
+## Quick Start
+
+```python
+import numpy as np
+from rat import RATranslator
+
+# 1. Prepare anchor embeddings from both models (K anchors, L2-normalized)
+anchor_a = ...  # (K, D_a) from model A
+anchor_b = ...  # (K, D_b) from model B
+
+# 2. Fit the translator
+translator = RATranslator(kernel="poly").fit_embeddings(anchor_a, anchor_b)
+
+# 3. Transform & retrieve
+query_emb = ...  # (N, D_a) from model A
+db_emb = ...     # (M, D_b) from model B
+results = translator.retrieve(query_emb, db_emb, top_k=10)
+# results["indices"]  → (N, 10) nearest neighbor indices
+# results["scores"]   → (N, 10) cosine similarity scores
+```
+
+Or transform individually for more control:
+
+```python
+q_rel = translator.transform(query_emb, "a")              # query side (no z-score)
+d_rel = translator.transform(db_emb, "b")                 # db side (z-score applied)
+d_rel = translator.transform(db_emb, "b", role="query")   # override: skip z-score
+```
+
+## Advanced: RATHub (multi-model)
+
+```python
+from rat import RATHub
+
+hub = RATHub(kernel="poly")
+hub.set_anchors("minilm", anchor_minilm)      # (K, 384)
+hub.set_anchors("e5", anchor_e5)              # (K, 1024)
+hub.set_anchors("bge", anchor_bge)            # (K, 384)
+
+# Transform from any model
+q = hub.transform("minilm", query_emb, role="query")
+d = hub.transform("e5", db_emb, role="db")
+
+# Or use retrieve directly
+results = hub.retrieve(query_emb, db_emb, "minilm", "e5", top_k=10)
+```
+
+## Paper
+
+See the [Zenodo record](https://zenodo.org/records/19401277) for the full experiment report.
+
+Experiment reproduction code is in `experiments/` (unchanged from the original research).
+
+## License
+
+MIT

rat_embed-0.1.0/pyproject.toml

@@ -0,0 +1,70 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "rat-embed"
+version = "0.1.0"
+description = "Relative Anchor Translation: zero-shot embedding space translation via similarity profiles"
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.10"
+authors = [
+    { name = "So Jiro" },
+]
+keywords = [
+    "embedding",
+    "zero-shot",
+    "model translation",
+    "relative representation",
+    "cross-model retrieval",
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "numpy>=1.24",
+]
+
+[project.optional-dependencies]
+models = ["sentence-transformers>=2.2"]
+dev = ["pytest>=7.0", "ruff>=0.1"]
+
+[project.urls]
+Homepage = "https://github.com/jiro-prog/rat-experiment"
+Repository = "https://github.com/jiro-prog/rat-experiment"
+Paper = "https://zenodo.org/records/19401277"
+Issues = "https://github.com/jiro-prog/rat-experiment/issues"
+
+[tool.hatch.build.targets.sdist]
+exclude = [
+    "data/",
+    "experiments/",
+    "examples/",
+    "results/",
+    "figures/",
+    "paper/",
+    "scripts/",
+    "src/",
+    "*.npy",
+    "config.py",
+    "requirements.txt",
+    "RAT-phase0-experiment.md",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["rat"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.ruff]
+target-version = "py310"
+line-length = 100

rat_embed-0.1.0/rat/__init__.py

@@ -0,0 +1,7 @@
+"""RAT: Relative Anchor Translation for zero-shot embedding space translation."""
+
+from rat.translator import RATranslator
+from rat.hub import RATHub
+
+__all__ = ["RATranslator", "RATHub"]
+__version__ = "0.1.0"

rat_embed-0.1.0/rat/hub.py

@@ -0,0 +1,293 @@
+"""RATHub: multi-model relative anchor translation."""
+
+from __future__ import annotations
+
+import json
+import warnings
+
+import numpy as np
+
+from rat.kernels import get_kernel
+from rat.normalize import compute_sim_mean, normalize_zscore, recommend_zscore
+
+
+class RATHub:
+    """Manage N embedding models and translate between any pair via RAT.
+
+    Each model's anchor embeddings are registered once. After that,
+    any model's embeddings can be transformed to the shared relative space.
+    """
+
+    def __init__(
+        self,
+        kernel: str = "poly",
+        kernel_params: dict | None = None,
+        normalize: str = "auto",
+        normalize_threshold: float = 0.15,
+        normalize_harmful_threshold: float = 0.65,
+        verbose: bool = False,
+    ):
+        self._kernel_name = kernel
+        self._kernel_params = kernel_params or {}
+        self._kernel_fn = get_kernel(kernel, kernel_params)
+        self._normalize = normalize
+        self._normalize_threshold = normalize_threshold
+        self._normalize_harmful_threshold = normalize_harmful_threshold
+        self._verbose = verbose
+
+        # Per-model state: {model_name: {...}}
+        self._models: dict[str, dict] = {}
+        self._anchor_k: int | None = None  # enforced anchor count
+
+    # ── Registration ──
+
+    def set_anchors(
+        self,
+        model_name: str,
+        anchor_embeddings: np.ndarray,
+    ) -> None:
+        """Register anchor embeddings for a model.
+
+        Parameters
+        ----------
+        model_name : identifier for this model
+        anchor_embeddings : (K, D) L2-normalized
+        """
+        _check_l2_normalized(anchor_embeddings)
+
+        K = anchor_embeddings.shape[0]
+        if self._anchor_k is None:
+            self._anchor_k = K
+        elif K != self._anchor_k:
+            raise ValueError(
+                f"Anchor count mismatch: expected {self._anchor_k}, got {K} for '{model_name}'"
+            )
+
+        sim_mean = compute_sim_mean(anchor_embeddings)
+
+        self._models[model_name] = {
+            "anchors": anchor_embeddings,
+            "sim_mean": sim_mean,
+        }
+
+        if self._verbose:
+            rec = recommend_zscore(sim_mean, self._normalize_threshold, self._normalize_harmful_threshold)
+            print(f"[RAT] {model_name} sim_mean={sim_mean:.3f} → z-score: {rec}")
+            print(
+                f"[RAT] Anchor shape: {model_name}="
+                f"({anchor_embeddings.shape[0]}, {anchor_embeddings.shape[1]})"
+            )
+
+    # ── Transform ──
+
+    def transform(
+        self,
+        model_name: str,
+        embeddings: np.ndarray,
+        role: str = "auto",
+    ) -> np.ndarray:
+        """Transform embeddings to relative representation.
+
+        Parameters
+        ----------
+        model_name : which model produced these embeddings
+        embeddings : (N, D) L2-normalized
+        role : "query" (skip z-score), "db" (apply z-score per normalize setting),
+               "auto" (treat as db)
+
+        Returns
+        -------
+        (N, K) relative representation
+        """
+        if model_name not in self._models:
+            raise RuntimeError(
+                f"Model '{model_name}' not registered. Call set_anchors() first."
+            )
+        if role not in ("query", "db", "auto"):
+            raise ValueError(f"role must be 'query', 'db', or 'auto', got '{role}'")
+
+        _check_l2_normalized(embeddings)
+
+        state = self._models[model_name]
+        anchors = state["anchors"]
+        rel = self._kernel_fn(embeddings, anchors)  # (N, K)
+
+        if self._should_apply_zscore(model_name, role):
+            rel = normalize_zscore(rel)
+
+        return rel
+
+    def _should_apply_zscore(self, model_name: str, role: str) -> bool:
+        """Decide whether to apply z-score for this model/role combination."""
+        if role == "query":
+            return False
+
+        # role is "db" or "auto" → check normalize setting
+        state = self._models[model_name]
+        sim_mean = state["sim_mean"]
+
+        if self._normalize == "never":
+            return False
+        if self._normalize == "always":
+            return True
+
+        # "auto": use recommendation
+        rec = recommend_zscore(sim_mean, self._normalize_threshold, self._normalize_harmful_threshold)
+        should = rec == "recommended"
+
+        if self._verbose:
+            if rec == "harmful":
+                print(
+                    f"[RAT] {model_name} sim_mean={sim_mean:.3f} >= "
+                    f"harmful_threshold {self._normalize_harmful_threshold} "
+                    f"→ z-score: harmful (skipping)"
+                )
+            elif should:
+                print(
+                    f"[RAT] {model_name} sim_mean={sim_mean:.3f} >= "
+                    f"{self._normalize_threshold} → applying z-score"
+                )
+            else:
+                print(
+                    f"[RAT] {model_name} sim_mean={sim_mean:.3f} < "
+                    f"{self._normalize_threshold} → z-score: not needed"
+                )
+
+        return should
+
+    # ── Retrieval ──
+
+    def retrieve(
+        self,
+        query_emb: np.ndarray,
+        db_emb: np.ndarray,
+        from_model: str,
+        to_model: str,
+        top_k: int = 10,
+    ) -> dict:
+        """Transform + cosine similarity retrieval.
+
+        Parameters
+        ----------
+        query_emb : (N, D_from) query embeddings from from_model
+        db_emb : (M, D_to) database embeddings from to_model
+        from_model : model name for queries
+        to_model : model name for database
+        top_k : number of results per query
+
+        Returns
+        -------
+        {"indices": (N, top_k), "scores": (N, top_k)}
+        """
+        q_rel = self.transform(from_model, query_emb, role="query")
+        d_rel = self.transform(to_model, db_emb, role="db")
+
+        # L2-normalize relative representations for cosine similarity
+        q_norm = q_rel / (np.linalg.norm(q_rel, axis=1, keepdims=True) + 1e-10)
+        d_norm = d_rel / (np.linalg.norm(d_rel, axis=1, keepdims=True) + 1e-10)
+
+        scores = q_norm @ d_norm.T  # (N, M)
+        top_k = min(top_k, scores.shape[1])
+        indices = np.argsort(-scores, axis=1)[:, :top_k]
+        sorted_scores = np.take_along_axis(scores, indices, axis=1)
+
+        return {"indices": indices, "scores": sorted_scores}
+
+    # ── Compatibility ──
+
+    def estimate_compatibility(
+        self,
+        model_a: str,
+        model_b: str,
+    ) -> dict:
+        """Estimate RAT compatibility between two registered models.
+
+        Returns
+        -------
+        dict with sim_mean_a, sim_mean_b, max_sim_mean, estimated_recall_at_1,
+        z_score_recommendation, warnings
+        """
+        if model_a not in self._models:
+            raise RuntimeError(f"Model '{model_a}' not registered.")
+        if model_b not in self._models:
+            raise RuntimeError(f"Model '{model_b}' not registered.")
+
+        sm_a = self._models[model_a]["sim_mean"]
+        sm_b = self._models[model_b]["sim_mean"]
+        max_sm = max(sm_a, sm_b)
+        rec = recommend_zscore(max_sm, self._normalize_threshold, self._normalize_harmful_threshold)
+
+        warn_list: list[str] = []
+        if max_sm > self._normalize_harmful_threshold:
+            warn_list.append(
+                f"High sim_mean ({max_sm:.3f}) detected. "
+                "Z-score normalization may be harmful for this pair."
+            )
+
+        return {
+            "sim_mean_a": sm_a,
+            "sim_mean_b": sm_b,
+            "max_sim_mean": max_sm,
+            "estimated_recall_at_1": None,  # Phase 2
+            "z_score_recommendation": rec,
+            "warnings": warn_list,
+        }
+
+    # ── Persistence ──
+
+    def save(self, path: str) -> None:
+        """Save hub state to .npz file."""
+        save_dict: dict[str, object] = {}
+
+        model_names = list(self._models.keys())
+        for name in model_names:
+            state = self._models[name]
+            save_dict[f"anchor_{name}"] = state["anchors"]
+
+        config = {
+            "kernel": self._kernel_name,
+            "kernel_params": self._kernel_params,
+            "normalize": self._normalize,
+            "normalize_threshold": self._normalize_threshold,
+            "normalize_harmful_threshold": self._normalize_harmful_threshold,
+            "model_names": model_names,
+            "sim_means": {n: self._models[n]["sim_mean"] for n in model_names},
+            "version": "0.1.0",
+        }
+        save_dict["config"] = np.array(json.dumps(config))
+
+        np.savez_compressed(path, **save_dict)
+
+    @classmethod
+    def load(cls, path: str) -> "RATHub":
+        """Load hub state from .npz file."""
+        if not path.endswith(".npz"):
+            path = path + ".npz"
+        data = np.load(path, allow_pickle=False)
+        config = json.loads(str(data["config"]))
+
+        hub = cls(
+            kernel=config["kernel"],
+            kernel_params=config.get("kernel_params"),
+            normalize=config["normalize"],
+            normalize_threshold=config["normalize_threshold"],
+            normalize_harmful_threshold=config.get("normalize_harmful_threshold", 0.65),
+        )
+
+        for name in config["model_names"]:
+            anchors = data[f"anchor_{name}"]
+            hub.set_anchors(name, anchors)
+            # Override sim_mean from saved config (avoids recomputation drift)
+            hub._models[name]["sim_mean"] = config["sim_means"][name]
+
+        return hub
+
+
+def _check_l2_normalized(embeddings: np.ndarray) -> None:
+    """Warn if embeddings don't appear L2-normalized."""
+    norms = np.linalg.norm(embeddings, axis=1)
+    if not np.allclose(norms, 1.0, atol=0.01):
+        warnings.warn(
+            "Input embeddings may not be L2-normalized. RAT expects normalized embeddings.",
+            stacklevel=3,
+        )

rat_embed-0.1.0/rat/kernels.py

@@ -0,0 +1,97 @@
+"""Kernel functions for computing relative representations."""
+
+from __future__ import annotations
+
+from functools import partial
+
+import numpy as np
+
+
+def poly_kernel(
+    X: np.ndarray,
+    A: np.ndarray,
+    degree: int = 2,
+    coef0: float = 1.0,
+) -> np.ndarray:
+    """Polynomial kernel: (X @ A.T + coef0) ** degree.
+
+    Parameters
+    ----------
+    X : (N, D) input embeddings
+    A : (K, D) anchor embeddings
+    degree : polynomial degree
+    coef0 : constant term
+
+    Returns
+    -------
+    (N, K) kernel matrix
+    """
+    return (X @ A.T + coef0) ** degree
+
+
+def cosine_kernel(
+    X: np.ndarray,
+    A: np.ndarray,
+) -> np.ndarray:
+    """Cosine similarity kernel: X @ A.T (assumes L2-normalized inputs).
+
+    Parameters
+    ----------
+    X : (N, D) input embeddings, L2-normalized
+    A : (K, D) anchor embeddings, L2-normalized
+
+    Returns
+    -------
+    (N, K) similarity matrix
+    """
+    return X @ A.T
+
+
+def rbf_kernel(
+    X: np.ndarray,
+    A: np.ndarray,
+    gamma: float | None = None,
+) -> np.ndarray:
+    """RBF (Gaussian) kernel: exp(-gamma * ||x - a||^2).
+
+    Parameters
+    ----------
+    X : (N, D) input embeddings
+    A : (K, D) anchor embeddings
+    gamma : kernel width. None uses 1/D
+
+    Returns
+    -------
+    (N, K) kernel matrix
+    """
+    if gamma is None:
+        gamma = 1.0 / X.shape[1]
+    # ||x - a||^2 = ||x||^2 + ||a||^2 - 2 x·a
+    X_sq = np.sum(X ** 2, axis=1, keepdims=True)  # (N, 1)
+    A_sq = np.sum(A ** 2, axis=1, keepdims=True)  # (K, 1)
+    dist_sq = X_sq + A_sq.T - 2.0 * (X @ A.T)    # (N, K)
+    dist_sq = np.maximum(dist_sq, 0.0)  # numerical safety
+    return np.exp(-gamma * dist_sq)
+
+
+def get_kernel(name: str, params: dict | None = None):
+    """Factory: return a callable(X, A) -> np.ndarray for the named kernel.
+
+    Parameters
+    ----------
+    name : "poly", "cosine", or "rbf"
+    params : optional kwargs forwarded to the kernel function
+
+    Returns
+    -------
+    callable(X, A) -> np.ndarray
+    """
+    params = params or {}
+    kernels = {
+        "poly": poly_kernel,
+        "cosine": cosine_kernel,
+        "rbf": rbf_kernel,
+    }
+    if name not in kernels:
+        raise ValueError(f"Unknown kernel '{name}'. Choose from: {list(kernels.keys())}")
+    return partial(kernels[name], **params)