Metadata-Version: 2.4
Name: neembed-geoopt
Version: 0.11.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>=2.3.0
Requires-Dist: sentence-transformers>=5.4.1
Requires-Dist: geoopt>=0.5.1
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:** Package version v0.11.0 adds reproducible encoder loading, differentiable query/document input roles, batched encoding, versioned checkpoint metadata, and documented public API/benchmark contracts. It retains v0.10.0 ProductManifold/Scaled mixed-curvature support and all v0.4-v0.9 single-manifold contracts. The API remains 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 constant-curvature 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

The current API supports Euclidean, Poincaré, Lorentz / Hyperboloid, SphereProjection, and generic signed-curvature Stereographic geometry through the same sentence-model, objective, evaluator, retrieval, and save/load surface. Geometry choice is task-dependent; neembed does not claim that curvature alone guarantees quality gains.

## Current scope

v0.4 keeps the fixed-curvature v0.3 path backward-compatible:

- Sentence Transformers as the pretrained encoder backend
- Poincaré-ball and Lorentz / Hyperboloid embeddings through Geoopt
- optional lower-dimensional tangent-space projection
- shared public curvature semantics across the two hyperbolic models
- geodesic distance and manifold-aware multiple-negatives ranking loss
- ordinary `AdamW` fine-tuning for the model-only path
- `encode()` / `distance()`, `ManifoldEmbeddingEvaluator`, DataLoader interoperability, and local save/load
- geometry-consistency regressions and a matched Euclidean-vs-Poincaré-vs-Lorentz engineering benchmark

v0.4 also adds:

- opt-in fixed vs learnable curvature for Poincaré and Lorentz
- true trainable manifold-valued `ManifoldPrototypes`
- `ManifoldPrototypeHierarchyLoss` for sentence assignments plus parent-child structure
- an explicit caller-supplied Geoopt Riemannian optimizer path for manifold parameters
- joint learnable-curvature / prototype training through Geoopt stabilization
- a compact fixed-vs-learnable structure regression example

v0.5 adds a focused retrieval workflow without introducing a retrieval framework:

- optional caller-supplied explicit hard negatives while preserving the original `(anchors, positives)` training contract
- Recall@K and MRR on the aligned retrieval evaluator
- `model.rank()` for small in-memory geodesic reranking
- nearest-prototype assignment evaluation for learned manifold structure
- one reproducible v0.5 retrieval regression example

v0.6 extends that workflow to exact caller-owned corpora without adding ANN infrastructure:

- chunked exact geodesic distance evaluation without materializing the full query-by-corpus distance matrix
- `exact_corpus_search()` for exact text-corpus top-k retrieval
- `ManifoldCorpusRetrievalEvaluator` with explicit query/corpus IDs and multi-positive Recall@K / MRR
- deterministic caller-invoked `mine_hard_negatives()` with positive, self, and additional exclusions
- a reproducible v0.6 exact-search → evaluation → mining → explicit-negative training regression example

v0.7 expands the retrieval-objective and evaluation surface without changing the v0.4-v0.6 contracts:

- `ManifoldTripletLoss` for aligned geodesic margin triplets
- `ManifoldMarginMSELoss` for teacher positive-minus-negative margin regression
- `ManifoldDistanceMSELoss` for direct geodesic-distance regression
- opt-in `ManifoldSymmetricMultipleNegativesRankingLoss` while keeping the original one-directional MNRL unchanged
- `ManifoldGradedCorpusRetrievalEvaluator` with nDCG@K while preserving the existing binary/multi-positive evaluator output
- a deterministic Poincaré objective-comparison example covering MNRL, Triplet, MarginMSE, DistanceMSE, MRR, Recall@K, and nDCG@K

v0.8 adds explicit hierarchy-native learning without introducing a graph framework:

- caller-owned node IDs, parent-child edges, optional depths, and directed negatives
- `ManifoldRadialOrderLoss`, `ManifoldDepthLoss`, and `ManifoldHierarchyTripletLoss`
- `ManifoldRetrievalHierarchyLoss` for one retrieval objective plus one hierarchy objective
- `ManifoldHierarchyEvaluator` for radial-order and depth-vs-radius structure diagnostics
- a deterministic Poincaré retrieval-only vs hierarchy-aware regression example

v0.9 adds constant-curvature geometry selection without changing the legacy hyperbolic curvature meaning:

- Euclidean as the explicit flat `K=0` backend
- SphereProjection for fixed positive sectional curvature
- generic Stereographic for fixed signed sectional curvature (`K<0`, `K=0`, or `K>0`)
- legacy Poincaré/Lorentz `curvature=c>0` remains the magnitude of `K=-c`
- float64 geometry operations for SphereProjection/Stereographic, including CPU geometry fallback when the encoder runs on Apple MPS
- a deterministic matched constant-curvature comparison example for engineering regression, not geometry ranking

v0.10 adds flat mixed-curvature products with fixed component scales, full-product retrieval, named distance diagnostics, and optional hierarchy supervision on a designated Poincaré/Lorentz component. See the [Mixed-curvature product guide](https://neembed.readthedocs.io/en/latest/user_guide/mixed_curvature.html) for configuration, compatibility, optimizer/persistence behavior, and the regression example. More geometry does not guarantee better task metrics.

Nested products, learnable product curvature/scales, learnable signed curvature crossing zero, advanced SPD/Siegel/Stiefel/Birkhoff geometry, ontology parsing, graph-database integration, ANN/vector-database integration, and distributed retrieval remain outside this scope.

A manifold-valued **output** does not by itself require Riemannian optimization: encoder/projection parameters and learnable curvature are not manifold-valued points. Detailed parameter, optimizer, persistence, and numerical behavior lives in the [Learnable structure guide](https://neembed.readthedocs.io/en/latest/user_guide/learnable_structure.html). The end-to-end retrieval composition and the distinction between small in-memory reranking, exact corpus search, and external ANN systems are documented in the [Retrieval workflow guide](https://neembed.readthedocs.io/en/latest/user_guide/retrieval.html). Objective and graded-evaluation selection for v0.7 is documented in the [Retrieval objectives guide](https://neembed.readthedocs.io/en/latest/user_guide/retrieval_objectives.html). Explicit v0.8 hierarchy supervision, origin/radius semantics, composition, and structure metrics are documented in the [Hierarchy-native learning guide](https://neembed.readthedocs.io/en/latest/user_guide/hierarchy.html). v0.9 geometry roles, curvature naming, compatibility, dtype/device behavior, and matched-comparison guidance are documented in the [Constant-curvature geometry guide](https://neembed.readthedocs.io/en/latest/user_guide/constant_curvature_semantics.html).

## Installation

```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. This model-only path keeps the ordinary AdamW behavior even though its outputs lie on a manifold. See the [Training guide](https://neembed.readthedocs.io/en/latest/user_guide/training.html) for the objective and batching details, the [Retrieval workflow guide](https://neembed.readthedocs.io/en/latest/user_guide/retrieval.html) for optional explicit negatives and retrieval evaluation, the [Retrieval objectives guide](https://neembed.readthedocs.io/en/latest/user_guide/retrieval_objectives.html) for v0.7 objective/metric selection, the [Hierarchy-native learning guide](https://neembed.readthedocs.io/en/latest/user_guide/hierarchy.html) for v0.8 explicit hierarchy supervision, the [Constant-curvature geometry guide](https://neembed.readthedocs.io/en/latest/user_guide/constant_curvature_semantics.html) before switching v0.9 geometry/curvature semantics, and the [Learnable structure guide](https://neembed.readthedocs.io/en/latest/user_guide/learnable_structure.html) before adding trainable manifold prototypes.

## Role-aware training and search

After installing the development checkout, download the two pinned encoders once
with `HF_HUB_DISABLE_XET=1 python scripts/preload_role_models.py`. The following
example then runs from the repository root with cached files only. E5-small-v2
uses English text with explicit query/passage instructions; pass raw texts.

<!-- v11-role-example -->
```python
import torch
from neembed import ManifoldMultipleNegativesRankingLoss, ManifoldSentenceTransformer, exact_corpus_search

model = ManifoldSentenceTransformer(
    "intfloat/e5-small-v2", revision="ffb93f3bd4047442299a41ebb6fa998a38507c52",
    local_files_only=True, device="cpu", embedding_dim=16,
)
model.encoder.prompts.update({"query": "query: ", "document": "passage: "})
queries = ["What is a dog?", "Which animals can fly?"]
documents = ["A dog is a mammal.", "Birds can fly.", "Trees are plants."]

loss_fn = ManifoldMultipleNegativesRankingLoss(model, input_options={
    "anchors": {"task": "query"}, "positives": {"task": "document"},
})
model.train()
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
optimizer.zero_grad()
loss = loss_fn(queries, documents[:2])
loss.backward()
optimizer.step()

embeddings = model.encode_query(queries, batch_size=1, convert_to_tensor=True)
results = exact_corpus_search(model, queries, documents, top_k=2, input_options={
    "queries": {"task": "query"}, "corpus": {"task": "document"},
})
model.save_pretrained("./saved_role_model")
restored = ManifoldSentenceTransformer.from_pretrained("./saved_role_model", local_files_only=True, device="cpu")
torch.testing.assert_close(restored.encode_query(queries, batch_size=1, convert_to_tensor=True), embeddings)
print(results)
```
<!-- /v11-role-example -->

`forward` and losses retain the caller's gradient context and model mode.
`encode` switches to eval mode and returns CPU NumPy arrays by default, or
geometry-device tensors with `convert_to_tensor=True`; call `model.train()`
when resuming training. `distance` is a no-grad helper with broadcast Tensor
output and leaves mode unchanged. Use `model.manifold.dist` on forward outputs
for differentiable distances. Rankings put the smallest geodesic distance first;
equal distances keep candidate/corpus order.

For intrinsic dimension `D`, single geometry outputs have width `D`, except
Lorentz's `D + 1`. Product width is `product_config.ambient_dim`, the sum of
component intrinsic dimensions plus one for each Lorentz component. Lorentz,
SphereProjection, Stereographic, and fixed-double products map/output float64;
Poincare/Euclidean retain ordinary dtype. Tensors follow the geometry device
(SphereProjection/Stereographic and fixed-double products fall back to CPU on
MPS); NumPy outputs use CPU. NumPy conversion must support the output dtype:
bfloat16 requires Tensor output. These policies do not qualify every device/dtype.

Unspecified roles keep raw-text behavior even with a saved default prompt.
Versionless v0.10-and-earlier checkpoints still load; new format-version 1
metadata preserves prompts, geometry, dtype and available provenance. Per-call
options and batch size, optimizer state, and external prototypes are not saved
as defaults. Float64 encoders include supplemental exact weights. See the
[public API contracts](https://neembed.readthedocs.io/en/latest/user_guide/api_contracts.html)
and [compatibility guide](https://neembed.readthedocs.io/en/latest/user_guide/encoder_compatibility.html).

## Documentation

Local saves now use versioned metadata for geometry, prompts, floating dtypes,
package version, and available base-model provenance. Versionless v0.10 and
earlier configurations remain loadable; restoration uses the saved local encoder.

The minimum qualified dependencies are PyTorch 2.3.0, Sentence Transformers
5.4.1, and Geoopt 0.5.1. A two-profile CPU CI verifies pinned MiniLM and E5-small-v2
encoders, including role-aware training, retrieval, and offline save/load. See
the [compatibility guide](https://neembed.readthedocs.io/en/latest/user_guide/encoder_compatibility.html)
for exact stacks, model revisions, and tested scope.

`encode()`, `encode_query()`, and `encode_document()` accept `batch_size=32`
by default, preserving input order and geometry dtype/device. Set a smaller
positive integer to bound each encoder call; empty sequences return an empty
matrix with the configured output width. Encoding leaves the model in eval mode.

`forward()` and `encode()` accept `task="query"` / `"document"`, `prompt_name`,
or an explicit `prompt`. `encode_query()` and `encode_document()` share the same
prompt selection with the differentiable training path. Calls without these
options retain raw-text behavior. See the [Inference guide](https://neembed.readthedocs.io/en/latest/user_guide/inference.html)
for precedence, prompt-aware pooling, routing support, and current integration limits.

Built-in losses accept `input_options`, mapping their text input names to the
same task/prompt options. Configure anchors as queries and positives/explicit
negatives as documents while keeping existing training batches. See the
[Retrieval objectives guide](https://neembed.readthedocs.io/en/latest/user_guide/retrieval_objectives.html)
for examples and independent hierarchy/composite input settings.

`rank()`, exact corpus search, evaluators, and offline hard-negative mining also
accept `input_options`. Reuse the same query/document settings as training,
using each API's text input names. See the
[Retrieval workflow guide](https://neembed.readthedocs.io/en/latest/user_guide/retrieval.html)
for a complete example and the input-name mapping.

Encoder loading accepts `revision`, `local_files_only`, `cache_folder`, and
`device` in both the constructor and local `from_pretrained()`. Pin a remote
encoder with a commit ID, or use `local_files_only=True` for local/cached files.
See the [Saving and Loading guide](https://neembed.readthedocs.io/en/latest/user_guide/saving_loading.html)
for examples and the distinction between encoder revision and local checkpoints.

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)
- [Constant-curvature geometry](https://neembed.readthedocs.io/en/latest/user_guide/constant_curvature_semantics.html)
- [Learnable structure](https://neembed.readthedocs.io/en/latest/user_guide/learnable_structure.html)
- [Training](https://neembed.readthedocs.io/en/latest/user_guide/training.html)
- [Retrieval workflow](https://neembed.readthedocs.io/en/latest/user_guide/retrieval.html)
- [Retrieval objectives](https://neembed.readthedocs.io/en/latest/user_guide/retrieval_objectives.html)
- [Hierarchy-native learning](https://neembed.readthedocs.io/en/latest/user_guide/hierarchy.html)
- [Evaluation](https://neembed.readthedocs.io/en/latest/user_guide/evaluation.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 main references from the repository root:

```bash
python examples/train_poincare.py
python examples/train_lorentz.py
python examples/v04_learnable_structure.py
python examples/v05_retrieval_workflow.py
python examples/v06_exact_retrieval_workflow.py
python examples/v07_objective_comparison.py
python examples/v08_hierarchy_learning.py
python examples/v09_constant_curvature_comparison.py
python examples/v10_mixed_curvature_workflow.py
```

- [examples/train_poincare.py](examples/train_poincare.py) — minimal Poincaré workflow
- [examples/train_lorentz.py](examples/train_lorentz.py) — Lorentz train/evaluate/inference with intrinsic-vs-ambient dimensions
- [examples/train_dataloader.py](examples/train_dataloader.py) — ordinary PyTorch `DataLoader` training and epoch validation
- [examples/train_hierarchy.py](examples/train_hierarchy.py) — minimal hierarchy-aware prototype objective
- [examples/v04_learnable_structure.py](examples/v04_learnable_structure.py) — fixed-vs-learnable structure regression diagnostics; not a superiority benchmark
- [examples/v05_retrieval_workflow.py](examples/v05_retrieval_workflow.py) — explicit hard negatives, Recall@K / MRR, in-memory geodesic reranking, and prototype assignment in one Poincaré regression workflow; not a research benchmark
- [examples/v06_exact_retrieval_workflow.py](examples/v06_exact_retrieval_workflow.py) — exact corpus search, explicit-ID corpus evaluation, offline hard-negative mining, and the existing three-sequence trainer in one Poincaré regression workflow; not a research benchmark
- [examples/v07_objective_comparison.py](examples/v07_objective_comparison.py) — deterministic comparison of MNRL, Triplet, MarginMSE, and DistanceMSE under fixed data/initialization with MRR, Recall@K, and nDCG@K diagnostics; not a research benchmark or superiority claim
- [examples/v08_hierarchy_learning.py](examples/v08_hierarchy_learning.py) — deterministic retrieval-only vs hierarchy-aware Poincaré regression with explicit caller-owned hierarchy supervision and separate retrieval/structure diagnostics; not a benchmark or superiority claim
- [examples/v09_constant_curvature_comparison.py](examples/v09_constant_curvature_comparison.py) — deterministic matched comparison of Euclidean, Poincaré, Lorentz, SphereProjection, and signed Stereographic geometry with retrieval and dtype/device diagnostics; engineering regression only, not a benchmark or geometry-superiority claim
- [examples/v10_mixed_curvature_workflow.py](examples/v10_mixed_curvature_workflow.py) — fixed-scale Poincaré × SphereProjection × Euclidean retrieval plus named hyperbolic-component hierarchy supervision, component distances, and save/load validation; not a benchmark or mixed-curvature superiority claim
- [experiments/README.md](experiments/README.md) — reproducible Euclidean-vs-Poincaré-vs-Lorentz engineering benchmark and interpretation limits

## License

`neembed` is released under the [MIT License](LICENSE). Third-party dependencies retain their own licenses.
