neembed-geoopt 0.1.0__tar.gz

neembed_geoopt-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 taishi-yamasaki
+
+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.

neembed_geoopt-0.1.0/PKG-INFO

@@ -0,0 +1,163 @@
+Metadata-Version: 2.4
+Name: neembed-geoopt
+Version: 0.1.0
+Summary: Fine-tune pretrained sentence embedding models in non-Euclidean spaces with Geoopt.
+Author: taishi-yamasaki
+Maintainer: taishi-yamasaki
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/t-yamsaki/neembed
+Project-URL: Documentation, https://neembed.readthedocs.io/en/latest/
+Project-URL: Repository, https://github.com/t-yamsaki/neembed
+Project-URL: Issues, https://github.com/t-yamsaki/neembed/issues
+Keywords: sentence-embeddings,hyperbolic-embeddings,non-euclidean,poincare,geoopt,sentence-transformers
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+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
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: sentence-transformers
+Requires-Dist: geoopt
+Provides-Extra: dev
+Requires-Dist: pytest>=8; extra == "dev"
+Dynamic: license-file
+
+# neembed
+
+**Fine-tune pretrained sentence embedding models in non-Euclidean spaces with Geoopt.**
+
+[Documentation](https://neembed.readthedocs.io/en/latest/) · [日本語](docs/README_ja.md)
+
+> **Status:** v0.1 is implemented and being prepared for its first public PyPI release. The API is intentionally small and may still evolve before a stable 1.0 release.
+
+`neembed` is a lightweight integration layer between pretrained Sentence Transformer models and manifold-valued representations. It keeps the pretrained encoder intact, optionally projects its Euclidean output, and delegates Poincaré-ball geometry to Geoopt.
+
+```text
+Pretrained Sentence Encoder
+        ↓
+Euclidean sentence embedding
+        ↓
+Projection head (optional)
+        ↓
+Tangent-space representation
+        ↓
+Geoopt manifold map
+        ↓
+Non-Euclidean embedding
+```
+
+## Why non-Euclidean embeddings?
+
+Hierarchical and tree-like relations can be awkward to represent in a flat Euclidean space. Hyperbolic spaces are a natural fit for rapidly expanding structures such as:
+
+- taxonomies and concept hierarchies
+- knowledge graphs
+- hierarchical labels
+- tree-like semantic relations
+
+v0.1 deliberately focuses on one geometry: the Poincaré ball.
+
+## v0.1
+
+The first release includes:
+
+- Sentence Transformers as the pretrained encoder backend
+- Poincaré-ball embeddings through Geoopt
+- optional lower-dimensional tangent-space projection
+- geodesic distance
+- manifold-aware multiple-negatives ranking / InfoNCE-style loss
+- ordinary `AdamW` fine-tuning
+- `encode()` and `distance()` inference helpers
+- `save_pretrained()` / `from_pretrained()` local persistence
+- numerical-stability tests, a runnable example, and a Euclidean baseline experiment
+
+Detailed behavior, assumptions, and API signatures live in the [documentation](https://neembed.readthedocs.io/en/latest/).
+
+## Installation
+
+After `v0.1.0` is published to PyPI:
+
+```bash
+pip install neembed-geoopt
+```
+
+The PyPI distribution is named `neembed-geoopt`; the Python import package remains `neembed`.
+
+For development:
+
+```bash
+git clone https://github.com/t-yamsaki/neembed.git
+cd neembed
+pip install -e ".[dev]"
+```
+
+## Quick start
+
+```python
+from neembed import (
+    ManifoldSentenceTransformer,
+    ManifoldMultipleNegativesRankingLoss,
+    ManifoldTrainer,
+)
+
+model = ManifoldSentenceTransformer(
+    "sentence-transformers/all-MiniLM-L6-v2",
+    manifold="poincare",
+    embedding_dim=64,
+    curvature=1.0,
+)
+
+loss = ManifoldMultipleNegativesRankingLoss(
+    model=model,
+    temperature=0.1,
+)
+
+trainer = ManifoldTrainer(model=model, loss=loss)
+
+train_batches = [
+    (["Shiba Inu", "Siamese cat"], ["dog", "cat"]),
+    (["dog", "cat"], ["mammal", "feline"]),
+]
+
+trainer.fit(train_batches, epochs=1)
+
+embeddings = model.encode(["Shiba Inu", "dog", "mammal"])
+distance = model.distance(embeddings[0], embeddings[1])
+print(float(distance))
+```
+
+Each anchor is paired with the positive at the same batch index. Because off-diagonal candidates become in-batch negatives, avoid duplicate positives within one batch. See the [Training guide](https://neembed.readthedocs.io/en/latest/user_guide/training.html) for the objective and batching details.
+
+## Documentation
+
+The full guide is hosted on Read the Docs:
+
+- [Installation](https://neembed.readthedocs.io/en/latest/getting_started/installation.html)
+- [Quick Start](https://neembed.readthedocs.io/en/latest/getting_started/quickstart.html)
+- [Architecture](https://neembed.readthedocs.io/en/latest/user_guide/architecture.html)
+- [Training](https://neembed.readthedocs.io/en/latest/user_guide/training.html)
+- [Inference](https://neembed.readthedocs.io/en/latest/user_guide/inference.html)
+- [Saving and Loading](https://neembed.readthedocs.io/en/latest/user_guide/saving_loading.html)
+- [API Reference](https://neembed.readthedocs.io/en/latest/#api-reference)
+
+## Examples and validation
+
+Run the end-to-end example:
+
+```bash
+python examples/train_poincare.py
+```
+
+- [examples/train_poincare.py](examples/train_poincare.py) contains the workflow.
+- [experiments/README.md](experiments/README.md) documents the fixed Euclidean-vs-Poincaré comparison experiment and its interpretation limits.
+
+## License
+
+`neembed` is released under the [MIT License](LICENSE). Third-party dependencies retain their own licenses.

neembed_geoopt-0.1.0/README.md

@@ -0,0 +1,132 @@
+# neembed
+
+**Fine-tune pretrained sentence embedding models in non-Euclidean spaces with Geoopt.**
+
+[Documentation](https://neembed.readthedocs.io/en/latest/) · [日本語](docs/README_ja.md)
+
+> **Status:** v0.1 is implemented and being prepared for its first public PyPI release. The API is intentionally small and may still evolve before a stable 1.0 release.
+
+`neembed` is a lightweight integration layer between pretrained Sentence Transformer models and manifold-valued representations. It keeps the pretrained encoder intact, optionally projects its Euclidean output, and delegates Poincaré-ball geometry to Geoopt.
+
+```text
+Pretrained Sentence Encoder
+        ↓
+Euclidean sentence embedding
+        ↓
+Projection head (optional)
+        ↓
+Tangent-space representation
+        ↓
+Geoopt manifold map
+        ↓
+Non-Euclidean embedding
+```
+
+## Why non-Euclidean embeddings?
+
+Hierarchical and tree-like relations can be awkward to represent in a flat Euclidean space. Hyperbolic spaces are a natural fit for rapidly expanding structures such as:
+
+- taxonomies and concept hierarchies
+- knowledge graphs
+- hierarchical labels
+- tree-like semantic relations
+
+v0.1 deliberately focuses on one geometry: the Poincaré ball.
+
+## v0.1
+
+The first release includes:
+
+- Sentence Transformers as the pretrained encoder backend
+- Poincaré-ball embeddings through Geoopt
+- optional lower-dimensional tangent-space projection
+- geodesic distance
+- manifold-aware multiple-negatives ranking / InfoNCE-style loss
+- ordinary `AdamW` fine-tuning
+- `encode()` and `distance()` inference helpers
+- `save_pretrained()` / `from_pretrained()` local persistence
+- numerical-stability tests, a runnable example, and a Euclidean baseline experiment
+
+Detailed behavior, assumptions, and API signatures live in the [documentation](https://neembed.readthedocs.io/en/latest/).
+
+## Installation
+
+After `v0.1.0` is published to PyPI:
+
+```bash
+pip install neembed-geoopt
+```
+
+The PyPI distribution is named `neembed-geoopt`; the Python import package remains `neembed`.
+
+For development:
+
+```bash
+git clone https://github.com/t-yamsaki/neembed.git
+cd neembed
+pip install -e ".[dev]"
+```
+
+## Quick start
+
+```python
+from neembed import (
+    ManifoldSentenceTransformer,
+    ManifoldMultipleNegativesRankingLoss,
+    ManifoldTrainer,
+)
+
+model = ManifoldSentenceTransformer(
+    "sentence-transformers/all-MiniLM-L6-v2",
+    manifold="poincare",
+    embedding_dim=64,
+    curvature=1.0,
+)
+
+loss = ManifoldMultipleNegativesRankingLoss(
+    model=model,
+    temperature=0.1,
+)
+
+trainer = ManifoldTrainer(model=model, loss=loss)
+
+train_batches = [
+    (["Shiba Inu", "Siamese cat"], ["dog", "cat"]),
+    (["dog", "cat"], ["mammal", "feline"]),
+]
+
+trainer.fit(train_batches, epochs=1)
+
+embeddings = model.encode(["Shiba Inu", "dog", "mammal"])
+distance = model.distance(embeddings[0], embeddings[1])
+print(float(distance))
+```
+
+Each anchor is paired with the positive at the same batch index. Because off-diagonal candidates become in-batch negatives, avoid duplicate positives within one batch. See the [Training guide](https://neembed.readthedocs.io/en/latest/user_guide/training.html) for the objective and batching details.
+
+## Documentation
+
+The full guide is hosted on Read the Docs:
+
+- [Installation](https://neembed.readthedocs.io/en/latest/getting_started/installation.html)
+- [Quick Start](https://neembed.readthedocs.io/en/latest/getting_started/quickstart.html)
+- [Architecture](https://neembed.readthedocs.io/en/latest/user_guide/architecture.html)
+- [Training](https://neembed.readthedocs.io/en/latest/user_guide/training.html)
+- [Inference](https://neembed.readthedocs.io/en/latest/user_guide/inference.html)
+- [Saving and Loading](https://neembed.readthedocs.io/en/latest/user_guide/saving_loading.html)
+- [API Reference](https://neembed.readthedocs.io/en/latest/#api-reference)
+
+## Examples and validation
+
+Run the end-to-end example:
+
+```bash
+python examples/train_poincare.py
+```
+
+- [examples/train_poincare.py](examples/train_poincare.py) contains the workflow.
+- [experiments/README.md](experiments/README.md) documents the fixed Euclidean-vs-Poincaré comparison experiment and its interpretation limits.
+
+## License
+
+`neembed` is released under the [MIT License](LICENSE). Third-party dependencies retain their own licenses.

neembed_geoopt-0.1.0/pyproject.toml

@@ -0,0 +1,62 @@
+[build-system]
+requires = ["setuptools>=77.0.3"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "neembed-geoopt"
+version = "0.1.0"
+description = "Fine-tune pretrained sentence embedding models in non-Euclidean spaces with Geoopt."
+readme = "README.md"
+requires-python = ">=3.10"
+license = "MIT"
+license-files = ["LICENSE"]
+authors = [
+    { name = "taishi-yamasaki" },
+]
+maintainers = [
+    { name = "taishi-yamasaki" },
+]
+keywords = [
+    "sentence-embeddings",
+    "hyperbolic-embeddings",
+    "non-euclidean",
+    "poincare",
+    "geoopt",
+    "sentence-transformers",
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Operating System :: OS Independent",
+    "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 = [
+    "torch",
+    "sentence-transformers",
+    "geoopt",
+]
+
+[project.urls]
+Homepage = "https://github.com/t-yamsaki/neembed"
+Documentation = "https://neembed.readthedocs.io/en/latest/"
+Repository = "https://github.com/t-yamsaki/neembed"
+Issues = "https://github.com/t-yamsaki/neembed/issues"
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8",
+]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+markers = [
+    "real_stack: loads the real Sentence Transformer and runs slow integration gates",
+]

neembed_geoopt-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

neembed_geoopt-0.1.0/src/neembed/__init__.py

@@ -0,0 +1,11 @@
+"""neembed package."""
+
+from neembed.losses import ManifoldMultipleNegativesRankingLoss
+from neembed.model import ManifoldSentenceTransformer
+from neembed.trainer import ManifoldTrainer
+
+__all__ = [
+    "ManifoldMultipleNegativesRankingLoss",
+    "ManifoldSentenceTransformer",
+    "ManifoldTrainer",
+]

neembed_geoopt-0.1.0/src/neembed/losses.py

@@ -0,0 +1,60 @@
+"""Losses for manifold-valued sentence embeddings."""
+
+from collections.abc import Sequence
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from neembed.model import ManifoldSentenceTransformer
+
+
+class ManifoldMultipleNegativesRankingLoss(nn.Module):
+    """In-batch ranking loss based on manifold geodesic distance.
+
+    Each anchor is paired with the positive at the same batch index. All
+    off-diagonal positive candidates are treated as negatives, so duplicate
+    positives within one batch should be avoided.
+
+    Args:
+        model: Manifold sentence model used to encode anchors and positives.
+        temperature: Positive, finite temperature used to scale distance logits.
+    """
+
+    def __init__(
+        self,
+        model: ManifoldSentenceTransformer,
+        temperature: float = 0.1,
+    ) -> None:
+        super().__init__()
+        if temperature <= 0 or not math.isfinite(temperature):
+            raise ValueError("temperature must be positive and finite")
+
+        self.model = model
+        self.temperature = float(temperature)
+
+    def forward(
+        self,
+        anchors: Sequence[str],
+        positives: Sequence[str],
+    ) -> torch.Tensor:
+        """Return the contrastive loss for aligned anchor-positive pairs.
+
+        Args:
+            anchors: Batch of anchor texts.
+            positives: Batch of positive texts aligned by index with ``anchors``.
+
+        Returns:
+            A scalar cross-entropy loss built from pairwise geodesic distances.
+        """
+        anchor_embeddings = self.model(anchors)
+        positive_embeddings = self.model(positives)
+
+        distances = self.model.manifold.dist(
+            anchor_embeddings[:, None, :],
+            positive_embeddings[None, :, :],
+        )
+        logits = -distances / self.temperature
+        targets = torch.arange(logits.shape[0], device=logits.device)
+        return F.cross_entropy(logits, targets)

neembed_geoopt-0.1.0/src/neembed/manifolds.py

@@ -0,0 +1,15 @@
+"""Manifold integration points for neembed."""
+
+import math
+
+import geoopt
+
+
+def get_manifold(name: str, curvature: float = 1.0) -> geoopt.PoincareBall:
+    """Return the Geoopt manifold used by the minimal neembed v0.1 API."""
+    if name != "poincare":
+        raise ValueError(f"Unsupported manifold: {name}")
+    if curvature <= 0 or not math.isfinite(curvature):
+        raise ValueError("curvature must be positive and finite")
+
+    return geoopt.PoincareBall(c=curvature)

neembed_geoopt-0.1.0/src/neembed/model.py

@@ -0,0 +1,182 @@
+"""Sentence Transformer integration for manifold-valued embeddings."""
+
+from collections.abc import Sequence
+import json
+from pathlib import Path
+from typing import Any
+
+import torch
+from sentence_transformers import SentenceTransformer
+from torch import nn
+
+from neembed.manifolds import get_manifold
+
+
+class ManifoldSentenceTransformer(nn.Module):
+    """Map pretrained sentence embeddings onto a configured manifold.
+
+    Args:
+        model_name_or_path: Sentence Transformer model name or local model path.
+        manifold: Manifold backend name. v0.1 supports only ``"poincare"``.
+        embedding_dim: Optional output dimension for a learned linear projection.
+            If omitted, the encoder embedding dimension is preserved.
+        curvature: Positive, finite Poincaré-ball curvature parameter.
+    """
+
+    def __init__(
+        self,
+        model_name_or_path: str,
+        *,
+        manifold: str = "poincare",
+        embedding_dim: int | None = None,
+        curvature: float = 1.0,
+    ) -> None:
+        super().__init__()
+        self.encoder = SentenceTransformer(model_name_or_path)
+
+        encoder_dim = self.encoder.get_embedding_dimension()
+        if encoder_dim is None:
+            raise ValueError("Sentence Transformer embedding dimension is unknown")
+
+        self._projection_dim = embedding_dim
+        self.projection: nn.Module
+        if embedding_dim is None:
+            self.projection = nn.Identity()
+            self.embedding_dim = encoder_dim
+        else:
+            self.projection = nn.Linear(encoder_dim, embedding_dim)
+            self.embedding_dim = embedding_dim
+        self.projection.to(self.encoder.device)
+
+        self.manifold_name = manifold
+        self.curvature = float(curvature)
+        self.manifold = get_manifold(self.manifold_name, self.curvature)
+        self.manifold.to(self.encoder.device)
+
+    def forward(self, sentences: Sequence[str]) -> torch.Tensor:
+        """Encode a batch and map the embeddings from the origin tangent space.
+
+        Args:
+            sentences: Batch of input texts.
+
+        Returns:
+            A tensor of manifold-valued embeddings with shape
+            ``(batch_size, embedding_dim)``.
+        """
+        features = self.encoder.preprocess(list(sentences))
+        features = {
+            key: value.to(self.encoder.device) if torch.is_tensor(value) else value
+            for key, value in features.items()
+        }
+        encoder_output: dict[str, Any] = self.encoder(features)
+        tangent = self.projection(encoder_output["sentence_embedding"])
+        return self.manifold.expmap0(tangent)
+
+    def encode(
+        self,
+        sentences: str | Sequence[str],
+        *,
+        convert_to_tensor: bool = False,
+    ) -> Any:
+        """Encode text as manifold-valued embeddings for inference.
+
+        Args:
+            sentences: A single text or a sequence of texts.
+            convert_to_tensor: Return a ``torch.Tensor`` instead of a NumPy array.
+
+        Returns:
+            A single embedding with shape ``(embedding_dim,)`` for string input,
+            or a batch with shape ``(batch_size, embedding_dim)`` for sequence
+            input. NumPy arrays are returned by default; tensors are returned when
+            ``convert_to_tensor=True``.
+
+        Notes:
+            Encoding switches the model to evaluation mode and runs under
+            ``torch.inference_mode()``, so returned embeddings do not track
+            gradients.
+        """
+        single_input = isinstance(sentences, str)
+        batch = [sentences] if single_input else list(sentences)
+
+        self.eval()
+        with torch.inference_mode():
+            embeddings = self(batch)
+
+        if single_input:
+            embeddings = embeddings[0]
+        if convert_to_tensor:
+            return embeddings
+        return embeddings.cpu().numpy()
+
+    def distance(self, a: Any, b: Any) -> torch.Tensor:
+        """Return the geodesic distance between two manifold embeddings.
+
+        Args:
+            a: First manifold embedding or array-like value.
+            b: Second manifold embedding or array-like value.
+
+        Returns:
+            A tensor containing the manifold geodesic distance.
+
+        Notes:
+            This is an inference helper. Inputs are moved to the model device and
+            dtype, and the distance is computed under ``torch.no_grad()``.
+        """
+        reference = next(self.parameters())
+        a_tensor = torch.as_tensor(a, device=reference.device, dtype=reference.dtype)
+        b_tensor = torch.as_tensor(b, device=reference.device, dtype=reference.dtype)
+
+        with torch.no_grad():
+            return self.manifold.dist(a_tensor, b_tensor)
+
+    def save_pretrained(self, output_path: str | Path) -> None:
+        """Save the encoder, projection weights, and neembed configuration.
+
+        Args:
+            output_path: Directory in which to save the model.
+        """
+        output_path = Path(output_path)
+        output_path.mkdir(parents=True, exist_ok=True)
+
+        self.encoder.save_pretrained(str(output_path / "encoder"))
+        config = {
+            "embedding_dim": self._projection_dim,
+            "manifold": self.manifold_name,
+            "curvature": self.curvature,
+        }
+        (output_path / "neembed_config.json").write_text(
+            json.dumps(config, indent=2) + "\n",
+            encoding="utf-8",
+        )
+        torch.save(self.projection.state_dict(), output_path / "projection.pt")
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_path: str | Path,
+    ) -> "ManifoldSentenceTransformer":
+        """Load a model previously saved with :meth:`save_pretrained`.
+
+        Args:
+            model_path: Directory containing a saved neembed model.
+
+        Returns:
+            The reconstructed manifold sentence model.
+        """
+        model_path = Path(model_path)
+        config = json.loads(
+            (model_path / "neembed_config.json").read_text(encoding="utf-8")
+        )
+        model = cls(
+            str(model_path / "encoder"),
+            manifold=config["manifold"],
+            embedding_dim=config["embedding_dim"],
+            curvature=config["curvature"],
+        )
+        projection_state = torch.load(
+            model_path / "projection.pt",
+            map_location="cpu",
+            weights_only=True,
+        )
+        model.projection.load_state_dict(projection_state)
+        return model